diff --git a/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py b/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py index 95f40dbbc..470bcec5d 100644 --- a/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py +++ b/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py @@ -851,22 +851,22 @@ def _pad_graph_req_pool_indices( pad = padded_bs - active_req_pool_indices.shape[0] if pad <= 0: return active_req_pool_indices - if self.config.spec_algo == "DFLASH": - # Route padding rows to the sentinel req-pool slot - # (max_req_pool_size), not slot 0. The DFLASH draft derives each - # row's block seq_len from valid_cache_lengths[req_pool], so - # padding rows pointing at slot 0 would grow unbounded with - # request 0's context and hang the draft block-decode kernel. - # The sentinel row stays zero-init (length 0, dummy page 0). - sentinel = int(self.config.max_req_pool_size) - return torch.cat( - [ - active_req_pool_indices, - active_req_pool_indices.new_full((pad,), sentinel), - ] - ) + # Route padding rows to the sentinel req-pool slot (max_req_pool_size), + # not slot 0 -- for EVERY spec algorithm. Slot 0 aliases the live first + # request: harmless for the plain verify decode (seq_len == 1), but any + # DRAFTER in the captured graph re-derives per-row state from the pool + # index -- DFLASH grows its block seq_len with request 0's context and + # hangs; Eagle/MTP reads request 0's committed length and resolves + # draft-KV write slots through request 0's live frontier (a write-write + # race + wild paged-indexer accesses at ragged batch sizes -> IMA on + # GB10 MTP2 conc8). The sentinel row stays zero-init (length 0, dummy + # page 0), keeping the derived lengths inert for every drafter backend. + sentinel = int(self.config.max_req_pool_size) return torch.cat( - [active_req_pool_indices, active_req_pool_indices.new_zeros(pad)] + [ + active_req_pool_indices, + active_req_pool_indices.new_full((pad,), sentinel), + ] ) def _set_graph_state_write_indices( diff --git a/python/tokenspeed/runtime/execution/model_executor.py b/python/tokenspeed/runtime/execution/model_executor.py index 7599d12e1..51e5c8ac4 100644 --- a/python/tokenspeed/runtime/execution/model_executor.py +++ b/python/tokenspeed/runtime/execution/model_executor.py @@ -252,10 +252,11 @@ def __init__( # ~1 page past context_len + spec_num_tokens. Without headroom # req_to_page overflows and the next draft block's page write goes # out of bounds, hanging the attention kernel. Pad generously; a few - # int32 columns per request. Non-DFLASH algorithms do not need this. - draft_block_reservation_slack = ( - config.spec_num_tokens * 64 if config.spec_algo == "DFLASH" else 0 - ) + # int32 columns per request. Kept for EVERY drafter (not just + # DFLASH): Eagle/MTP draft steps also write KV at the reservation + # frontier, and running out degrades to a clamped KV-drop near + # context_len. + draft_block_reservation_slack = config.spec_num_tokens * 64 max_num_pages_per_req = ( config.context_len + config.spec_num_tokens @@ -727,10 +728,9 @@ def _run_sampling( accept_lengths = self._apply_force_single_token_verify( accept_lengths, 0, num_decodes, ctx.decode_input_ids ) - if self.config.spec_algo == "DFLASH": - accept_lengths = self._cap_accept_to_context_len( - accept_lengths, sampling_info.req_pool_indices[:num_decodes] - ) + accept_lengths = self._cap_accept_to_context_len( + accept_lengths, sampling_info.req_pool_indices[:num_decodes] + ) return output_tokens, accept_lengths logits = logits_output.next_token_logits @@ -745,10 +745,9 @@ def _run_sampling( decode_accept = self._apply_force_single_token_verify( decode_accept, num_extends, num_decodes, ctx.decode_input_ids ) - if self.config.spec_algo == "DFLASH": - decode_accept = self._cap_accept_to_context_len( - decode_accept, sampling_info.req_pool_indices[num_extends:] - ) + decode_accept = self._cap_accept_to_context_len( + decode_accept, sampling_info.req_pool_indices[num_extends:] + ) if ( prefill_out.next_token_logprobs is not None and decode_out.next_token_logprobs is not None @@ -1816,9 +1815,10 @@ def execute_forward_op( time.perf_counter() - forward_step_start ) * 1000.0 - if self.config.spec_algo == "DFLASH": + if self.config.spec_algo is not None: # Clamp the committed-length delta so no request grows past - # context_len. Done here (outside the graph) so it reaches + # context_len (every drafter can overshoot near the limit). + # Done here (outside the graph) so it reaches # both _update_runtime_state and the scheduler page # reservation; see _clamp_committed_to_context_len. output_lengths = self._clamp_committed_to_context_len(