Internal Wiki

主题

RFC-0035: UnifiedRadixCache SWA ownership, lifecycle, and validation ​

Overview ​

This RFC specifies how to add sliding-window attention (SWA) to UnifiedRadixCache safely for hybrid full-attention + SWA models. It defines the ownership, lifecycle, routing, and validation contracts required to move this model class from the legacy SWARadixCache path to one FULL-backed unified radix tree with an independent SWA component.

The existing SWARadixCache remains the explicit rollback path. Production routing changes only after the allocator, component, request-lifecycle, scheduler, and validation contracts in this document pass their gates.

Tracking and relationship to existing designs ​

This RFC extends RFC-3 and sgl-project/sglang-jax#1341; it does not replace their overall Unified Hybrid Radix Cache direction. Its one deliberate refinement to that issue is that every SWA-capable cache, including SWAChunkCache, owns request-tail reclamation through evict_req_swa(). The scheduler therefore never decides whether direct pool freeing is safe.

Background and motivation ​

The legacy SWARadixCache manages hybrid SWA models today. UnifiedRadixCache already provides a component model and supports FULL and Recurrent components, but safely adding SWA requires more than registering a new component:

  • FULL and SWA use separate pools connected by a rank-local mapping.
  • SWA may be evicted while the FULL-backed tree node remains alive.
  • Healing may transfer freshly allocated request slots into an existing tree node.
  • Request completion, retraction, abort, and unfinished relock paths must preserve the same lock and ownership boundaries.
  • DP ranks may reuse the same local index values and therefore require explicit rank-local operations.
  • A flush can hide a pre-existing leak unless ownership is checked before state is cleared.

Without explicit contracts at these boundaries, basic cache-hit tests can pass while resources are leaked, released twice, or released from the wrong rank.

Goals ​

  1. Add a device-only SWAComponent to one FULL-backed UnifiedRadixCache tree while preserving tombstone, healing, page-aligned window, LRU, and lock semantics.
  2. Make FULL/SWA ownership transfer and freeing explicit, rank-local, and exactly once across allocator, tree, and request lifecycles.
  3. Put request reclamation behind cache ownership APIs, keep scheduler dispatch capability-based, and preserve the legacy radix and chunk routes.
  4. Validate the new path with per-rank ownership ledgers and deterministic baseline/Unified comparisons.

Non-goals ​

The first version does not support:

  • SWA HiCache L2/L3 transfer or host backup;
  • PD/disaggregation SWA transfer;
  • speculative decoding or EAGLE with Unified SWA;
  • FULL+SWA+RECURRENT in one tree;
  • performance tuning;
  • removal of SWARadixCache.

Explicitly requesting Unified SWA with an unsupported runtime combination must fail before cache construction or allocator mutation. It must not silently fall back.

Proposal ​

The design uses one radix tree whose base value is the FULL component. SWA is an independent component on the same nodes, with its own pool indices, LRU state, lock boundary, and accounting. The core coordinates component ownership without learning SWA-specific window policy.

Request-side SWA reclamation belongs to each SWA-capable cache. The scheduler uses capability and ownership APIs, passes the scheduler-observed DP rank, and does not branch on concrete cache types or infer whether a range is tree-owned.

The production route remains disabled until the guarded route, CPU regressions, evidence tooling, and paired TPU gates are complete.

Alternatives considered ​

  1. Keep the scheduler tree/chunk branching proposed in sgl-project/sglang-jax#1341. This is smaller, but leaves ownership policy in the scheduler and risks direct-freeing a tree-protected prefix when another SWA-capable cache is added.
  2. Create a second SWA radix tree. This reuses the legacy structure but defeats the UnifiedRadixCache single-tree component model and duplicates prefix ownership.
  3. Keep the issue body as the only design document. This is easy to edit, but does not provide the versioned RFC path, index, and review history used by primatrix/wiki.

Constraints ​

  • All allocator mapping, translation, and freeing operations are explicit about dp_rank.
  • Logical token boundaries and physical page boundaries remain distinct.
  • Internal SWA eviction must preserve FULL/tree ownership; leaf deletion may cascade across both pools.
  • The internal ledger is an observability and test surface, not a stable public API.
  • CPU or PR-gate success does not authorize TPU allocation. Each TPU gate requires a separate resource and immutable-runtime preflight followed by explicit approval.

Design details ​

1. Tree and ownership model ​

There is one radix tree whose base value is the FULL component. For each node:

  • ComponentType.FULL stores FULL-pool indices.
  • ComponentType.SWA stores the corresponding SWA-pool indices, or None after independent SWA eviction.
  • the node's dp_rank selects the mapping and allocator state.

Outside an active free group, every physical FULL or SWA slot has exactly one owner: free, request-owned, or tree-owned. Inside a free group, surrendered slots enter a transient pending-free state: they are no longer request/tree-owned, are not yet available for allocation, and become free only when the group is flushed.

An internal SWA eviction frees only SWA slots and clears their mappings. It leaves the FULL value and tree structure intact, creating an SWA tombstone. Recomputing the prefix may heal that tombstone by transferring freshly allocated request FULL+SWA slots to the existing node.

2. Dual-pool allocator contract ​

SWATokenToKVPoolAllocator provides explicit rank-local operations:

  • translate_full_to_swa(full_indices, dp_rank, require_mapped=True) returns a copy of mapped SWA indices and fails loudly on an unexpected missing mapping.
  • free_full(full_indices, dp_rank) frees only FULL slots while preserving SWA slots and their mapping for an ordered cascade.
  • free_swa(full_indices, dp_rank) accepts FULL indices, captures mapped SWA slots, clears mapping immediately, and releases those physical SWA slots. Repeated calls are no-ops.
  • the existing composite free() releases both ownership domains.

Within a free group, composite and ownership-specific frees share per-rank, per-pool deduplication. Mapping is cleared as soon as SWA ownership is surrendered, while physical reuse remains delayed until free_group_end(). Group completion flushes the two physical queues directly and never reinterprets an already-cleared mapping. A failed two-pool allocation rolls back FULL slots, SWA slots, and mapping atomically.

3. Unified core and component contract ​

Every component receives validated, non-null cache initialization parameters. Constructing SWA without a positive sliding-window size fails.

Each component's overlap hook returns a consumed boundary: the slice offset from which the component/tree takes ownership of request slots. The core takes the minimum boundary across components and composite-frees only the duplicate interval that no component consumed. The request caller must not free that interval again. Logical matched-prefix length remains separate from ownership transfer.

When fresh request slots unevict a FULL tombstone, the core invokes auxiliary-component recovery exactly once before committing replacement values. Component metadata is node-local, and auxiliary components use independent LRU timestamps and candidates without changing existing FULL or Recurrent ordering.

4. SWA lifecycle ​

A match is valid only when the path has a continuous live SWA suffix covering the sliding window. A tombstone resets the continuous-length counter. LRU refresh protects the smallest page-aligned physical suffix sufficient for the window; when a page is larger than the logical window, the final physical page remains protected.

Overlap healing has three cases:

  1. The live request boundary is before the tombstone node: the node adopts the entire new slice.
  2. The boundary is inside the node: the node is split at a page-aligned boundary and the suffix adopts the request slots.
  3. The boundary is after the node: the tombstone remains and consumes no request slots.

Healing frees only the tombstone's old FULL slots. The adopted request FULL indices already carry valid SWA mappings.

Internal SWA eviction calls free_swa() with FULL indices, marks SWA as a tombstone, and preserves FULL/tree ownership. Leaf removal uses the fixed cascade free_full() then free_swa(), followed by clearing the FULL value, so each pool is released exactly once.

Locks protect only the final physical window. Acquisition records both the component UUID boundary and skipped tombstone nodes; release consumes the same data. A split moves the UUID to the new parent rather than copying it.

5. Request lifecycle and cache capabilities ​

BasePrefixCache.supports_swa() defaults to False; legacy SWA radix, SWA chunk, and Unified FULL+SWA return True.

Every SWA-capable cache implements:

python
evict_req_swa(req, pre_len, dp_rank)

The method may release only request-owned SWA slots and must never release a tree-owned prefix. Unified derives the protected boundary from the request's last tree node; legacy radix and chunk caches provide equivalent ownership-specific implementations.

All three implementations use the same monotonic, page-safe reclaim contract:

text
old = max(req.swa_evicted_seqlen, tree_protected_len)
new = floor_to_page(max(old, pre_len - sliding_window_size - page_size))
free only request_row[old:new]
req.swa_evicted_seqlen = new

The protected and reclaimed boundaries never move backward, and the final active physical page is not released early.

The request stores the complete lock-release parameters, including the SWA UUID and skipped-node set. Finished, unfinished, retracted, aborted, and reset paths release or replace those parameters symmetrically.

Scheduler admission, reclaim, sanity checks, and size queries use capabilities and ownership APIs rather than concrete cache classes. Existing pool-layout checks on SWATokenToKVPoolAllocator remain allocator checks, not cache dispatch.

6. Internal observability and evidence ​

Each SWA-capable cache exposes the same internal/test-only per-rank ledger schema. request-owned values are computed from live request rows; they are not inferred as a residual. Separate fields detect duplicate request owners, request/tree overlap, invalid mapping, and duplicate SWA mapping.

For each pool and rank, an idle snapshot must satisfy:

text
available + tree_evictable + tree_protected + request_owned == capacity

The ledger endpoint is read-only, enabled only by the cache-report flag, and is not a stable public API. The scheduler supplies observed DP rank and route information; the client cannot forge them. Every required pre-flush snapshot is idle and has a non-zero cold capacity that remains constant across all phases for that rank; corresponding baseline/Unified cold capacities must match. Post-flush availability must equal the same cold capacity, with request/tree owners and mappings cleared.

An online test collector stores raw input/output ID sidecars, per-case ledger snapshots, launch metadata, and hash-verified logs. A separate offline checker uses a versioned, checked-in GATE_SCHEMA and pinned fixture hashes to recompute absolute route/configuration, request completeness, exact outputs, ledger invariants, and forbidden-error counts instead of trusting artifact-reported expectations or pass fields. The schema fixes the required workload and parallelism, including v6 TP4/DP1 and v7 TP8/DP2 with attention tensor axis size 4 and EP8.

7. Routing and compatibility ​

ConfigurationCache implementation
Hybrid SWA + Radix disabledSWAChunkCache
Hybrid SWA + Radix enabled + Unified enabledUnifiedRadixCache(FULL, SWA)
Hybrid SWA + Radix enabled + Unified disabledlegacy SWARadixCache
Non-hybrid modelexisting behavior

Radix-disabled routing takes precedence over the Unified flag. The new-route predicate is exactly hybrid_swa && unified && !disable_radix_cache; unsupported-combination gates apply only to that predicate. The new route is activated only after allocator, core, component, request, scheduler, and test contracts pass. Unsupported Unified SWA combinations fail before construction. With Radix enabled, disabling the Unified flag remains the rollback mechanism.

Test plan ​

CPU correctness gate ​

  • allocator ownership, mapping, grouped free, rollback, and DP isolation;
  • component overlap/recovery order and independent metadata/LRU;
  • SWA match, tombstone, three healing cases, locks, internal/leaf eviction;
  • page sizes 1/128/256 and DP1/DP2;
  • finished/unfinished/retract/abort request paths and per-rank ledger;
  • capability-based scheduler behavior, route matrix, fail-fast behavior, collector, and checker;
  • legacy SWA, Unified FULL, Recurrent, and HiCache-FULL regressions.

The gate runs a versioned, checked-in CPU test manifest. Existing environment-specific skips must be documented in the allowlist, and new or modified tests must not introduce skip or xfail.

v6e PR gate ​

Run google/gemma-4-31B-it on v6e 2x2, four JAX devices, TP4/DP1, page size 128. Baseline and Unified use the same clean commit, immutable model/tokenizer revision, runtime, hardware, inputs, and sampling; the Unified flag is the only route difference. Route, prefix reuse, SWA eviction/tombstone/healing, pre/post-flush ledger, and every raw output_ids sequence must pass the offline checker.

After correctness passes, run a separate legacy-versus-Unified performance A/B on the same clean commit. Use shared-prefix and random/no-sharing workloads; keep the runtime, inputs, sampling, and launch parameters identical except for the Unified route flag. Store performance artifacts separately and report TTFT, throughput, hit rate, ratio/delta, and variability. A performance regression does not invalidate a correctness PASS and does not expand the first version into performance tuning.

A v6 PASS makes the change eligible for PR review or merge; it does not close the implementation task.

v7x final gate ​

Run XiaomiMiMo/MiMo-V2-Flash on v7x-8 (four chips/eight JAX devices), TP8/DP2/EP8, page size 256. Verify scheduler-reported rank assignment, isolation and reuse on both ranks, dual-pool pressure, per-rank ledger, and exact raw outputs. If either baseline or Unified fails to trigger either the FULL or SWA pool path at 64 requests, rerun both routes from fresh servers with the paired 128-request workload; never mix rounds.

After the CPU gate passes, run each TPU gate only after its resource, model revision, image, and launch command have been reviewed and approved. The implementation closes only when complete paired raw artifacts produce a v7 checker PASS and the CPU, v6e, and v7x evidence all refer to the same final clean PR head commit.

Impact ​

  • SGL-JAX cache stack: allocator, unified tree/component seams, request lifecycle, scheduler capability use, route construction, and test-only observability change together.
  • Compatibility: non-hybrid routing and flag-off hybrid behavior remain unchanged. SWARadixCache remains available for rollback.
  • Operational behavior: unsupported Unified SWA combinations fail during startup instead of entering an unvalidated path.
  • Public API: no stable external API is added. The ledger endpoint is internal/test-only and flag-gated.
  • Resources: this RFC and its CPU work authorize no TPU allocation. Hardware runs remain separately approved gates.

Implementation plan ​

The detailed task checklist and acceptance state live in sgl-project/sglang-jax#1531. The implementation order is:

PhaseTaskDependencyVerifiable deliverable
Baseline and ownership foundation0Frozen source baselineSaved baseline logs, critical call-path notes, and a FULL/SWA ownership oracle
Baseline and ownership foundation1Task 0Rank-local dual-pool allocator APIs plus mapping, grouped-free, rollback, and DP-isolation tests
Baseline and ownership foundation2Task 1Unified overlap/recovery seams plus component metadata/LRU contract tests; production route unchanged
SWA correctness closure3Tasks 1–2SWAComponent plus page 1/128/256 and DP1/DP2 lifecycle tests
SWA correctness closure4Task 3Request reclaim/lock lifecycle implementations and per-rank ledger tests
System integration and CPU closure5Task 4Capability-based scheduler integration, online collector, offline checker, fixtures, and focused tests
System integration and CPU closure6Task 5Guarded production route, unsupported-combination fail-fast checks, and route-matrix tests
System integration and CPU closure7Tasks 1–6Versioned CPU test manifest, lint output, and complete regression evidence
v6e PR delivery8Task 7 and separate resource approvalPaired DP1 correctness artifacts and a separate performance A/B report
v6e PR delivery9Task 8Reviewable implementation PR with rollback instructions, evidence links, and tested head SHA
v7x final closure10Task 9 and separate resource approvalPaired DP2 raw artifacts, dual-pool pressure evidence, exact outputs, and final checker PASS

Routing is activated last. A phase stops on an ownership ambiguity, regression, silent fallback, missing evidence, or failed gate; later phases do not compensate for an earlier failure.

Risks and mitigations ​

RiskMitigation
Double-free or stale mapping during overlap healing and grouped freesExplicit consumed-boundary ownership, mapping-clear versus physical-flush timing, cross-API deduplication, and focused allocator/core tests
Releasing live or tree-owned SWA pagesCache-owned reclaim API with a monotonic protected boundary and page-safe formula
Cross-rank contamination from equal local indicesMandatory dp_rank on mapping/free APIs, scheduler-origin rank reporting, and DP2 isolation tests
False PASS caused by flush or self-reported artifactsRequired pre-flush ledger snapshots, constant cold capacities, raw sidecars/logs, pinned schemas, and offline recomputation
Production route regression or unsupported combinationMutually exclusive route table, exact activation predicate, fail-before-mutation checks, and explicit legacy rollback
Page-size-dependent window or lock errorsPhysical tail-window rules and page 1/128/256 coverage across DP1/DP2
TPU evidence cannot be reproducedSame-clean-commit paired A/B runs, immutable revisions, exact launch metadata, hash readback, and no mixed rounds

Review focus ​

Review should focus on the allocator/free-group ownership contract, overlap recovery, cache-owned request reclamation, lock symmetry, the internal ledger boundary, route/fail-fast behavior, and the CPU/v6e/v7x acceptance sequence.