Fix nonpad_kv_seqlen in static-cache generation example (before-scatter -> after-scatter) - #365
Merged
Merged
Conversation
…eration The static-cache generation example fed the BEFORE-scatter valid KV count as nonpad_kv_seqlen. Both nonpad_kv_seqlen and write_indices were initialized to 0 and only incremented AFTER session.run, so the value fed each step always equaled write_indices instead of the spec-correct count. This is wrong in prefill (feeds 0 instead of prompt_len) and off-by-one on every decode step. The model scatters K/V into the cache before the maskless is_causal Attention reads nonpad_kv_seqlen, so under the bottom-right contract the valid count must include the just-written chunk: write_indices + cur_seq_len. Feed that directly and remove the now-redundant nonpad_kv_seqlen accumulator (it duplicated write_indices). Matches the authoritative reference in tests/static_cache_parity_test.py (prefill nonpad = prompt_len = write_indices(0) + query_len; decode nonpad = valid_len + 1 = write_indices(valid_len) + 1) and the consumer src/mobius/components/_attention.py (scatter precedes the nonpad read). Verified by running the example on CPU with Qwen/Qwen2.5-0.5B: 'The capital of France is' -> ' Paris. It is the largest city in' (coherent; pre-fix the all-zero prefill nonpad produced wrong logits). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
Update the module-docstring bullet to match the fix's after-scatter semantics: nonpad_kv_seqlen is the number of valid KV entries after the current chunk is scattered (write_indices + cur_seq_len), not a generic running count. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Fixes the canonical static-cache generation example to feed nonpad_kv_seqlen as the after-scatter valid-KV count (write_indices + cur_seq_len), matching the bottom-right causal is_causal contract.
Changes:
- Correct
nonpad_kv_seqlenfed intosession.run()towrite_indices + cur_seq_len(after-scatter valid KV count). - Remove the redundant
nonpad_kv_seqlenaccumulator that previously duplicatedwrite_indices. - Add docstring and inline comments clarifying the after-scatter
nonpad_kv_seqlenpattern for prefill and decode.
Performance Comparison
|
|
The author of this PR, titaiwangms, is not an activated member of this organization on Codecov. |
justinchuby
approved these changes
Jun 19, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
examples/static_cache_generation.pyfednonpad_kv_seqlenas the before-scatter KV count. A separatenonpad_kv_seqlenaccumulator stayed always-equal towrite_indices(both initialized to 0, both incremented only aftersession.run), so the example passednonpad == write_indices. The spec-correct value iswrite_indices + cur_seq_len— the valid KV count after the current chunk is scattered into the cache, per the bottom-rightis_causalcontract (onnx/onnx#8068).Impact
Wrong in every phase:
0instead ofprompt_len, so the kernel saw ~0 valid keys and produced wrong logits.On the CPU EP the maskless graph runs-but-wrong silently (no error raised). This is example-only — it is the canonical usage reference and does not affect library/graph code or PR #364.
Fix
nonpad_kv_seqlen = write_indices + cur_seq_lenat the run inputs.nonpad_kv_seqlenaccumulator (it merely duplicatedwrite_indices).Verified
tests/static_cache_parity_test.py(prefill nonpad = prompt_len = write_indices(0) + query_len, lines 322-327; decode nonpad = valid_len + 1 = write_indices(valid_len) + 1, lines 361-362) and the consumersrc/mobius/components/_attention.py:169-183(the scatter precedes the masklessis_causalAttention'snonpadread).Qwen/Qwen2.5-0.5B-> coherent output:The capital of France is->Paris. It is the largest city in. Pre-fix the all-zero prefill nonpad silently produced wrong logits.Reference