refactor(unified-memory): translate the KV write location once, at ForwardBatch construction - #35245
Merged
Merged
Conversation
caihuali95
requested review from
Fridge003,
HaiShaw,
Qiaolin-Yu,
Ying1123,
alphabetc1,
hanming-lu,
hebiao064,
hnyls2002,
huangtingwei9988,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
August 18, 2026 04:24
1 task
This was referenced Aug 18, 2026
ch-wan
reviewed
Aug 22, 2026
caihuali95
force-pushed
the
mainline/write-loc-rebind
branch
from
August 24, 2026 20:46
4701729 to
b6488bb
Compare
caihuali95
requested review from
BBuf,
DarkSharpness,
HydraQYH,
celve and
yuan-luo
as code owners
August 24, 2026 20:46
5 tasks
…onstexpr `create_flashinfer_kv_indices_triton` declares `req_to_token_ptr_stride` as `tl.constexpr`, which was free while the only caller passed `req_to_token.stride(0)` -- one value for the life of the process. It is not free now: under the unified pool the argument is `index_table.row_stride`, i.e. the eager table's width, which `index_table_for_batch` recomputes per batch as `ceil(max_seq_len / page_size)`. A constexpr is part of Triton's specialization key, so every distinct batch max-seq-len compiles a fresh kernel -- a recompile every few decode steps at small page sizes. Make it a runtime argument. `ENTRY_PAGE_SIZE` stays constexpr: it is the page size, fixed for the run, and it selects between two different address computations. Also widen `req_pool_index` to int64 before it multiplies the stride. The product indexes `req_to_token`, whose element count exceeds int32 on a large `max_num_reqs x max_context_len` table. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… not just the reference The module contract says a `-1` in `req_to_token` lands on entry 0, the reserved sink, and the CPU reference implements it (`torch.where(tok < 0, 0, tok // page_size)`). The Triton kernel did not, so the two halves of one function disagreed -- and CI only exercises the CPU half. At `page_size > 1` truncating division hid it: `-1 // ps` is 0, and `MultiEndedAllocator.clear()` pins `virtual_to_physical[0] = 0`, so the lane reached the sink anyway. At `page_size == 1` there is nothing to truncate: `-1 // 1` is -1, so the load reads `v2p[-1]` -- one element before the table -- and `tl.maximum(phys * mult, 0)` passes whatever was there through as a live KV id. Measured on an H200 by backing `v2p` with a tensor whose preceding element was set to 123456789: the kernel emitted 987654312 (= 123456789 * 8) where the sink was expected. No mainline path puts `-1` inside a row's live prefix today, so this is the guard the module documents rather than a live miscompute. One `tl.where` makes both page sizes and both implementations agree. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ns_cpu
`index_table_for_batch` took `seq_lens_cpu is not None` as "the CPU mirror is
live". It is not that signal, and the code this series replaces said so in as
many words:
seq_lens_sum is the reliable "mirror present" signal: it is
None-preserving into the replay view, unlike seq_lens_cpu (always a
non-None but stale slice for gpu_only batches).
Here the consequence is worse than the one that comment guarded. The value
sets the table's WIDTH, so a stale-and-smaller max silently under-sizes it:
the columns past the built prefix keep the fresh zeros -- entry 0, the sink --
for tokens the kernel genuinely wants to read. Wrong attention output, no
error anywhere.
Only the eager path reaches this today (every replay path goes through
`build_index_table(captured=True)`), so this is disarming the trap rather than
fixing a live miscompute. Use `seq_lens_sum` as the guard and keep the table's
own width as the fallback, which is always safe.
The unit fake grows the field, defaulting to the real sum, so a case that
wants the gpu_only regime asks for it explicitly.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Same pass as the parent PR, over what this one added.
- `kv_index_translator` module docstring: the three id spaces, the table shape
and the affine rebuild law stay -- that is the subsystem's contract and a
module docstring is where design rationale belongs. The argument for having
a choke point at all ("each backend would be a place to get it wrong") is
PR-body material and goes, along with the WHY/WHAT section headings that
only make sense to someone deciding whether to approve it.
- `kv_read_table` module docstring: same, condensed.
- The `-1` guard keeps the fact that makes it necessary -- Triton's `//`
truncates, so `-1 // 1` is -1 and would read before the table -- and drops
the retelling of how the two implementations came to disagree.
- `index_table_for_batch` keeps the trap (`seq_lens_cpu` is non-None but stale
on a gpu_only batch, and a stale max under-sizes the table); drops the walk
through the consequence.
- Triton backend's write hook and two test fixtures: keep the constraint, drop
the derivation.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Leftovers from the rename. The module is the translator, what it hands out is the read table, and the per-forward write target is the write loc -- a few comments and one test class still said choke point / canonical / rail. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…, not a mode `build_index_table(captured=True)` said the wrong thing. Nothing about the call is captured -- all six call sites are in `init_forward_metadata_out_graph` (or `_apply_cuda_graph_metadata`, which it calls), so the build kernel runs out of graph on both arms. What the flag actually selected was the DESTINATION: the module's own pre-allocated buffers, because a captured graph will later bake their pointers, versus a fresh allocation. That is a property of the buffer, not of the caller's graph state, and a reader who takes `captured=True` to mean "this runs inside a graph" is misled. It was not really a choice either: every explicit caller passed `True`. `False` was only ever reached through `index_table_for_batch`. So pass the buffer. `out=` / `out_sliding_window=` replace `captured`, and `out=None` allocates as before -- which is the shape `fill_read_table(out=...)` already had for trtllm_mla / flashmla, so the two are now one operation and `fill_read_table` is a three-line wrapper that returns the tensor instead of the table. The width rule comes with it: a caller-owned table may be padded wider than req_to_token's span, so both paths clamp. `ensure_capture_buffers` goes; `make_read_table_buffers` gives the same zero-filled shape and the BACKEND owns what it allocates, next to every other capture-stable buffer it already holds -- which is where cuda-graph loc management belongs. The implicit "call ensure_capture_buffers first" ordering, enforced by a runtime assert, becomes unrepresentable: you cannot pass a buffer you have not made. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ch-wan
force-pushed
the
mainline/write-loc-rebind
branch
from
August 30, 2026 22:14
87b2c19 to
ff658bf
Compare
Every caller owns the block table it passes in and discards the return value; the docstring also still named the pre-into= parameter.
ch-wan
force-pushed
the
mainline/write-loc-rebind
branch
from
August 30, 2026 22:31
19af95a to
29fb983
Compare
The comments this PR adds that narrate their own call site, and the em-dashes it introduces in comments and docstrings (.claude/rules/comment-style.md requires ASCII). Only lines this PR owns; pre-existing text is left alone.
Give `KVIndexTranslator` the WRITE half of the id-space contract, in two
phases, inert until the next commit wires the callers. One rule carried
over from the read path: per-forward translation RESULTS live on per-batch
products, never as translator state — the write half holds NO per-forward
state at all.
Phase 1 — `rebind_write_loc(forward_batch)` converts `out_cache_loc` to
FULL-side kernel-facing ids exactly once, at ForwardBatch construction.
REBIND, never mutate: the convert returns a FRESH tensor, so the
ScheduleBatch's aliased tensor stays virtual for the radix/accept/inflight
machinery. Construction is the only legal moment: the plan stream may
snapshot the tensor into capture-stable buffers immediately after, and a
captured graph carries no convert nodes to fix it later.
Phase 2 — the sliding-window write loc becomes a field on the per-batch
`KVIndexTable` (`sliding_window_write_loc`), computed by `build_index_table`
at the same moment as the index table. It derives POINTWISE from the kernel-facing
full-side VALUES, by inverting the full-side page scaling through the
allocator's physical->virtual table and mapping the virtual page through
the swa side's own table:
offset = kernel-facing % (page_size * full_multiplier)
virt_page = full_p2v[kernel-facing // (page_size * full_multiplier)]
swa_loc = swa_v2p[virt_page] * (page_size * swa_multiplier) + offset
Deriving from values is the point: the DP pad appends zeros (slot 0 is the
reserved padding slot in every id space), TBO slices, and runner buffer
copies all PRESERVE values — so whatever tensor a batch carries at build
time derives correctly, with no identity tracking, no handover calls, and
no ordering constraint between the two phases. Static SWA pools get the
same field from their own legacy full->swa translate, at the same build.
The fused draft region has its own page stride and its own disposition;
its locs never flow through this derivation.
- python/sglang/srt/mem_cache/multi_ended_allocator.py: expose
`full_p2v_page_table` on the SWA composite (the inverse of
`full_v2p_page_table`).
- python/sglang/srt/mem_cache/kv_index_translator.py: the field, the
two-phase surface, and the derivation.
- test/registered/unit/mem_cache/test_kv_index_translator.py: pin the
fresh rebind + untouched virtual alias; the round-trip property
`field(full(t)) == swa(t)` across page sizes {1, 4, 64} and real
multipliers; pad lanes derive to slot 0 with no bookkeeping; slices
crossing the pad boundary and fresh equal-value copies derive pointwise
with no handover; tombstoned swa pages clamp to slot 0; the static-pool
field equals the pool's legacy translate; no-loc / no-swa-side builds
carry None; a rebind retires the previous forward's read-table memo.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…Batch prep Flip the unified-memory write path onto the translator's two-phase contract: `init_new` rebinds `out_cache_loc` to full-side kernel-facing ids exactly once, every downstream full-side consumer becomes a passthrough, and the sliding-window side is read off the per-batch build (`KVIndexTable.sliding_window_write_loc`) — metadata-time sites read the field, and the per-layer forward sites read backend metadata, whose captured variant is the capture-stable buffer a recorded kernel must keep reading. No batch transform needs carry code: pads, TBO slices, and buffer copies preserve the values phase 2 derives from. - python/sglang/srt/model_executor/forward_batch_info.py: `init_new` calls `kv_index_translator.rebind_write_loc` (before the idle early-return, so idle batches also retire the read memo); `_pad_inputs_to_size` pads with zeros and hands nothing over — pad lanes derive to slot 0 by value. - python/sglang/srt/model_executor/runner/eager_runner.py: the registry rebuild needs no handover either; the rebuilt loc carries the same values. - python/sglang/srt/layers/attention/triton_backend.py: the write-side duck-typed `_translate_kv_loc` hook is deleted; the eager full write loc and the capture-stable refill consume `forward_batch.out_cache_loc` directly (the refill's convert becomes a copy — valid because capture batches are zero-filled and slot 0 is reserved as padding in every id space, pinned in test_multi_ended_allocator). The sliding-window write sites split by when they run: the eager metadata site and the cuda-graph refill read the build's field (the captured build now takes `out_cache_loc` and returns the table to the refill), and the two per-layer forward sites (window-layer read-back and extend) read `forward_metadata.swa_out_cache_loc` — the capture-safe delivery, mirroring the full side's `out_cache_loc_full_physical`. - python/sglang/srt/layers/attention/trtllm_mla_backend.py: the kernel-facing write-loc refill becomes a copy of the already-rebound loc; `loc_is_kernel_facing` retires (every loc reaching the door is kernel-facing). - python/sglang/srt/mem_cache/memory_pool.py: the MLA WRITE door (`set_mla_kv_buffer`) forwards `loc` untouched; the READ door keeps its convert — its indices are req_to_token-produced and the read-side migration is the next PR's topic. `KVWriteLoc` docs updated: `loc` is kernel-facing on every pool. - python/sglang/srt/mem_cache/unified_memory_pool.py: the factory's `_full_translate` wiring is now read-door-only; say so. - test/registered/unit/mem_cache/test_full_loc_fast_path.py: the MLA write door's pin inverts with the contract — it now asserts the door forwards `loc` UNTOUCHED, so a re-added door convert (which would double-convert every unified MLA write) turns it red. The read door keeps its convert-exactly-once pin. - test/registered/unit/mem_cache/test_multi_ended_allocator.py: pin `virtual 0 <-> physical 0` across alloc/free/compaction churn, on BOTH sub-pools — the zero-fill/copy equivalences and the pad-lane derivation are only valid while slot 0 stays the reserved padding slot everywhere. - test/registered/unit/mem_cache/test_unified_mha_views.py: end-to-end over the real factory — the rebind emits the full-attention kernel-facing loc and the build derives the sliding-window one, both checked against the formulas over VIRTUAL ids. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…cewise NO prefill capture backend is wired for the unified pool: the prefill graph runner builds its ForwardBatch directly, so it never runs the write-loc rebind (`rebind_write_loc`) and a captured prefill batch would carry VIRTUAL ids — the captured store would silently write wrong slots. The gate only rejected TC_PIECEWISE while the generic prefill default is BREAKABLE, so the DEFAULT unified invocation was broken out of the box; it only ever worked when --disable-piecewise-cuda-graph (a deprecated alias for --cuda-graph-backend-prefill=disabled) happened to be passed. - python/sglang/srt/server_args.py: `_handle_unified_memory_pool` turns prefill capture OFF (with a warning) whenever the user did not ask for a prefill backend explicitly, and raises an actionable error when they did -- an explicit request is never silently overridden. Decode capture, the wired path, is untouched. - test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py: the default is auto-disabled for every non-disabled backend, an explicit request raises, an already-disabled config is a no-op, and decode capture survives all three. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The write contract's SEMANTICS live with the translator and are pinned in test_kv_index_translator.py; what nothing pinned yet is the WIRING — the one ForwardBatch call site the contract hangs on. Phase 2 needs no ForwardBatch-side wiring at all (pads, slices, and buffer copies preserve the values it derives from), so exactly one call site remains to pin: - test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py: `init_new` must call `kv_index_translator.rebind_write_loc` — a construction path that skipped it would ship VIRTUAL write ids to the kernels, a silent wrong-slot store under the unified pool (checked structurally); end-to-end over the REAL `_pad_inputs_to_size` with a live translator — pad lanes are zeros and zeros derive to slot 0, and a slice of the padded tensor (the TBO-child shape) derives pointwise — plus the empty-loc rebind edge. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…r does Moving the write translate to ForwardBatch construction removed `_full_translate` from `set_mla_kv_buffer` and, with it, the comment that explained what `_full_translate` is for. What is left is a lone `self._full_translate = lambda ids: ids`, one surviving caller, and a set/get pair that now disagree about the id space of `loc` for no stated reason. The disagreement is deliberate and load-bearing: - `set_mla_kv_buffer` gets `forward_batch.out_cache_loc`, already rebound to kernel-facing ids, so it must NOT translate. - `get_mla_kv_buffer` gets `prefix_chunk_kv_indices[i]` / `fetch_mha_one_shot_kv_indices()`, which ForwardBatch builds straight out of req_to_token and nothing in this series touches, so it is still VIRTUAL and must. Write that down where the attribute is defined. Finishing the cleanup by deleting it would send the chunked-prefix MLA read through `v2p[kernel_id // ps]`, past the end of the table, on exactly the models the unified pool supports. Also clamp the intermediate virtual page in `_swa_write_loc_unified`. An unmapped physical page reads back as -1, and the next gather relied on torch wrapping that onto the v2p table's last element, which is the -1 trailing sentinel, so the final clamp still reached the sink. Two unrelated invariants held that up; one clamp replaces both. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`rebind_write_loc` runs in `ForwardBatch.init_new`, and that is the only place it runs. Every direct `ForwardBatch(...)` construction skips it -- both graph runners, dspark_draft, dflash_worker_v2, the eagle runners, two_batch_overlap -- and a skipped rebind is silent: the store lands on the wrong slots and only the model output is wrong. Nothing could see it. `maybe_detect_oob(loc, 0, size + page_size)` bounds locs by `_view_rows`, and a VIRTUAL id is `blocks_per_page` times smaller than the kernel-facing id for the same slot, so it is comfortably inside those bounds. This series already paid for that once: prefill cuda-graph capture had to be turned off after the captured batch was found carrying virtual ids, and it was found by reading the code. Give the property a probe. A kernel-facing id is `phys_page * (page_size * blocks_per_page) + offset` with `offset < page_size`, so its remainder modulo the page stride is always below page_size; a virtual id satisfies that only if it happens to land in the first block, so a batch of them trips it essentially always. Checked with `torch._assert_async` under the existing SGLANG_ENABLE_ASYNC_ASSERT gate -- same facility and same no-sync cost profile as the OOB probe beside it, and no new environment variable. `KVCache.kernel_page_blocks` carries the count (1 = no kernel-facing space, so the probe is vacuous); the two unified pools set it from their spec. Verified on an H200 that kernel-facing locs pass and virtual locs fire with the intended message. `UnifiedMLATokenToKVPool` also stops open-coding `spec.layer_num` for its view row count, matching `_build_mla_views` -- same value, one source. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Same pass, over what this PR added. The `_full_translate` block was the worst of them: eleven lines, most of them addressed to whoever might delete the attribute next. The fact underneath is one sentence -- the two MLA doors take different id spaces, and which is which -- so that is what stays. `maybe_detect_kernel_facing_loc`, the `_swa_write_loc_unified` clamp and the rebind test's module docstring get the same treatment: keep the property being checked and the trap it exists for, drop the walk through why it works and what it was a response to. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…r, not the read table The sliding-window WRITE loc rode along on KVIndexTable, so the one backend that needed it had to keep an out_cache_loc= argument threaded through build_index_table and _apply_cuda_graph_metadata, and the latter had to return the table purely so the refill could read one field off it. Ask the translator for the loc where it is needed instead: the field, the argument, the coupling between into= and out_cache_loc=, and the return value all go away, and a new backend has one method to call rather than a table-construction argument to remember.
The comments this PR adds that narrate their own call site, and the em-dashes it introduces in comments and docstrings (.claude/rules/comment-style.md requires ASCII). Only lines this PR owns; pre-existing text is left alone.
ch-wan
force-pushed
the
mainline/write-loc-rebind
branch
from
August 31, 2026 06:49
29fb983 to
54a59c2
Compare
5 tasks
nzr-niu
pushed a commit
to nzr-niu/sglang
that referenced
this pull request
Sep 1, 2026
…rwardBatch construction (sgl-project#35245) Co-authored-by: Caihua Li <caihua.li@bytedance.com> Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
RolaoDenthu
pushed a commit
to RolaoDenthu/sglang
that referenced
this pull request
Sep 1, 2026
…rwardBatch construction (sgl-project#35245) Co-authored-by: Caihua Li <caihua.li@bytedance.com> Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This was referenced Sep 1, 2026
This was referenced Sep 2, 2026
StevenChenSE
pushed a commit
to StevenChenSE/sglang
that referenced
this pull request
Sep 6, 2026
…rwardBatch construction (sgl-project#35245) Co-authored-by: Caihua Li <caihua.li@bytedance.com> Co-authored-by: Cheng Wan <cheng.wan@radixark.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
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
Under the unified memory pool a KV location has two forms: the virtual id that ScheduleBatch-side
machinery (radix cache, accept bookkeeping, in-flight tracking) reads, and the kernel-facing id the
store kernels need. Today each attention backend and each pool door performs that translation
itself. The same location is translated a different number of times depending on which path a forward
takes, and every new consumer has to remember to translate — a silent-corruption failure mode,
since writing a virtual id as though it were physical lands KV in the wrong slot rather than raising.
Modifications
Establishes a single rule for the write path: ScheduleBatch-side tensors stay virtual always, and
each ForwardBatch is translated exactly once, at its construction. Every downstream consumer
— backend metadata, cuda-graph refill, the write door — becomes a passthrough, and the
backends' own write-side translations are removed. The rebind produces a fresh tensor rather
than mutating in place, because the scheduler-side machinery reads the same tensors and
requires virtual ids; for hybrid-SWA models the two rails are derived in the one order that is
correct, since a single virtual id maps to two different kernel-facing ids.
The translation results live on the per-runner
KVIndexTranslator(the pool subsystem), not on theForwardBatch — with more hybrid attention kinds a forward carries more kernel-facing id spaces,
and one ForwardBatch field per id space does not scale. Backends fetch the swa-side rail into
their own attention metadata through
resolve_swa_write_loc, which recognizes the preparedloc or any torch view of it by address-range containment. That is what removes every piece of
carry code: TBO child slices and cuda-graph replay views resolve as-is (this PR touches neither
two_batch_overlap.pynor the graph runners), and the two transforms that genuinely replacethe tensor — the eager input registry's static-buffer rebuild and DP-sync padding — hand the
new tensor over explicitly.
Because the contract is now positional rather than local, the choke point enforces it: fetching
the rail for a loc that was never prepared refuses loudly instead of silently writing virtual ids as
physical. The PR also turns prefill cuda-graph capture off under the unified pool (no prefill
capture backend is wired for the rebind; the default invocation was broken without this, masked
wherever
--disable-piecewise-cuda-graphhappened to be passed).Accuracy Tests
GSM8K, unified vs. baseline.
Speed Tests and Profiling
Serving benchmark (heavy-decode, radix-retract), unified vs. baseline at matched attention backend.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #33365731844
Latest PR Test (Extra): ❌ Run #33365731572
Latest PR Test (AMD ROCm 7.2): ❌ Run #33365731652