[AMD] DeepSeek-V4: add aiter fused mHC post+pre with cross-layer boundary dispatch (#32577)
Co-authored-by: 1am9trash <1am9trash@gmail.com> Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
karverma-amd
co-authored by
1am9trash
HAI
1am9trash
HAI
parent
db570fe619
commit
d315eb7250
@@ -4,13 +4,55 @@ from typing import Optional, Tuple
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import get_bool_env_var, is_gfx95_supported, is_hip
|
||||
from sglang.srt.utils.common import is_sm120_supported
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_is_hip = is_hip()
|
||||
_is_gfx95_supported = is_gfx95_supported()
|
||||
|
||||
_FUSED_HC_POST_PRE_M_THRESHOLD = 64
|
||||
_FUSED_HC_POST_PRE_CACHE: dict[tuple, dict[str, torch.Tensor]] = {}
|
||||
_TRITON_MHC_POST_PRE_OPS = None
|
||||
_TRITON_MHC_POST_PRE_RUNTIME_DISABLED = False
|
||||
|
||||
# aiter fused mHC state. The kernel self-gates the fused-vs-unfused decision per
|
||||
# arch internally (see aiter.ops.mhc.mhc_fused_post_pre), so this wrapper only
|
||||
# tracks import/runtime availability and never re-implements the token threshold.
|
||||
_AITER_MHC_FUSED_POST_PRE_RUNTIME_DISABLED = False
|
||||
_AITER_MHC_IMPORT_WARNED = False
|
||||
|
||||
|
||||
def _is_fused_mhc_post_pre_enabled() -> bool:
|
||||
# SM120 disables the standalone TileLang pre path. mhc_fused_post_pre does
|
||||
# not read that flag and dispatches independently for both small and large
|
||||
# token batches, so the standalone pre flag must not veto the fused opt-in.
|
||||
return (
|
||||
envs.SGLANG_OPT_FUSE_MHC_POST_PRE.get()
|
||||
and envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get()
|
||||
and (envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get() or is_sm120_supported())
|
||||
)
|
||||
|
||||
|
||||
def _is_aiter_gfx95_mhc_available() -> bool:
|
||||
return _is_hip and get_bool_env_var("SGLANG_USE_AITER") and is_gfx95_supported()
|
||||
|
||||
|
||||
def _is_production_mhc_enabled() -> bool:
|
||||
return _is_fused_mhc_post_pre_enabled() or _is_aiter_gfx95_mhc_available()
|
||||
|
||||
|
||||
def is_cross_layer_mhc_fusion_enabled() -> bool:
|
||||
"""Whether DeepSeek-V4 may defer mHC post across the attn/MoE boundary.
|
||||
|
||||
Cross-layer fusion requires a fused post+pre kernel to be available: either
|
||||
the TileLang path (``SGLANG_OPT_FUSE_MHC_POST_PRE`` + TileLang pre/post) or
|
||||
the aiter HIP path on a supported gfx95 device.
|
||||
"""
|
||||
return _is_production_mhc_enabled()
|
||||
|
||||
|
||||
def _get_triton_mhc_post_pre_ops():
|
||||
global _TRITON_MHC_POST_PRE_OPS
|
||||
@@ -153,3 +195,220 @@ def try_fused_hc_post_pre(
|
||||
return None
|
||||
|
||||
return new_residual, layer_input_out, bufs["h_post"], bufs["h_res"], False
|
||||
|
||||
|
||||
def try_aiter_fused_mhc_post_pre(
|
||||
layer_input: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
post: torch.Tensor,
|
||||
comb: torch.Tensor,
|
||||
hc_fn: torch.Tensor,
|
||||
hc_scale: torch.Tensor,
|
||||
hc_base: torch.Tensor,
|
||||
rms_eps: float,
|
||||
hc_eps: float,
|
||||
hc_post_mult: float,
|
||||
sinkhorn_iters: int,
|
||||
norm_weight: Optional[torch.Tensor],
|
||||
norm_eps: Optional[float],
|
||||
) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool]]:
|
||||
"""Fused mhc_post + next-layer mhc_pre via the aiter HIP kernel.
|
||||
|
||||
Returns ``(next_residual, layer_input, post_mix, comb_mix, norm_applied)`` or
|
||||
``None`` to let the caller fall back. The aiter kernel internally chooses
|
||||
between its fused and unfused (mhc_post + mhc_pre) implementations based on the
|
||||
token count and detected arch, so no token threshold is applied here.
|
||||
"""
|
||||
global _AITER_MHC_FUSED_POST_PRE_RUNTIME_DISABLED, _AITER_MHC_IMPORT_WARNED
|
||||
|
||||
if (
|
||||
_AITER_MHC_FUSED_POST_PRE_RUNTIME_DISABLED
|
||||
or not _is_aiter_gfx95_mhc_available()
|
||||
or layer_input.shape[0] == 0
|
||||
or layer_input.dim() != 2
|
||||
or residual.dim() != 3
|
||||
):
|
||||
return None
|
||||
|
||||
try:
|
||||
from aiter.ops.mhc import mhc_fused_post_pre
|
||||
except Exception as err:
|
||||
if not _AITER_MHC_IMPORT_WARNED:
|
||||
logger.warning("aiter fused mHC is unavailable, falling back: %s", err)
|
||||
_AITER_MHC_IMPORT_WARNED = True
|
||||
_AITER_MHC_FUSED_POST_PRE_RUNTIME_DISABLED = True
|
||||
return None
|
||||
|
||||
norm_kwargs = {}
|
||||
if norm_weight is not None:
|
||||
norm_kwargs["norm_weight"] = norm_weight
|
||||
norm_kwargs["norm_eps"] = norm_eps if norm_eps is not None else rms_eps
|
||||
|
||||
try:
|
||||
post_mix, comb_mix, layer_input_out, next_residual = mhc_fused_post_pre(
|
||||
layer_input=layer_input,
|
||||
residual_in=residual,
|
||||
post_layer_mix=post,
|
||||
comb_res_mix=comb,
|
||||
fn=hc_fn,
|
||||
hc_scale=hc_scale,
|
||||
hc_base=hc_base,
|
||||
rms_eps=rms_eps,
|
||||
hc_pre_eps=hc_eps,
|
||||
hc_sinkhorn_eps=hc_eps,
|
||||
hc_post_mult_value=hc_post_mult,
|
||||
sinkhorn_repeat=sinkhorn_iters,
|
||||
**norm_kwargs,
|
||||
)
|
||||
except Exception as err:
|
||||
logger.warning(
|
||||
"aiter fused mHC kernel failed, disabling fallback path: %s", err
|
||||
)
|
||||
_AITER_MHC_FUSED_POST_PRE_RUNTIME_DISABLED = True
|
||||
return None
|
||||
|
||||
post_out = post_mix.squeeze(-1) if post_mix.ndim == 3 else post_mix
|
||||
return (
|
||||
next_residual,
|
||||
layer_input_out,
|
||||
post_out,
|
||||
comb_mix,
|
||||
norm_weight is not None,
|
||||
)
|
||||
|
||||
|
||||
def try_mhc_fused_post_pre_boundary(
|
||||
layer_input: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
post: torch.Tensor,
|
||||
comb: torch.Tensor,
|
||||
hc_fn: torch.Tensor,
|
||||
hc_scale: torch.Tensor,
|
||||
hc_base: torch.Tensor,
|
||||
hc_mult: int,
|
||||
rms_eps: float,
|
||||
hc_eps: float,
|
||||
hc_post_mult: float,
|
||||
sinkhorn_iters: int,
|
||||
norm_weight: Optional[torch.Tensor],
|
||||
norm_eps: Optional[float],
|
||||
fn_transpose: bool,
|
||||
is_gfx95_supported_flag: bool,
|
||||
) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool]]:
|
||||
"""Dispatch the fused mHC post+pre across the attn/MoE boundary.
|
||||
|
||||
Preference order (first available wins): aiter HIP kernel, then the Triton
|
||||
kernel. Returns ``None`` when neither fires so the caller can fall back to the
|
||||
TileLang path or the unfused ``hc_post`` + ``hc_pre`` sequence.
|
||||
|
||||
The aiter and Triton kernels expect opposite ``fn`` orientations: the Triton
|
||||
kernel takes ``hc_fn`` transposed (``fn_transpose``), while aiter consumes it
|
||||
in the native ``mhc_pre`` layout.
|
||||
"""
|
||||
aiter_result = try_aiter_fused_mhc_post_pre(
|
||||
layer_input,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
hc_fn,
|
||||
hc_scale,
|
||||
hc_base,
|
||||
rms_eps,
|
||||
hc_eps,
|
||||
hc_post_mult,
|
||||
sinkhorn_iters,
|
||||
norm_weight,
|
||||
norm_eps,
|
||||
)
|
||||
if aiter_result is not None:
|
||||
return aiter_result
|
||||
|
||||
triton_fn = hc_fn.T if fn_transpose else hc_fn
|
||||
return try_fused_hc_post_pre(
|
||||
layer_input,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
triton_fn,
|
||||
hc_scale,
|
||||
hc_base,
|
||||
hc_mult,
|
||||
rms_eps,
|
||||
hc_eps,
|
||||
hc_post_mult,
|
||||
sinkhorn_iters,
|
||||
is_gfx95_supported_flag,
|
||||
)
|
||||
|
||||
|
||||
def apply_mhc_post_pre_boundary(
|
||||
layer_input: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
post: torch.Tensor,
|
||||
comb: torch.Tensor,
|
||||
hc_fn: torch.Tensor,
|
||||
hc_scale: torch.Tensor,
|
||||
hc_base: torch.Tensor,
|
||||
hc_mult: int,
|
||||
rms_eps: float,
|
||||
hc_eps: float,
|
||||
hc_post_mult: float,
|
||||
sinkhorn_iters: int,
|
||||
norm_weight: Optional[torch.Tensor],
|
||||
norm_eps: Optional[float],
|
||||
*,
|
||||
fn_transpose: bool,
|
||||
) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool]]:
|
||||
# Try the aiter/Triton fused post+pre kernels first; if neither fires,
|
||||
# fall back to the TileLang fused kernel, else return None so the caller
|
||||
# runs the unfused hc_post + hc_pre sequence.
|
||||
fused = try_mhc_fused_post_pre_boundary(
|
||||
layer_input,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
hc_fn,
|
||||
hc_scale,
|
||||
hc_base,
|
||||
hc_mult,
|
||||
rms_eps,
|
||||
hc_eps,
|
||||
hc_post_mult,
|
||||
sinkhorn_iters,
|
||||
norm_weight,
|
||||
norm_eps,
|
||||
fn_transpose,
|
||||
_is_gfx95_supported,
|
||||
)
|
||||
if fused is not None:
|
||||
return fused
|
||||
|
||||
if not _is_fused_mhc_post_pre_enabled():
|
||||
return None
|
||||
|
||||
from sglang.srt.models.deepseek_v4 import _get_mhc_ops
|
||||
|
||||
post_in = post.unsqueeze(-1) if post.ndim == 2 else post
|
||||
(
|
||||
residual,
|
||||
post_out,
|
||||
comb_out,
|
||||
layer_input_out,
|
||||
) = _get_mhc_ops().mhc_fused_post_pre(
|
||||
layer_input,
|
||||
residual,
|
||||
post_in,
|
||||
comb,
|
||||
hc_fn,
|
||||
hc_scale,
|
||||
hc_base,
|
||||
rms_eps,
|
||||
hc_eps,
|
||||
hc_eps,
|
||||
hc_post_mult,
|
||||
sinkhorn_iters,
|
||||
norm_weight=norm_weight,
|
||||
norm_eps=norm_eps,
|
||||
)
|
||||
post_out = post_out.squeeze(-1) if post_out.ndim == 3 else post_out
|
||||
return residual, layer_input_out, post_out, comb_out, True
|
||||
|
||||
@@ -138,7 +138,8 @@ from sglang.srt.model_loader.weight_utils import (
|
||||
)
|
||||
from sglang.srt.models.dbrx import ReplicatedLinear
|
||||
from sglang.srt.models.deepseek_common.amd.deepseek_v4_fused_mhc import (
|
||||
try_fused_hc_post_pre,
|
||||
apply_mhc_post_pre_boundary,
|
||||
is_cross_layer_mhc_fusion_enabled,
|
||||
)
|
||||
from sglang.srt.models.deepseek_common.utils import (
|
||||
_use_aiter_bpreshuffle_gfx95,
|
||||
@@ -168,7 +169,6 @@ from sglang.srt.utils import (
|
||||
log_info_on_rank0,
|
||||
make_layers,
|
||||
)
|
||||
from sglang.srt.utils.common import is_sm120_supported
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
from sglang.srt.utils.hf_transformers_utils import get_rope_config
|
||||
|
||||
@@ -219,17 +219,6 @@ DEEPSEEK_V4_STACKED_PARAMS_MAPPING: List[Tuple[str, str, int]] = [
|
||||
]
|
||||
|
||||
|
||||
def _is_fused_mhc_post_pre_enabled() -> bool:
|
||||
# SM120 disables the standalone TileLang pre path. mhc_fused_post_pre does
|
||||
# not read that flag and dispatches independently for both small and large
|
||||
# token batches, so the standalone pre flag must not veto the fused opt-in.
|
||||
return (
|
||||
envs.SGLANG_OPT_FUSE_MHC_POST_PRE.get()
|
||||
and envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get()
|
||||
and (envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get() or is_sm120_supported())
|
||||
)
|
||||
|
||||
|
||||
# FlashInfer's mhc_pre_big_fuse only accepts these split-K counts.
|
||||
_FLASHINFER_MHC_PRE_SPLITS = (1, 2, 4, 8, 16)
|
||||
|
||||
@@ -1792,7 +1781,7 @@ class DeepseekV4DecoderLayer(nn.Module):
|
||||
) = make_hc_mixing_params(hc_mult, config.hidden_size)
|
||||
self.rms_norm_eps = config.rms_norm_eps
|
||||
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
|
||||
self.use_fused_mhc_post_pre = _is_fused_mhc_post_pre_enabled()
|
||||
self.use_fused_mhc_post_pre = is_cross_layer_mhc_fusion_enabled()
|
||||
self._input_layernorm_weight_bf16 = None
|
||||
self._post_attention_layernorm_weight_bf16 = None
|
||||
|
||||
@@ -2022,7 +2011,12 @@ class DeepseekV4DecoderLayer(nn.Module):
|
||||
use_fused = self.use_fused_mhc_post_pre
|
||||
|
||||
if prev_residual is not None and use_fused:
|
||||
residual, post, comb, hidden_states = _get_mhc_ops().mhc_fused_post_pre(
|
||||
input_norm_weight = (
|
||||
self._input_layernorm_weight_bf16
|
||||
if self._input_layernorm_weight_bf16 is not None
|
||||
else self.input_layernorm.weight.data
|
||||
)
|
||||
fused = apply_mhc_post_pre_boundary(
|
||||
hidden_states,
|
||||
prev_residual,
|
||||
prev_post,
|
||||
@@ -2030,19 +2024,63 @@ class DeepseekV4DecoderLayer(nn.Module):
|
||||
self.hc_attn_fn,
|
||||
self.hc_attn_scale,
|
||||
self.hc_attn_base,
|
||||
self.hc_mult,
|
||||
self.rms_norm_eps,
|
||||
self.hc_eps,
|
||||
self.hc_eps,
|
||||
_MHC_POST_MULT_VALUE,
|
||||
self.hc_sinkhorn_iters,
|
||||
norm_weight=(
|
||||
self._input_layernorm_weight_bf16
|
||||
if self._input_layernorm_weight_bf16 is not None
|
||||
else self.input_layernorm.weight.data
|
||||
),
|
||||
norm_eps=self.input_layernorm.variance_epsilon,
|
||||
input_norm_weight,
|
||||
self.input_layernorm.variance_epsilon,
|
||||
fn_transpose=False,
|
||||
)
|
||||
x_quant = None
|
||||
if fused is not None:
|
||||
residual, hidden_states, post, comb, norm_fused = fused
|
||||
if not norm_fused:
|
||||
# The Triton fused post+pre returns the layer input WITHOUT
|
||||
# the input layernorm applied (norm_fused=False). Apply it
|
||||
# (fp8-quant on aiter gfx95) before attention, exactly as the
|
||||
# unfused hc_pre path below does; otherwise unnormalized
|
||||
# activations reach self_attn.
|
||||
if _use_aiter and _is_gfx95_supported:
|
||||
x_quant, hidden_states = _fused_rmsnorm_fp8_quant(
|
||||
hidden_states,
|
||||
self.input_layernorm.weight,
|
||||
self.rms_norm_eps,
|
||||
)
|
||||
else:
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
x_quant = None
|
||||
else:
|
||||
x_quant = None
|
||||
else:
|
||||
# Fused dispatch declined: close the previous layer's deferred
|
||||
# mHC post (prev_residual/prev_post/prev_comb) before opening this
|
||||
# layer's pre. Skipping hc_post here would drop the previous-layer
|
||||
# post state and corrupt all subsequent layers.
|
||||
hidden_states = self.hc_post(
|
||||
hidden_states, prev_residual, prev_post, prev_comb
|
||||
)
|
||||
residual = hidden_states
|
||||
hidden_states, post, comb, norm_fused = self.hc_pre(
|
||||
hidden_states,
|
||||
self.hc_attn_fn,
|
||||
self.hc_attn_scale,
|
||||
self.hc_attn_base,
|
||||
norm=self.input_layernorm,
|
||||
forward_batch=forward_batch,
|
||||
)
|
||||
if not norm_fused:
|
||||
if _use_aiter and _is_gfx95_supported:
|
||||
x_quant, hidden_states = _fused_rmsnorm_fp8_quant(
|
||||
hidden_states,
|
||||
self.input_layernorm.weight,
|
||||
self.rms_norm_eps,
|
||||
)
|
||||
else:
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
x_quant = None
|
||||
else:
|
||||
x_quant = None
|
||||
else:
|
||||
residual = hidden_states
|
||||
hidden_states, post, comb, norm_fused = self.hc_pre(
|
||||
@@ -2075,12 +2113,17 @@ class DeepseekV4DecoderLayer(nn.Module):
|
||||
)
|
||||
|
||||
if use_fused:
|
||||
fused_mhc = try_fused_hc_post_pre(
|
||||
post_attn_norm_weight = (
|
||||
self._post_attention_layernorm_weight_bf16
|
||||
if self._post_attention_layernorm_weight_bf16 is not None
|
||||
else self.post_attention_layernorm.weight.data
|
||||
)
|
||||
fused = apply_mhc_post_pre_boundary(
|
||||
hidden_states,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
self.hc_ffn_fn.T,
|
||||
self.hc_ffn_fn,
|
||||
self.hc_ffn_scale,
|
||||
self.hc_ffn_base,
|
||||
self.hc_mult,
|
||||
@@ -2088,32 +2131,30 @@ class DeepseekV4DecoderLayer(nn.Module):
|
||||
self.hc_eps,
|
||||
_MHC_POST_MULT_VALUE,
|
||||
self.hc_sinkhorn_iters,
|
||||
_is_gfx95_supported,
|
||||
post_attn_norm_weight,
|
||||
self.post_attention_layernorm.variance_epsilon,
|
||||
fn_transpose=True,
|
||||
)
|
||||
if fused_mhc is not None:
|
||||
residual, hidden_states, post, comb, norm_fused = fused_mhc
|
||||
if fused is not None:
|
||||
residual, hidden_states, post, comb, norm_fused = fused
|
||||
if not norm_fused:
|
||||
# The Triton fused post+pre skips the post-attention
|
||||
# layernorm (norm_fused=False); apply it before the MoE,
|
||||
# matching the unfused hc_pre path below.
|
||||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
else:
|
||||
residual, post, comb, hidden_states = _get_mhc_ops().mhc_fused_post_pre(
|
||||
hidden_states = self.hc_post(hidden_states, residual, post, comb)
|
||||
residual = hidden_states
|
||||
hidden_states, post, comb, norm_fused = self.hc_pre(
|
||||
hidden_states,
|
||||
residual,
|
||||
post.unsqueeze(-1) if post.ndim == 2 else post,
|
||||
comb,
|
||||
self.hc_ffn_fn,
|
||||
self.hc_ffn_scale,
|
||||
self.hc_ffn_base,
|
||||
self.rms_norm_eps,
|
||||
self.hc_eps,
|
||||
self.hc_eps,
|
||||
_MHC_POST_MULT_VALUE,
|
||||
self.hc_sinkhorn_iters,
|
||||
norm_weight=(
|
||||
self._post_attention_layernorm_weight_bf16
|
||||
if self._post_attention_layernorm_weight_bf16 is not None
|
||||
else self.post_attention_layernorm.weight.data
|
||||
),
|
||||
norm_eps=self.post_attention_layernorm.variance_epsilon,
|
||||
norm=self.post_attention_layernorm,
|
||||
forward_batch=forward_batch,
|
||||
)
|
||||
norm_fused = True
|
||||
if not norm_fused:
|
||||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
else:
|
||||
hidden_states = self.hc_post(hidden_states, residual, post, comb)
|
||||
residual = hidden_states
|
||||
@@ -2352,11 +2393,63 @@ class DeepseekV4DecoderLayer(nn.Module):
|
||||
def op_mhc_post_attn_pre_mlp(self, state):
|
||||
# Close the attention mHC (hc_post), then open the FFN-side mHC pre +
|
||||
# post-attention layernorm. Produces the 2D MoE input.
|
||||
#
|
||||
# Pop each boundary tensor from the state EXACTLY ONCE, up front, and
|
||||
# reuse the locals for both the fused attempt and the non-fused
|
||||
# fallback. apply_mhc_post_pre_boundary() returns None when it declines
|
||||
# to fuse -- most importantly for the 0-token DP two-batch-overlap idle
|
||||
# ubatch -- in which case control must fall through to the unfused
|
||||
# hc_post. Popping in the fused call's arguments and again in the
|
||||
# fallback would double-pop -> KeyError: 'hidden_states_after_attn' on
|
||||
# every idle DP rank. use_fused_mhc_post_pre is on whenever the aiter
|
||||
# gfx95 mHC path is available (is_cross_layer_mhc_fusion_enabled), so
|
||||
# this fallback is reached under DP regardless of the TileLang env.
|
||||
hidden_states_after_attn = state.pop("hidden_states_after_attn")
|
||||
attn_residual = state.pop("attn_residual")
|
||||
attn_post = state.pop("attn_post")
|
||||
attn_comb = state.pop("attn_comb")
|
||||
|
||||
if self.use_fused_mhc_post_pre:
|
||||
post_attn_norm_weight = (
|
||||
self._post_attention_layernorm_weight_bf16
|
||||
if self._post_attention_layernorm_weight_bf16 is not None
|
||||
else self.post_attention_layernorm.weight.data
|
||||
)
|
||||
fused = apply_mhc_post_pre_boundary(
|
||||
hidden_states_after_attn,
|
||||
attn_residual,
|
||||
attn_post,
|
||||
attn_comb,
|
||||
self.hc_ffn_fn,
|
||||
self.hc_ffn_scale,
|
||||
self.hc_ffn_base,
|
||||
self.hc_mult,
|
||||
self.rms_norm_eps,
|
||||
self.hc_eps,
|
||||
_MHC_POST_MULT_VALUE,
|
||||
self.hc_sinkhorn_iters,
|
||||
post_attn_norm_weight,
|
||||
self.post_attention_layernorm.variance_epsilon,
|
||||
fn_transpose=True,
|
||||
)
|
||||
if fused is not None:
|
||||
ffn_residual, hidden_states, post, comb, norm_fused = fused
|
||||
if not norm_fused:
|
||||
# The Triton fused post+pre skips the post-attention
|
||||
# layernorm (norm_fused=False); apply it before the MoE,
|
||||
# matching the unfused hc_pre path below.
|
||||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
state.ffn_residual = ffn_residual
|
||||
state.ffn_post = post
|
||||
state.ffn_comb = comb
|
||||
state.hidden_states_mlp_input = hidden_states
|
||||
return
|
||||
|
||||
hidden_states = self.hc_post(
|
||||
state.pop("hidden_states_after_attn"),
|
||||
state.pop("attn_residual"),
|
||||
state.pop("attn_post"),
|
||||
state.pop("attn_comb"),
|
||||
hidden_states_after_attn,
|
||||
attn_residual,
|
||||
attn_post,
|
||||
attn_comb,
|
||||
)
|
||||
ffn_residual = hidden_states
|
||||
hidden_states, post, comb, norm_fused = self.hc_pre(
|
||||
@@ -2630,7 +2723,7 @@ class DeepseekV4Model(nn.Module):
|
||||
) = make_hc_head_params(hc_mult, config.hidden_size)
|
||||
|
||||
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
|
||||
self.use_fused_mhc_post_pre = _is_fused_mhc_post_pre_enabled()
|
||||
self.use_fused_mhc_post_pre = is_cross_layer_mhc_fusion_enabled()
|
||||
if self.dsa_enable_prefill_cp:
|
||||
self.cp_size = get_parallel().attn_cp_size
|
||||
|
||||
|
||||
@@ -0,0 +1,400 @@
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.models.deepseek_common.amd import deepseek_v4_fused_mhc
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=4, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestAmdFusedMhcCrossLayerGating(unittest.TestCase):
|
||||
"""Gating and dispatch-preference tests (CPU, no kernels required)."""
|
||||
|
||||
def test_tilelang_fuse_flag_enables_cross_layer_fusion(self):
|
||||
with (
|
||||
envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(True),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(True),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(True),
|
||||
):
|
||||
self.assertTrue(deepseek_v4_fused_mhc.is_cross_layer_mhc_fusion_enabled())
|
||||
|
||||
@mock.patch.object(deepseek_v4_fused_mhc, "is_sm120_supported", return_value=True)
|
||||
def test_sm120_enables_fusion_with_tilelang_pre_disabled(self, _mock_sm120):
|
||||
# Regression (PR review): consolidating _is_fused_mhc_post_pre_enabled into
|
||||
# this module must preserve the SM120 special case. SM120 disables the
|
||||
# standalone TileLang pre path, but mhc_fused_post_pre dispatches
|
||||
# independently, so fuse+post enabled with the pre flag OFF must still
|
||||
# enable fusion when SM120 is supported. The pre-fix consolidation
|
||||
# required the pre flag unconditionally and silently disabled fusion on
|
||||
# SM120.
|
||||
with (
|
||||
envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(True),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(True),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(False),
|
||||
):
|
||||
self.assertTrue(deepseek_v4_fused_mhc._is_fused_mhc_post_pre_enabled())
|
||||
|
||||
@mock.patch.object(deepseek_v4_fused_mhc, "is_sm120_supported", return_value=False)
|
||||
def test_no_sm120_still_requires_tilelang_pre(self, _mock_sm120):
|
||||
# Negative branch: the (pre OR sm120) clause must not degrade to
|
||||
# always-true. With SM120 unsupported and the pre flag off, fuse+post
|
||||
# alone must not enable the standalone TileLang fused path.
|
||||
with (
|
||||
envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(True),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(True),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(False),
|
||||
):
|
||||
self.assertFalse(deepseek_v4_fused_mhc._is_fused_mhc_post_pre_enabled())
|
||||
|
||||
def test_is_fused_mhc_post_pre_enabled_policy(self):
|
||||
# Full gating table for _is_fused_mhc_post_pre_enabled, migrated from the
|
||||
# removed test_deepseek_v4_fused_mhc_policy.py now that the helper lives
|
||||
# in this module (it used to patch deepseek_v4.is_sm120_supported /
|
||||
# deepseek_v4._is_fused_mhc_post_pre_enabled, both gone after the
|
||||
# consolidation -> the registered CPU test AttributeError'd). Fusion
|
||||
# requires the opt-in flag AND TileLang post AND (TileLang pre OR SM120).
|
||||
cases = [
|
||||
# (fuse, pre, post, sm120, expected)
|
||||
(True, False, True, True, True), # SM120 waives the standalone pre flag
|
||||
(True, False, True, False, False), # non-SM120 still needs the pre flag
|
||||
(True, True, True, False, True), # non-SM120 with the pre flag on
|
||||
(False, False, True, True, False), # fusion opt-in is required
|
||||
(True, False, False, True, False), # TileLang post is required
|
||||
]
|
||||
for fuse, pre, post, sm120, expected in cases:
|
||||
with self.subTest(fuse=fuse, pre=pre, post=post, sm120=sm120):
|
||||
with (
|
||||
envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(fuse),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(pre),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(post),
|
||||
mock.patch.object(
|
||||
deepseek_v4_fused_mhc,
|
||||
"is_sm120_supported",
|
||||
return_value=sm120,
|
||||
),
|
||||
):
|
||||
self.assertEqual(
|
||||
deepseek_v4_fused_mhc._is_fused_mhc_post_pre_enabled(),
|
||||
expected,
|
||||
)
|
||||
|
||||
@mock.patch.object(deepseek_v4_fused_mhc, "is_gfx95_supported", return_value=True)
|
||||
@mock.patch.object(deepseek_v4_fused_mhc, "get_bool_env_var", return_value=True)
|
||||
@mock.patch.object(deepseek_v4_fused_mhc, "_is_hip", True)
|
||||
def test_aiter_gfx95_enables_cross_layer_fusion(self, _mock_aiter, _mock_gfx95):
|
||||
# TileLang flags off: fusion must still enable via the aiter gfx95 path.
|
||||
with (
|
||||
envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(False),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(False),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(False),
|
||||
):
|
||||
self.assertTrue(deepseek_v4_fused_mhc.is_cross_layer_mhc_fusion_enabled())
|
||||
|
||||
@mock.patch.object(deepseek_v4_fused_mhc, "is_gfx95_supported", return_value=False)
|
||||
@mock.patch.object(deepseek_v4_fused_mhc, "get_bool_env_var", return_value=True)
|
||||
@mock.patch.object(deepseek_v4_fused_mhc, "_is_hip", True)
|
||||
def test_aiter_cross_layer_disabled_without_gfx95(self, _mock_aiter, _mock_gfx95):
|
||||
with (
|
||||
envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(False),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(False),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(False),
|
||||
):
|
||||
self.assertFalse(deepseek_v4_fused_mhc.is_cross_layer_mhc_fusion_enabled())
|
||||
|
||||
@mock.patch.object(deepseek_v4_fused_mhc, "is_gfx95_supported", return_value=True)
|
||||
@mock.patch.object(deepseek_v4_fused_mhc, "get_bool_env_var", return_value=False)
|
||||
@mock.patch.object(deepseek_v4_fused_mhc, "_is_hip", True)
|
||||
def test_aiter_path_skips_without_sglang_use_aiter(self, _mock_aiter, _mock_gfx95):
|
||||
result = deepseek_v4_fused_mhc.try_aiter_fused_mhc_post_pre(
|
||||
layer_input=mock.Mock(shape=(32, 7168), dim=2, device="cpu"),
|
||||
residual=mock.Mock(dim=3),
|
||||
post=mock.Mock(),
|
||||
comb=mock.Mock(),
|
||||
hc_fn=mock.Mock(),
|
||||
hc_scale=mock.Mock(),
|
||||
hc_base=mock.Mock(),
|
||||
rms_eps=1e-6,
|
||||
hc_eps=1e-6,
|
||||
hc_post_mult=2.0,
|
||||
sinkhorn_iters=20,
|
||||
norm_weight=mock.Mock(),
|
||||
norm_eps=1e-6,
|
||||
)
|
||||
self.assertIsNone(result)
|
||||
|
||||
@mock.patch.object(
|
||||
deepseek_v4_fused_mhc,
|
||||
"try_aiter_fused_mhc_post_pre",
|
||||
return_value=("res", "hs", "post", "comb", True),
|
||||
)
|
||||
@mock.patch.object(deepseek_v4_fused_mhc, "try_fused_hc_post_pre")
|
||||
def test_boundary_prefers_aiter_over_triton(self, mock_triton, mock_aiter):
|
||||
result = deepseek_v4_fused_mhc.try_mhc_fused_post_pre_boundary(
|
||||
layer_input=mock.Mock(shape=(32, 7168), dim=2),
|
||||
residual=mock.Mock(dim=3),
|
||||
post=mock.Mock(),
|
||||
comb=mock.Mock(),
|
||||
hc_fn=mock.Mock(),
|
||||
hc_scale=mock.Mock(),
|
||||
hc_base=mock.Mock(),
|
||||
hc_mult=4,
|
||||
rms_eps=1e-6,
|
||||
hc_eps=1e-6,
|
||||
hc_post_mult=2.0,
|
||||
sinkhorn_iters=20,
|
||||
norm_weight=mock.Mock(),
|
||||
norm_eps=1e-6,
|
||||
fn_transpose=True,
|
||||
is_gfx95_supported_flag=True,
|
||||
)
|
||||
self.assertEqual(result, ("res", "hs", "post", "comb", True))
|
||||
mock_triton.assert_not_called()
|
||||
mock_aiter.assert_called_once()
|
||||
|
||||
@mock.patch.object(
|
||||
deepseek_v4_fused_mhc, "try_aiter_fused_mhc_post_pre", return_value=None
|
||||
)
|
||||
@mock.patch.object(
|
||||
deepseek_v4_fused_mhc,
|
||||
"try_fused_hc_post_pre",
|
||||
return_value=("res", "hs", "post", "comb", False),
|
||||
)
|
||||
def test_boundary_falls_back_to_triton(self, mock_triton, mock_aiter):
|
||||
hc_fn = mock.Mock()
|
||||
hc_fn.T = "transposed_fn"
|
||||
result = deepseek_v4_fused_mhc.try_mhc_fused_post_pre_boundary(
|
||||
layer_input=mock.Mock(shape=(32, 7168), dim=2),
|
||||
residual=mock.Mock(dim=3),
|
||||
post=mock.Mock(),
|
||||
comb=mock.Mock(),
|
||||
hc_fn=hc_fn,
|
||||
hc_scale=mock.Mock(),
|
||||
hc_base=mock.Mock(),
|
||||
hc_mult=4,
|
||||
rms_eps=1e-6,
|
||||
hc_eps=1e-6,
|
||||
hc_post_mult=2.0,
|
||||
sinkhorn_iters=20,
|
||||
norm_weight=mock.Mock(),
|
||||
norm_eps=1e-6,
|
||||
fn_transpose=True,
|
||||
is_gfx95_supported_flag=True,
|
||||
)
|
||||
self.assertEqual(result, ("res", "hs", "post", "comb", False))
|
||||
mock_aiter.assert_called_once()
|
||||
# fn_transpose=True must hand the Triton kernel the transposed fn.
|
||||
self.assertEqual(mock_triton.call_args.args[4], "transposed_fn")
|
||||
|
||||
|
||||
class TestAmdFusedMhcAttnBoundaryFallback(unittest.TestCase):
|
||||
"""Regression: the attn-side boundary fallback must close the previous
|
||||
layer's deferred mHC post before opening the current layer's pre.
|
||||
|
||||
When ``apply_mhc_post_pre_boundary`` declines to fuse (returns ``None``) --
|
||||
reachable after an aiter import/kernel failure permanently disables the fused
|
||||
path -- the fallback must call
|
||||
``hc_post(hidden_states, prev_residual, prev_post, prev_comb)`` before
|
||||
``hc_pre``. The pre-fix code ran ``hc_pre`` directly on the raw input and
|
||||
dropped ``prev_residual``/``prev_post``/``prev_comb``, so the previous
|
||||
layer's deferred post was never applied and every subsequent layer computed
|
||||
on corrupted activations.
|
||||
|
||||
Drives the real ``DeepseekV4DecoderLayer.forward`` on a mocked layer with the
|
||||
fused dispatcher forced to ``None`` and halts at ``self_attn`` via a sentinel,
|
||||
so only the boundary fallback executes.
|
||||
"""
|
||||
|
||||
def test_fallback_closes_previous_post_before_pre(self):
|
||||
try:
|
||||
from sglang.srt.models.deepseek_v4 import DeepseekV4DecoderLayer
|
||||
except Exception as e: # pragma: no cover - env without full model deps
|
||||
self.skipTest(f"deepseek_v4 import unavailable: {e}")
|
||||
|
||||
class _StopForward(Exception):
|
||||
pass
|
||||
|
||||
layer = mock.Mock()
|
||||
layer.use_fused_mhc_post_pre = True
|
||||
layer._input_layernorm_weight_bf16 = None
|
||||
closed_post = object()
|
||||
layer.hc_post.return_value = closed_post
|
||||
# norm_fused=True keeps the fallback off the fp8-quant / layernorm branch.
|
||||
layer.hc_pre.return_value = (object(), object(), object(), True)
|
||||
layer.self_attn.maybe_use_decode_attn_tp.side_effect = _StopForward
|
||||
|
||||
hs_in = object()
|
||||
prev_residual, prev_post, prev_comb = object(), object(), object()
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"sglang.srt.models.deepseek_v4.apply_mhc_post_pre_boundary",
|
||||
return_value=None,
|
||||
),
|
||||
self.assertRaises(_StopForward),
|
||||
):
|
||||
DeepseekV4DecoderLayer.forward(
|
||||
layer,
|
||||
positions=object(),
|
||||
hidden_states=hs_in,
|
||||
input_ids=object(),
|
||||
forward_batch=object(),
|
||||
input_ids_global=object(),
|
||||
prev_residual=prev_residual,
|
||||
prev_post=prev_post,
|
||||
prev_comb=prev_comb,
|
||||
)
|
||||
|
||||
# The deferred previous-layer post must be closed with exactly the
|
||||
# prev_* tensors, and hc_pre must then run on the closed result.
|
||||
layer.hc_post.assert_called_once_with(
|
||||
hs_in, prev_residual, prev_post, prev_comb
|
||||
)
|
||||
layer.hc_pre.assert_called_once()
|
||||
self.assertIs(layer.hc_pre.call_args.args[0], closed_post)
|
||||
|
||||
|
||||
class TestAmdFusedMhcNormFusedHandling(unittest.TestCase):
|
||||
"""Regression: a fused-success result with ``norm_fused=False`` must have its
|
||||
layernorm applied at the call site before the activation reaches attention.
|
||||
|
||||
``try_fused_hc_post_pre`` (the Triton fused post+pre) always returns
|
||||
``norm_fused=False`` -- it does not apply the input/post-attention layernorm.
|
||||
The boundary dispatcher reaches it whenever the aiter kernel declines
|
||||
(notably after an aiter import/kernel failure permanently disables the aiter
|
||||
path). The pre-fix fused-success branch unpacked the tuple and fed the raw
|
||||
(unnormalized) hidden_states straight into ``self_attn`` with ``x_quant=None``,
|
||||
silently corrupting every subsequent layer. The fix mirrors the unfused
|
||||
``hc_pre`` branch: apply the input layernorm when ``norm_fused`` is False.
|
||||
|
||||
Drives the real ``DeepseekV4DecoderLayer.forward`` with the boundary forced to
|
||||
return ``norm_fused=False`` and halts at ``self_attn`` via a sentinel.
|
||||
"""
|
||||
|
||||
def test_fused_success_applies_input_layernorm_when_not_norm_fused(self):
|
||||
try:
|
||||
import sglang.srt.models.deepseek_v4 as deepseek_v4
|
||||
from sglang.srt.models.deepseek_v4 import DeepseekV4DecoderLayer
|
||||
except Exception as e: # pragma: no cover - env without full model deps
|
||||
self.skipTest(f"deepseek_v4 import unavailable: {e}")
|
||||
|
||||
class _StopForward(Exception):
|
||||
pass
|
||||
|
||||
layer = mock.Mock()
|
||||
layer.use_fused_mhc_post_pre = True
|
||||
layer._input_layernorm_weight_bf16 = None
|
||||
|
||||
fused_hs = object()
|
||||
residual, post, comb = object(), object(), object()
|
||||
# Fused dispatch SUCCEEDS but reports the input layernorm was NOT applied
|
||||
# (norm_fused=False) -- the Triton fused post+pre contract.
|
||||
normed = object()
|
||||
layer.input_layernorm.return_value = normed
|
||||
layer.self_attn.maybe_use_decode_attn_tp.side_effect = _StopForward
|
||||
|
||||
# Force the non-aiter (torch layernorm) branch deterministically so the
|
||||
# test does not depend on the runner arch and needs no real tensors.
|
||||
with (
|
||||
mock.patch.object(deepseek_v4, "_use_aiter", False),
|
||||
mock.patch.object(deepseek_v4, "_is_gfx95_supported", False),
|
||||
mock.patch(
|
||||
"sglang.srt.models.deepseek_v4.apply_mhc_post_pre_boundary",
|
||||
return_value=(residual, fused_hs, post, comb, False),
|
||||
),
|
||||
self.assertRaises(_StopForward),
|
||||
):
|
||||
DeepseekV4DecoderLayer.forward(
|
||||
layer,
|
||||
positions=object(),
|
||||
hidden_states=object(),
|
||||
input_ids=object(),
|
||||
forward_batch=object(),
|
||||
input_ids_global=object(),
|
||||
prev_residual=object(),
|
||||
prev_post=object(),
|
||||
prev_comb=object(),
|
||||
)
|
||||
|
||||
# The fused (unnormalized) layer input must be run through the input
|
||||
# layernorm before attention. Pre-fix this was never called on the
|
||||
# fused-success path.
|
||||
layer.input_layernorm.assert_called_once_with(fused_hs)
|
||||
|
||||
|
||||
def _hardware_available() -> bool:
|
||||
try:
|
||||
import torch
|
||||
|
||||
if not (
|
||||
torch.cuda.is_available() and deepseek_v4_fused_mhc.is_gfx95_supported()
|
||||
):
|
||||
return False
|
||||
from aiter.ops.mhc import mhc_fused_post_pre # noqa: F401
|
||||
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
_hardware_available(), "requires a gfx95 device with aiter mHC kernels"
|
||||
)
|
||||
class TestAmdFusedMhcNumerical(unittest.TestCase):
|
||||
"""On-device equivalence of the aiter fused kernel vs unfused mhc_post+mhc_pre.
|
||||
|
||||
Asserts the proven invariants: ``next_residual`` is bit-exact and
|
||||
``layer_input``/``post_mix`` match within bf16 tolerance. ``comb_mix`` is
|
||||
intentionally not asserted here -- its raw-tensor value differs between the
|
||||
fused and unfused kernels at the production Sinkhorn setting, and correctness
|
||||
is established end-to-end (fused-on vs fused-off token match). See the PR
|
||||
description; extend this test once the end-to-end sign-off pins the expected
|
||||
comb_mix convention.
|
||||
"""
|
||||
|
||||
def _run(self, m, hc_mult=4, hidden=7168, sinkhorn_iters=20):
|
||||
import torch
|
||||
from aiter.ops import mhc
|
||||
|
||||
dev = "cuda:0"
|
||||
torch.manual_seed(0)
|
||||
hc_mult3 = hc_mult * 2 + hc_mult * hc_mult
|
||||
li = (torch.randn(m, hidden, device=dev) * 0.02).bfloat16()
|
||||
res = (torch.randn(m, hc_mult, hidden, device=dev) * 0.02).bfloat16()
|
||||
post = torch.randn(m, hc_mult, device=dev) * 0.02
|
||||
comb = torch.randn(m, hc_mult, hc_mult, device=dev) * 0.02
|
||||
fn = (torch.randn(hc_mult3, hc_mult * hidden, device=dev) * 0.02).bfloat16()
|
||||
scl = torch.ones(hc_mult3, device=dev)
|
||||
base = torch.zeros(hc_mult3, device=dev)
|
||||
nw = torch.ones(hidden, device=dev).bfloat16()
|
||||
kw = dict(
|
||||
rms_eps=1e-6,
|
||||
hc_pre_eps=1e-6,
|
||||
hc_sinkhorn_eps=1e-6,
|
||||
hc_post_mult_value=2.0,
|
||||
sinkhorn_repeat=sinkhorn_iters,
|
||||
norm_weight=nw,
|
||||
norm_eps=1e-6,
|
||||
)
|
||||
post_mix, _comb_mix, li_out, next_res = mhc.mhc_fused_post_pre(
|
||||
li, res, post, comb, fn, scl, base, force_fused=True, **kw
|
||||
)
|
||||
ref_next = torch.empty_like(res)
|
||||
mhc.mhc_post(ref_next, li, res, post, comb)
|
||||
ref_post, _ref_comb, ref_li = mhc.mhc_pre(ref_next, fn, scl, base, **kw)
|
||||
|
||||
torch.testing.assert_close(next_res, ref_next, rtol=0, atol=0)
|
||||
torch.testing.assert_close(li_out.float(), ref_li.float(), rtol=3e-2, atol=3e-2)
|
||||
torch.testing.assert_close(
|
||||
post_mix.float(), ref_post.float(), rtol=3e-2, atol=3e-2
|
||||
)
|
||||
|
||||
def test_equivalence_decode(self):
|
||||
self._run(m=32)
|
||||
|
||||
def test_equivalence_prefill(self):
|
||||
self._run(m=96)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,79 +0,0 @@
|
||||
"""Unit tests for the DeepSeek-V4 fused-MHC enable policy."""
|
||||
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import sglang.srt.models.deepseek_v4 as deepseek_v4
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestDeepseekV4FusedMHCPolicy(CustomTestCase):
|
||||
def _is_enabled(
|
||||
self,
|
||||
*,
|
||||
fuse: bool,
|
||||
tilelang_pre: bool,
|
||||
tilelang_post: bool,
|
||||
sm120: bool,
|
||||
) -> bool:
|
||||
with (
|
||||
envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(fuse),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(tilelang_pre),
|
||||
envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(tilelang_post),
|
||||
patch.object(deepseek_v4, "is_sm120_supported", return_value=sm120),
|
||||
):
|
||||
return deepseek_v4._is_fused_mhc_post_pre_enabled()
|
||||
|
||||
def test_sm120_allows_fused_opt_in_with_standalone_pre_disabled(self):
|
||||
self.assertTrue(
|
||||
self._is_enabled(
|
||||
fuse=True,
|
||||
tilelang_pre=False,
|
||||
tilelang_post=True,
|
||||
sm120=True,
|
||||
)
|
||||
)
|
||||
|
||||
def test_other_platform_still_requires_tilelang_pre(self):
|
||||
self.assertFalse(
|
||||
self._is_enabled(
|
||||
fuse=True,
|
||||
tilelang_pre=False,
|
||||
tilelang_post=True,
|
||||
sm120=False,
|
||||
)
|
||||
)
|
||||
self.assertTrue(
|
||||
self._is_enabled(
|
||||
fuse=True,
|
||||
tilelang_pre=True,
|
||||
tilelang_post=True,
|
||||
sm120=False,
|
||||
)
|
||||
)
|
||||
|
||||
def test_fusion_opt_in_and_tilelang_post_remain_required(self):
|
||||
self.assertFalse(
|
||||
self._is_enabled(
|
||||
fuse=False,
|
||||
tilelang_pre=False,
|
||||
tilelang_post=True,
|
||||
sm120=True,
|
||||
)
|
||||
)
|
||||
self.assertFalse(
|
||||
self._is_enabled(
|
||||
fuse=True,
|
||||
tilelang_pre=False,
|
||||
tilelang_post=False,
|
||||
sm120=True,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user