fix(unified-memory): forward the KV-index translator through every wrapper backend - #37307
Merged
ch-wan merged 3 commits intoSep 1, 2026
Merged
Conversation
caihuali95
requested review from
Fridge003,
HaiShaw,
Qiaolin-Yu,
Ying1123,
alphabetc1,
hanming-lu,
hebiao064,
hnyls2002,
huangtingwei9988,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann,
yizhang2077 and
yuan-luo
as code owners
September 1, 2026 00:43
…apper backend
`AttentionBackend.kv_index_translator` is a class attribute defaulting to
None, so a backend that WRAPS another and does not re-expose the inner
copy answers "this backend needs no translation". Every producer that
reaches the translator through the live backend rather than through a
runner then skips translation silently. The MLA chunked-prefix-cache path
is one: `ForwardBatchDeepseekMHAMixin.prepare_chunked_kv_indices` reads
`get_attn_backend().kv_index_translator` and translates only
`if src is not None`, so it hands the kernel raw VIRTUAL ids for the
shared prefix. Nothing raises -- virtual and physical ids share a value
range, so the gather simply reads the wrong rows.
The damage needs the unified pool (translation actually required), an MLA
backend (only fa3 / flashmla take the chunked-prefix path; Triton does
not), and a radix cache (only then is there a shared prefix to read back).
On Kimi-Linear that is gsm8k 0.89-0.92 -> 0.05-0.31 on four of four such
cells, while the same configuration on Triton, and every chunked-prefill
cell, stays correct.
Four wrappers need the forward, not one. They spell the wrapped field
`full_attn_backend`, `prefill_backend`/`decode_backend`, `backend`/
`swa_backend` and `dense`, so nothing keyed on the attribute name finds
them all:
* layers/attention/hybrid_linear_attn_backend.py: from the full-attention
backend, beside the pool and req_to_token it already forwards.
* layers/attention/hybrid_attn_backend.py: from the model runner, which is
where it already takes the pool.
* layers/attention/dots_hybrid_backend.py: twice -- the DSA wrapper from
its inner backend, the SWA/MLA wrapper from its swa backend.
* layers/attention/minimax_sparse_backend.py: the hybrid wrapper, from its
dense backend.
mem_cache/kv_index_translator.py: `assert_backends_carry_translator`, which
refuses at boot when a reachable backend carries someone else's translator
or none at all, and is inert on a pool that needs no translation.
model_executor/model_runner.py: run that guard over the built backends,
right after they are constructed. Adding a forward is easy to forget in a
NEW wrapper and the failure is silent by construction, so the runtime
check is what makes it loud; the source scan below only covers shapes it
can recognise.
test/registered/unit/layers/attention/test_kv_translate_ownership.py: a
wrapper-forwarding case in the file that already owns translator
ownership. Derived, not enumerated -- a wrapper is any AttentionBackend
subclass TAKING another AttentionBackend as a constructor parameter, so a
new one is covered when it is written. Keying on the parameter type rather
than the attribute name is what makes it find all four. Red on the pre-fix
tree naming every offender, green after.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
caihuali95
force-pushed
the
mainline/wrapper-translator-fix
branch
from
September 1, 2026 01:56
0edaa1e to
d71d21b
Compare
…per-PR Replaces the source-scan wrapper test with an e2e cell that exercises the defect. The scan was file-scoped: `_iter_sources()` yields whole files and all three regexes `.search()` them, so one forwarding class anywhere in a file covers every wrapper in it. `hybrid_linear_attn_backend.py` holds four AttentionBackend subclasses and one forward -- putting the bug back while leaving any other `self.kv_index_translator =` in that file keeps the test green. It also required the wrapped parameter to carry a type annotation, and it checked the shape rather than the damage. `ModelRunner`'s `assert_backends_carry_translator` already covers the shape, over the live object graph rather than the source text. What was missing is an e2e cell: no per-PR test paired the unified pool with an MLA model, so nothing reached the MLA chunked-prefix / MHA-one-shot path that reads its translator off `get_attn_backend()`. The unified Kimi-Linear cell that would have caught it exists but is nightly, so it never gated the PR that introduced the bug. test_kimi_linear_models.py: add the resolved-default unified cell. This file is already base-b on 2-gpu-large and already launches this model twice, so the cost is one more server start in an allocated job, not the new 4-gpu job that promoting the nightly file would need. est_time 600 -> 900. The nightly file keeps its copy: it resolves a different default on its H100 runner, so the two cover different backends. Measured on Kimi-Linear-48B TP2, flashinfer, gsm8k 200q 5-shot, CUDA graphs on: 0.365 before the fix, 0.895 after (threshold 0.88). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ch-wan
force-pushed
the
mainline/wrapper-translator-fix
branch
from
September 1, 2026 02:06
2a5bdfd to
c2e6458
Compare
…ects Review follow-up on the two suggestions. The e2e cell only builds HybridLinearAttnBackend, via Kimi-Linear. The other six wrappers had nothing per-PR: `assert_backends_carry_translator` fires only when the wrapper is actually constructed, and no per-PR unified job builds Dots or MiniMax. This adds an object-graph test in the file that already owns translator ownership. Only the inner that must supply the translator carries it; every other inner carries None. A wrapper that copies from the linear / sparse / DSA side -- which leave the class default None -- ends up with None and fails, which a source scan cannot see. The wrapper set is derived per CLASS (an AttentionBackend subclass whose own __init__ takes another backend), not per file, and a completeness test asserts the derived set equals the set with instances. That is what the deleted scan got wrong: it searched whole files, so one forwarding class anywhere in a file covered every wrapper in it. Seven classes come out, including ShortConvHybridAttnBackend and TboAttnBackend. HybridAttnBackend takes its translator from the runner rather than an inner and reads the spec bag in __init__, so it is built under override_server_args. Verified by injection, each red on the new test and green on the old scan: drop the HybridLinear forward while leaving an unrelated `self.kv_index_translator = None` in that file (2 failed); point MiniMax at sparse_backend instead of dense_backend (1 failed); add a wrapper with no instance (1 failed). Clean tree: 5 passed, 7 subtests. test_kimi_linear_models.py: cut the class docstring to the black-box constraint and drop the restated measurements, per comment-style for a bug-regression test. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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
ch-wan
added a commit
that referenced
this pull request
Sep 2, 2026
Nothing in test/registered covered unified memory on a hybrid sliding-window
model. The existing unified e2e tests are MLA (Kimi-Linear) or ShortConv
(Inkling); the existing SWA e2e tests run the static pool. That missing cell is
why a config that dies inside one eval shipped.
google/gemma-4-E2B-it, 1 GPU, triton, --enable-unified-memory
--disable-radix-cache --mem-fraction-static 0.8 --max-total-tokens 60000
GSM8K over 200
narrow table 0.05 <- red on main
fixed passed <- 65 s wall clock
Every argument is load-bearing, and it took several wrong turns to establish
that:
- `--max-total-tokens 60000` is the trigger, not a convenience. The failure
needs cumulative churn past `swa.num_pages`; at this model's default budget
that table is 3.8M entries and it would take dozens of runs.
- the eval has to be this one. Synthetic short prompts at 3.7x the nominal
churn passed on main; the 5-shot GSM8K traffic is what reaches it.
- it cannot use `GSM8KMixin`: that path scores this model 0.155 whatever the
pool, because its chat template does not fit the model's reasoning config.
`run_eval` directly gives 0.87 on a healthy server.
Registered at `base-b` / `1-gpu-large` -- per-PR, not nightly. The lesson from
the last regression here (#37307) was that the guard existed but was
nightly-gated, so it never blocked the PR that broke it.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ch-wan
added a commit
that referenced
this pull request
Sep 2, 2026
Nothing in test/registered covered unified memory on a hybrid sliding-window
model. The existing unified e2e tests are MLA (Kimi-Linear) or ShortConv
(Inkling); the existing SWA e2e tests run the static pool. That missing cell is
why a config that dies inside one eval shipped.
google/gemma-4-E2B-it, 1 GPU, triton, --enable-unified-memory
--disable-radix-cache --mem-fraction-static 0.8 --max-total-tokens 60000
GSM8K over 200
narrow table 0.05 <- red on main
fixed passed <- 65 s wall clock
Every argument is load-bearing, and it took several wrong turns to establish
that:
- `--max-total-tokens 60000` is the trigger, not a convenience. The failure
needs cumulative churn past `swa.num_pages`; at this model's default budget
that table is 3.8M entries and it would take dozens of runs.
- the eval has to be this one. Synthetic short prompts at 3.7x the nominal
churn passed on main; the 5-shot GSM8K traffic is what reaches it.
- it cannot use `GSM8KMixin`: that path scores this model 0.155 whatever the
pool, because its chat template does not fit the model's reasoning config.
`run_eval` directly gives 0.87 on a healthy server.
Registered at `base-b` / `1-gpu-large` -- per-PR, not nightly. The lesson from
the last regression here (#37307) was that the guard existed but was
nightly-gated, so it never blocked the PR that broke it.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ch-wan
added a commit
that referenced
this pull request
Sep 2, 2026
Nothing in test/registered covered unified memory on a hybrid sliding-window
model. The existing unified e2e tests are MLA (Kimi-Linear) or ShortConv
(Inkling); the existing SWA e2e tests run the static pool. That missing cell is
why a config that dies inside one eval shipped.
google/gemma-4-E2B-it, 1 GPU, triton, --enable-unified-memory
--disable-radix-cache --mem-fraction-static 0.8 --max-total-tokens 60000
GSM8K over 200
narrow table 0.05 <- red on main
fixed passed <- 65 s wall clock
Every argument is load-bearing, and it took several wrong turns to establish
that:
- `--max-total-tokens 60000` is the trigger, not a convenience. The failure
needs cumulative churn past `swa.num_pages`; at this model's default budget
that table is 3.8M entries and it would take dozens of runs.
- the eval has to be this one. Synthetic short prompts at 3.7x the nominal
churn passed on main; the 5-shot GSM8K traffic is what reaches it.
- it cannot use `GSM8KMixin`: that path scores this model 0.155 whatever the
pool, because its chat template does not fit the model's reasoning config.
`run_eval` directly gives 0.87 on a healthy server.
Registered at `base-b` / `1-gpu-large` -- per-PR, not nightly. The lesson from
the last regression here (#37307) was that the guard existed but was
nightly-gated, so it never blocked the PR that broke it.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
5 tasks
ch-wan
added a commit
that referenced
this pull request
Sep 2, 2026
Nothing in test/registered covered unified memory on a hybrid sliding-window
model. The existing unified e2e tests are MLA (Kimi-Linear) or ShortConv
(Inkling); the existing SWA e2e tests run the static pool. That missing cell is
why a config that dies inside one eval shipped.
google/gemma-4-E2B-it, 1 GPU, triton, --enable-unified-memory
--disable-radix-cache --mem-fraction-static 0.8 --max-total-tokens 60000
GSM8K over 200
narrow table 0.05 <- red on main
fixed passed <- 65 s wall clock
Every argument is load-bearing, and it took several wrong turns to establish
that:
- `--max-total-tokens 60000` is the trigger, not a convenience. The failure
needs cumulative churn past `swa.num_pages`; at this model's default budget
that table is 3.8M entries and it would take dozens of runs.
- the eval has to be this one. Synthetic short prompts at 3.7x the nominal
churn passed on main; the 5-shot GSM8K traffic is what reaches it.
- it cannot use `GSM8KMixin`: that path scores this model 0.155 whatever the
pool, because its chat template does not fit the model's reasoning config.
`run_eval` directly gives 0.87 on a healthy server.
Registered at `base-b` / `1-gpu-large` -- per-PR, not nightly. The lesson from
the last regression here (#37307) was that the guard existed but was
nightly-gated, so it never blocked the PR that broke it.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ch-wan
added a commit
that referenced
this pull request
Sep 2, 2026
Nothing in test/registered covered unified memory on a hybrid sliding-window
model. The existing unified e2e tests are MLA (Kimi-Linear) or ShortConv
(Inkling); the existing SWA e2e tests run the static pool. That missing cell is
why a config that dies inside one eval shipped.
google/gemma-4-E2B-it, 1 GPU, triton, --enable-unified-memory
--disable-radix-cache --mem-fraction-static 0.8 --max-total-tokens 60000
GSM8K over 200
narrow table 0.05 <- red on main
fixed passed <- 65 s wall clock
Every argument is load-bearing, and it took several wrong turns to establish
that:
- `--max-total-tokens 60000` is the trigger, not a convenience. The failure
needs cumulative churn past `swa.num_pages`; at this model's default budget
that table is 3.8M entries and it would take dozens of runs.
- the eval has to be this one. Synthetic short prompts at 3.7x the nominal
churn passed on main; the 5-shot GSM8K traffic is what reaches it.
- it cannot use `GSM8KMixin`: that path scores this model 0.155 whatever the
pool, because its chat template does not fit the model's reasoning config.
`run_eval` directly gives 0.87 on a healthy server.
Registered at `base-b` / `1-gpu-large` -- per-PR, not nightly. The lesson from
the last regression here (#37307) was that the guard existed but was
nightly-gated, so it never blocked the PR that broke it.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ch-wan
added a commit
that referenced
this pull request
Sep 2, 2026
Nothing in test/registered covered unified memory on a hybrid sliding-window
model. The existing unified e2e tests are MLA (Kimi-Linear) or ShortConv
(Inkling); the existing SWA e2e tests run the static pool. That missing cell is
why a config that dies inside one eval shipped.
google/gemma-4-E2B-it, 1 GPU, triton, --enable-unified-memory
--disable-radix-cache --mem-fraction-static 0.8 --max-total-tokens 60000
GSM8K over 200
narrow table 0.05 <- red on main
fixed passed <- 65 s wall clock
Every argument is load-bearing, and it took several wrong turns to establish
that:
- `--max-total-tokens 60000` is the trigger, not a convenience. The failure
needs cumulative churn past `swa.num_pages`; at this model's default budget
that table is 3.8M entries and it would take dozens of runs.
- the eval has to be this one. Synthetic short prompts at 3.7x the nominal
churn passed on main; the 5-shot GSM8K traffic is what reaches it.
- it cannot use `GSM8KMixin`: that path scores this model 0.155 whatever the
pool, because its chat template does not fit the model's reasoning config.
`run_eval` directly gives 0.87 on a healthy server.
Registered at `base-b` / `1-gpu-large` -- per-PR, not nightly. The lesson from
the last regression here (#37307) was that the guard existed but was
nightly-gated, so it never blocked the PR that broke it.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Leoyzen
added a commit
to Leoyzen/sglang
that referenced
this pull request
Sep 2, 2026
Wires the Mooncake direct external linker (PR sgl-project#37205 backend, sgl-project#37307 translator passthrough) into the serving stack as an opt-in direct-L3 mode for UnifiedRadixCache, with no host cache tier: - server_args: --enable-unified-cache-external-linker (default FALSE, runtime stays off unless explicitly enabled) + --unified-cache-external-linker-backend {mooncake}. - arg_groups/hicache_hook: mutual-exclusion guard against --enable-hierarchical-cache and --hicache-storage-backend; linker mode skips hicache normalization entirely. - registry: route to _create_unified_radix_cache under the flag and attach MooncakeDirectLinker via init_cache_linker, registering its layer-done counter with the kv cache and tp worker. - scheduler: consolidate aborted-request cache-state release into _release_aborted_request (fires for hicache storage or the external linker); poll linker events, and gate NO_TOKEN / hicache-consumer batch behavior under the linker flag. - kv_cache_builder: forward attn_cp_rank / attn_cp_size into CacheInitParams (consumed by the storage config suffix). - kv_cache_configurator: keep the DSA indexer-K cache under the linker. - full_component: honor the PREPARE phase of ExternalLinkerLoadPhase (return the transfer unmodified; assert COMMIT after it). - tests: MooncakeTestServices lifecycle helper, DSV4-Flash and GLM-5.2 direct-linker KL E2E suites (registered as extra-b CUDA CI), plus the unit-test fixtures for the new scheduler/registry fields. Applied cleanly over the local branch (no textual conflicts; scheduler abort consolidation coexists with the local chunked-prefill abort re-dispatch and fail-soft changes). Default-off runtime behavior is preserved per the port-but-don't-switch decision. (cherry picked from commit d2ade53, PR sgl-project#37381)
ch-wan
added a commit
that referenced
this pull request
Sep 2, 2026
Nothing in test/registered covered unified memory on a hybrid sliding-window
model. The existing unified e2e tests are MLA (Kimi-Linear) or ShortConv
(Inkling); the existing SWA e2e tests run the static pool. That missing cell is
why a config that dies inside one eval shipped.
google/gemma-4-E2B-it, 1 GPU, triton, --enable-unified-memory
--disable-radix-cache --mem-fraction-static 0.8 --max-total-tokens 60000
GSM8K over 200
narrow table 0.05 <- red on main
fixed passed <- 65 s wall clock
Every argument is load-bearing, and it took several wrong turns to establish
that:
- `--max-total-tokens 60000` is the trigger, not a convenience. The failure
needs cumulative churn past `swa.num_pages`; at this model's default budget
that table is 3.8M entries and it would take dozens of runs.
- the eval has to be this one. Synthetic short prompts at 3.7x the nominal
churn passed on main; the 5-shot GSM8K traffic is what reaches it.
- it cannot use `GSM8KMixin`: that path scores this model 0.155 whatever the
pool, because its chat template does not fit the model's reasoning config.
`run_eval` directly gives 0.87 on a healthy server.
Registered at `base-b` / `1-gpu-large` -- per-PR, not nightly. The lesson from
the last regression here (#37307) was that the guard existed but was
nightly-gated, so it never blocked the PR that broke it.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ch-wan
added a commit
that referenced
this pull request
Sep 2, 2026
Nothing in test/registered covered unified memory on a hybrid sliding-window
model. The existing unified e2e tests are MLA (Kimi-Linear) or ShortConv
(Inkling); the existing SWA e2e tests run the static pool. That missing cell is
why a config that dies inside one eval shipped.
google/gemma-4-E2B-it, 1 GPU, triton, --enable-unified-memory
--disable-radix-cache --mem-fraction-static 0.8 --max-total-tokens 60000
GSM8K over 200
narrow table 0.05 <- red on main
fixed passed <- 65 s wall clock
Every argument is load-bearing, and it took several wrong turns to establish
that:
- `--max-total-tokens 60000` is the trigger, not a convenience. The failure
needs cumulative churn past `swa.num_pages`; at this model's default budget
that table is 3.8M entries and it would take dozens of runs.
- the eval has to be this one. Synthetic short prompts at 3.7x the nominal
churn passed on main; the 5-shot GSM8K traffic is what reaches it.
- it cannot use `GSM8KMixin`: that path scores this model 0.155 whatever the
pool, because its chat template does not fit the model's reasoning config.
`run_eval` directly gives 0.87 on a healthy server.
Registered at `base-b` / `1-gpu-large` -- per-PR, not nightly. The lesson from
the last regression here (#37307) was that the guard existed but was
nightly-gated, so it never blocked the PR that broke it.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ch-wan
added a commit
that referenced
this pull request
Sep 2, 2026
Nothing in test/registered covered unified memory on a hybrid sliding-window
model. The existing unified e2e tests are MLA (Kimi-Linear) or ShortConv
(Inkling); the existing SWA e2e tests run the static pool. That missing cell is
why a config that dies inside one eval shipped.
google/gemma-4-E2B-it, 1 GPU, triton, --enable-unified-memory
--disable-radix-cache --mem-fraction-static 0.8 --max-total-tokens 60000
GSM8K over 200
narrow table 0.05 <- red on main
fixed passed <- 65 s wall clock
Every argument is load-bearing, and it took several wrong turns to establish
that:
- `--max-total-tokens 60000` is the trigger, not a convenience. The failure
needs cumulative churn past `swa.num_pages`; at this model's default budget
that table is 3.8M entries and it would take dozens of runs.
- the eval has to be this one. Synthetic short prompts at 3.7x the nominal
churn passed on main; the 5-shot GSM8K traffic is what reaches it.
- it cannot use `GSM8KMixin`: that path scores this model 0.155 whatever the
pool, because its chat template does not fit the model's reasoning config.
`run_eval` directly gives 0.87 on a healthy server.
Registered at `base-b` / `1-gpu-large` -- per-PR, not nightly. The lesson from
the last regression here (#37307) was that the guard existed but was
nightly-gated, so it never blocked the PR that broke it.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
whybeyoung
added a commit
that referenced
this pull request
Sep 3, 2026
The upstream PR made part of its changes inside its merge-with-main commits (chunk_size -> state_chunk_size in _init_track_ssm_indices, and the track_ssm_h_batch_src integer index for the fp32 snapshot copy). Cherry-picking only the non-merge commits dropped them, which raised NameError: name 'chunk_size' is not defined at prefill time. The file now matches refs/pull/34820/head minus the unrelated kv_index_translator line from main #37307 that 07c8f72 lacks. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
StevenChenSE
pushed a commit
to StevenChenSE/sglang
that referenced
this pull request
Sep 6, 2026
…apper backend (sgl-project#37307) 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
AttentionBackend.kv_index_translatoris a class attribute that defaults toNone. A backend that wraps another and does not re-expose the innerbackend's copy therefore answers "this backend needs no translation".
Every producer that reaches the translator through the live backend rather
than through a runner then skips translation silently. The MLA
chunked-prefix-cache path is one:
ForwardBatchDeepseekMHAMixin.prepare_chunked_kv_indicesreadsget_attn_backend().kv_index_translatorand translates onlyif src is not None, so it hands the kernel raw virtual ids for the shared prefix.Nothing raises — virtual and physical ids share a value range, so the gather
simply reads the wrong rows.
Hybrid models (mamba/GDN + full attention) serve every forward through
HybridLinearAttnBackend, which forwards the pool andreq_to_tokenbut notthe translator. Three other wrappers have the same omission.
The damage needs all three of:
--enable-unified-memory(translation actuallyrequired), an MLA backend (only fa3 / flashinfer / flashmla take the
chunked-prefix path; Triton does not), and a radix cache (only then is there a
shared prefix to read back). On Kimi-Linear-48B that is gsm8k 0.89–0.92 →
0.05–0.31 on six of six such cells, while the same configuration on Triton,
and every chunked-prefill cell, stays correct.
Modifications
Four wrappers need the forward, not one. They spell the wrapped field
full_attn_backend,prefill_backend/decode_backend,backend/swa_backendand
dense, so nothing keyed on the attribute name finds them all:layers/attention/hybrid_linear_attn_backend.py— from the full-attentionbackend, beside the pool and
req_to_tokenit already forwards.layers/attention/hybrid_attn_backend.py— from the model runner, which iswhere it already takes the pool.
layers/attention/dots_hybrid_backend.py— twice: the DSA wrapper from itsinner backend, the SWA/MLA wrapper from its swa backend.
layers/attention/minimax_sparse_backend.py— the hybrid wrapper, from itsdense backend.
mem_cache/kv_index_translator.py—assert_backends_carry_translator, whichrefuses at boot when a reachable backend carries someone else's translator or
none at all, and is inert on a pool that needs no translation.
model_executor/model_runner.py— run that guard over the built backends, rightafter they are constructed. Adding a forward is easy to forget in a new
wrapper and the failure is silent by construction, so the runtime check is what
makes it loud.
test/registered/unit/layers/attention/test_kv_translate_ownership.py— anobject-graph test in the file that already owns translator ownership. It
builds a live instance of every wrapper and gives only the inner that must
supply the translator a copy; every other inner carries
None. A wrapper thatcopies from the linear / sparse / DSA side ends up with
Noneand fails — asource scan cannot see that.
Derived, not enumerated: a wrapper is any
AttentionBackendsubclass takinganother
AttentionBackendas a constructor parameter, so a new one is coveredthe day it is written, and a completeness test asserts the derived set equals
the set with instances. The derivation is per class, not per file: an
earlier revision of this PR scanned whole files, which passes as soon as any one
class in a file forwards — and
hybrid_linear_attn_backend.pyholds fourAttentionBackendsubclasses and one forward, so it covered the very file thebug lived in. Seven wrappers come out, including
ShortConvHybridAttnBackendand
TboAttnBackend.Verified by injection, each red on the new test and green on the old scan: drop
the HybridLinear forward while leaving an unrelated
self.kv_index_translator = Nonein that file; point MiniMax atsparse_backend; add a wrapper with no instance.test/registered/models_e2e/test_kimi_linear_models.py— the damage guard.assert_backends_carry_translatoronly fires when a wrapper is actually built,and no per-PR job pairs the unified pool with an MLA model, so nothing reached
the chunked-prefix / MHA-one-shot producer. This file is already
base-bon2-gpu-largeand already launches Kimi-Linear twice, so the cost is one moreserver start in an allocated job rather than a new 4-GPU job.
est_time600 → 900. The nightly
test_kimi_linear_unified_memory.pykeeps its copy: itresolves a different default backend on its H100 runner.
Red-then-green on that cell:
0.37 not >= 0.88on main without the fix,passing on this head.
Accuracy Tests
Kimi-Linear-48B-A3B-Instruct, tp2,
--enable-unified-memory, radix cache,gsm8k (n=200 per cell), unified vs baseline:
All six back within ±0.015 of baseline. The always-correct cells (chunked
prefill, Triton) are unchanged.
Whole matrix, 68 cells / 240 gsm8k runs across Kimi, Qwen3.5-9B, gpt-oss-20b,
Falcon-H1-7B and Llama-3.1-8B on triton / fa3 / flashinfer / flashmla:
112 baseline-vs-unified pairs, median +0.005, zero Kimi outliers, 240 ok /
0 crashed. The residual outliers are all gpt-oss, whose scores sit near 0.5
where the binomial sigma at n=200 is ~0.05.
Checklist
CI States
Latest PR Test (Base): ❌ Run #33463751335
Latest PR Test (Extra): ❌ Run #33463751148
Latest PR Test (AMD ROCm 7.2): ❌ Run #33463751216