[ROCm/gfx95] Fix fp8 per-channel attention for Kimi-K2.7-code-mxfp4 o… (#31105)

Co-authored-by: Hung <Emmanuel0612@users.noreply.github.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>

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

@@ -28,6 +28,7 @@ from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mha imp
     resolve_attn_backend,
 )
 from sglang.srt.models.deepseek_common.utils import (
+    _is_block_scale_fp8,
     _use_aiter_bpreshuffle_gfx95,
     _use_aiter_gfx95,
 )
@@ -74,10 +75,7 @@ class DeepseekMHARocmForwardMixin:
                 # on gfx95, we can still use fused RMSNorm+FP8 quant, but MUST request
                 # the unquantized output for q_lora; otherwise q_lora becomes the (fp8,scale)
                 # tuple.
-                if (
-                    _use_aiter_gfx95
-                    and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
-                ):
+                if _use_aiter_gfx95 and _is_block_scale_fp8(self.q_b_proj):
                     q_quanted, q_lora, _, _ = fused_rms_fp8_group_quant(
                         q,
                         self.q_a_layernorm.weight,
@@ -121,7 +119,7 @@ class DeepseekMHARocmForwardMixin:
                     None,
                 )
                 q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
-            elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
+            elif _use_aiter_gfx95 and _is_block_scale_fp8(self.q_b_proj):
                 q, _, _, _ = fused_rms_fp8_group_quant(
                     q,
                     self.q_a_layernorm.weight,
@@ -152,7 +150,7 @@ class DeepseekMHARocmForwardMixin:
         kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
         latent_cache = latent_cache.unsqueeze(1)
 
-        if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
+        if _use_aiter_gfx95 and _is_block_scale_fp8(self.kv_b_proj):
             kv_a_quanted, kv_a, _, _ = fused_rms_fp8_group_quant(
                 kv_a,
                 self.kv_a_layernorm.weight,
@@ -243,7 +241,7 @@ class DeepseekMHARocmForwardMixin:
                 )
             )[0]
         else:
-            if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
+            if _use_aiter_gfx95 and _is_block_scale_fp8(self.kv_b_proj):
                 kv = self.kv_b_proj(kv_a_quanted)[0]
             else:
                 kv = self.kv_b_proj(kv_a)[0]

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

@@ -53,6 +53,7 @@ from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mla imp
 )
 from sglang.srt.models.deepseek_common.utils import (
     FORWARD_ABSORB_CORE_ATTENTION_BACKENDS,
+    _is_block_scale_fp8,
     _is_gfx95_supported,
     _use_aiter,
     _use_aiter_bpreshuffle_gfx95,
@@ -224,7 +225,7 @@ def rocm_absorb_v_bmm(
         # _bmm_buf is already (batch, heads, dim) contiguous
         if attn.o_proj.weight.dtype == torch.uint8:
             attn_bmm_output = fused_flatten_mxfp4_quant(_bmm_buf)
-        elif attn.o_proj.weight.dtype == torch.float8_e4m3fn:
+        elif _is_block_scale_fp8(attn.o_proj):
             attn_bmm_output = fused_flatten_fp8_group_quant(
                 _bmm_buf,
                 group_size=128,
@@ -240,7 +241,7 @@ def rocm_absorb_v_bmm(
     elif attn.o_proj.weight.dtype == torch.uint8:
         attn_bmm_output = attn_bmm_output.transpose(0, 1)
         attn_bmm_output = fused_flatten_mxfp4_quant(attn_bmm_output)
-    elif attn.o_proj.weight.dtype == torch.float8_e4m3fn:
+    elif _is_block_scale_fp8(attn.o_proj):
         attn_bmm_output = attn_bmm_output.transpose(0, 1)
         attn_bmm_output = fused_flatten_fp8_group_quant(
             attn_bmm_output,
@@ -335,7 +336,7 @@ class DeepseekMLARocmForwardMixin:
                     self.kv_a_layernorm.weight,
                     self.kv_a_layernorm.variance_epsilon,
                 )
-            elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
+            elif _use_aiter_gfx95 and _is_block_scale_fp8(self.q_b_proj):
                 if self.use_dsa:
                     q_quanted, q_lora, k_nope, _ = fused_rms_fp8_group_quant(
                         q,

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

@@ -606,6 +606,10 @@ class DeepseekV2WeightLoaderMixin:
                     else:
                         weight = w
                         weight_scale = self_attn.kv_b_proj.weight_scale
+                        # Per-channel scale is 1D [out]; reshape to [out, 1] so it
+                        # broadcasts correctly against weight [out, in].
+                        if weight_scale.dim() == 1:
+                            weight_scale = weight_scale.view(-1, 1)
 
                     w, scale = channel_quant_to_tensor_quant(weight, weight_scale)
                     self_attn.w_scale = scale
@@ -638,6 +642,7 @@ class DeepseekV2WeightLoaderMixin:
                 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 = (
                     quark_post_load_weights(self_attn, w, "mxfp4")

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

@@ -72,6 +72,23 @@ FORWARD_ABSORB_CORE_ATTENTION_BACKENDS = [
 ]
 
 
+def _is_block_scale_fp8(proj: torch.nn.Module) -> bool:
+    """Return True if proj uses block-scale fp8 quantization.
+
+    Per-channel fp8 has weight_scale shape [N, 1] (one scale per output row).
+    Block-scale fp8 has weight_scale shape [N, K/block_size] (multiple columns).
+    The fused gfx95 kernels (fused_rms_fp8_group_quant, fused_flatten_fp8_group_quant)
+    are only compatible with block-scale layouts — per-channel layers must fall
+    through to the plain bf16 path instead.
+    """
+    if not hasattr(proj, "weight") or proj.weight.dtype != torch.float8_e4m3fn:
+        return False
+    weight_scale = getattr(proj, "weight_scale", None)
+    if weight_scale is None or weight_scale.dim() != 2:
+        return False
+    return weight_scale.shape[-1] > 1
+
+
 def awq_dequantize_func():
     """
     Get the AWQ dequantize function for the current device

python/sglang/srt/models/deepseek_v2.py

@@ -179,6 +179,7 @@ from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
 from sglang.srt.models.deepseek_common.utils import (
     _device_sm,
     _get_llama_4_scaling,
+    _is_block_scale_fp8,
     _is_cpu,
     _is_cpu_amx_available,
     _is_cuda,
@@ -2402,17 +2403,35 @@ class DeepseekV2DecoderLayer(nn.Module):
     def _detect_gfx95_quant_format(self) -> str:
         if not _is_gfx95_supported:
             return ""
-        weight = getattr(
-            getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None), "weight", None
-        )
+        proj = getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None)
+        weight = getattr(proj, "weight", None)
         if weight is None:
             return ""
         if weight.dtype == torch.uint8:
             return "mxfp4"
         if weight.dtype == getattr(torch, "float8_e4m3fn", None):
-            return "fp8"
+            # Use _is_block_scale_fp8 to distinguish block-scale fp8 (K/128 scale
+            # cols, compatible with fused_rms_fp8_group_quant) from per-channel fp8
+            # ([N, 1] scale, must use the plain bf16 path).
+            # weight_scale may not be reshaped yet at __init__ time — return
+            # "fp8_pending" so _resolve_gfx95_quant_format re-checks on first forward.
+            weight_scale = getattr(proj, "weight_scale", None)
+            if weight_scale is None:
+                return "fp8_pending"
+            return "fp8" if _is_block_scale_fp8(proj) else ""
         return ""
 
+    def _resolve_gfx95_quant_format(self) -> str:
+        """Re-evaluate after weights are loaded if still pending."""
+        fmt = getattr(self, "_gfx95_quant_format", "")
+        if fmt == "fp8_pending":
+            fmt = self._detect_gfx95_quant_format()
+            if fmt == "fp8_pending":
+                # weight_scale still unavailable — default to bf16 (safe fallback).
+                fmt = ""
+            self._gfx95_quant_format = fmt
+        return fmt
+
     def _is_layer_sparse(self, layer_id: int, is_nextn: bool) -> bool:
         return is_nextn or (
             self.config.n_routed_experts is not None
@@ -2440,7 +2459,7 @@ class DeepseekV2DecoderLayer(nn.Module):
                 residual,
                 forward_batch,
                 captured_last_layer_outputs=captured_last_layer_outputs,
-                quant_format=getattr(self, "_gfx95_quant_format", ""),
+                quant_format=self._resolve_gfx95_quant_format(),
             )
         )
 

test/registered/amd/accuracy/mi35x/test_kimi_k27_code_mxfp4_eval_mi35x.py

@@ -0,0 +1,184 @@
+"""MI35x Kimi-K2.7-Code-MXFP4 aiter MLA backend accuracy tests (4-GPU)
+
+Tests Kimi-K2.7-Code-MXFP4 with the aiter unified attention backend on MI35x.
+This model uses mixed quantization: mxfp4 for MoE layers and fp8 per-channel
+for attention projections (q_a_proj, q_b_proj, kv_a_proj_with_mqa, kv_b_proj,
+o_proj). The per-channel fp8 detection fix ensures the correct kernel path is
+selected for each layer type.
+
+Registry: nightly-amd-4-gpu-mi35x-kimi-k27-code-mxfp4-aiter-mla suite
+"""
+
+import os
+import unittest
+from dataclasses import dataclass
+from typing import List, Optional
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    is_in_ci,
+    popen_launch_server,
+    write_github_step_summary,
+)
+
+register_amd_ci(
+    est_time=7200,
+    suite="nightly-amd-4-gpu-mi35x-kimi-k27-code-mxfp4-aiter-mla",
+    nightly=True,
+)
+
+
+@dataclass
+class ModelConfig:
+    """Configuration for a model variant to test."""
+
+    model_path: str
+    tp_size: int = 4
+    accuracy_threshold: float = 0.94
+    other_args: Optional[List[str]] = None
+    env_vars: Optional[dict] = None
+    timeout: Optional[int] = None
+    variant: Optional[str] = None
+
+    def __post_init__(self):
+        if self.other_args is None:
+            self.other_args = []
+        if self.env_vars is None:
+            self.env_vars = {}
+
+    def get_display_name(self) -> str:
+        if self.variant:
+            return f"{self.model_path} ({self.variant})"
+        return self.model_path
+
+
+def get_kimi_k27_code_mxfp4_models() -> List[ModelConfig]:
+    """Get Kimi-K2.7-Code-MXFP4 model configurations for MI35x."""
+    common_kwargs = {
+        "model_path": "amd/Kimi-K2.7-Code-MXFP4",
+        "tp_size": 4,
+        "accuracy_threshold": 0.94,
+        "timeout": 3600,
+    }
+    common_args = [
+        "--attention-backend",
+        "aiter",
+        "--disable-radix-cache",
+        "--mem-fraction-static",
+        "0.90",
+        "--kv-cache-dtype",
+        "fp8_e4m3",
+        "--trust-remote-code",
+        "--watchdog-timeout",
+        "1200",
+        "--enable-aiter-allreduce-fusion",
+    ]
+
+    return [
+        ModelConfig(
+            **common_kwargs,
+            variant="default",
+            other_args=common_args,
+        ),
+    ]
+
+
+class TestKimiK27CodeMXFP4AiterMlaEvalMI35x(unittest.TestCase):
+    """Kimi-K2.7-Code-MXFP4 aiter MLA backend accuracy tests on MI35x."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.models = get_kimi_k27_code_mxfp4_models()
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.num_questions = int(os.environ.get("GSM8K_NUM_QUESTIONS", "1319"))
+
+    def test_kimi_k27_code_mxfp4_accuracy(self):
+        """Test Kimi-K2.7-Code-MXFP4 with GSM8K completion benchmark."""
+        from types import SimpleNamespace
+
+        from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
+
+        all_results = []
+        summary = "### Kimi-K2.7-Code-MXFP4 aiter MLA (MI35x)\n\n"
+        summary += "| Model | Variant | TP | Accuracy | Threshold | Status |\n"
+        summary += "| ----- | ------- | -- | -------- | --------- | ------ |\n"
+
+        for config in self.models:
+            display_name = config.get_display_name()
+            with self.subTest(model=display_name):
+                print(f"\n{'='*60}")
+                print(f"Testing: {display_name}")
+                print(f"{'='*60}")
+
+                env = os.environ.copy()
+                for key, value in config.env_vars.items():
+                    env[key] = value
+
+                other_args = list(config.other_args)
+                other_args.extend(["--tp", str(config.tp_size)])
+                timeout = config.timeout or DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
+
+                try:
+                    process = popen_launch_server(
+                        model=config.model_path,
+                        base_url=self.base_url,
+                        timeout=timeout,
+                        other_args=other_args,
+                        env=env,
+                    )
+
+                    try:
+                        args = SimpleNamespace(
+                            num_shots=8,
+                            data_path=None,
+                            num_questions=self.num_questions,
+                            parallel=self.num_questions,
+                            max_new_tokens=512,
+                            host="http://127.0.0.1",
+                            port=int(self.base_url.split(":")[-1]),
+                        )
+                        metrics = run_eval_few_shot_gsm8k(args)
+                        acc = metrics["accuracy"]
+
+                        passed = acc >= config.accuracy_threshold
+                        status = "PASS" if passed else "FAIL"
+                        print(
+                            f"  accuracy={acc:.3f} threshold={config.accuracy_threshold} {status}"
+                        )
+
+                        all_results.append(
+                            {
+                                "model": display_name,
+                                "accuracy": acc,
+                                "passed": passed,
+                            }
+                        )
+                        summary += f"| {config.model_path} | {config.variant or 'N/A'} | {config.tp_size} | {acc:.3f} | {config.accuracy_threshold} | {status} |\n"
+
+                    finally:
+                        kill_process_tree(process.pid)
+
+                except Exception as e:
+                    summary += f"| {config.model_path} | {config.variant or 'N/A'} | {config.tp_size} | N/A | {config.accuracy_threshold} | ERROR |\n"
+                    all_results.append(
+                        {
+                            "model": display_name,
+                            "accuracy": None,
+                            "passed": False,
+                            "error": str(e),
+                        }
+                    )
+
+        if is_in_ci():
+            write_github_step_summary(summary)
+
+        failed = [r for r in all_results if not r["passed"]]
+        if failed:
+            raise AssertionError(f"Failed models: {[r['model'] for r in failed]}")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/amd/test_fp8_per_channel_detection.py

@@ -0,0 +1,75 @@
+"""Unit tests for _is_block_scale_fp8 per-channel vs block-scale fp8 detection.
+
+Tests the helper that distinguishes block-scale fp8 (weight_scale [N, K/128],
+compatible with fused gfx95 group-quant kernels) from per-channel fp8
+(weight_scale [N, 1], must use the plain bf16 path).
+
+These tests run on CPU and require no GPU, guarding the regression surface
+cheaply without waiting for a full nightly accuracy run.
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.test.ci.ci_register import register_amd_ci
+
+register_amd_ci(est_time=10, suite="stage-a-test-1-gpu-small-amd")
+
+
+def _make_proj(weight_dtype, weight_scale_shape=None):
+    """Create a fake projection module with the given weight/scale configuration."""
+    proj = SimpleNamespace()
+    proj.weight = torch.empty(64, 512, dtype=weight_dtype)
+    if weight_scale_shape is not None:
+        proj.weight_scale = torch.empty(*weight_scale_shape, dtype=torch.float32)
+    return proj
+
+
+class TestIsBlockScaleFp8(unittest.TestCase):
+    """Unit tests for _is_block_scale_fp8 detection helper."""
+
+    def setUp(self):
+        from sglang.srt.models.deepseek_common.utils import _is_block_scale_fp8
+
+        self.fn = _is_block_scale_fp8
+
+    def test_block_scale_fp8_returns_true(self):
+        """Block-scale fp8: weight_scale [N, K/128] — should return True."""
+        proj = _make_proj(torch.float8_e4m3fn, weight_scale_shape=(64, 4))
+        self.assertTrue(self.fn(proj))
+
+    def test_per_channel_fp8_returns_false(self):
+        """Per-channel fp8: weight_scale [N, 1] — should return False."""
+        proj = _make_proj(torch.float8_e4m3fn, weight_scale_shape=(64, 1))
+        self.assertFalse(self.fn(proj))
+
+    def test_non_fp8_weight_returns_false(self):
+        """bf16 weight is not fp8 at all — should return False."""
+        proj = _make_proj(torch.bfloat16, weight_scale_shape=(64, 4))
+        self.assertFalse(self.fn(proj))
+
+    def test_uint8_mxfp4_returns_false(self):
+        """uint8 mxfp4 weight — should return False (handled separately)."""
+        proj = _make_proj(torch.uint8, weight_scale_shape=(64, 4))
+        self.assertFalse(self.fn(proj))
+
+    def test_no_weight_scale_returns_false(self):
+        """No weight_scale attribute — should return False gracefully."""
+        proj = _make_proj(torch.float8_e4m3fn)  # no weight_scale
+        self.assertFalse(self.fn(proj))
+
+    def test_1d_weight_scale_returns_false(self):
+        """1D weight_scale [N] (not yet reshaped) — should return False."""
+        proj = _make_proj(torch.float8_e4m3fn, weight_scale_shape=(64,))
+        self.assertFalse(self.fn(proj))
+
+    def test_no_weight_attribute_returns_false(self):
+        """No weight attribute — should return False gracefully."""
+        proj = SimpleNamespace()
+        self.assertFalse(self.fn(proj))
+
+
+if __name__ == "__main__":
+    unittest.main()