Fix NaN in triton EAGLE spec-v2 draft-extend CUDA graph at topk>1 (wrong qo_indptr stride)#27545
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the _update_draft_extend_buffers method in triton_backend.py to use self.num_draft_tokens instead of self.speculative_num_steps + 1 when the V2 draft-extend forward mode is active. The reviewer identified a critical omission where max_extend_len in _build_cuda_graph_forward_metadata is still hardcoded to self.speculative_num_steps + 1, which could lead to an underestimated Triton attention kernel grid size and incorrect outputs when topk > 1.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| num_tokens_per_bs = ( | ||
| self.num_draft_tokens | ||
| if forward_mode.is_draft_extend_v2() | ||
| else self.speculative_num_steps + 1 | ||
| ) |
There was a problem hiding this comment.
While updating num_tokens_per_bs to use self.num_draft_tokens for the V2 draft-extend path is correct, there is a corresponding omission in _build_cuda_graph_forward_metadata (around line 903).\n\nCurrently, max_extend_len is still hardcoded to self.speculative_num_steps + 1 for all draft-extend modes:\npython\n elif forward_mode.is_draft_extend(include_v2=True):\n return ForwardMetadata(\n attn_logits=None,\n attn_lse=None,\n max_extend_len=self.speculative_num_steps + 1,\n\nWhen forward_mode.is_draft_extend_v2() is active and topk > 1, self.num_draft_tokens can be larger than self.speculative_num_steps + 1. If max_extend_len is too small, the Triton attention kernel grid size (which is computed as triton.cdiv(max_len_extend, BLOCK_M)) will be underestimated, causing some query blocks to not be launched and resulting in unprocessed tokens or incorrect outputs.\n\nPlease update _build_cuda_graph_forward_metadata to also use self.num_draft_tokens for the V2 draft-extend path:\npython\n elif forward_mode.is_draft_extend(include_v2=True):\n return ForwardMetadata(\n attn_logits=None,\n attn_lse=None,\n max_extend_len=(\n self.num_draft_tokens\n if forward_mode.is_draft_extend_v2()\n else self.speculative_num_steps + 1\n ),\n
|
/rerun-test test_spec_eagle_triton.py test_spec_eagle_topk.py test_spec_eagle_topk_page.py |
|
Results for 🚀 |
1 participant
The draft-extend CUDA-graph metadata built
qo_indptrwith a per-req stride ofnum_steps+1, but the runner lays outnum_draft_tokenstokens/req. For topk>1 these differ, so at bs>1 each req's queries read from the wrong offset -> NaN (bs==1 and eager are unaffected). Usenum_draft_tokensfor the v2 draft-extend path.CI States
Latest PR Test (Base): ❌ Run #27126544910
Latest PR Test (Extra): ❌ Run #27126543998