Skip to content

[flashinfer] Pass window_left at plan time for DFlash verify - #35208

Open
SubSir wants to merge 1 commit into
sgl-project:mainfrom
SubSir:jianchen/dflash-swa-window-plan-time
Open

SubSir wants to merge 1 commit into
sgl-project:mainfrom
SubSir:jianchen/dflash-swa-window-plan-time

Conversation

@SubSir

@SubSir SubSir commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

Motivation

#31501 added plan-time window_left for the SWA paged prefill wrapper and gated it on
spec_info is None. The same commit noted, on the index trim right above it, that the trim
is request-granular and "exactness comes from plan-time window_left".

DFlash verify therefore still runs with the window compiled out. The whole block shares one
left bound anchored at the block start, so the query at block offset j attends to j keys
that a window ending at its own position excludes. DFlash verify is a linear block with no
tree mask -- the worker passes custom_mask=None -- so nothing else supplies the window.

Modifications

  • flashinfer_backend.py: allow the plan-time window for SpecInputType.DFLASH_VERIFY.
    EAGLE tree verify keeps the previous behaviour.
  • speculative_target_verify_runner.py: correct the DFlash reference mask from
    prefix_len - window to q_pos - window.

The reference mask was built from the same request-granular rule as the code under test,
which is why the pre-existing runner_dflash_verify_swa_chain case passed. With the
reference corrected, that case becomes a regression test for this fix: at
prefix_lens=(3, 5) and sliding_window_size=4, the last query of the second request sees
two keys too many without the backend change.

Accuracy Tests

Verify output changes on SWA layers whenever the committed prefix exceeds the window. The
window it removes is at most draft_token_num - 1 keys at the oldest end of the window, so
the effect is small in absolute terms; the point is that verify and decode now use the same
window definition.

Measured on H200 with a sliding-window DFlash draft (window 2047, block 16) on flashinfer,
24 gsm8k prompts at T=0, this branch against the same tree with the backend hunk reverted:

prompt this branch reverted
~3070 tokens, so the prefix passes the window accept 6.050 accept 6.041
gsm8k as-is, ~400 tokens with the completion accept 6.545 accept 6.545

The short row is the control: below the window the two masks are elementwise equal, and the
arms agree to the digit, which is also what makes the long row readable -- at T=0 with fixed
prompts the pipeline is deterministic, so 0.009 is the effect and not noise. It is that small
because the keys at stake are at most 15 out of 2047, at the oldest end of the window.

test/registered/attention/unittests/swa/test_flashinfer.py::runner_dflash_verify_swa_chain
fails before the backend change and passes after it.

Speed Tests and Profiling

None. Passing window_left >= 0 selects the flashinfer module with the window predicate
compiled in; no additional work per step.

Checklist


CI States

Latest PR Test (Base): ❌ Run #32075240918
Latest PR Test (Extra): ❌ Run #32075240770

sgl-project#31501 added plan-time `window_left` for the SWA paged prefill wrapper and
gated it on `spec_info is None`, noting in the same commit that the index
trim is request-granular and "exactness comes from plan-time window_left".
DFlash verify therefore still runs with the window compiled out: every query
in the block shares one left bound anchored at the block start, so a query at
block offset j sees j keys that a window ending at its own position excludes.

DFlash verify is a linear block with no tree mask (the worker passes
custom_mask=None), so nothing else supplies the window. Whitelist
DFLASH_VERIFY; EAGLE tree verify keeps the old behaviour.

The existing `runner_dflash_verify_swa_chain` case did not catch this because
its reference mask was built from the same request-granular rule as the code
under test. Correct the reference to the per-query window, which turns that
pre-existing case into a regression test for this fix: with prefix_lens=(3,5)
and window=4 the last query of the second request sees two keys too many
without the change.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants