From 6795840de680260b21788adce9be16bcbd8af1ac Mon Sep 17 00:00:00 2001 From: Jacob Date: Mon, 6 Jul 2026 00:48:30 -0500 Subject: [PATCH 1/4] [AMD] [GLM5] Use fused fp8 MLA absorbed bmm on gfx950 by quantizing w_kc/w_vc On gfx950, GLM (GlmMoeDsaForCausalLM) MLA absorbed weights (w_kc/w_vc) are kept bf16, so forward_mla falls to the slow per-batched torch.bmm (rocBLAS) path. This quantizes them to per-tensor fp8_e4m3fn at load (mirroring the DeepSeek fp8 flow), which routes the q_nope/attn-output projections through the fused aiter batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant kernel. GLM-5.2-MXFP4, MI355X TP4, i1024/o1024 conc4 (median): TPOT 15.66 -> 14.89 ms (-4.9%); GSM8K(1319) 0.940 -> 0.929. --- .../deepseek_common/deepseek_weight_loader.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index ad83266b7753..3375598f4472 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -625,6 +625,22 @@ def post_load_weights( torch.bfloat16 ) + # GLM (GlmMoeDsa) MLA absorbed weights (w_kc/w_vc) load as bf16, which + # forces the slow per-batched torch.bmm path in forward_mla on ROCm. Quantize + # to per-tensor fp8_e4m3fn (mirroring the DeepSeek fp8 flow) so the fused + # batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant + # kernel is used instead. Must be e4m3fn (not fnuz) so forward_mla's dtype + # gate matches on gfx950. + if ( + _use_aiter_gfx95 + and self.config.architectures + and self.config.architectures[0] == "GlmMoeDsaForCausalLM" + and w.dtype == torch.bfloat16 + ): + from sglang.srt.layers.quantization.fp8_utils import input_to_float8 + + w, self_attn.w_scale = input_to_float8(w, dtype=torch.float8_e4m3fn) + w_kc, w_vc = w.unflatten( 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) From 7e0d137f4707ccd5e50eb4a79506deff157de2b8 Mon Sep 17 00:00:00 2001 From: Jacob Date: Tue, 7 Jul 2026 01:31:49 -0500 Subject: [PATCH 2/4] [AMD] [GLM5] Remove per-layer direct_copy in fp8 MLA absorbed v_up bmm on gfx950 The fp8 absorbed-bmm path ran the aiter batched_gemm_a8w8 with transpose_bm=False, producing a (heads, tokens, v_head_dim) output that the downstream o_proj prep turned into (tokens, heads*v) via transpose(0,1).flatten, which is non-contiguous and forces an at::native direct_copy (elementwise_manual_unroll) once per decode layer (~5us/layer on MI355X). The MXFP4 path avoids this by writing the bmm output batch-major. Allocate the bmm output as (tokens, heads, v_head_dim) and write it via transpose_bm=True, so the existing _bmm_buf path flattens it as a free view. GLM-5.2-MXFP4 MI355X TP4 i1024/o512 conc4: median TPOT 14.53 -> 14.30 ms; GSM8K(200) 0.940, 0 invalid. --- .../attention_forward_methods/forward_mla.py | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 3de047b36251..6143eb3532eb 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -836,16 +836,31 @@ def forward_absorb_core( else: _bmm_buf = None if _use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn: - attn_bmm_output = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( + # Write the absorbed v_up bmm output batch-major + # (tokens, heads, v_head_dim) via transpose_bm=True so the + # downstream flatten(1, 2) in the `_bmm_buf is not None` block is + # a free view. With the old transpose_bm=False the output was + # (heads, tokens, v) and needed a transpose(0,1).flatten copy = + # a per-layer direct_copy (elementwise_manual_unroll ~5us/layer), + # which ATOM / the MXFP4 path do not pay. + _bmm_buf = torch.empty( + attn_output.shape[0], + self.num_local_heads, + self.w_vc.shape[-1], + device=attn_output.device, + dtype=torch.bfloat16, + ) + batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( X=attn_output, WQ=self.w_vc.transpose(-1, -2), w_scale=self.w_scale, group_size=128, - YQ=None, - transpose_bm=False, + YQ=_bmm_buf, + transpose_bm=True, transpose_bm_in=True, dtype=torch.bfloat16, ) + attn_bmm_output = _bmm_buf else: attn_bmm_output = torch.bmm( attn_output.to(torch.bfloat16).transpose(0, 1), From eed73c43f16a7df309f4a8af1ff8ee5ca99ef635 Mon Sep 17 00:00:00 2001 From: Jacob Date: Sun, 16 Aug 2026 23:53:59 -0500 Subject: [PATCH 3/4] Shrink diff: condense comments, drop dead assignment, reuse arch check Reduce the injected code to the two changes that carry the feature: the per-tensor e4m3fn quantization of GLM's bf16 kv_b_proj, and the batch-major output layout for the fp8 absorbed V bmm. - Condense the two rationale comments; the measured per-layer copy cost lives in the PR description instead of the source. - Drop `attn_bmm_output = _bmm_buf`, which every branch below the `_bmm_buf is not None` check already reassigns. - Import input_to_float8 at module level alongside the other fp8_utils helpers instead of inside the branch. - Hoist `arch` so the new gate and the adjacent quark gate share one architectures lookup. Co-authored-by: Cursor --- .../forward_mla_rocm.py | 10 ++------- .../deepseek_common/deepseek_weight_loader.py | 21 +++++++------------ 2 files changed, 10 insertions(+), 21 deletions(-) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py index 0bbc30266c2d..928eb6352fc3 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py @@ -203,13 +203,8 @@ def rocm_absorb_v_bmm( else: _bmm_buf = None if _use_aiter_gfx95 and attn.w_kc.dtype == torch.float8_e4m3fn: - # Write the absorbed v_up bmm output batch-major - # (tokens, heads, v_head_dim) via transpose_bm=True so the - # downstream flatten(1, 2) in the `_bmm_buf is not None` block is - # a free view. With the old transpose_bm=False the output was - # (heads, tokens, v) and needed a transpose(0,1).flatten copy = - # a per-layer direct_copy (elementwise_manual_unroll ~5us/layer), - # which ATOM / the MXFP4 path do not pay. + # As in the mxfp4 path above, write (batch, heads, dim) so the + # post-GEMM flatten is a free view instead of a copy. _bmm_buf = torch.empty( attn_output.shape[0], attn.num_local_heads, @@ -227,7 +222,6 @@ def rocm_absorb_v_bmm( transpose_bm_in=True, dtype=torch.bfloat16, ) - attn_bmm_output = _bmm_buf else: attn_bmm_output = torch.bmm( attn_output.to(torch.bfloat16).transpose(0, 1), diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index 9ab655bccceb..724f02b25f06 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -31,6 +31,7 @@ block_quant_dequant, block_quant_to_tensor_quant, channel_quant_to_tensor_quant, + input_to_float8, inverse_transform_scale_ue8m0, normalize_e4m3fn_to_e4m3fnuz, quant_weight_ue8m0, @@ -500,6 +501,8 @@ def post_load_weights( if layer_id < self.config.num_hidden_layers: layer_ids.add(layer_id) + arch = self.config.architectures[0] if self.config.architectures else None + for layer_id in layer_ids: self_attn = ( self.model.layers[layer_id].self_attn @@ -631,20 +634,13 @@ def post_load_weights( torch.bfloat16 ) - # GLM (GlmMoeDsa) MLA absorbed weights (w_kc/w_vc) load as bf16, which - # forces the slow per-batched torch.bmm path in forward_mla on ROCm. Quantize - # to per-tensor fp8_e4m3fn (mirroring the DeepSeek fp8 flow) so the fused - # batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant - # kernel is used instead. Must be e4m3fn (not fnuz) so forward_mla's dtype - # gate matches on gfx950. + # GLM ships kv_b_proj as bf16, which falls back to torch.bmm. Quantize to + # per-tensor e4m3fn (not fnuz) to match forward_mla_rocm's dtype gate. if ( _use_aiter_gfx95 - and self.config.architectures - and self.config.architectures[0] == "GlmMoeDsaForCausalLM" + and arch == "GlmMoeDsaForCausalLM" and w.dtype == torch.bfloat16 ): - from sglang.srt.layers.quantization.fp8_utils import input_to_float8 - w, self_attn.w_scale = input_to_float8(w, dtype=torch.float8_e4m3fn) w_kc, w_vc = w.unflatten( @@ -655,9 +651,8 @@ def post_load_weights( _use_aiter_gfx95 and self.quant_config is not None and self.quant_config.get_name() == "quark" - and self.config.architectures - and self.config.architectures[0] - == "DeepseekV3ForCausalLM" # Avoid processing other models like GlmMoeDsaForCausalLM + # Avoid processing other models like GlmMoeDsaForCausalLM + and arch == "DeepseekV3ForCausalLM" and w.dtype not in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) ): w_kc, self_attn.w_scale_k, w_vc, self_attn.w_scale_v = ( From 8d82ccecce943db85a71e68106c94f419a9961c5 Mon Sep 17 00:00:00 2001 From: Jacob Date: Mon, 17 Aug 2026 01:00:28 -0500 Subject: [PATCH 4/4] Keep the diff inside this PR's own scope Restore the adjacent quark gate and drop the hoisted `arch` local, so the change no longer touches code outside the feature. The GLM gate reads self.config.architectures directly again, matching the surrounding style. Co-authored-by: Cursor --- .../models/deepseek_common/deepseek_weight_loader.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index 724f02b25f06..88ce9196de5a 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -501,8 +501,6 @@ def post_load_weights( if layer_id < self.config.num_hidden_layers: layer_ids.add(layer_id) - arch = self.config.architectures[0] if self.config.architectures else None - for layer_id in layer_ids: self_attn = ( self.model.layers[layer_id].self_attn @@ -638,7 +636,8 @@ def post_load_weights( # per-tensor e4m3fn (not fnuz) to match forward_mla_rocm's dtype gate. if ( _use_aiter_gfx95 - and arch == "GlmMoeDsaForCausalLM" + and self.config.architectures + and self.config.architectures[0] == "GlmMoeDsaForCausalLM" and w.dtype == torch.bfloat16 ): w, self_attn.w_scale = input_to_float8(w, dtype=torch.float8_e4m3fn) @@ -651,8 +650,9 @@ def post_load_weights( _use_aiter_gfx95 and self.quant_config is not None and self.quant_config.get_name() == "quark" - # Avoid processing other models like GlmMoeDsaForCausalLM - and arch == "DeepseekV3ForCausalLM" + and self.config.architectures + and self.config.architectures[0] + == "DeepseekV3ForCausalLM" # Avoid processing other models like GlmMoeDsaForCausalLM and w.dtype not in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) ): w_kc, self_attn.w_scale_k, w_vc, self_attn.w_scale_v = (