Skip to content

feat(unified-memory): read unified pool from attention backends fa3/flashinfer/trtllm_mha/flashmla - #34613

Merged
ch-wan merged 12 commits into
sgl-project:mainfrom
caihuali95:mainline/read-choke-point
Aug 31, 2026
Merged

ch-wan merged 12 commits into
sgl-project:mainfrom
caihuali95:mainline/read-choke-point

Conversation

@caihuali95

@caihuali95 caihuali95 commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

Note on scope — this PR was split. It previously contained both the read-path refactor
and this backend migration. The refactor is now a separate PR that this one stacks on, so
the mechanism can be reviewed independently of the per-backend work.
Stack order: #34602#35245#35247 → this PR.

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

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.

model class before after
unified MHA/SWA (uniform rows) triton triton, fa3, fa4, flashinfer, trtllm_mha
unified MLA triton, fa3, trtllm_mla, flashinfer, cutedsl_mla, tokenspeed_mla + flashmla
asymmetric K/V, or env-forced strided triton triton (unchanged)

unified_mem_hooks.py is deleted. After the fold, two translate residents remain
outside 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, outside
the forward path). A registered source scan enforces this — it fails if any
file under layers/attention/ calls a unified translate, probes for one, or
imports the deleted hooks module.

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 at matched backend. Every newly enabled backend is at parity.

Backend Median Δ
Triton +0.00 pt
fa3 +0.00 pt
FlashInfer +0.00 pt
flashmla +0.00 pt

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.

model backend ITL
Qwen3.5-9B (GDN) triton +0.55% / −0.29%
fa3 +0.10% / −0.20%
flashinfer +0.35% / −0.22%
gpt-oss-20b (SWA) triton −1.13% / −0.15%
fa3 +0.35% / +0.04%
Falcon-H1-7B (Mamba) triton +1.47% / −0.38%

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 #33366081537
Latest PR Test (Extra): ❌ Run #33366158884
Latest PR Test (AMD ROCm 7.2): ❌ Run #33366081659

@caihuali95

Copy link
Copy Markdown
Contributor Author

/tag-and-rerun-ci extra

Caihua Li and others added 12 commits August 31, 2026 06:54
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
ch-wan force-pushed the mainline/read-choke-point branch from 2a3e960 to 2be1787 Compare August 31, 2026 06:54
@ch-wan ch-wan added the run-ci label Aug 31, 2026
@ch-wan
ch-wan merged commit 8bb776d into sgl-project:main Aug 31, 2026
129 of 150 checks passed
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>
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>
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