[Bugfix][Model] Fix Inkling NVIDIA sconv cache block-size mismatch - #51951
Open
fjosw wants to merge 2 commits into
Open
[Bugfix][Model] Fix Inkling NVIDIA sconv cache block-size mismatch#51951fjosw wants to merge 2 commits into
fjosw wants to merge 2 commits into
Conversation
The KV-cache planner may enlarge the conv group's block size (W=4 -> e.g. 8) when the attention page is a multiple of the conv page, and slot mappings are then built with the enlarged size. NVIDIA call sites must index with the runtime block size; the AMD variant already does. Fails today: InklingConvState has no cache_block_size accessor. Co-authored-by: Kimi-K3 Signed-off-by: Fabian Joswig <fjosw@users.noreply.github.com>
unify_kv_cache_spec_page_size enlarges the conv group's block size whenever the attention page is a multiple of the conv page (e.g. Inkling-Small-NVFP4 with --block-size=64: attention page ends up 2x the conv page and the conv block grows 4 -> 8). Slot mappings and block tables are then built with the enlarged size, while all NVIDIA call sites still indexed with the kernel window size W=4, so conv state reads/writes landed in wrong pages: corrupted outputs, greedy nondeterminism, and cross-request contamination. Fix it the same way the AMD variant already does: expose InklingConvState.cache_block_size (runtime token dim of the bound cache, falling back to the kernel window size before binding) and use it when launching fused_sconv, fused_qkvr_prep, and the Lamport RS/conv/AG kernels. All kernels already take the conv block size as a runtime argument. Co-authored-by: Kimi-K3 Signed-off-by: Fabian Joswig <fjosw@users.noreply.github.com>
fjosw
requested review from
AndreasKaratzas,
DarkLight1337 and
ywang96
as code owners
August 12, 2026 08:14
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
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
1 participant
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.
Purpose
While benchmarking Inkling-Small on Blackwell I ran into a bug when tuning the KV cache
block size: with
--block-size 64,unify_kv_cache_spec_page_sizeenlargesthe sconv conv group's block size to match the attention page (4 -> 8 here),
but every NVIDIA call site still indexed the cache with the static kernel
window size W=4. Slot mappings use the enlarged size, so this landed conv
reads/writes in the wrong pages: greedy generation degenerates into garbage
within a few tokens.
AMD already handles this correctly via
InklingConvState.cache_block_size(#48841). This ports that property to NVIDIA and switches all call sites to it.
Tagging @WoosukKwon (who contributed these call sites in #48799) and @tjtanaa (AMD implementation which handles this correctly).
AI-assisted, every line reviewed and tested by the submitter.
Test Plan
pytest tests/models/inkling/test_sconv_cache_layout.py -v
E2E: Inkling-Small-NVFP4, TP=4,
--block-size 64, greedy completion of"The capital of France is Paris. The capital of Germany is", before/after
on top of nightly main.
Test Result
Unit: fails on main (no
cache_block_sizeaccessor), passes with the fix.E2E: unpatched output collapses into a
they they they...loop; patchedoutput is coherent: "Berlin. The capital of Italy is Rome. The capital of
Spain is Madrid...".