Skip to content

refactor(unified-memory): translate the KV write location once, at ForwardBatch construction - #35245

Merged
ch-wan merged 19 commits into
sgl-project:mainfrom
caihuali95:mainline/write-loc-rebind
Aug 31, 2026
Merged

ch-wan merged 19 commits into
sgl-project:mainfrom
caihuali95:mainline/write-loc-rebind

Conversation

@caihuali95

@caihuali95 caihuali95 commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Note on scope — this PR was split out. It was previously part of #34602 alongside the
dense per-layer view feature. Separating them lets the layout change and this refactor be
reviewed on their own terms. The stack was also reordered per review: this PR now sits ON
the read-path choke point (#35247) and keeps its state there, so nothing is stored on the
ForwardBatch. Stack order: #34602#35247 → this PR → #34613.

Only the last 4 commits are for this PR. The other commits are from the stacked PRs.

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 the
ForwardBatch — 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 prepared
loc 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.py nor the graph runners), and the two transforms that genuinely replace
the 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-graph happened to be passed).

Note: the accuracy and speed tables below were measured on the previous revision of this stack. The series has since been restructured per review and rebased onto current main; a revalidation pass is in progress and the tables will be refreshed.

Accuracy Tests

GSM8K, unified vs. baseline.

Model family Median Δ Range
Qwen3.5-9B +0.00 pt −1.5 … +1.5
Kimi-Linear-48B +0.00 pt −2.5 … +1.0
Falcon-H1-7B +0.50 pt −3.0 … +3.0
gpt-oss-20b −1.75 pt −13.0 … +10.5

Speed Tests and Profiling

Serving benchmark (heavy-decode, radix-retract), unified vs. baseline at matched attention backend.

Model family Median Δ ITL
Qwen3.5-9B +0.38%
Kimi-Linear-48B +1.15%
Falcon-H1-7B −0.39%
gpt-oss-20b +0.46%

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ❌ Run #33365731844
Latest PR Test (Extra): ❌ Run #33365731572
Latest PR Test (AMD ROCm 7.2): ❌ Run #33365731652

Comment thread python/sglang/srt/batch_overlap/two_batch_overlap.py Outdated
Comment thread python/sglang/srt/model_executor/forward_batch_info.py Outdated
ch-wan and others added 6 commits August 30, 2026 22:11
…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
ch-wan force-pushed the mainline/write-loc-rebind branch from 87b2c19 to ff658bf Compare August 30, 2026 22:14
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
ch-wan force-pushed the mainline/write-loc-rebind branch from 19af95a to 29fb983 Compare August 30, 2026 22:31
ch-wan and others added 10 commits August 31, 2026 06:46
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
ch-wan force-pushed the mainline/write-loc-rebind branch from 29fb983 to 54a59c2 Compare August 31, 2026 06:49
@ch-wan
ch-wan merged commit 29578d5 into sgl-project:main Aug 31, 2026
106 of 123 checks passed
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>
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants