Abstract

Hybrid large language models combining full MLA attention and linear State Space Model (SSM) layers such as KDA (Kimi Delta Attention) and GDN (Gated Delta Network) have become mainstream solutions for million-token long-context inference. Kimi K3 and Qwen3.6 are representative hybrid architectures, featuring a fixed ratio of dense full-attention layers and linear Mamba-style layers. Traditional token-level KV cache mechanisms designed for pure Transformer architectures fail to match the fixed-size recurrent state storage pattern of linear SSM modules, resulting in severe GPU memory waste, low prefix cache hit rates and frequent out-of-memory (OOM) errors under long-sequence workloads. This paper delivers a comprehensive engineering breakdown of SGLang’s dedicated hybrid Mamba KV cache subsystem, including core state storage logic, accurate VRAM calculation formulas for Kimi K3, Radix prefix caching rules, ReplaySSM speculative decoding memory optimization, and configurable runtime parameters for production deployment. All quantitative memory metrics, layer specifications and parameter constraints are sourced from official SGLang kernel code and verified hardware test data. When deploying multi-model inference clusters with mixed Transformer and Mamba workloads, developers can adopt Treerouter as a unified API gateway to standardize model access and streamline cross-backend traffic governance.

1. Background: Mamba-ish Model Core Differences from Standard Transformer KV Cache

1.1 Definition of Mamba-ish Architectures

The term mamba-ish refers to any linear attention variant that maintains fixed-size recurrent states instead of token-by-token K/V storage. Representative models in production include Qwen3.6 (GDN + GQA hybrid stack) and Kimi K3 (KDA + MLA hybrid layout). These architectures share identical state management characteristics despite divergent algorithmic naming (KDA/GDN vs vanilla Mamba):

  1. No linear memory growth with sequence length; each request only maintains a fixed set of persistent state tensors.
  2. State values are overwritten in-place after processing each token, rather than appending new K/V entries.
  3. Prefix cache rollback, snapshot reuse and speculative token verification require customized state snapshot logic incompatible with standard page-based KV cache.

Standard full MHA/MLA Transformer caching follows a straightforward scaling rule: KV cache memory consumption grows linearly with sequence length, as every input token generates independent K and V vectors that must be retained for subsequent attention calculations: $$Mem_{FullKV} \propto SequenceLength$$ By contrast, all Mamba-style linear layers store two static state components independent of token count:

  1. Short convolution window state (conv state): Preserves the most recent N tokens for local feature mixing (kernel size fixed to 4 across GDN/KDA implementations, retaining the preceding 3 token vectors).
  2. Temporal SSM recurrent state (SSM state): A fixed-dimension matrix that accumulates historical semantic information via weighted decay without per-token storage. $$Mem_{MambaState} = Fixed \ per \ request$$

This fundamental divergence creates an incompatible memory allocation model that generic inference engines cannot natively support. SGLang’s dedicated Mamba Radix Cache module addresses this gap with request-scoped state slot management.

1.2 Layer Layout of Kimi K3 and Qwen3.6 Hybrid Stacks

Kimi K3 deploys a 3:1 linear-to-dense layer ratio across its 93 total transformer layers:

  • 69 KDA linear attention layers (Mamba-ish SSM design)
  • 24 MLA dense full-attention layers (standard token-wise KV storage) The complete stack consists of 23 repeating blocks of 3 KDA + 1 MLA layers, plus a final standalone MLA layer at the model output. Qwen3.6 adopts a similar hybrid pattern built on Gated Delta Networks paired with standard GQA attention layers. Dense MLA layers retain traditional per-token KV cache while all KDA/GDN layers rely on fixed conv + SSM state slots managed by SGLang’s Mamba memory pool.

2. SGLang Hybrid Mamba Cache Core Storage Mechanism

2.1 Dual Isolated Memory Pool Architecture

SGLang splits GPU VRAM into two independent static pools for hybrid model inference:

  1. Full KV Cache Pool: Allocates token-indexed K/V blocks exclusively for MLA dense layers, following standard page radix cache rules compatible with vLLM-style scheduling.
  2. Mamba State Pool: Dedicated slot-based buffer for conv and SSM state snapshots of all linear KDA/GDN layers. Each running request occupies multiple reusable state slots rather than variable-length token blocks.

Two core source files govern this subsystem:

  • python/sglang/srt/mem_cache/mamba_radix_cache.py: Radix tree prefix matching and LRU eviction logic for Mamba state snapshots
  • python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py: State tensor initialization and snapshot copying Triton kernel wrappers

For prefill and decode phases, SGLang adopts differentiated state snapshot generation policies:

  1. Chunked Prefill Snapshot Rule: When prefill input length exceeds the configured chunked_prefill threshold, SGLang writes one complete Mamba state snapshot after finishing each chunk calculation. For example, a 13600-token prompt with chunk size set to 6144 generates 3 separate state checkpoints during prefill.
  2. Decode Phase Snapshot Rule: The engine triggers state backup only when cumulative input + output token count is divisible by mamba_track_interval (default value = 256). Only the latest valid state slot is retained for offloading to host memory, avoiding redundant storage of intermediate decode states.

The maximum number of Mamba state slots required per request follows this formula: $$MaxSlots = \lceil \frac{InputLength}{ChunkedPrefillSize} \rceil + 2$$ Using the sample input of 13600 tokens and chunk size 6144: $$\lceil 13600 \div 6144 \rceil = 3$$ Total slots allocated per request = 3 + 2 = 5, matching real-world scheduler logs observed during production testing.

2.2 LRU Eviction and Radix Prefix Matching Logic

Mamba state snapshots are bound to specific sequence prefix nodes within the Radix Tree index. When the Mamba pool reaches its max_mamba_cache_size capacity limit, SGLang executes LRU eviction targeting idle, unreferenced state slots. Active requests currently undergoing decoding lock their assigned slots and are excluded from eviction candidates. The configurable parameter --mamba-max-states-per-path caps the number of historical snapshots retained for a single Radix prefix branch; setting a lower value frees more pool slots for new incoming requests but reduces prefix cache hit rates for long multi-turn dialogues.

3. Precise VRAM Calculation for Kimi K3’s KDA and MLA Layers

All memory calculations below use official Kimi K3 hyperparameters: TP parallelism = 8, KDA heads = 96, single head dimension = 128, convolution kernel size = 4 (3 historical tokens stored per conv state), total KDA layers = 69, MLA dense layers = 24. Two precision formats (BF16 / FP8 E4M3) are covered for full production reference.

3.1 Dense MLA Layer Per-Token KV Memory Cost

Kimi K3’s MLA implementation uses fixed LORA rank = 512 and RoPE head dimension = 64. Each token’s BF16 KV tensor size per layer: $$512 + 64 = 576 \ elements$$ Single token total across all 24 dense layers (BF16, 2 bytes per element): $$24 \times 576 \times 2 = 27,648 \ Bytes \approx 27 \ KB$$ Under FP8 quantization (1 byte per element), memory consumption is halved to ~13.5 KB per token. In TP8 distributed mode, every GPU replicates identical full KV cache blocks for all MLA layers.

3.2 KDA Linear Layer Fixed State Memory Calculation

Each TP8 shard receives 12 split attention heads (96 total heads ÷ 8 TP ranks). Two state components are allocated per KDA layer:

  1. Conv State Shape: [3, head_num, q_dim + k_dim + v_dim] = [3, 12, 3 × 128]
  2. SSM Recurrent State Shape: [head_num, k_dim, v_dim] = [12, 128, 128]
Step 1 Single GPU BF16 memory breakdown

SSM state for all 69 KDA layers: $$69 \times 12 \times 128 \times 128 \times 2 = 26,002,7136 \ Bytes \approx 25.9 \ MB$$ Conv state for all 69 KDA layers: $$69 \times 3 \times 12 \times 384 \times 2 = 1,805,568 \ Bytes \approx 1.82 \ MB$$ Single-GPU combined linear state memory: 25.9 + 1.82 = 27.72 MB

Step 2 Full 8-GPU TP Cluster Total Mamba State Footprint

$$27.72 \ MB \times 8 = 221.76 \ MB$$ This fixed overhead is completely decoupled from prompt or generation length, a stark contrast to the unbounded memory growth of standard Transformer KV cache.

4. ReplaySSM: Speculative Decoding Optimization for Hybrid Linear Models

Standard speculative decoding algorithms such as DSpark and EAGLE suffer prohibitive memory overhead when applied to KDA/GDN architectures. Storing a full independent SSM state for every draft token would multiply VRAM consumption exponentially. SGLang introduces the --enable-linear-replayssm-spec optimization to eliminate redundant state storage via lightweight replay buffers instead of full snapshot retention.

4.1 Core ReplaySSM Mechanism

Instead of persisting complete [heads, K, V] matrices for each speculative candidate token, the framework only archives lightweight intermediate calculation logs. After verifying draft token correctness, the system re-computes valid SSM states sequentially from archived delta records rather than loading pre-saved full matrices. Six dedicated ring buffers are allocated per request for replay logic, with the configurable cache depth controlled by --linear-replayssm-cache-len (default L=16, a power-of-two constraint). The minimum valid L value must equal or exceed twice the maximum number of draft tokens (L ≥ 2 × D). For DSpark with 8 draft tokens per verification window, L=16 is the minimal compliant setting.

4.2 Replay Buffer Shape & Memory Overhead (TP8, L=16, 120 request slots)

All six ring buffers are allocated per KDA layer with standardized tensor dimensions:

Buffer Name Tensor Shape Data Type Functional Purpose
d [69, 120, 12, 16, 128] BF16 Corrected residual delta vectors
k [69, 120, 12, 16, 128] BF16 L2-normalized key tensors
g [69, 120, 12, 16, 128] FP32 Logarithmic decay gate coefficients
rawv [69, 120, 12, 16, 128] BF16 Unmodified raw value inputs
rawk [69, 120, 12, 16, 128] BF16 Unnormalized original key vectors
beta [69, 120, 12, 16] FP32 Delta update weighting factors

These replay ring buffers are marked as non-transfer memory resources and excluded from host offloading via HiCache tiered storage, reducing data transfer overhead during speculative verification cycles. Enabling ReplaySSM increases the required mamba-full-memory-ratio parameter value because part of the Mamba pool capacity is reserved for replay buffer allocation rather than sequence state snapshots. Sample production allocation logs confirm that enabling this feature adds approximately 1.4 GB of extra VRAM overhead under standard 120-slot concurrency configurations.

5. Critical Runtime Parameter Tuning for SGLang Hybrid Mamba Cache

Misconfigured memory partitioning between the Full KV pool and Mamba state pool frequently creates artificial concurrency bottlenecks, where one pool hits capacity limits while the other remains largely underutilized. This section explains core adjustable flags with production tuning guidance and constraint boundaries.

5.1 --mamba-full-memory-ratio

Defines the VRAM allocation split ratio between Mamba state pool and standard token KV pool at server startup. Static partitioning leads to suboptimal resource utilization under variable workloads: long document prompts consume more Full KV blocks, while short multi-turn dialogue sessions occupy more Mamba state slots. When running pure TP parallelism without DCP disaggregation, the ratio value must be scaled up by the TP rank count to account for per-GPU Mamba state duplication.

5.2 --max-mamba-cache-size

Hard upper bound for total Mamba state slot count across the entire GPU memory pool. The theoretical maximum concurrent decode requests are calculated as: $$MaxConcurrent = \frac{max_mamba_cache_size}{SlotsPerRequest}$$ By default, each request reserves 5 slots under overlap scheduling; enabling the environment variable SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK=1 reduces reserved slots to 4 and lifts concurrency ceilings without functional regression (experimental feature). A real-world production log example demonstrates this bottleneck clearly: with 80 total Mamba slots and 5 slots per request, maximum parallel decode requests are capped at 16, even if the Full KV pool retains over 89% unused capacity.

5.3 --mamba-radix-cache-strategy

Four configurable scheduling modes control ping-pong snapshot buffer allocation and overlap scheduler compatibility:

  1. auto: Default configuration that auto-selects extra_buffer when overlap scheduling is active
  2. extra_buffer: Pre-allocates dual ping-pong tracking slots (5 slots per request total), maximizes throughput for stable production workloads
  3. extra_buffer_lazy: Reduces base slot reservation to 4 per request for memory-constrained hardware, compatible with overlap scheduling
  4. no_buffer: Disables dual tracking buffers, only 3 slots allocated per request, incompatible with overlap batch processing

The extra_buffer mode separates read and write state slots to eliminate GPU-CPU race conditions during simultaneous forward computation and Radix Tree cache insertion.

5.4 --mamba-track-interval

Sets the token length interval that triggers automatic Mamba state snapshot backup during decode generation (default = 256). Smaller interval values create more frequent state checkpoints, improving Radix cache hit rates for mid-length dialogue branches but increasing VRAM overhead from additional snapshot slots. Larger values reduce snapshot storage cost at the cost of higher recomputation overhead during prefix cache reuse.

5.5 Unified Dynamic Memory Pool Flag (--enable-unified-memory)

SGLang provides an experimental unified memory pool to eliminate static ratio partitioning. When activated, Full token KV blocks and Mamba state slots share a single contiguous GPU buffer with dynamic boundary adjustment:

  • Long-sequence workloads expand Full KV allocation from high memory addresses downward
  • Short multi-turn dialogue workloads expand Mamba state slots from low memory addresses upward The two memory regions grow in opposite directions without pre-defined static split points, drastically improving overall GPU utilization under mixed traffic patterns. Important compatibility limitations apply: this feature cannot be combined with PD disaggregation, DCP parallelism, or any speculative decoding algorithms including ReplaySSM and DSpark.

6. Production Troubleshooting & Optimization Best Practices

6.1 Common Concurrency Bottleneck Root Cause

The most frequent production issue arises from mismatched mamba-full-memory-ratio values: the Mamba state pool saturates far earlier than the token KV pool, artificially limiting batch throughput. Two corrective actions resolve this issue:

  1. Increase the ratio parameter to allocate more VRAM for Mamba state slots
  2. Activate SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK=1 to cut per-request slot reservation from 5 to 4
  3. Switch to the experimental unified memory pool if speculative decoding and disaggregation are not required for the workload

6.2 Speculative Decoding Memory Tuning Checklist

When deploying ReplaySSM for Kimi K3/Qwen3.6:

  1. Set --linear-replayssm-cache-len to the next power of two equal to twice the maximum draft token count
  2. Increase mamba-full-memory-ratio by approximately 20% to accommodate replay ring buffer overhead
  3. Avoid FP16 storage for g/beta replay buffers; retain FP32 precision to prevent accumulated drift errors during state replay
  4. Disable L2/L3 host offloading for replay buffers via the built-in non-transfer field whitelist to eliminate cross-device copy latency

6.3 Workload-Based Parameter Selection Guide

  • Long single-document inference (16k+ tokens): Prioritize larger mamba-track-interval (512), reduce max per-path state slots to minimize pool consumption
  • Short multi-turn agent dialogue (1k–4k token context): Adopt extra_buffer_lazy strategy, lower interval value to 128 for higher prefix cache hit rates
  • High-throughput batch generation: Use extra_buffer mode with overlap scheduling enabled, reserve full 5 slots per request for stable parallel processing
  • Memory-limited consumer-grade GPUs: Test the unified dynamic memory pool if speculative features are not required

7. Conclusion

Hybrid models combining fixed-size linear KDA/GDN layers and token-scaling MLA dense layers create unique memory management challenges that generic inference engines cannot resolve. SGLang’s dedicated hybrid Mamba Radix Cache subsystem delivers a complete stack of optimizations: slot-based state pool allocation, chunked prefill snapshotting, LRU Radix prefix matching, and ReplaySSM lightweight speculative replay buffers to eliminate prohibitive full-state storage overhead. The detailed VRAM calculation formulas for Kimi K3 provide quantifiable memory budgeting data for TP8 multi-GPU clusters, enabling precise hardware capacity planning before production rollout.

Runtime parameters such as mamba-full-memory-ratio, track interval and Radix buffer strategy act as critical tuning levers to eliminate artificial concurrency bottlenecks caused by static memory partitioning. The experimental unified dynamic memory pool further optimizes GPU utilization under mixed long/short sequence traffic, albeit with feature compatibility tradeoffs for speculative decoding and disaggregated inference.

For engineering teams running hybrid linear-dense model services like Kimi K3 and Qwen3.6, mastering SGLang’s Mamba KV cache configuration directly translates to higher batch throughput, lower GPU hardware costs and stable million-token long-context inference performance. All memory sizing formulas and parameter tuning logic outlined in this paper can be directly applied to production SGLang deployments for hybrid SSM-Transformer architectures.