[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>
This commit is contained in:
Hsiu-Chun, Hung
co-authored by
Hung
HaiShaw
Hung
HaiShaw
parent
1c06c160f9
commit
e74ea5b1d7
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user