[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>

python/sglang/srt/model_executor/model_runner.py

@@ -2485,6 +2485,22 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 f"Unsupported kv_cache_dtype: {self.server_args.kv_cache_dtype}."
             )
 
+        # DFLASH: fa4 draft attention can't read the target's fp8 KV (needs K.dtype == Q.dtype),
+        # so give the fa4 draft its own compute-dtype KV. fp8-capable backends keep the target dtype.
+        if (
+            self.is_draft_worker
+            and self.spec_algorithm.is_dflash()
+            and self.server_args.speculative_draft_attention_backend == "fa4"
+            and self.kv_cache_dtype != self.dtype
+        ):
+            logger.info(
+                "DFLASH fa4 draft: overriding KV cache dtype %s -> %s "
+                "(fa4 needs K.dtype == Q.dtype; cannot read the target's quantized KV).",
+                self.kv_cache_dtype,
+                self.dtype,
+            )
+            self.kv_cache_dtype = self.dtype
+
     def init_cublas(self):
         """We need to run a small matmul to init cublas. Otherwise, it will raise some errors later."""
         dtype = torch.float16

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -117,6 +117,12 @@ class DFlashWorkerV2(BaseSpecWorker):
         self.nccl_port = nccl_port
         self._target_worker = target_worker
         self.model_runner = target_worker.model_runner
+        self._need_mamba_verify_commit = (
+            self.model_runner.mambaish_config is not None
+            and hasattr(
+                self.model_runner.attn_backend, "update_mamba_state_after_mtp_verify"
+            )
+        )
         self.page_size = server_args.page_size
         # Normalized in arg_groups.speculative_hook.handle_speculative_decoding.
         self.draft_window_size: Optional[int] = (
@@ -1099,9 +1105,9 @@ class DFlashWorkerV2(BaseSpecWorker):
         cache per-step intermediate states. After acceptance, we need to commit the
         state corresponding to each request's last accepted step.
         """
-        attn_backend = self.target_worker.model_runner.attn_backend
-        if not hasattr(attn_backend, "update_mamba_state_after_mtp_verify"):
+        if not self._need_mamba_verify_commit:
             return
+        attn_backend = self.target_worker.model_runner.attn_backend
 
         last_correct_step_indices = commit_lens.to(torch.int64) - 1
         mamba_steps_to_track = None
@@ -1530,12 +1536,8 @@ class DFlashWorkerV2(BaseSpecWorker):
         batch.out_cache_loc = verify_out_cache_loc
         sampling_info = batch.sampling_info
 
-        need_mamba_verify_commit = hasattr(
-            self.target_worker.model_runner.attn_backend,
-            "update_mamba_state_after_mtp_verify",
-        )
         seq_lens_pre_verify = (
-            batch.seq_lens.clone() if need_mamba_verify_commit else None
+            batch.seq_lens.clone() if self._need_mamba_verify_commit else None
         )
         seq_lens_cpu_backup = batch.seq_lens_cpu
         seq_lens_sum_backup = batch.seq_lens_sum
@@ -1655,7 +1657,7 @@ class DFlashWorkerV2(BaseSpecWorker):
                     1, accept_len.to(torch.int64)[:, None], bonus[:, None]
                 )
 
-        if need_mamba_verify_commit:
+        if self._need_mamba_verify_commit:
             assert seq_lens_pre_verify is not None
             self._update_target_mamba_state_after_verify(
                 batch=batch,

test/registered/quant/test_kimi_k26_nvfp4_dflash.py

@@ -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()