[isytwu]

Draft. This is an initial cut rebased onto current main; it has not been
re-run end to end since the rebase (the version this history was measured on
was validated on Kimi K3, TP8, MI355X — see the linked results below). Opening
now to get early feedback on the approach; will follow up with fresh
measurements before this is ready to merge.

Motivation

Builds on #35687 / #37578 (the external-cache linker and its UMBP/mori backend),
which are verified on GLM-5.2 and DeepSeek-V4 but not on any hybrid mamba model.
Kimi K3 is one (MLA full attention + KDA linear attention, backed by
HybridLinearKVPool), and two separate things stopped it from working:

• MambaComponent.build_external_linker_transfer was a stub that raised
  unconditionally.
• _update_load's COMMIT phase assumed every component's device_indices are
  page-aligned token slots, which crashes a slot-addressed component (one state
  per node, not a token range) with shape '[-1, 64]' is invalid for input of size 1 at page_size > 1.

Modifications

Hybrid mamba support (mamba.py, linker_pool_assembler.py,
hybrid_pool_assembler.py)

A mamba checkpoint is one whole state slot per node, not a token range, so the
MAMBA DevicePoolEntry is addressed by slot id over the pool's conv and
temporal buffers, alongside the ordinary paged KV entry for the full-attention
layers. Its transfer carries the single key of the node's end boundary under
TRAILING_PAGES, which is the shape the lookup's boundary filter already
expects. The group is deliberately not rank_replicated: the MLA KV is
TP-replicated but the KDA state is not.

update_external_linker_load returns the slot allocated during LOAD when a
load is called off; int8 mamba checkpoints are excluded, since they live in a
pool the linker's device group does not map.

Slot-addressed indices at page_size > 1 (components/base.py,
unified_cache_linker.py)

TreeComponent.linker_indices_are_paged now declares whether a component's
device_indices are page-aligned, and the COMMIT-phase slicing is skipped for
components that are not.

The COMMIT-phase test for "is this slot-addressed checkpoint still ours to
write into" is insert_result.mamba_exist — that is exactly the condition
load_back() itself acts on: it hands the slot to tree_core.insert() as
InsertParams.mamba_value and frees it there iff mamba_exist. An
adopted-range test looks like the natural fit here and is wrong:
record_adopted_range() is only ever called for BASE_COMPONENT_TYPE and
SWA, never for MAMBA, so that test is always false and drops every mamba
load — and in the mamba_exist == False branch the slot is already attached to
the tree node by that point, so the dropped transfer leaves the node holding a
mamba slot whose contents were never written, later used as a valid
checkpoint. Caught by per-pool load-byte counters (mamba load bytes 0 against
nonzero mamba host hits), not by any end-to-end metric, since the
restored-token count is a KV measure and says nothing about whether the
recurrent state moved.

A checkpoint that fails the test is dropped without being freed: releasing
it again through the component's ABORT contract would double-free a slot the
tree already owns (the idle-pool invariant catches this as a mamba count one
larger than the pool, leaked_mamba_pages=None). ABORT is for a load that
never reached insert.

UMBPDirectLinker load-completion tracking (umbp_direct_linker.py)

Thread the request ids and a ready event through the load queue and publish
them to a completed-loads queue once a batch's layer-wise reads land,
mirroring the existing offload path.

Configuration requirement

Use this path with page_size >= 64, not the page_size=1 that
MambaComponent otherwise forces. At page_size=1 every token becomes its
own stored object and the restore stops paying for itself. page_size > 1
requires the mamba track buffer, i.e. --mamba-radix-cache-strategy extra_buffer (or extra_buffer_lazy), which is what the Kimi K3 cookbook
already recommends.

Known limitation

The reachable prefix is capped at floor((L-1)/grid)*grid, where grid = lcm(mamba_cache_chunk_size, page_size) — 64 for Kimi K3, whose config carries
no mamba_chunk_size so FLA_CHUNK_SIZE applies. Two rules combine:

• A checkpoint only lands on a multiple of grid.
• match_prefix probes 0..L-1, one token short of the prompt, because a
  request must keep a token to forward.

So a 16384-token prompt can restore at most 16320 tokens, and a prompt replayed
at exactly its cached length restores nothing: its only checkpoint sits at
16384, one grid step past the probe ceiling.

This all-or-nothing behaviour is correct rather than a degradation bug.
Restored KV without the matching recurrent state at that boundary is unusable,
and the state cannot be recomputed without a forward pass over the prefix, so
MAMBA's TRAILING_PAGES filter vetoes the boundary instead of returning KV
that cannot be used. Partial-prefix reuse, the case prefix caching exists for,
is unaffected. Raising the ceiling means changing where MambaComponent
places checkpoints, which is out of scope here.

Checklist

• Format your code according to the Code Formatting with Pre-Commit — not yet run in this environment; will run before merge.
• Add unit tests as outlined in the Running Unit Tests — the sibling direct-linker pool builders (_build_dsa_device_pool_group, _build_deepseek_v4_device_pool_group) have none either; exercising this path needs a live hybrid KV pool. Will validate end to end before this leaves draft.
• Update documentation / docstrings / example tutorials as needed — no user-facing surface changed; the configuration requirement and the reachable-prefix ceiling are documented above.
• Provide throughput / latency benchmark results and accuracy evaluation results — pending re-run on this rebased branch; a prior revision was validated end to end on Kimi K3 (TP8, MI355X) and restored the expected token counts with zero leak/OOM/reshape errors, but those numbers predate this rebase and are not re-posted here.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

---

CI States

Latest PR Test (Base): ❌ Run #38022058748
Latest PR Test (Extra): ❌ Run #38022058719
Latest PR Test (AMD ROCm 10): ❌ Run #38022058670

[isytwu]

Superseded by #45016, which rebuilds this work on current main now that #41343 provides the Mamba linker hooks (LOAD / PREPARE / ABORT). It adds the FULL + Mamba pool group for Kimi K3, the UMBP side, DCP and the DSPARK draft KV, with unit, KL and end-to-end accuracy tests and benchmarks against HiCache. Leaving this PR open until #45016 lands.

[isytwu]

Closing: hybrid Mamba support for the direct linker landed in #44977, and the remaining part, storing the replicated MLA KV once on UMBP, continues in #45016.