[Spec] DFlash: support pure-MLA targets with an fp8 KV cache (Kimi-K2.x-NVFP4) (#29218)

Co-authored-by: Hao Phan <htphan@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-07 19:52:14 -07:00
parent db40fd83d2
commit 7bc343470f
3 changed files with 96 additions and 8 deletions
@@ -0,0 +1,70 @@
+import unittest
+
+from sglang.test.accuracy_test_runner import AccuracyTestParams
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.performance_test_runner import PerformanceTestParams
+from sglang.test.run_combined_tests import run_combined_tests
+from sglang.test.test_utils import ModelLaunchSettings
+
+# Kimi-K2.6 NVFP4 (pure-MLA target, fp8 KV) + DFlash speculative decoding on 8x B200, tp=8.
+register_cuda_ci(est_time=3600, suite="nightly-8-gpu-b200", nightly=True)
+
+MODEL_PATH = "nvidia/Kimi-K2.6-NVFP4"
+DRAFT_MODEL_PATH = "nvidia/Kimi-K2.6-DFlash"
+
+# trtllm_mla verify only; cuteDSL fold verify depends on the flashinfer version.
+EXTRA_ARGS = [
+    "--trust-remote-code",
+    "--quantization=modelopt_fp4",
+    "--moe-runner-backend=flashinfer_trtllm",
+    "--fp4-gemm-backend=flashinfer_cutlass",
+    "--attention-backend=trtllm_mla",
+    "--kv-cache-dtype=fp8_e4m3",
+    "--mem-fraction-static=0.85",
+    "--max-running-requests=16",
+    "--speculative-algorithm=DFLASH",
+    f"--speculative-draft-model-path={DRAFT_MODEL_PATH}",
+    "--speculative-num-draft-tokens=8",
+    "--speculative-draft-attention-backend=fa4",
+    "--speculative-draft-model-quantization=unquant",
+    "--speculative-draft-window-size=4096",
+]
+
+
+class TestKimiK26Nvfp4Dflash(unittest.TestCase):
+    """Kimi-K2.6 NVFP4 (pure-MLA, fp8 KV) with DFlash speculative decoding on 8x B200 (tp=8).
+
+    Runs both an accuracy test (gsm8k) and a performance test (bs=1/8/16), and gates the
+    speculative-decoding accept length. Guards the pure-MLA fp8-KV DFlash path.
+    """
+
+    def test_kimi_k26_nvfp4_dflash(self):
+        variants = [
+            ModelLaunchSettings(
+                MODEL_PATH,
+                tp_size=8,
+                extra_args=EXTRA_ARGS,
+                variant="TP8+DFLASH",
+            ),
+        ]
+
+        run_combined_tests(
+            models=variants,
+            test_name="Kimi-K2.6-NVFP4 DFlash",
+            # Thresholds from a measured tp=8 run: gsm8k 0.936 (full set), accept length ~2.66.
+            accuracy_params=AccuracyTestParams(
+                dataset="gsm8k",
+                baseline_accuracy=0.92,
+                num_examples=200,
+                api="completion",
+            ),
+            performance_params=PerformanceTestParams(
+                batch_sizes=[1, 8, 16],
+                spec_accept_length_threshold=2.0,
+                profile_dir="performance_profiles_kimi_k26_nvfp4_dflash",
+            ),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()