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

[ch-wan]

Stack — first of four. Nothing depends on it upstream; #37512 → #37550 → #37560 sit on top, in that order.

Motivation

Two independent costs in the unified memory pool's hot path.

The read table's cuda-graph grid was sized by max_context_len. build_kv_read_table launched (bs, cdiv(max_pages, BLOCK)) blocks. The eager path passes the batch's live maximum, but the captured path passes max_context_len — and a cuda-graph capture bakes the grid. Every decode replay ran 4096 column blocks per row for sequences 8k long: 99.2% of the 262144 blocks loaded two scalars, failed the cols < n_pages mask, and exited.

The tell is that the cost did not vary with the sequence length — 158.5us at seq=1024 and 160.2us at seq=131072, one launch per step. It was pure block-scheduling overhead, and it grows with cuda_graph_max_bs: 631us per step at bs=256.

The allocator's free path issued five kernels where one suffices. _free_lazy's v2p gather and both tombstone scatters are one lane's work per page, so they collapse into a single kernel — the inverse of the alloc_bind_inplace the alloc fast path already uses.

Modifications

• build_kv_read_table's grid comes from bs alone; each program strides over its columns with the loop bound read from the device-side seq_lens. The grid stays static so a capture bakes it, and the work tracks the real lengths at replay. BLOCK 512 / 8 warps / 1024 blocks is the joint optimum over bs 1..256 × seq 1k..128k, flat within ~10%.
• _free_lazy and bind in MultiEndedAllocator use two new fused Triton kernels (free_unbind_inplace, bind_inplace), each with a pure-torch CPU path because the allocator's unit tests run on CPU.
• The two torch.cats stay: _free_phys_pages is a _CapacityField whose contract is rebind-never-mutate, and the rebind is what invalidates the available_size memos.
• test_free_paths_actually_use_index_fill asserted the free path's source contains index_fill_. The fusion writes the tombstones inside the kernel instead, which serves that file's real invariant — the -1 never crosses the bus — strictly better. It now checks for any of the sanctioned no-sync forms.
• virtual_slot.py had no test in the tree at all, and on CPU both launchers take their reference path, so the kernel doing the tombstoning was never executed by the suite. Added CUDA-gated coverage over 20 randomized bind/free rounds.

Accuracy Tests

Kimi-Linear-48B TP2 + --enable-unified-memory, GSM8K over 400: 0.9125 on both sides.

Treat that as a sanity check, not the argument: the same unchanged server returns 0.9175 / 0.9175 / 0.9125 across three runs, so 400-example GSM8K on this box carries a ±2 question band.

The correctness argument is the kernel-level equality. The page table is verified against an independent CPU model over 30 randomized cases — page_size 1..4, mixed seq_lens, -1 holes in req_to_token, tombstoned v2p rows, and a non-zero out so the prefix-only contract has to leave the tail intact — plus both degenerate shapes. A captured graph replays bit-identically to a fresh launch as seq_lens moves over 1 / 64 / 999 / 4096, which is the property the static grid exists for.

The fused allocator kernels are checked against the torch implementation over 20 randomized page sets (both tables plus the returned physical pages), empty input, and CPU/CUDA agreement. Dropping tl.store(p2v_ptr + p, -1) from the kernel turns two of the new tests red.

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

Speed Tests and Profiling

Read-table grid, H100, under graph replay:

shape	before	after
bs=64 seq=8224	158.7us	2.6us	61x
bs=64 seq=1024	158.9us	1.7us	92x
bs=256 seq=8224	631.3us	4.9us	129x
bs=256 seq=131072	642.8us	81.4us	8x
bs=1 seq=8224	3.4us	1.6us	2x

Allocator fusion, Kimi-Linear TP2 + unified, bs=64 / 8k in / 10 decode steps:

_free_lazy  305 launches / 430us -> 183 / 259us
bind        140 launches / 220us ->  70 /  70us

End to end, Kimi-Linear TP2, bs=64 / 8k input / 64 output, three runs a side on the same idle GPU pair:

main    5164.96 / 5166.33 / 5167.17 tok/s
branch  5239.63 / 5241.60 / 5250.28 tok/s      +1.50%

That is 172us per decode step against the 156us the kernel saves — the same number within run-to-run spread. A fresh decode profile confirms it: build_kv_read_table_kernel 1606us → 47us over the same 10 launches, 1.18% → 0.03% of decode kernel time, with every large kernel (moe, all_reduce, MLA attention) inside noise.

Both sides ran sequentially on the same idle GPU pair. That is not a formality on a shared box: a TP pool is sized from the minimum free memory across ranks, so a side-by-side run on two GPU sets measures whoever else is resident.

Kernel timings were taken under torch.cuda.CUDAGraph replay. A plain Python launch loop floors every Triton launch at ~16us regardless of the work, which made a 60-config sweep read as uniformly optimal.

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.

Not in this PR

The out buffer these kernels fill is still max_bs x max_context_len int32 — 1.0 GiB at 1M context — for the same reason the grid was. That is #37512.

---

CI States

Latest PR Test (Base): ❌ Run #33682076564
Latest PR Test (Extra): ❌ Run #33682076343
Latest PR Test (AMD ROCm 7.2): ❌ Run #33682076492

---

CI States

Latest PR Test (Base): 🚫 Run #33697177974
Latest PR Test (Extra): 🚫 Run #33697177892
Latest PR Test (AMD ROCm 7.2): 🚫 Run #33697178027

[chatgpt-codex-connector]

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-02T02:46:04.955504Z	f9feea3	PR opened

ℹ️ 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.