feat(unified-memory): read unified pool from attention backends fa3/flashinfer/trtllm_mha/flashmla - #34613
Merged
Conversation
caihuali95
requested review from
BBuf,
DarkSharpness,
Fridge003,
HaiShaw,
HydraQYH,
Qiaolin-Yu,
Ying1123,
alphabetc1,
celve,
hanming-lu,
hebiao064,
hnyls2002,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann,
yizhang2077 and
yuan-luo
as code owners
August 12, 2026 18:06
Contributor
Author
|
/tag-and-rerun-ci extra |
caihuali95
force-pushed
the
mainline/read-choke-point
branch
from
August 18, 2026 03:52
d90771d to
f869768
Compare
This was referenced Aug 18, 2026
Merged
ch-wan
force-pushed
the
mainline/read-choke-point
branch
from
August 30, 2026 22:31
08e38cf to
01455d8
Compare
ch-wan
force-pushed
the
mainline/read-choke-point
branch
from
August 31, 2026 00:19
01455d8 to
2ae6c76
Compare
ch-wan
force-pushed
the
mainline/read-choke-point
branch
from
August 31, 2026 06:49
0a7a84f to
2a3e960
Compare
Point the FlashAttention backend (fa3 and fa4, which is the same class) at the batch's `KVIndexTable` for its page tables, so it reads the unified pool's per-layer views without any id-space logic of its own. Byte-identical for non-unified pools, where the view is `req_to_token` + `req_pool_indices`. - python/sglang/srt/layers/attention/flashattention_backend.py: hold the runner's `KVIndexTranslator`, take the per-batch view where page tables are built, and assert the two shapes it cannot serve yet (unified + local attention). The metadata builders gather from the view instead of `req_to_token`, and the swa write-target refill derives the write loc through the translator (`sliding_window_write_loc_for`) with the capture arm zero-filling — a runner-built capture batch has no prepared write loc, and slot 0 is the sink in every id space. The FB-carried write-loc parameter is gone; capture/replay is an explicit `in_capture` flag, the convention this backend already uses for its prefill graph path. - python/sglang/kernels/ops/attention/metadata.py: thread the read table and its page size through the decode metadata kernel so it can reconstruct token ids from a page-granular read table. - python/sglang/srt/mem_cache/memory_pool.py, python/sglang/srt/mem_cache/unified_memory_pool.py: `KVWriteLoc.full_loc` becomes an optional same-space alias — since the construction-time rebind the generic `loc` is already kernel-facing, so the composite falls back to it when no alias is passed (fa3 passes none; only triton's captured path still carries its capture-stable buffer). - test/registered/unit/layers/attention/test_flashattention_graph_metadata.py, test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py: give the fake runner a real `KVIndexTranslator` over its non-unified pool (the probe disables it), since the backend now dereferences that attribute during __init__. - test/registered/attention/test_normal_decode_set_metadata.py, test/registered/unit/mem_cache/test_unified_mla_block_table.py: cover the page-granular source table in the decode-metadata and block-table cases. - test/registered/unit/mem_cache/test_full_loc_fast_path.py: follow the door signature change above. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ranslator Point both flashinfer backends at the batch's `KVIndexTable` for their KV index buffers, so they read the unified pool's per-layer views without any id-space logic of their own. Byte-identical for non-unified pools, where the view is `req_to_token` + `req_pool_indices`. - python/sglang/srt/layers/attention/flashinfer_backend.py: build the KV indices from the view on the decode and extend paths, and resolve the swa write rail through the source (`sliding_window_write_loc_for`) at the eager metadata build and the cuda-graph refill — the capture arm zero-fills (a runner-built capture batch has no prepared write loc; slot 0 is the sink in every id space). The static-pool read translate keeps its legacy in-place map, guarded on the view not being kernel-facing. - python/sglang/srt/layers/attention/flashinfer_mla_backend.py: thread the view through the plan/replay chain alongside the CPU-length arguments upstream added to the same signatures (sgl-project#34614), so both sets of parameters reach `call_begin_forward` together. - python/sglang/kernels/ops/kvcache/kv_indices.py: the shared builder takes the view's source page size, so one read table serves token-level consumers. - test/registered/unit/mem_cache/test_unified_mla_block_table.py: extend the block-table cases to the flashinfer path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Build the MLA block tables through the translator's `fill_read_table` instead of each backend gathering `req_to_token` and translating, so trtllm_mla and the two backends that subclass it read the unified pool's per-layer views with no id-space logic of their own. Unchanged for non-unified pools, which keep their existing builder. - python/sglang/srt/layers/attention/trtllm_mla_backend.py: fill the padded block table's live prefix via `fill_read_table` when the translator is translating, and collapse the fused fp8 write-loc selection into a named `_resolve_fused_write_loc` helper (captured decode reads the refilled kernel-facing buffer; eager decode on a unified pool is simply not covered by the fused path). - python/sglang/srt/layers/attention/cutedsl_mla_backend.py, python/sglang/srt/layers/attention/flashmla_backend.py: same dispatch at their own table builders, which is what brings both backends onto the unified pool. - python/sglang/kernels/ops/kvcache/kv_indices.py: the flashmla index kernel takes the view's source page size, matching the flashinfer one. - test/registered/unit/mem_cache/test_unified_mla_block_table.py: cover the block-table path for each of the three backends, and the prefix-only fill (columns past the live pages keep the backend's own -1 sentinel). - test/registered/unit/mem_cache/test_multi_ended_allocator.py: cover the page-granular entries the block tables are built from. `fill_read_table`'s cap at the widest legal column — a padded table may extend past `req_to_token` — is what makes the padded block tables here safe, but it lands with the translator in the previous PR, not in this commit. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Bind the trtllm_mha page tables to the translator's capture-stable read-table buffers, so the recorded graph reads tables that are refreshed out of the graph rather than rebuilding them from `req_to_token` inside it. Unchanged for non-unified pools, which keep building their own tables. - python/sglang/srt/layers/attention/trtllm_mha_backend.py: take the captured view at metadata setup and point the graph metadata at its tables; the in-graph builder is told to skip the page-table and sliding-window write-loc work it must not redo, under one named `unified` predicate applied uniformly to those arguments. The swa write loc derives through the translator (`sliding_window_write_loc_for`) at the eager metadata build and the cuda-graph refill; the capture arm zero-fills (a runner-built capture batch has no prepared write loc, and slot 0 is the sink in every id space). - python/sglang/kernels/ops/kvcache/trtllm_mha_graph_metadata.py: accept the skip flag and leave the caller-owned tables untouched when it is set. - test/registered/attention/test_trtllm_mha_graph_metadata.py: pin that the skip path leaves the read tables byte-identical while still filling the sequence-length outputs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ot at the door Finish the read-path consolidation for the MLA model doors: the deepseek MHA mixin's req_to_token-derived indices (chunked-prefix and one-shot) are translated where they are PRODUCED, through the runner's translator, and both pool doors become passthrough. Every id space is now translated at exactly one place — writes at ForwardBatch construction, block-table reads inside the read table build, and flat-index reads at their production sites — so the `_full_translate` hook has no caller left and is deleted. - python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py: `prepare_chunked_kv_indices` / `fetch_mha_one_shot_kv_indices` run `kv_index_translator.translate_full` on their freshly built indices (after the DCP filter — an elementwise translate commutes with subsetting); the memo then holds the kernel-facing result. The source is reached through `get_attn_backend()`, the mixin's existing accessor pattern. - python/sglang/srt/layers/attention/base_attn_backend.py: class-level `kv_index_translator = None`, so the mixin None-checks instead of probing — a backend that never set it cannot serve the unified pool (the server-args allow-list enforces that). - python/sglang/srt/layers/attention/tbo_backend.py: forward the primary's source, next to the two pool attributes already forwarded. - python/sglang/srt/mem_cache/memory_pool.py: `get_mla_kv_buffer` forwards `loc` untouched, matching the write door; the `_full_translate` attribute is gone. - python/sglang/srt/mem_cache/unified_memory_pool.py: the factory wires no translate hook; a note states where each id space is kernel-facing instead. - test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py: the production-site pins — one-shot indices translated exactly once with the memo holding the translated result, the raw producer output passing through on an unmigrated backend (base-class None), and the read door forwarding `loc` untouched over a real HybridLinearKVPool. - test/registered/unit/mem_cache/test_full_loc_fast_path.py: the read door's pin inverts — it asserted translate-exactly-once, it now asserts UNTOUCHED forwarding, so a re-added door translate (a double translate for every unified MLA prefix read) turns it red. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…e boundary With every backend on the translator, no attention backend calls a unified translate surface or probes for one. Delete the capability shim they used to go through, and add a scan that keeps it that way. - python/sglang/srt/layers/attention/unified_mem_hooks.py: deleted. Its last importer was its own detection test; its guarded semantics live on in `KVIndexTranslator`'s isinstance probe. - test/registered/unit/layers/attention/test_kv_translate_choke_point.py: new. Scans python/sglang/srt/layers/attention/ and fails on a unified translate call, a translate-capability getattr probe, or an import of the deleted module. All three are red on the pre-migration backends. Two translate residents remain outside the translator by design, neither reachable from a backend: the allocator-internal v2p implementations, which are the mechanism the translator calls, and the PD transfer-plane translate, which stages for RDMA outside the forward path. - test/registered/unit/mem_cache/test_unified_mla_block_table.py: drop the hook-detection case along with the module it tested; the multiplier-1-is-still-unified regression it also carried is kept. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
With the read path behind the translator, unified-memory models no longer
have to run on Triton. Widen the page-major layout gate from one arm to two,
so each pool shape admits exactly the backends that can now read it.
- python/sglang/srt/server_args.py: `_handle_page_major_kv_layout` gains an
MHA/SWA arm and grows its MLA one. Unified MLA admits {triton, fa3,
trtllm_mla, flashinfer, cutedsl_mla, tokenspeed_mla, flashmla} — cutlass_mla
stays rejected, never having been exercised. Unified MHA/SWA admits {triton,
fa3, fa4, flashinfer, trtllm_mha}. Plain page-major without the unified pool
keeps the envelope-strided views only Triton reads.
- test/registered/unit/server_args/test_page_major_backend_allowlist.py:
rewritten for the two arms — the allowed set per arm, MLA-only kernels
barred from the MHA arm, and the per-layer-view families rejected without the unified
pool.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… gate Add GPU-CI cells covering every arm the backend allow-list can now take, so each pool shape is exercised by a real boot rather than only by unit pins. All cells are label-gated or nightly, so there is no per-commit cost. - test/registered/page_major/test_page_major_gpt_oss.py: unified cells for triton and fa3 on a uniform-row hybrid-SWA model, plus a pin-free resolved-default boot (the pattern that caught both review defects in sgl-project#32972). No flashinfer cell — gpt-oss uses attention sinks. - test/registered/page_major/test_page_major_qwen_hybrid.py: unified cells for triton, fa3, flashinfer and the resolved default on a GDN hybrid; the conv/SSM state stays envelope-strided in every cell, with linear and mamba pinned to triton as the gate enforces. - test/registered/models_e2e/test_kimi_linear_unified_memory.py: a flashmla cell at its page-size-64 snap, exercising the read table block-table builder together with the 64-token sink floor and the view tail pad end to end. `est_time` registrations are bumped for the added boots. DSPARK with --enable-unified-memory has no cell here: it needs a DSPARK-capable draft checkpoint that no registered tiny model provides today, so its verify-shape contracts stay unit-pinned in test_kv_index_translator.py and test_triton_unified_verify_sources.py. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The rename left the prose behind. Three words were still in use for things that already have names in this series: "choke point" for the translator, "canonical" for the read table it hands out, and "rail" for the per-forward write loc. `test_kv_translate_choke_point.py` -> `test_kv_translate_ownership.py`, which is what it actually scans for: id-space translation is owned by exactly two places and no attention backend may re-grow its own. Its docstring also named `apply_unified_kv_loc_rebind`, which does not exist -- the entry point is `rebind_write_loc`. Untouched, different meanings: `canonicalize_stride`, "non-canonical stride", and "the canonical in-place form" (as in "the usual form"). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Same pass, over the backend wirings. These carried the most of it: each one narrated what the eager branches above had produced, what was being swapped in for what, and why the swap was safe -- a description of the change rather than of the code. What stays is what a reader cannot get from the code: that the translator's tables are PAGE-granular and already kernel-facing (so the `// page_size` reduction and the static full->swa map do not apply to them), that on replay `kv_indices` IS the buffer the captured wrapper reads and must never be rebound, that only the live prefix is refreshed and the stale tail is bounded by seq_lens_k, that a capture batch never went through `init_new`, and that `None` on the backend attribute means "no translate". The allow-list keeps its per-family bullets -- that is the decision the block exists to record -- and loses the paragraph explaining the id surface. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ecode updaters The decode gathers now address rows through `kv_view.row_ids`, so req_pool_indices survived only as `bs = len(req_pool_indices)` at the bottom of a four-level thread: _apply_cuda_graph_metadata -> update -> update_single_wrapper / update_sliding_window / update_cross_attention -> call_begin_forward. Derive bs from paged_kernel_lens instead and delete the parameter from all five, in both the MHA and the MLA decode updater. The prefill updaters keep it: `spec_info.generate_attn_arg_prefill` still indexes req_to_token by pool row, which is also why FlashInferMLAAttnBackend._apply_cuda_graph_metadata keeps the argument -- its target-verify arm is the one remaining consumer.
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/read-choke-point
branch
from
August 31, 2026 06:54
2a3e960 to
2be1787
Compare
nzr-niu
pushed a commit
to nzr-niu/sglang
that referenced
this pull request
Sep 1, 2026
…lashinfer/trtllm_mha/flashmla (sgl-project#34613) Co-authored-by: Caihua Li <caihua.li@bytedance.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.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
…lashinfer/trtllm_mha/flashmla (sgl-project#34613) Co-authored-by: Caihua Li <caihua.li@bytedance.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
ch-wan
added a commit
that referenced
this pull request
Sep 1, 2026
`get_mla_kv_buffer` is a read door with the caller-translates contract: the pool never translates, and its docstring names the two production sites the DeepSeek MHA mixin owns. DCP adds a third that nobody updated. `prepare_decode_context_parallel_metadata` builds `dcp_local_prefix_kv_indices` straight off `req_to_token`, collapses `loc // dcp_size`, and hands the result to `get_mla_kv_buffer` to gather the shared prefix. On the unified pool those are VIRTUAL ids read as kernel-facing ones -- the two spaces share a value range, so nothing raises; the gather just returns another request's KV. Fail-silent by construction, the same family as #34613 -> #37307: the hook is `Optional` and identity-defaulted, so "forgot to wire it" and "no translation needed" are the same value. Kimi-Linear-48B, TP2, flashinfer, GSM8K 5-shot 200q, --enable-unified-memory --dcp-size 2, CUDA graph on: 0.000 (Invalid 1.000) -> 0.930 (Invalid 0.000). Static-pool DCP is 0.910 and unified without DCP is 0.905, so this lands the DCP cell back in the same band rather than merely improving it. The damage tracked prefix reuse, which is what made it look like a graph bug: radix on + graph 0.000, radix on + eager 0.650, radix off 0.825 (chunked prefill still produces prefixes). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ch-wan
added a commit
that referenced
this pull request
Sep 1, 2026
Follow-up to the loc-space unification on this branch. Both items are the same shape as the bugs that unification removed: a translation whose absence reads as "no translation needed". `MLATokenToKVPool.set_kv_buffer` selected `loc % dcp_size == dcp_rank` and then wrote WITHOUT dividing -- widened ids straight into a rank-local buffer. It has been there since #14194, which gave the two MLA write doors different treatment for the same input: `set_mla_kv_buffer`'s kernel divides, this one never did. It cannot be repaired into a correct write, because its two possible callers disagree on the loc space: flashinfer-MLA's `k_rope is None` branch passes a WIDENED loc, while the Triton backend passes one it already collapsed. So the door refuses under DCP instead. Probed for reachability first -- DeepSeek-V2- Lite TP2/DCP2 on flashinfer over decode, GSM8K, 9k chunked prefill and batched prefix-cache reuse never enters it; every MLA write goes through `set_mla_kv_buffer`. Nothing reachable changes behaviour. `HybridLinearKVPool.mamba_translate` defaulted to identity. The unified pool holds VIRTUAL mamba slot ids and installs its translate after construction (the pool is one hop of a cycle ending at the allocator that owns it), so a dropped install silently offloaded the wrong slots -- the fail-silent shape that cost a released regression in the KV read path (#34613 -> #37307). The default now refuses; `mamba_slot_identity` is the static pool's explicit answer, and the one static construction site that exercises HiCache offload now says so. Both guards verified red on the pre-fix code and green on the fix. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
5 tasks
StevenChenSE
pushed a commit
to StevenChenSE/sglang
that referenced
this pull request
Sep 6, 2026
…lashinfer/trtllm_mha/flashmla (sgl-project#34613) Co-authored-by: Caihua Li <caihua.li@bytedance.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.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
With the read path behind a single choke point, the remaining attention backends can read
the unified memory pool without carrying any id-space logic of their own. This PR does that
migration and widens the server's allow-list accordingly, which is what lifts the unified pool
from Triton-only to the backends users actually run.
Modifications
Migrates FlashAttention (fa3/fa4), FlashInfer's MHA and MLA backends, trtllm_mla together
with the two backends that subclass it, and trtllm_mha onto the choke point's canonical page
tables. Each backend keeps its existing builder for non-unified pools and dispatches to the
shared one only when the source is active, so the change is inert off the unified path. The MLA
model doors' read indices are translated at their production sites (the deepseek MHA mixin),
which lets both pool doors become passthrough and deletes the last pool-resident translate
hook. The now-unused backend hook shim is deleted and replaced by an enforcement scan, so
the boundary cannot be re-crossed silently. The startup allow-list is widened to match what the backends
can now serve — fa3/fa4/FlashInfer/trtllm_mha for the MHA and SWA families, plus flashmla
for MLA — and end-to-end cells pin the resulting matrix.
tritontriton, fa3, fa4, flashinfer, trtllm_mhatriton, fa3, trtllm_mla, flashinfer, cutedsl_mla, tokenspeed_mlaflashmlatritontriton(unchanged)unified_mem_hooks.pyis deleted. After the fold, two translate residents remainoutside the choke point, each by design and neither reachable from a backend: the
allocator-internal v2p implementations (the mechanism the choke point calls) and
the PD transfer-plane
translate_kv_indices_for_transfer(RDMA staging, outsidethe forward path). A registered source scan enforces this — it fails if any
file under
layers/attention/calls a unified translate, probes for one, orimports the deleted hooks module.
Accuracy Tests
GSM8K, unified vs. baseline at matched backend. Every newly enabled backend is at parity.
Speed Tests and Profiling
Serving benchmark at matched batch, production configuration: radix
cache + overlap scheduler + cuda graphs + page_size 256.
Two workloads:
heavy-decode (512-token prompts, 128 concurrent) and radix-retract
(4k shared-prefix prompts, 192 concurrent, forcing eviction and retraction).
Positive = unified is worse. Each row is heavy-decode / radix-retract.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #33366081537
Latest PR Test (Extra): ❌ Run #33366158884
Latest PR Test (AMD ROCm 7.2): ❌ Run #33366081659