[spec] Misc defensive guards for EAGLE draft KV indexing#27486
Conversation
…e + hidden-pool bounds
|
/rerun-test test_spec_eagle_topk_page.py test_spec_eagle_fa3.py test_spec_eagle_triton.py |
|
Results for 🚀 🚀 |
![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 introduces bounds checking and validation to prevent illegal memory accesses (IMA) and integer overflows in KV cache operations and speculative decoding. Specifically, it adds a pool_size parameter to bound slot values during KV cache stores and moves, asserts proper draft token configurations for multi-layer EAGLE workers, and guards against int32 overflows in flat offset calculations. Feedback on these changes suggests casting pool_size to the source tensor's data type in Triton to avoid potential type mismatches, and wrapping an inner loop in assign_hidden_states_pool_kernel with a conditional block to optimize performance and simplify masking.
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.
…maybe_oob/A2; pool_size -> size_limit
|
/rerun-test test_spec_eagle_topk_page.py test_spec_eagle_fa3.py test_spec_eagle_triton.py test_deepseek_v3_fp4_mtp_small.py |
|
Results for 🚀 🚀 🚀 |
|
/tag-and-rerun-ci extra |
1 participant
Three independent defensive guards for the EAGLE draft KV-indexing paths. Behavior-preserving for valid configs -- each assert fires only on a real bug or misconfiguration.
Guards
MultiLayerEagle*Worker.__init__): assertnum_draft_tokens == num_steps + 1(mirrors the existingeagle_worker_v2check) so a mismatched config fails loud at startup instead of a later cryptic shape mismatch. This also subsumes theassign_hidden_states_poolseq-offset underflow (which requiresextend_len < pool_size, impossible once this invariant holds), so no kernel-level guard is added there.kv_indicesint32-overflow guard (draft_kv_indices_buffer_width): assertnum_seqs * topk * max_context_len < 2**31-- the draftkv_indicesflat offset is computed in int32, so an extremebatch * topk * contextwould silently wrap.store_kvcacheslot-value fast-fail (JIT CUDA, the hotset_kv_bufferwrite path): a stale / out-of-range slot id would otherwise cause an illegal-address crash displaced into a victim kernel (hard to attribute). The kernel nowassert(0 <= index < size_limit)before the write -- always-on (the kvcache JIT compiles withoutNDEBUG), no host sync, and it fails at the culprit with a clear message instead of a displaced IMA. (pool_sizearg renamed tosize_limit;size_limit = self.size + self.page_size= cache rows = real slots + the reserved padding slot.)move_kv_cache, the other slot-value consumer, keeps relying on its existing asyncmaybe_detect_oobprobe (no in-kernel change).CI States
Latest PR Test (Base): ✅ Run #27108922233
Latest PR Test (Extra): ✅ Run #27108922191