From 2c2a4f243a6289fbf3d7b4530226003aea568123 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Fri, 5 Jun 2026 16:22:04 -0400 Subject: [PATCH] [Bug] Fix EAGLE draft CUDA-graph `kv_indices` under-allocation for `topk > 1` (#27338) --- .../srt/layers/attention/flashinfer_backend.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 567ffadad..70fd29d65 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -1584,6 +1584,19 @@ class FlashInferMultiStepDraftBackend: 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, only sometimes surfacing as a crash. + 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) ]( @@ -1641,8 +1654,11 @@ class FlashInferMultiStepDraftBackend: self.common_template(forward_batch, kv_indices, call_fn) def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + # generate_draft_decode_kv_indices packs topk per-branch sequences per row, + # so the row needs the topk factor -- same as the eager init_forward_metadata + # (batch_size * topk * max_context_len). Dropping it overflows the buffer. 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", )