From 5160f7914ebf8e23e9827c7c683675ed9e8373e2 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 6 Jun 2026 16:28:34 -0700 Subject: [PATCH] Fix MLA EAGLE draft CUDA-graph `kv_indices` under-allocation for `topk > 1` (#27460) --- python/sglang/srt/debug_utils/pr_fix_toggle.py | 12 ++++++++++++ .../layers/attention/flashinfer_mla_backend.py | 17 ++++++++++++++++- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/debug_utils/pr_fix_toggle.py b/python/sglang/srt/debug_utils/pr_fix_toggle.py index 1f1edc361..31c59c6f0 100644 --- a/python/sglang/srt/debug_utils/pr_fix_toggle.py +++ b/python/sglang/srt/debug_utils/pr_fix_toggle.py @@ -102,12 +102,24 @@ patches: """ +_PR_REVERT_YAML_27460 = """ +patches: + - target: sglang.srt.layers.attention.flashinfer_mla_backend.FlashInferMLAMultiStepDraftBackend.init_cuda_graph_state + edits: + - match: | + (self.speculative_num_steps, max_bs * self.topk * self.max_context_len), + replacement: | + (self.speculative_num_steps, max_bs * self.max_context_len), +""" + + _PR_FIX_REVERT_YAML: Dict[int, str] = { 25015: _PR_REVERT_YAML_25015, 26329: _PR_REVERT_YAML_26329, 27338: _PR_REVERT_YAML_27338, 27360: _PR_REVERT_YAML_27360, 26972: _PR_REVERT_YAML_26972, + 27460: _PR_REVERT_YAML_27460, } diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index a1b6f2981..da4881e88 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -943,6 +943,19 @@ class FlashInferMLAMultiStepDraftBackend: bs = self.topk * num_seqs seq_lens_sum = forward_batch.seq_lens_sum + # Fail fast on an undersized kv_indices row: the kernel would otherwise + # write OOB and silently corrupt memory. + required_kv_indices_len = ( + seq_lens_sum * self.topk + bs * self.speculative_num_steps + ) + assert required_kv_indices_len <= kv_indices_buffer.shape[1], ( + f"EAGLE draft kv_indices row too small: need {required_kv_indices_len} " + f"but row width is {kv_indices_buffer.shape[1]} (topk={self.topk}, " + f"num_seqs={num_seqs}, seq_lens_sum={seq_lens_sum}, " + f"num_steps={self.speculative_num_steps}); the buffer must be sized " + f"max_bs * topk * max_context_len." + ) + self.generate_draft_decode_kv_indices[ (self.speculative_num_steps, num_seqs, self.topk) ]( @@ -993,8 +1006,10 @@ class FlashInferMLAMultiStepDraftBackend: self.common_template(forward_batch, kv_indices, call_fn) def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + # Row holds topk per-branch sequences (generate_draft_decode_kv_indices), so + # it needs the topk factor, matching the eager init_forward_metadata. self.cuda_graph_kv_indices = torch.zeros( - (self.speculative_num_steps, max_bs * self.max_context_len), + (self.speculative_num_steps, max_bs * self.topk * self.max_context_len), dtype=torch.int32, device="cuda", )