Conversation
ormandj
requested review from
Fridge003,
HaiShaw,
Qiaolin-Yu,
hanming-lu,
hebiao064,
ispobock,
merrymercy,
yizhang2077 and
yuan-luo
as code owners
September 2, 2026 04:37
Open
2 tasks
ormandj
changed the title
ormandj
changed the title
ormandj
force-pushed
the
pr/kda-blocked-extend
branch
from
September 8, 2026 02:45
27e648e to
aa05881
Compare
…_kda calls TritonKDAKernel.extend ran chunk_kda over the whole extend in one call, so the chunk kernels' workspace (l2-normed q/k, gate cumsums, w/u/kg projections, per-chunk state buffer, intra-chunk A matrices) scaled with the extend length. The scheduler bounds that length (max_prefill_tokens, chunked_prefill_size), but the workspace is a per-layer transient on top of the weights and pools, and on a memory-tight KDA deployment the prefill chunk that keeps throughput up was too large for these kernels. Add SGLANG_KDA_EXTEND_BLOCK_TOKENS (opt-in; default 0 and any value below the 64-token chunk disable, larger values round down to a multiple of 64). When set, run chunk_kda over blocks of at most that many tokens. Consecutive whole sequences that fit within the block share one call, so short-sequence batches stay batched; only longer sequences are split, one call per block, chaining their recurrent state through a per-sequence fp32 scratch slot (gathered from the pool, zero for the padded-row sentinel -1, written back once and masked for -1; the kernel loads initial_state[index] and stores the final state back, INPLACE_UPDATE, skipping both for -1). The one cast back to a bf16/fp16 pool is the single call's one final rounding. Block boundaries fall on chunk boundaries; the output assembles in place in one contiguous v buffer (o=v); intermediate states go into one preallocated packed tensor. Each call's cu_seqlens comes from a small content-keyed LRU cache so repeated layouts reuse one tensor and hit fla's identity cache. Blocking is skipped when extend_seq_lens_cpu is missing or does not match the device metadata in count or total (the boundaries agree by construction: query_start_loc is the device cumsum of the same host list), and under graph capture. Blocking is off by default because each block is a full kernel pipeline and each distinct layout costs one metadata sync; that latency trade has not been measured. Numerical contract: bitwise equal to the single call when every call dispatches the same intra-chunk kernel variant (B * num_chunks * H <= 256 in fla/kda.py); rounding-level differences when blocking moves a call across that threshold. Add test_kda_extend_blocked.py: GPU tests compare the blocked path against the single call (exact where the kernel variant is the same, including bf16/fp16 pools, ragged final blocks, grouped short sequences and a padded -1 row with a nonzero slot 0; a spied threshold-crossing 8-head case with a tolerance), with a spy proving the call plan. CPU tests cover the plan, grouping, metadata fallback, scratch carry, padded-row zero state and the cu_seqlens cache bound with a stand-in kernel whose results depend on the loaded state. GPU tests were not executed for this revision.
ormandj
force-pushed
the
pr/kda-blocked-extend
branch
from
September 8, 2026 21:43
aa05881 to
e96e6f0
Compare
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.
Motivation
Long or batched prefills can run out of GPU memory inside KDA attention because its temporary workspace grows with every token passed to
chunk_kda. Lowering the global prefill chunk size also changes batching for the rest of the model.This PR lets KDA process a large extend in smaller internal blocks while preserving the global prefill limits. It bounds the workspace allocated by each KDA call, not total model or request memory. An earlier per-sequence revision allowed author-observed GLM-5.3-Flash long/image prefills to complete after OOMs; no full-model receipt for the current grouped implementation is attached.
Modifications
SGLANG_KDA_EXTEND_BLOCK_TOKENS, disabled by default. Values below 64 retain the original call; larger values round down to a 64-token boundary.Accuracy Tests
Author CPU validation at
290173ceef: 8 tests and 14 subtests passed, with CUDA hidden. Its Python runtime and test trees are unchanged at refreshed heade96e6f0d12on mainafe90a8bc9. Earlier GPU and serving results retain their stated source scope.The real-kernel regression starts the blocked and unblocked paths from identical pool copies, spies on every KDA call to enforce the block limit, and compares outputs, final states and untouched slots. Cases include ragged/grouped sequences, partial final blocks, reduced-precision pools, strided inputs, padded rows and intermediate states. CPU fake-kernel cases check call planning, state carry, metadata caching and fallback.
Author-reported SM120 result: 15 passed and 20 subtests. CPU-only run: 8 passed and 14 subtests. Results are bitwise equal when blocking preserves the intra-chunk kernel variant; the dispatch-threshold-crossing case is checked at
atol=rtol=1e-2, not exact equality.Speed Tests and Profiling
Blocking trades extra kernel calls and possible metadata synchronization for lower per-call workspace. Author-reported launch-bound microbenchmarks on one RTX PRO 6000 Blackwell took 0.50-0.65 ms unblocked versus 3.67-4.45 ms with eight 2048-token calls. Those one-head shapes measure call overhead, not model throughput. Peak memory of the current grouped path was not measured.
Checklist
Developed with AI assistance.
CI States
Latest PR Test (Base): ❌ Run #34282201461
Latest PR Test (Extra): ❌ Run #34282201310
Latest PR Test (AMD ROCm 7.2): ❌ Run #34282201526