[AMD] [GLM5] fp8 MLA absorbed bmm for GLM-5.2 on gfx950 (#30519)

Co-authored-by: Thomas Wang <thomawan@amd.com>
Co-authored-by: sogalin_codegen <39478626+sogalin@users.noreply.github.com>

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py

@@ -203,17 +203,24 @@ def rocm_absorb_v_bmm(
     else:
         _bmm_buf = None
         if _use_aiter_gfx95 and attn.w_kc.dtype == torch.float8_e4m3fn:
-            attn_bmm_output = (
+            # As in the mxfp4 path above, write (batch, heads, dim) so the
-                batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
+            # post-GEMM flatten is a free view instead of a copy.
-                    X=attn_output,
+            _bmm_buf = torch.empty(
-                    WQ=attn.w_vc.transpose(-1, -2),
+                attn_output.shape[0],
-                    w_scale=attn.w_scale,
+                attn.num_local_heads,
-                    group_size=128,
+                attn.w_vc.shape[-1],
-                    YQ=None,
+                device=attn_output.device,
-                    transpose_bm=False,
+                dtype=torch.bfloat16,
-                    transpose_bm_in=True,
+            )
-                    dtype=torch.bfloat16,
+            batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
-                )
+                X=attn_output,
+                WQ=attn.w_vc.transpose(-1, -2),
+                w_scale=attn.w_scale,
+                group_size=128,
+                YQ=_bmm_buf,
+                transpose_bm=True,
+                transpose_bm_in=True,
+                dtype=torch.bfloat16,
             )
         else:
             attn_bmm_output = torch.bmm(

python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py

@@ -31,6 +31,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
     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,
@@ -639,6 +640,16 @@ class DeepseekV2WeightLoaderMixin:
                         torch.bfloat16
                     )
 
+            # 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 w.dtype == torch.bfloat16
+            ):
+                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)