minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[AMD] [GLM-5.3-Flash Day 0] Route mHC through AITER on gfx950 (#38545)

Browse Source 
Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com>
Co-authored-by: Thomas Wang <thomawan@amd.com>
Co-authored-by: Kevin Mi <45493463+kevin-mii@users.noreply.github.com>
Co-authored-by: Kevin Mi <mikevin920@yahoo.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Raiden Makoto
2026-09-21 22:24:29 -07:00
committed by GitHub
co-authored by Raiden-Makoto Thomas Wang Kevin Mi Kevin Mi Claude Fable 5.1
parent 264da63319
commit 095e45100b
2 changed files with 352 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/layernorm/mhc.py
+139 -4
View File
@@ -18,10 +18,99 @@ from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_interleave
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.utils.common import strict_contiguous
from sglang.srt.runtime_context import get_parallel, get_platform
from sglang.srt.utils import is_gfx95_supported, is_hip
from sglang.srt.utils.common import is_gfx1250_supported
logger = logging.getLogger(__name__)
_AITER_MHC_RUNTIME_DISABLED = False
_AITER_MHC_ACTIVE_LOGGED = False
def _use_aiter_mhc() -> bool:
return (
not _AITER_MHC_RUNTIME_DISABLED
and is_gfx95_supported()
and envs.SGLANG_USE_AITER.get()
)
def _try_aiter_mhc_pre(
residual: torch.Tensor,
fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
rms_eps: float,
hc_pre_eps: float,
hc_sinkhorn_eps: float,
hc_post_mult_value: float,
sinkhorn_repeat: int,
norm_weight: torch.Tensor | None,
norm_eps: float | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
global _AITER_MHC_RUNTIME_DISABLED, _AITER_MHC_ACTIVE_LOGGED
try:
from aiter.ops.mhc import mhc_pre as aiter_mhc_pre
except Exception as err:
logger.warning("AITER mHC pre is unavailable, falling back: %s", err)
_AITER_MHC_RUNTIME_DISABLED = True
return None
kwargs = {}
if norm_weight is not None:
kwargs["norm_weight"] = norm_weight
kwargs["norm_eps"] = norm_eps if norm_eps is not None else rms_eps
try:
result = aiter_mhc_pre(
residual,
fn,
hc_scale,
hc_base,
rms_eps,
hc_pre_eps,
hc_sinkhorn_eps,
hc_post_mult_value,
sinkhorn_repeat,
**kwargs,
)
except Exception as err:
logger.warning("AITER mHC pre failed, disabling fast path: %s", err)
_AITER_MHC_RUNTIME_DISABLED = True
return None
if not _AITER_MHC_ACTIVE_LOGGED:
logger.info("Using AITER gfx950 mHC pre/post kernels")
_AITER_MHC_ACTIVE_LOGGED = True
return result
def _try_aiter_mhc_post(
x: torch.Tensor,
residual: torch.Tensor,
post_layer_mix: torch.Tensor,
comb_res_mix: torch.Tensor,
) -> torch.Tensor | None:
global _AITER_MHC_RUNTIME_DISABLED
try:
from aiter.ops.mhc import mhc_post as aiter_mhc_post
except Exception as err:
logger.warning("AITER mHC post is unavailable, falling back: %s", err)
_AITER_MHC_RUNTIME_DISABLED = True
return None
out = torch.empty_like(residual)
try:
aiter_mhc_post(out, x, residual, post_layer_mix, comb_res_mix)
except Exception as err:
logger.warning("AITER mHC post failed, disabling fast path: %s", err)
_AITER_MHC_RUNTIME_DISABLED = True
return None
return out
# This module is imported during model-registry discovery. Do not import the real
# TileLang package here: it loads native CUDA stubs. The proxy below lets
# module-level @tilelang.jit declarations parse, then imports and applies real
@@ -119,6 +208,24 @@ pass_configs = {
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
}
def _use_deep_gemm_hc_prenorm() -> bool:
if is_hip() or not envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get():
return False
from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM
return ENABLE_JIT_DEEPGEMM
def _use_tilelang_mhc_pre() -> bool:
return envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get() and not is_hip()
def _use_tilelang_mhc_post() -> bool:
return envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get() and not is_hip()
FP8 = "float8_e4m3"
BF16 = "bfloat16"
FP32 = "float32"
@@ -1041,7 +1148,7 @@ def mhc_pre(
num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
)
if envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get():
if _use_deep_gemm_hc_prenorm():
n_splits = _compute_num_split_for_mhc_pre(num_tokens, hc_hidden_size)
gemm_out_mul = torch.empty(
@@ -1653,7 +1760,7 @@ def mhc_fused_post_pre(
hidden_size,
)
if envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get():
if _use_deep_gemm_hc_prenorm():
import deep_gemm
deep_gemm.tf32_hc_prenorm_gemm(
@@ -1847,7 +1954,25 @@ def _mhc_pre_dispatch(
norm_eps: float | None = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]:
assert residual.dim() == 3, f"residual must be (s, n, h); got {residual.shape}"
if not envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get():
if _use_aiter_mhc():
result = _try_aiter_mhc_pre(
residual=residual,
fn=fn,
hc_scale=hc_scale,
hc_base=hc_base,
rms_eps=rms_eps,
hc_pre_eps=hc_pre_eps,
hc_sinkhorn_eps=hc_sinkhorn_eps,
hc_post_mult_value=hc_post_mult_value,
sinkhorn_repeat=sinkhorn_repeat,
norm_weight=norm_weight,
norm_eps=norm_eps,
)
if result is not None:
post_mix, comb_mix, layer_input = result
return post_mix, comb_mix, layer_input, norm_weight is not None
if not _use_tilelang_mhc_pre():
post_mix, comb_mix, layer_input = _mhc_pre_torch(
residual=residual,
fn=fn,
@@ -1886,7 +2011,17 @@ def _mhc_post_dispatch(
) -> torch.Tensor:
assert x.dim() == 2 and residual.dim() == 3
assert post_layer_mix.dim() == 3 and comb_res_mix.dim() == 3
if not envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get():
if _use_aiter_mhc():
result = _try_aiter_mhc_post(
x=x,
residual=residual,
post_layer_mix=post_layer_mix,
comb_res_mix=comb_res_mix,
)
if result is not None:
return result
if not _use_tilelang_mhc_post():
return _mhc_post_torch(x, residual, post_layer_mix, comb_res_mix)
return mhc_post(x, residual, post_layer_mix, comb_res_mix)