[AMD] [GLM5] use tuned MXFP4 kernels for MLA absorbed BMM

python/sglang/srt/environ.py

@@ -1457,6 +1457,9 @@ class Envs:
     SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM = EnvBool(False)
     SGLANG_DSA_TOPK_BROADCAST = EnvBool(False)
     SGLANG_DISABLE_DSA_INDEXER_FUSION = EnvBool(False)
+    # Opt-in GLM MLA absorbed-BMM backend that keeps w_kc/w_vc in packed
+    # MXFP4 and dispatches the matching AITER FP4 kernels.
+    SGLANG_USE_MXFP4_MLA_BMM = EnvBool(False)
     # Opt-in perf path for --dsa-prefill-backend flashmla_sparse_q8: fuse the
     # absorbed q bmm with the nope/rope concat + fp8 cast so q is written
     # directly in fp8 ("born fp8") and the standalone concat-cast kernel

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

@@ -60,6 +60,7 @@
     _is_cuda,
     _is_hip,
     _is_musa,
+    _use_aiter_gfx95,
 )
 from sglang.srt.runtime_context import get_exec, get_parallel
 from sglang.srt.state_capturer.indexer_topk import (
@@ -110,6 +111,114 @@ def is_mla_dcp_lse_base_on_e(attention_backend: Optional[str]) -> bool:
     from sglang.kernels.ops.gemm import bmm_fp8
 
 
+if _use_aiter_gfx95:
+    from aiter.ops.triton._triton_kernels.gemm.batched.batched_gemm_a16wfp4 import (
+        _get_config as _get_mxfp4_bmm_config,
+    )
+    from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
+
+    from sglang.srt.layers.quantization.rocm_mxfp4_utils import (
+        batched_gemm_afp4wfp4_pre_quant,
+    )
+
+
+def _get_single_split_mxfp4_bmm_config(x: torch.Tensor, weight: torch.Tensor) -> dict:
+    config, _ = _get_mxfp4_bmm_config(x.shape[1], weight.shape[1], x.shape[2])
+    config = config.copy()
+    config["NUM_KSPLIT"] = 1
+    return config
+
+
+def _get_glm_mxfp4_k_bmm_config(x: torch.Tensor, weight: torch.Tensor) -> dict:
+    config = _get_single_split_mxfp4_bmm_config(x, weight)
+    # GLM's K-up has K=192. Larger blocks over-read its six E8M0 scale groups.
+    config["BLOCK_SIZE_K"] = 64
+    return config
+
+
+def _get_glm_mxfp4_v_bmm_config(x: torch.Tensor, weight: torch.Tensor) -> dict:
+    config = _get_single_split_mxfp4_bmm_config(x, weight)
+    # Keep AITER's small-M decode buckets; use the profiled prefill tiles.
+    if x.shape[1] > 256:
+        config["BLOCK_SIZE_M"] = 128
+        config["BLOCK_SIZE_K"] = 128
+    return config
+
+
+def _run_tuned_mxfp4_bmm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output: torch.Tensor,
+    config: dict,
+    *,
+    transpose_bm: bool,
+) -> torch.Tensor:
+    return batched_gemm_a16wfp4(
+        x,
+        weight,
+        weight_scale,
+        y=output,
+        config=config,
+        transpose_bm=transpose_bm,
+        prequant=True,
+        y_scale=None,
+        dtype=torch.bfloat16,
+    )
+
+
+def _run_mxfp4_k_bmm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    if x.shape[2] == 192 and weight.shape[1] == 512:
+        _run_tuned_mxfp4_bmm(
+            x,
+            weight,
+            weight_scale,
+            output,
+            _get_glm_mxfp4_k_bmm_config(x, weight),
+            transpose_bm=False,
+        )
+        return
+
+    batched_gemm_afp4wfp4_pre_quant(
+        x,
+        weight,
+        weight_scale,
+        torch.bfloat16,
+        output,
+    )
+
+
+def _run_mxfp4_v_bmm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output: torch.Tensor,
+) -> torch.Tensor:
+    if x.shape[2] == 512 and weight.shape[1] == 256:
+        return _run_tuned_mxfp4_bmm(
+            x,
+            weight,
+            weight_scale,
+            output,
+            _get_glm_mxfp4_v_bmm_config(x, weight),
+            transpose_bm=True,
+        )
+
+    batched_gemm_afp4wfp4_pre_quant(
+        x,
+        weight,
+        weight_scale,
+        torch.bfloat16,
+        output.transpose(0, 1),
+    )
+    return output
+
+
 def should_defer_dsa_cp_kv_gather(
     *,
     dsa_prefill_cp: bool,
@@ -529,6 +638,21 @@ def forward_absorb_prepare(
                     expected_m,
                 )
                 q_nope_out = q_nope_out[:, :expected_m, :]
+            elif _use_aiter_gfx95 and self.w_kc.dtype == torch.uint8:
+                x = q_nope.transpose(0, 1)
+                q_nope_out = torch.empty(
+                    x.shape[0],
+                    x.shape[1],
+                    self.w_kc.shape[2],
+                    device=x.device,
+                    dtype=torch.bfloat16,
+                )
+                _run_mxfp4_k_bmm(
+                    x,
+                    self.w_kc.transpose(-2, -1),
+                    self.w_scale_k.transpose(-2, -1),
+                    q_nope_out,
+                )
             elif self.w_kc.dtype == torch.float8_e4m3fn:
                 if _is_cpu:
                     q_nope_out = torch.bmm(
@@ -847,6 +971,21 @@ def forward_absorb_core(
             attn_bmm_output = (
                 attn_bmm_output[:, :expected_m, :].transpose(0, 1).flatten(1, 2)
             )
+        elif _use_aiter_gfx95 and self.w_vc.dtype == torch.uint8:
+            x = attn_output.transpose(0, 1)
+            bmm_output = torch.empty(
+                x.shape[1],
+                x.shape[0],
+                self.w_vc.shape[2],
+                device=x.device,
+                dtype=torch.bfloat16,
+            )
+            attn_bmm_output = _run_mxfp4_v_bmm(
+                x,
+                self.w_vc.transpose(-2, -1),
+                self.w_scale_v.transpose(-2, -1),
+                bmm_output,
+            ).flatten(1, 2)
         elif self.w_vc.dtype == torch.float8_e4m3fn:
             if _is_cpu:
                 attn_bmm_output = torch.bmm(

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

@@ -48,6 +48,8 @@
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
 from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mla import (
+    _run_mxfp4_k_bmm,
+    _run_mxfp4_v_bmm,
     _select_local_dcp_heads_for_autotune,
     is_dcp_mla_decode_phase,
     is_mla_dcp_lse_base_on_e,
@@ -134,7 +136,6 @@ def fused_qk_rmsnorm_bf16(q, q_weight, q_eps, k, k_weight, k_eps):
     )
 
     from sglang.srt.layers.quantization.rocm_mxfp4_utils import (
-        batched_gemm_afp4wfp4_pre_quant,
         fused_flatten_mxfp4_quant,
         fused_rms_mxfp4_quant,
     )
@@ -158,11 +159,10 @@ def rocm_absorb_q_bmm(
             device=x.device,
             dtype=torch.bfloat16,
         )
-        batched_gemm_afp4wfp4_pre_quant(
+        _run_mxfp4_k_bmm(
             x,
             attn.w_kc.transpose(-2, -1),
             attn.w_scale_k.transpose(-2, -1),
-            torch.bfloat16,
             q_nope_out,
         )
     else:
@@ -210,14 +210,13 @@ def rocm_absorb_v_bmm(
             device=x.device,
             dtype=torch.bfloat16,
         )
-        attn_bmm_output = _bmm_buf.transpose(0, 1)
-        batched_gemm_afp4wfp4_pre_quant(
+        _bmm_buf = _run_mxfp4_v_bmm(
             x,
             attn.w_vc.transpose(-2, -1),
             attn.w_scale_v.transpose(-2, -1),
-            torch.bfloat16,
-            attn_bmm_output,
+            _bmm_buf,
         )
+        attn_bmm_output = _bmm_buf
     else:
         _bmm_buf = None
         if _use_aiter_gfx95 and attn.w_kc.dtype == torch.float8_e4m3fn:

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

@@ -670,19 +670,29 @@ def post_load_weights(
                         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 (
+            # GLM absorbed weights load as bf16. Keep the current per-tensor FP8
+            # conversion as the default rollback, or preserve packed MXFP4 weights
+            # and per-head scales for the opt-in AITER absorbed-BMM backend.
+            is_glm_bf16_absorbed_weight = (
                 _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)
+            )
+            use_mxfp4_mla_bmm = (
+                is_glm_bf16_absorbed_weight and envs.SGLANG_USE_MXFP4_MLA_BMM.get()
+            )
+            if use_mxfp4_mla_bmm:
+                w_kc, self_attn.w_scale_k, w_vc, self_attn.w_scale_v = (
+                    quark_post_load_weights(self_attn, w, "mxfp4")
+                )
+            else:
+                if is_glm_bf16_absorbed_weight:
+                    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)
+                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)
 
             if (
                 _use_aiter_gfx95