Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
YazhiGao
marked this pull request as ready for review
July 27, 2026 05:15
YazhiGao
requested review from
BBuf,
DarkSharpness,
Ying1123,
alphabetc1,
celve,
hanming-lu,
hnyls2002,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
July 27, 2026 05:15
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
YazhiGao
force-pushed
the
agent/skip-kv-write-reserved-padding-slots
branch
from
July 27, 2026 05:16
f70477b to
cb1e044
Compare
21 tasks
adityavaid
pushed a commit
to adityavaid/sglang
that referenced
this pull request
Jul 29, 2026
Co-authored-by: Andrew Gu <andrew@thinkingmachines.ai>
jinzhenfan
pushed a commit
to jinzhenfan/sglang
that referenced
this pull request
Jul 29, 2026
Co-authored-by: Andrew Gu <andrew@thinkingmachines.ai>
This was referenced Jul 29, 2026
This was referenced Aug 1, 2026
jetd1
pushed a commit
to jetd1/sglang
that referenced
this pull request
Aug 3, 2026
Co-authored-by: Andrew Gu <andrew@thinkingmachines.ai>
hassellof
added a commit
to hassellof/sglang
that referenced
this pull request
Aug 6, 2026
fused_store_flashmla_cache and fused_store_indexer_cache compute
page = index >> kPageBits and store unconditionally, even though the
kernel comment says "always load the value from input (don't store if
invalid)" — the skip was never implemented, and unlike the generic pool
kernel there is no bounds check at all.
Two ways invalid indices reach these kernels in normal operation:
* CUDA-graph padding rows carry out_cache_loc = 0, the allocator's
reserved slot (PagedTokenToKVPoolAllocator.clear() starts free pages
at 1, and the SWA page LUT maps 0 -> 0). Undefined padding values are
silently stored into slot 0, poisoning the page that page-table
padding also points at; partially-valid tiles read the garbage back
before masking.
* The SWA page LUT deliberately keeps a live -1 ("no page") entry at
the end. Unguarded, page goes negative and the store is an
out-of-bounds write before the cache base.
Guard both kernels with store_valid = index > 0, wrapping only the
stores (not early-returning) so the PDL trigger still fires and warp
reductions see the full warp. The predicate is uniform per block /
per warp, so there is no added divergence.
Same class of fix as sgl-project#32477 for the generic store_cache kernel; the
regression test mirrors test_store_cache_reserved_skip_index. The -1
case is made deterministic without compute-sanitizer by placing a
sentinel-filled guard page immediately before the cache view: the
unguarded kernel's negative-page store lands in the guard page.
5 tasks
saturn-acc
pushed a commit
to saturn-acc/sglang
that referenced
this pull request
Aug 16, 2026
Co-authored-by: Andrew Gu <andrew@thinkingmachines.ai>
This was referenced Aug 22, 2026
jakki-amd
pushed a commit
to jakki-amd/sglang
that referenced
this pull request
Sep 9, 2026
Co-authored-by: Andrew Gu <andrew@thinkingmachines.ai>
Atituiset
pushed a commit
to Atituiset/sglang
that referenced
this pull request
Sep 10, 2026
Co-authored-by: Andrew Gu <andrew@thinkingmachines.ai>
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
5 participants
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.
Motivation
CUDA-graph padding rows can contain undefined K/V values while still targeting the reserved padding slot at index 0. In sliding-window attention, page-table padding also points at that reserved page. A partially valid tile may load the page before masking invalid lanes, allowing a NaN value to contaminate the result through
0 * NaN.Changes
reserved_skip_indexargument to the JIT KV-cache store kernel-1as an opt-outValidation
pre-commit run --files python/sglang/kernels/jit/csrc/elementwise/kvcache.cuh python/sglang/kernels/ops/kvcache/kvcache.py python/sglang/srt/mem_cache/memory_pool.py test/registered/kernels/ops/kvcache/test_store_cache.pypytest -q test/registered/kernels/ops/kvcache/test_store_cache.py -x— 340 passedCI States
Latest PR Test (Base): ✅ Run #30371139759
Latest PR Test (Extra): ✅ Run #30371134897