Build the unified read stream directly, without the page-table rectangle by ch-wan · Pull Request #37512 · sgl-project/sglang

ch-wan · 2026-09-02T02:43:34Z

Stack — second of four. Depends on #37511 (its base branch); review only the commits above Test the fused tombstone by what it must do. #37550 → #37560 sit on top.

Motivation

Three of the seven backends that read the unified pool never wanted a page table. They built a [max_bs, max_pages] rectangle and handed it straight to create_flashinfer_kv_indices_triton to repack into the indptr-addressed stream a paged wrapper plans over; no attention kernel ever read the rectangle.

A rectangle's width is bounded by max_context_len, because any single request could be that long and every row must be as wide as the widest. At ps=1, max_bs 256 and a 1M context that reserves 268M slots against a pool that holds 8.15M tokens — 32x more than any reachable state. The stream holds one id per resident token, so its length is sum(seq_lens), which the pool bounds exactly.

Modifications

One gather-and-translate, two delivery forms:

PAGE TABLE    out[b, c] = entry(b, c)                              rows a uniform stride apart
TOKEN STREAM  out[row_starts[b] + p] = entry(b, p//ps)*ps + p%ps   indptr-addressed

The rule that comes out of it is stateable: a consumer that plans over a stream calls fill_packed_read_stream (flashinfer_mla, flashinfer, triton); a consumer whose kernel reads a page table fills the table it already owns, via fill_read_table (flashattention, trtllm_mha, trtllm_mla, flashmla).

The two forms share one @triton.jit, differing by two constexpr: row_starts null means a uniform stride, and EMIT_PER_TOKEN selects page or token granularity. The CPU reference paths share _entries for the same reason.
flashinfer_mla, flashinfer and triton fill the stream directly. kv_view was threaded through five frames of each flashinfer updater chain; the decode chains never needed the real pool indices (a rectangle is indexed by batch lane, so its row ids were arange(bs)) and now take them instead.
The second half of the rule was already half-true: trtllm_mla and flashmla pass their own block table and consume it in place. flashattention and trtllm_mha each held a second rectangle — trtllm_mha's own page_table is explicitly not written under unified (skip_page_table=unified) while kv_read_tables is what gets bound, so one of its two identically-shaped buffers was dead. Both now match their siblings, and make_capture_tables is deleted.
normal_decode_set_metadata gains SKIP_PAGE_TABLE, a constexpr early-out after the prefix sum, so the translator can write fa3's page tables directly. Its now-unreachable src_is_read_table path goes with it: -41 lines, one mechanism in the file instead of two.
Two predicates get names of their own: KVIndexTranslator.reads_are_translated (unified, minus the DCP defer), which callers need before picking a source, and fill_packed_read_stream(sliding_window=...), selecting the swa sub-pool's own v2p.
Bug fix, pre-existing: the page table was built one length and read another. The builder got seq_lens; fa3's fused kernel and trtllm_mha's attention kernels bound their reads by cache_seqlens = seq_lens + delta on a draft decode or target verify. The columns in between were read and never translated. A comment claimed this was unreachable — "Spec is asserted off under the unified pool" — but kv_cache_hook allows --speculative-algorithm DSPARK, a chain draft with topk 1, which is exactly that branch. Both now fill to metadata.cache_seqlens_int32.
Bug fix, pre-existing: fill_read_table asserted is_translating, but build_index_table returns the passthrough without touching into when reads are deferred, which is what DCP does. trtllm_mla returns early on dcp_enabled; flashmla does not, and its block_kv_indices is a reused buffer, so unified + DCP + flashmla would have read the previous batch's indices. The precondition is now reads_are_translated, so that combination fails by name.

Accuracy Tests

Branch vs base, GSM8K over 400 unless noted:

config	base	branch
Kimi-Linear TP2, flashinfer_mla	0.9125	0.9175
Kimi-Linear TP2 + `--dcp-size 2`	0.9175 / 0.9225	0.9150 / 0.9175
Falcon-H1-1.5B, flashinfer	0.7725 / 0.7850 / 0.7825	0.7850 / 0.7775 / 0.7775
Falcon-H1-1.5B, triton	0.7850 / 0.7825 / 0.7775	0.7750 / 0.7850 / 0.7850
Falcon-H1-1.5B, fa3 (5 runs)	mean 0.7775	mean 0.7850
gpt-oss-20b, triton (SWA, 4 runs)	mean 0.9387	mean 0.9387

Means match on the Falcon and gpt-oss rows; the Kimi rows sit inside this box's ±2 question band, which three runs of one unchanged server reproduce. Two rows carry their own weight: the DCP row is the one path where fill_packed_read_stream's return value is load-bearing (the stream comes back deliberately virtual and the caller finishes it), and the gpt-oss row is the sliding-window stream, which nothing else here exercises.

On the fa3 row: at three runs a side the two bands looked disjoint (base 0.7625–0.7800, branch 0.7825–0.7975). They are not — the base's own spread is 11 questions over five runs.

Kernel level: the stream builder is bit-identical to the two steps it replaces over 40 randomized cases (page_size 1..4, mixed seq_lens, -1 holes, tombstoned v2p rows, half the trials driving kv_start_idx). A CPU test pins packed[indptr[b] + p] == ids[b, p // ps] * ps + p % ps, so a change to either builder alone turns red.

test/registered/unit/mem_cache/: 2127 passed.

Speed Tests and Profiling

This PR buys memory, not speed. Kimi-Linear TP2, 1M context, cuda_graph_max_bs 256:

decode CUDA graph memory   2.89 GB -> 1.89 GB        (2.97 -> 1.97 under --dcp-size 2)

Throughput is unchanged, and that is the honest framing: the repack this removes costs ~5.8us per step against ~780ms of decode, two orders of magnitude under the run-to-run spread. (At kernel level the fused build is 2.1–6.5x the two steps it replaces; that ratio does not survive contact with a decode loop.)

A caveat on the memory figures: they come from the CUDA graph end ... mem usage log line, which is empty_cache() + mem_get_info() — true device-free memory, and therefore a lower bound rather than the allocation delta. On gpt-oss-20b it reports 0.67 GB on both sides while torch.cuda.memory_allocated() across init_cuda_graph_state reports base +545.3 MiB vs tip +289.3 MiB — 256.0 MiB, exactly the two rectangles. The caching allocator satisfies the request out of the block it already holds for the KV pool, so device-free does not move. Non-zero deltas are real; zero ones prove nothing.

Also removed: one kernel launch per step per backend (the repack), and for fa3 the verbatim copy of the read table into page_table.

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations. — no user-facing surface changes.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

On #35245

That PR introduced KVIndexTranslator to give id-space translation a single owner, after every backend reached for the hook itself through getattr(allocator, "translate_kv_loc_for_kernel", None) or getattr(...). That design is kept in full — no backend here gains a v2p read or a translate hook; grep for v2p_page_table|kernel_page_multiplier|translate_kv_loc_for_kernel outside mem_cache/ still returns nothing. The converted call sites know strictly less than before: they used to drive the repack kernel themselves with four fields of the table's internal layout (ids, row_ids, row_stride, entry_page_size), and now hand over a buffer and get back a bool.

What this PR replaces is that PR's choice of a shared, memoized rectangle as the universal intermediate. Those two are in tension — a rectangle is shareable precisely because it is consumer-independent, while a stream's layout is that consumer's indptr — and the trade resolves against the cache at every batch size, because the thing being cached is cheaper than the thing the cache forces afterwards:

rectangle build (shared, once)     2.66 us
repack, per consumer               5.64 us
fused stream build, per consumer   2.83 us

consumers   cached rect + N repacks    N fused builds
    1              8.73 us                2.79 us
    2             14.23 us                5.46 us
    4             25.20 us               10.75 us

2.83N < 2.66 + 5.64N for every N.

Known gap

trtllm_mha has no end-to-end run behind it. No configuration on this box reaches trtllm_mha with unified memory: Falcon-H1 OOMs in forward_extend and gemma-4 is rejected for head_dim: 512. Its change is a mechanical swap onto the API its two MLA siblings already use and the unit suite passes, but it wants a reviewer with the hardware.

Cross-shape validation

Benchmarked against the static pool across four hybrid shapes -- same build, only
--enable-unified-memory differing. Schedulable KV, accuracy and throughput all
match; graph memory moves in both directions.

model	type	KV tokens	graph GB	GSM8K	tput
gemma-4-E2B-it	SWA	-0.00%	0.84 -> 0.84	0.8717 -> 0.8600	+2.00%
Qwen3.5-0.8B	18 linear + 6 full	+0.00%	1.20 -> 0.74	0.5400 -> 0.5350	+1.67%
Kimi-Linear-48B	MLA + KDA	0.00%	0.92 -> 1.93	0.9075 -> 0.9100	-0.63%
Inkling-Small	full + SWA + ShortConv	0.00%	0.57 -> 0.47	0.975 -> 0.975	+1.21%

Inkling at TP4 is identical down to the token: max_total_num_tokens=450605 and
max_mamba_cache_size=207 on both sides, from the same rest=4.962 GB budget.
Kimi at TP2 likewise, 4536315 both sides. At --mem-fraction-static 0.9 both
also boot with 74778 tokens, so unified needs no extra headroom.

Accuracy is clean throughout -- the worst gap is gemma's 2.3 questions per 200,
against 1 sigma of 0.0236, i.e. 0.49 sigma. Kimi's +1.01 GB of graph memory is
not the index buffers: probing init_cuda_graph_state reports +1025.0 MiB on
both sides, so that 1.0 GiB cuda_graph_kv_indices is a static-pool cost too.

One methodology note, because the first pass got it wrong: a multi-GPU A/B run
side by side on two GPU sets is not a valid comparison on a shared box. The TP
pool is sized from the minimum free memory across ranks, so one foreign 1.2 GB
process on one card of one set shrinks that whole side's KV -- it read as a 22.5%
Inkling regression and a 1.57% Kimi one, both of which vanish when the two sides
run sequentially on the same cards.

Follow-ups

cuda_graph_kv_indices is still sized max_bs * max_context_len; the stream's
length is bounded by the pool, so it can be ~32.6 MB instead of 1.07 GB.

CI States

Latest PR Test (Base): 🚫 Run #33697241821
Latest PR Test (Extra): 🚫 Run #33697241692
Latest PR Test (AMD ROCm 7.2): ❌ Run #33697241950

chatgpt-codex-connector · 2026-09-02T02:43:43Z

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-02T05:55:00.590271Z	`9540561`	New commits

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

ch-wan · 2026-09-02T05:07:21Z

/tag-and-rerun-ci

`build_kv_read_table` fills a `[bs, max_pages]` page table, which three of the four unified-pool consumers then hand straight to `create_flashinfer_kv_indices_triton` to repack into the indptr-addressed stream a wrapper plans over. The rectangle is never read by an attention kernel on those paths -- it exists only to be re-gathered. `build_kv_read_table_packed` does both halves in one pass, writing `out[indptr[b] + p] = entry * ps + p % ps` directly. This commit adds the kernel and the translator entry point; the backends move over next. Two things follow from the layout rather than from the fusion. The stream holds one id per resident token, so its length is `sum(seq_lens)` -- bounded by the pool, where a rectangle is bounded by `max_bs x max_context_len` and at ps=1 reserves 32x more slots than the pool can ever fill. And the repack turns out to cost more than the build it reads: bs= 64 seq= 8224 8.6us -> 2.8us 3.1x bs= 8 seq= 8224 7.6us -> 1.8us 4.2x bs=256 seq= 8224 12.3us -> 5.8us 2.1x bs= 64 seq=131072 162.8us -> 25.1us 6.5x Verified bit-identical to the two steps it replaces over 40 randomized cases (page_size 1..4, mixed seq_lens, `-1` holes, tombstoned v2p rows, and half the trials driving `kv_start_idx` so the windowed gather the non-MLA path uses is covered), plus CPU/CUDA agreement on the pure-torch reference the allocator's CPU tests run through. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Both indices updaters built a `[max_bs, max_pages]` page table and immediately handed it to `create_flashinfer_kv_indices_triton` to repack; nothing else read it. They now fill the stream directly and the rectangle is gone from this backend. The updaters took `kv_view` rather than `req_pool_indices` because the rectangle is indexed by batch lane -- the row ids were `arange(bs)` and the real pool indices were never needed. Filling from `req_to_token` needs them, so the decode updater takes `req_pool_indices` instead of the view; the prefill updater already had it. Kimi-Linear TP2 + unified, 1M context, cuda_graph_max_bs 256: decode CUDA graph memory 2.89 GB -> 1.89 GB which is the 1.0 GiB the rectangle held, to the reported precision. Throughput is unchanged: the repack this removes is ~5.8us per step against ~780ms of decode, two orders of magnitude under the run-to-run spread. GSM8K 0.9175 vs 0.9125 on main -- the same server returns 0.9175 / 0.9175 / 0.9125 across three runs, so 400-example GSM8K on this box carries a +-2 question band and neither number is evidence either way. The correctness argument is the randomized equality against the two steps this replaces, not the eval. `make_capture_tables` still serves flashattention, flashinfer, triton and trtllm_mha. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The packed builder arrived as a second `@triton.jit` alongside the page-table one, duplicating the gather-and-translate formula that is the load-bearing part of this file -- a fix to the entry derivation would have had to land twice. They now share one kernel. What differs is an output layout and a granularity, both constexpr: `row_starts` null means rows sit a uniform stride apart, and `EMIT_PER_TOKEN` selects whether an item is a page (emit the entry) or a token (emit `entry * ps + pos % ps`). The two public launchers stay, so a call site still names the form it wants; the CPU reference paths share `_entries` for the same reason the kernels now share a body. Deriving the per-row item count moved into the kernel. Computing `n_pages` on the host first cost the page-table path 2.66us -> 5.55us at bs=64 / seq=8224 -- a whole launch for one elementwise op, which a captured graph then replays every step. The kernel already loads `seq_lens`; the delivery form tells it what an item is. No behaviour change: both forms re-verified against what they were checked against before (the page table against an independent CPU model over 30 randomized cases plus graph replay at four sequence lengths; the stream against [page table + repack] over 40, half driving `kv_start_idx`), and the page table's timings are unchanged across bs 1..256 x seq 1k..128k. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Same move as the MLA backend, over a wider surface: `kv_view` was threaded through five frames of each updater chain. The prefill chain already carried `req_pool_indices`, so those frames just lose a parameter; the decode chain never needed the real pool indices (a rectangle is indexed by batch lane, so its row ids were `arange(bs)`) and now takes them instead. Two things the view was answering for get names of their own: - `use_swa_source` asked `kv_view.is_translated`, which has to be known before the fill picks a source. That predicate is now `KVIndexTranslator.reads_are_translated` -- unified, minus the DCP defer. - `fill_packed_read_stream` gained `sliding_window`, selecting the swa sub-pool's own v2p and multiplier. Built from VIRTUAL ids, never chained through full-physical, which is what the swa canonical array was for. `fill_packed_read_stream` now takes the batch size from `seq_lens` and slices `req_pool_indices` to it. Callers hold indices padded to the graph bucket while asking for fewer rows; the grid used to come from the indices, so the extra lanes would have run with no length to bound them. Falcon-H1-1.5B + unified + flashinfer, branch vs base, GSM8K over 400: base 0.7725 / 0.7850 / 0.7825 mean 0.7800 tip 0.7850 / 0.7775 / 0.7775 mean 0.7800 decode CUDA graph memory 0.89 GB -> 0.77 GB. (A 1.5B model's context makes the rectangle small; the 1.0 GiB figure is Kimi-Linear at 1M.) Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The last stream consumer. `_fill_kv_indptr_and_indices` and `update_sliding_window_buffer` both built a page table and immediately repacked it through `create_flashinfer_kv_indices_triton`; both now fill their CSR buffer directly. The SWA one already had a window offset per row, which the stream builder takes as `kv_start_idx`. Triton's index buffers are int64, so the builder emits either width, chosen from the destination's dtype. Validated on Falcon-H1-1.5B + unified + triton (hybrid mamba, so the full stream only), GSM8K over 400: base 0.7850 / 0.7825 / 0.7775 mean 0.7817 tip 0.7750 / 0.7850 / 0.7850 mean 0.7817 The SWA stream is NOT covered by that run and I could not cover it: the one hybrid-SWA model available here, gemma-4-31b-it, dies under unified memory on plain main -- a device-side index assert, GSM8K 0.0 and the scheduler gone, where the same model on the same commit with the static pool scores 0.98. This branch reproduces it identically, so it is not from this change, but it means the SWA wiring here rests on the kernel-level equality (which does cover `kv_start_idx` and both id spaces) and not on an end-to-end run. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

`make_capture_tables` is gone, and with it the last rectangle the unified pool allocated on anyone's behalf. trtllm_mla and flashmla were already doing this right: they call `fill_read_table` with their own block table and consume it in place. The other two page-table consumers each held a second rectangle instead. **trtllm_mha** allocated `kv_read_tables` and, at capture, rebound `metadata.page_table` to it -- while its own `page_table`, allocated in `init_cuda_graph_state`, is explicitly *not written* under unified (`page_table=None`, `skip_page_table=unified`). Two buffers of identical shape and dtype, one dead. It now fills its own, which also retires the `in_capture` rebinding special case: one buffer, refreshed in place on capture and on every replay-prep, so there is nothing to bind once. **flashattention** built the rectangle and had `normal_decode_set_metadata` copy it verbatim into `metadata.page_table`. That fused kernel gains `SKIP_PAGE_TABLE`, a constexpr early-out after the prefix sum, so the translator can write the page tables directly and the kernel is left with the work only it can do. The skip path launches one block, because one block is what the prefix sum uses. `fill_read_table` grew `sliding_window_out` to fill the swa twin in the same pass, which both of these need and the two MLA callers do not. Measured, Falcon-H1-1.5B + unified + fa3, decode CUDA graph memory 0.52 GB -> 0.39 GB, with GSM8K over 400 across five runs a side: base 0.7775 / 0.7625 / 0.7800 / 0.7900 / 0.7775 mean 0.7775 tip 0.7825 / 0.7975 / 0.7825 / 0.7825 / 0.7800 mean 0.7850 At three runs a side those bands looked disjoint; they are not, and the base's own spread is 11 questions. **trtllm_mha has no end-to-end run behind it.** No configuration on this box reaches trtllm_mha with unified memory: Falcon-H1 OOMs in `forward_extend` (215 GiB), and gemma-4 is rejected for `head_dim: 512`. The change is a mechanical swap onto the API its two MLA siblings already use, and the unit suite passes, but a reviewer with the hardware should exercise it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Two findings from auditing the read path, both predating this branch. **The page table was built one length and read another.** The builder was handed `seq_lens`; fa3's fused metadata kernel and trtllm_mha's attention kernels bound their page-table reads by `cache_seqlens`, which is `seq_lens + seq_len_delta` on a draft decode and `seq_lens + draft_token_num` on a target verify. The columns in between were read and never translated -- zeros (the page-0 sink) on a fresh capture buffer, and the previous batch's entries on a replay. A comment at fa3's draft-decode site asserted this could not happen: "Spec is asserted off under the unified pool." It is not. `kv_cache_hook` allows `--speculative-algorithm DSPARK` with unified memory, and DSPARK is a chain draft with topk 1, which is exactly the branch that comment sits in. Both now fill to `metadata.cache_seqlens_int32` -- the tensor the reader bounds by, already computed by the fused call above -- instead of to `seq_lens`. For fa3 that meant running the fused call first, which costs nothing: it was already going to run, and it is what produces the bound. **`fill_read_table` could silently fill nothing.** It asserted `is_translating`, but `build_index_table` returns the passthrough view without touching `into` when reads are deferred, which is what DCP does. trtllm_mla returns early on `dcp_enabled` before reaching it; flashmla does not, and its `block_kv_indices` is a reused buffer, so unified + DCP + flashmla would have read the previous batch's indices. The precondition is now `reads_are_translated`, so that combination fails by name instead. Also drops `src_is_read_table` / `swa_src_table` / `SRC_IS_KERNEL_PAGE_TABLE` from the fused metadata kernel. fa3 was the only caller and no longer takes that path, so the "copy the read table verbatim" branch is dead: -41 lines, and one mechanism in the file instead of two. `test_unified_mla_block_table.py` moves onto the new one; the invariant it pins -- the page table ends up holding kernel-facing ids -- is unchanged. Falcon-H1-1.5B + unified + fa3 after the reorder: 0.7825 / 0.7900 / 0.7900, against a base spanning 0.7625-0.7900 over five runs. Decode graph memory stays 0.39 GB. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

sgl-project#37511 and sgl-project#37512 took the unified pool's READ path off the per-step launch path: the read table is built by one kernel whose grid is sized from bs, and the flat consumers (fa3, flashinfer, triton, trtllm_mha) were converted to a packed stream. The WRITE path was left eager, and the block-table backends (trtllm_mla, cutedsl_mla, tokenspeed_mla, flashmla) were not converted at all. Hopper's MLA default is fa3, so it got both halves of that work. Blackwell runs the trtllm_mla family, which got neither -- which is why the Hopper fix did not carry over. `translate_kv_loc_for_kernel` was 6 torch ops (floor_divide, remainder, take, mul, add, clamp) and `translate_write_loc_for_kernel` layered 6 more on top (remainder, eq, floor_divide, zeros_like, where, copy). Both run per forward on `out_cache_loc`, OUTSIDE any cuda graph, so each op is a real launch on the critical path. Profiling a bs=8 decode on B300 (24 steps, 2 ranks): GPU kernel time unified 147.56 ms static 153.67 ms -6.11 ms wall span unified 81.10 ms static 77.92 ms +3.18 ms (+4.1%) cudaLaunchKernel unified 988 static 532 +456 cpu_op time unified 51.40 ms static 35.24 ms +16.16 ms GPU time was already LOWER; the loss was ~10 extra launches per step per rank. That is invisible next to a Hopper decode step and is not next to a Blackwell one. Route both translates through one Triton kernel in virtual_slot.py, the module sgl-project#37511 added for exactly this class of fusion. It carries the DCP owner rule, so the widened-loc path is the same single launch. Triton truncates division toward zero where torch floors it, so a negative loc is tested explicitly rather than relying on the division to reach the v2p sentinel row. Measured on B300, TP2 DCP2, cutedsl_mla, unified vs static decode throughput over 6 shape/batch points: mean gap -1.70% -> -0.23%. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The four spec-input CSR builders took a raw `(req_pool_indices, req_to_token)` pair and ran `create_flashinfer_kv_indices_triton` themselves, so a translating batch had no way to reach them: the ids in `req_to_token` are VIRTUAL under the unified pool, and nothing in that argument pair can say so. They now take the `KVIndexTranslator` and gather through `fill_packed_read_stream`, which is what upstream's non-spec branches in these same functions already do. One call replaces the hand-rolled launch: translator.fill_packed_read_stream( req_pool_indices=req_pool_indices, seq_lens=paged_kernel_lens, # widened just above indptr=cum_kv_seq_len, total_tokens=total_tokens, out=kv_indices, ) Three things fall out of using it rather than a page-table rectangle: * No rectangle. The old path built a `[bs, max_pages]` table and repacked it into the CSR stream the wrapper actually plans over -- the exact pattern sgl-project#37512 removed from the non-spec branches, kept alive here only because spec was refused on unified pools. At page size 1 that rectangle reserves far more slots than the pool has tokens. * No widening contract. The builders already widen their own lens (`paged_kernel_lens + draft_token_num`, per-row for DFlash's ragged layout) and now hand those widened lens straight to the translator. There is no separate table whose live prefix could fall short, so `seq_len_delta` is not part of this path at all. * No id-space plumbing. `sliding_window=` selects the swa sub-pool's own id space, mirroring the non-spec branches, instead of re-aiming a table at it. Behaviour on a plain pool is unchanged by construction: with `reads_are_translated` false, `fill_packed_read_stream` runs the same `create_flashinfer_kv_indices_triton` launch over the same `req_to_token` at `ENTRY_PAGE_SIZE=1` -- the call these builders made by hand. layers/attention/aiter_backend.py: its three updaters bind `kv_index_translator` beside the `req_to_token` they already hold. aiter is absent from the spec-verify allow-list, so it only ever sees the non-translating branch -- the bind lets it satisfy the signature, not a claim that it translates. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

sgl-project#37511 and sgl-project#37512 took the unified pool's READ path off the per-step launch path: the read table is built by one kernel whose grid is sized from bs, and the flat consumers (fa3, flashinfer, triton, trtllm_mha) were converted to a packed stream. The WRITE path was left eager, and the block-table backends (trtllm_mla, cutedsl_mla, tokenspeed_mla, flashmla) were not converted at all. Hopper's MLA default is fa3, so it got both halves of that work. Blackwell runs the trtllm_mla family, which got neither -- which is why the Hopper fix did not carry over. `translate_kv_loc_for_kernel` was 6 torch ops (floor_divide, remainder, take, mul, add, clamp) and `translate_write_loc_for_kernel` layered 6 more on top (remainder, eq, floor_divide, zeros_like, where, copy). Both run per forward on `out_cache_loc`, OUTSIDE any cuda graph, so each op is a real launch on the critical path. Profiling a bs=8 decode on B300 (24 steps, 2 ranks): GPU kernel time unified 147.56 ms static 153.67 ms -6.11 ms wall span unified 81.10 ms static 77.92 ms +3.18 ms (+4.1%) cudaLaunchKernel unified 988 static 532 +456 cpu_op time unified 51.40 ms static 35.24 ms +16.16 ms GPU time was already LOWER; the loss was ~10 extra launches per step per rank. That is invisible next to a Hopper decode step and is not next to a Blackwell one. Route both translates through one Triton kernel in virtual_slot.py, the module sgl-project#37511 added for exactly this class of fusion. It carries the DCP owner rule, so the widened-loc path is the same single launch. Triton truncates division toward zero where torch floors it, so a negative loc is tested explicitly rather than relying on the division to reach the v2p sentinel row. Measured on B300, TP2 DCP2, cutedsl_mla, unified vs static decode throughput over 6 shape/batch points: mean gap -1.70% -> -0.23%. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The four spec-input CSR builders took a raw `(req_pool_indices, req_to_token)` pair and ran `create_flashinfer_kv_indices_triton` themselves, so a translating batch had no way to reach them: the ids in `req_to_token` are VIRTUAL under the unified pool, and nothing in that argument pair can say so. They now take the `KVIndexTranslator` and gather through `fill_packed_read_stream`, which is what upstream's non-spec branches in these same functions already do. One call replaces the hand-rolled launch: translator.fill_packed_read_stream( req_pool_indices=req_pool_indices, seq_lens=paged_kernel_lens, # widened just above indptr=cum_kv_seq_len, total_tokens=total_tokens, out=kv_indices, ) Three things fall out of using it rather than a page-table rectangle: * No rectangle. The old path built a `[bs, max_pages]` table and repacked it into the CSR stream the wrapper actually plans over -- the exact pattern sgl-project#37512 removed from the non-spec branches, kept alive here only because spec was refused on unified pools. At page size 1 that rectangle reserves far more slots than the pool has tokens. * No widening contract. The builders already widen their own lens (`paged_kernel_lens + draft_token_num`, per-row for DFlash's ragged layout) and now hand those widened lens straight to the translator. There is no separate table whose live prefix could fall short, so `seq_len_delta` is not part of this path at all. * No id-space plumbing. `sliding_window=` selects the swa sub-pool's own id space, mirroring the non-spec branches, instead of re-aiming a table at it. Behaviour on a plain pool is unchanged by construction: with `reads_are_translated` false, `fill_packed_read_stream` runs the same `create_flashinfer_kv_indices_triton` launch over the same `req_to_token` at `ENTRY_PAGE_SIZE=1` -- the call these builders made by hand. layers/attention/aiter_backend.py: its three updaters bind `kv_index_translator` beside the `req_to_token` they already hold. aiter is absent from the spec-verify allow-list, so it only ever sees the non-translating branch -- the bind lets it satisfy the signature, not a claim that it translates. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…gle (sgl-project#37512) Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

sgl-project#37511 and sgl-project#37512 took the unified pool's READ path off the per-step launch path: the read table is built by one kernel whose grid is sized from bs, and the flat consumers (fa3, flashinfer, triton, trtllm_mha) were converted to a packed stream. The WRITE path was left eager, and the block-table backends (trtllm_mla, cutedsl_mla, tokenspeed_mla, flashmla) were not converted at all. Hopper's MLA default is fa3, so it got both halves of that work. Blackwell runs the trtllm_mla family, which got neither -- which is why the Hopper fix did not carry over. `translate_kv_loc_for_kernel` was 6 torch ops (floor_divide, remainder, take, mul, add, clamp) and `translate_write_loc_for_kernel` layered 6 more on top (remainder, eq, floor_divide, zeros_like, where, copy). Both run per forward on `out_cache_loc`, OUTSIDE any cuda graph, so each op is a real launch on the critical path. Profiling a bs=8 decode on B300 (24 steps, 2 ranks): GPU kernel time unified 147.56 ms static 153.67 ms -6.11 ms wall span unified 81.10 ms static 77.92 ms +3.18 ms (+4.1%) cudaLaunchKernel unified 988 static 532 +456 cpu_op time unified 51.40 ms static 35.24 ms +16.16 ms GPU time was already LOWER; the loss was ~10 extra launches per step per rank. That is invisible next to a Hopper decode step and is not next to a Blackwell one. Route both translates through one Triton kernel in virtual_slot.py, the module sgl-project#37511 added for exactly this class of fusion. It carries the DCP owner rule, so the widened-loc path is the same single launch. Triton truncates division toward zero where torch floors it, so a negative loc is tested explicitly rather than relying on the division to reach the v2p sentinel row. Measured on B300, TP2 DCP2, cutedsl_mla, unified vs static decode throughput over 6 shape/batch points: mean gap -1.70% -> -0.23%. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The four spec-input CSR builders took a raw `(req_pool_indices, req_to_token)` pair and ran `create_flashinfer_kv_indices_triton` themselves, so a translating batch had no way to reach them: the ids in `req_to_token` are VIRTUAL under the unified pool, and nothing in that argument pair can say so. They now take the `KVIndexTranslator` and gather through `fill_packed_read_stream`, which is what upstream's non-spec branches in these same functions already do. One call replaces the hand-rolled launch: translator.fill_packed_read_stream( req_pool_indices=req_pool_indices, seq_lens=paged_kernel_lens, # widened just above indptr=cum_kv_seq_len, total_tokens=total_tokens, out=kv_indices, ) Three things fall out of using it rather than a page-table rectangle: * No rectangle. The old path built a `[bs, max_pages]` table and repacked it into the CSR stream the wrapper actually plans over -- the exact pattern sgl-project#37512 removed from the non-spec branches, kept alive here only because spec was refused on unified pools. At page size 1 that rectangle reserves far more slots than the pool has tokens. * No widening contract. The builders already widen their own lens (`paged_kernel_lens + draft_token_num`, per-row for DFlash's ragged layout) and now hand those widened lens straight to the translator. There is no separate table whose live prefix could fall short, so `seq_len_delta` is not part of this path at all. * No id-space plumbing. `sliding_window=` selects the swa sub-pool's own id space, mirroring the non-spec branches, instead of re-aiming a table at it. Behaviour on a plain pool is unchanged by construction: with `reads_are_translated` false, `fill_packed_read_stream` runs the same `create_flashinfer_kv_indices_triton` launch over the same `req_to_token` at `ENTRY_PAGE_SIZE=1` -- the call these builders made by hand. layers/attention/aiter_backend.py: its three updaters bind `kv_index_translator` beside the `req_to_token` they already hold. aiter is absent from the spec-verify allow-list, so it only ever sees the non-translating branch -- the bind lets it satisfy the signature, not a claim that it translates. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

ch-wan requested review from BBuf, DarkSharpness, Fridge003, HaiShaw, HydraQYH, Qiaolin-Yu, Ying1123, alphabetc1, celve, hanming-lu, hebiao064, hnyls2002, huangtingwei9988, hzh0425, ispobock, merrymercy, xiezhq-hermann, yizhang2077 and yuan-luo as code owners September 2, 2026 02:43

github-actions Bot added the jit-kernel label Sep 2, 2026

github-actions Bot added run-ci blackwell SM100/SM120 labels Sep 2, 2026

ch-wan mentioned this pull request Sep 2, 2026

Converge the two SWA predicates, and stop conditioning the capture sink on the pool #37550

Merged

5 tasks

ch-wan removed the run-ci label Sep 2, 2026

This was referenced Sep 2, 2026

Fix unified SWA: size a non-owner's v2p by the id space it must address #37560

Merged

Size the unified read-table grid from bs, and fuse the allocator's tombstone scatters #37511

Merged

ch-wan force-pushed the cheng/unified-drop-capture-tables branch from 9540561 to d6b0123 Compare September 2, 2026 08:11

ch-wan force-pushed the cheng/unified-alloc-fusion branch from f9feea3 to 972729f Compare September 2, 2026 19:50

ch-wan force-pushed the cheng/unified-drop-capture-tables branch from d6b0123 to 308631e Compare September 2, 2026 19:50

ch-wan force-pushed the cheng/unified-alloc-fusion branch from 972729f to 7d39aaf Compare September 2, 2026 20:54

ch-wan force-pushed the cheng/unified-drop-capture-tables branch from 308631e to f315619 Compare September 2, 2026 20:54

ch-wan force-pushed the cheng/unified-alloc-fusion branch from 7d39aaf to 9a9fae4 Compare September 2, 2026 23:54

Base automatically changed from cheng/unified-alloc-fusion to main September 2, 2026 23:54

ch-wan and others added 7 commits September 2, 2026 23:54

ch-wan force-pushed the cheng/unified-drop-capture-tables branch from f315619 to 4841c24 Compare September 2, 2026 23:55

ch-wan merged commit d9848b9 into main Sep 2, 2026
6 of 17 checks passed

ch-wan deleted the cheng/unified-drop-capture-tables branch September 2, 2026 23:55

bingxche mentioned this pull request Sep 3, 2026

[CI Monitor] Daily Report - 2026-09-03 bingxche/sglang-ci-bot#184

Open

StevenChenSE pushed a commit to StevenChenSE/sglang that referenced this pull request Sep 6, 2026

Build the unified read stream directly, without the page-table rectan…

1c1ab05

…gle (sgl-project#37512) Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

Navigation Menu

Uh oh!

Build the unified read stream directly, without the page-table rectangle - #37512

Build the unified read stream directly, without the page-table rectangle#37512
ch-wan merged 7 commits into
mainfrom
cheng/unified-drop-capture-tables

ch-wan commented Sep 2, 2026 •

edited by github-actions Bot

Loading

Uh oh!

chatgpt-codex-connector Bot commented Sep 2, 2026 •

edited

Loading

Uh oh!

ch-wan commented Sep 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

ch-wan commented Sep 2, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

On #35245

Known gap

Cross-shape validation

Follow-ups

CI States

Uh oh!

chatgpt-codex-connector Bot commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codex Review Summary

Uh oh!

ch-wan commented Sep 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

ch-wan commented Sep 2, 2026 •

edited by github-actions Bot

Loading

chatgpt-codex-connector Bot commented Sep 2, 2026 •

edited

Loading