From 7bc343470f2973ebfc1106c18ecfa6f40d7c8eeb Mon Sep 17 00:00:00 2001 From: Thanhhao <31717833+thanhhao98@users.noreply.github.com> Date: Wed, 8 Jul 2026 09:52:14 +0700 Subject: [PATCH] [Spec] DFlash: support pure-MLA targets with an fp8 KV cache (Kimi-K2.x-NVFP4) (#29218) Co-authored-by: Hao Phan Co-authored-by: Claude Opus 4.8 --- .../sglang/srt/model_executor/model_runner.py | 16 +++++ .../srt/speculative/dflash_worker_v2.py | 18 ++--- .../quant/test_kimi_k26_nvfp4_dflash.py | 70 +++++++++++++++++++ 3 files changed, 96 insertions(+), 8 deletions(-) create mode 100644 test/registered/quant/test_kimi_k26_nvfp4_dflash.py diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 8509b3ad9..ef2190fa0 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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 diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 7c38edc35..29c540579 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/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, diff --git a/test/registered/quant/test_kimi_k26_nvfp4_dflash.py b/test/registered/quant/test_kimi_k26_nvfp4_dflash.py new file mode 100644 index 000000000..f8fd0aaa8 --- /dev/null +++ b/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()