Skip to content

fix(spec): restore sentinel padding + context clamps for every drafter #618

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jasl wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(spec): restore sentinel padding + context clamps for every drafter #618

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 15 additions & 15 deletions python/tokenspeed/runtime/execution/cuda_graph_wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
28 changes: 14 additions & 14 deletions python/tokenspeed/runtime/execution/model_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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(
Expand Down