From e74ea5b1d709c5476973fd383b6204398d0c31cf Mon Sep 17 00:00:00 2001 From: "Hsiu-Chun, Hung" <160560375+Emmanuel0612@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:51:57 +0800 Subject: [PATCH] =?UTF-8?q?[ROCm/gfx95]=20Fix=20fp8=20per-channel=20attent?= =?UTF-8?q?ion=20for=20Kimi-K2.7-code-mxfp4=20o=E2=80=A6=20(#31105)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Hung Co-authored-by: HaiShaw --- .../forward_mha_rocm.py | 12 +- .../forward_mla_rocm.py | 7 +- .../deepseek_common/deepseek_weight_loader.py | 5 + .../srt/models/deepseek_common/utils.py | 17 ++ python/sglang/srt/models/deepseek_v2.py | 29 ++- .../test_kimi_k27_code_mxfp4_eval_mi35x.py | 184 ++++++++++++++++++ .../amd/test_fp8_per_channel_detection.py | 75 +++++++ 7 files changed, 314 insertions(+), 15 deletions(-) create mode 100644 test/registered/amd/accuracy/mi35x/test_kimi_k27_code_mxfp4_eval_mi35x.py create mode 100644 test/registered/amd/test_fp8_per_channel_detection.py diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py index 00fa2407c..39fc8ab09 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py +++ b/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] 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 04f1a8f1b..21e6e6369 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 @@ -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, 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 165119e2d..ef0077d6b 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/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") diff --git a/python/sglang/srt/models/deepseek_common/utils.py b/python/sglang/srt/models/deepseek_common/utils.py index a7d374310..f56f47ece 100644 --- a/python/sglang/srt/models/deepseek_common/utils.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index ce3c75852..325bc9ebb 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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(), ) ) diff --git a/test/registered/amd/accuracy/mi35x/test_kimi_k27_code_mxfp4_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_kimi_k27_code_mxfp4_eval_mi35x.py new file mode 100644 index 000000000..42d0e5166 --- /dev/null +++ b/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() diff --git a/test/registered/amd/test_fp8_per_channel_detection.py b/test/registered/amd/test_fp8_per_channel_detection.py new file mode 100644 index 000000000..5a94bdd5a --- /dev/null +++ b/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()