[shiyang814-cpu]

Summary

This PR introduces QSA cache sharding for QSA-enabled Qwen3.8 hybrid models. It shards the sequence-length-dependent QSA states:

• QSA raw K/V history
• QSA compressed cache

The implementation reuses SGLang's page-interleave placement and allocator infrastructure, and adds:

• allocator-visible logical QSA cache sharding
• decode allocation over the widened logical slot space
• exact distributed QSA Top-K
• owner-side sparse attention
• FP32 log-sum-exp output merging
• CUDA Graph-compatible communication and metadata buffers
• one shared CUDA Graph collective arena across all QSA layers

The implementation is functionally validated and benchmarked with the Qwen3.8-Flash-Next-FP8 checkpoint on NVIDIA H20 GPUs.

Performance

The benchmark uses theQwen3.8-Flash-Next-FP8 model on 8 NVIDIA H20 GPUs with TP8 and EP8. It uses random token IDs with an output length of 1, a maximum concurrency of 1, and a chunked prefill size of 8192, with Decode CUDA Graph enabled and Radix Cache disabled.

Input length	Mode	TTFT	Input throughput	Logical capacity
256K	Unsharded	16,798.12 ms	15,605.55 tok/s	2,692,032
256K	QSA shard 4	15,471.77 ms	16,943.38 tok/s	10,218,240
512K	Unsharded	58,526.52 ms	8,958.13 tok/s	2,692,032
512K	QSA shard 4	35,682.70 ms	14,693.06 tok/s	10,218,240
1M	Unsharded	279,878.39 ms	3,572.98 tok/s	2,692,032
1M	QSA shard 4	90,356.03 ms	11,067.33 tok/s	10,218,240

Input length	TTFT reduction	Speedup	Throughput improvement
256K	7.90%	1.09x	+8.57%
512K	39.03%	1.64x	+64.02%
1M	67.72%	3.10x	+209.75%

QSA cache sharding increases the logical QSA cache capacity from 2,692,032 to 10,218,240 tokens, approximately 3.80x.

Concurrent Long-Context Workload

With an input length of 1M tokens, 8 requests, and a maximum concurrency of 8:

Metric	Unsharded	QSA shard 4	Improvement
Input throughput	12,135.16 tok/s	14,351.96 tok/s	+18.27%
Mean TTFT	372.84 s	314.44 s	-15.66%
P99 TTFT	653.09 s	552.14 s	-15.46%
Benchmark duration	659.24 s	557.42 s	-15.45%

All 8 requests completed successfully in both configurations.

Accuracy

The standard SGLang GSM8K evaluation used 1,314 examples. The LongBench-v2 evaluation used 120 examples filtered to 10K-200K context lengths, with temperature 0.6, top-p 0.95, and up to 16,384 output tokens.

Both evaluations used chunked_prefill_size=8192.

Dataset	Baseline	QSA shard 4	Delta
GSM8K	0.977	0.978	+0.001
LongBench-v2	0.554	0.549	-0.005

The sharded results remain close to the unsharded baseline on both short-context reasoning and long-context evaluation.

Motivation

The QSA raw K/V history and compressed cache grow linearly with sequence length.

Before this change, every TP rank stored the complete QSA history. This caused two long-context bottlenecks:

1. the replicated QSA history limited the total cache capacity to the physical capacity of one rank;
2. every rank scanned the full compressed history during QSA token selection.

The impact becomes increasingly significant as the context length grows. On Qwen3.8-Flash-Next, the unsharded implementation spends substantially more time scanning and maintaining QSA history at 512K and 1M context lengths.

This PR assigns each QSA cache page to one owner rank. Each rank stores and scans only its local shard while preserving the exact unsharded QSA semantics.

Implementation

QSA cache placement

QSA raw K/V history and compressed cache reuse SGLang's existing page-interleave placement and allocator infrastructure.

Each QSA cache page is assigned to one owner rank. The scheduler, request token mapping, cache allocator, and QSA KV pools operate on the same widened logical slot space, while each rank stores only its local physical shard.

The allocator supports both prefill and decode allocation, keeping cache placement consistent throughout the request lifecycle.

Distributed QSA execution

Each rank scans only its local compressed-cache shard and produces local candidates.

The local candidates are merged into an exact global Top-K using deterministic global token-ID tie-breaking. The selected tokens are then routed to their raw K/V owners for sparse-attention computation.

Owner-side partial outputs are merged using FP32 log-sum-exp normalization, preserving the original unsharded QSA semantics.

The distributed path supports both eager execution and CUDA Graph replay.

CUDA Graph buffers

QSA communication uses fixed-address buffers for CUDA Graph compatibility.

All QSA layers share one collective buffer arena because the layers execute serially and do not require concurrent per-layer storage.

For the largest tested graph bucket with 48 QSA layers, this reduces the collective arena from approximately 4.42 GiB to 94.35 MiB.

CUDA Graph metadata

QSA CUDA Graph metadata construction is limited to compressed pages visible to the current sequence.

Unused page-table tails are not cleared during every replay because downstream consumers are bounded by graph_compressed_lengths.

This preserves correctness while reducing metadata overhead for long contexts.

Memory accounting

The KV cache configurator reserves NCCL CUDA Graph private pools and PLE prefill workspace before sizing the KV cache.

The QSA collective arena additionally uses exact byte accounting verified against its allocated tensor storage.

Usage

python3 -m sglang.launch_server \
  --model-path /models/Qwen3.8-Flash-Next-FP8 \
  --tp-size 8 \
  --ep-size 8 \
  --qsa-cache-sharding-size 4 \
  --chunked-prefill-size 8192 \
  --max-prefill-tokens 8192 \
  --mem-fraction-static 0.85 \
  --cuda-graph-max-bs-decode 16 \
  --context-length 1048576 \
  --linear-attn-prefill-backend flashinfer

--qsa-cache-sharding-size 4 creates a four-rank QSA cache-sharding group inside the TP group.

Testing

The implementation is validated with the Qwen3.8-Flash-Next QSA configuration.

Correctness coverage includes:

• unsharded and sharded QSA token selection;
• exact distributed Top-K;
• deterministic global token-ID tie-breaking;
• raw K/V owner mapping;
• owner-side sparse attention;
• FP32 LSE merging;
• prefill and decode paths;
• eager execution;
• CUDA Graph capture and replay;
• shrinking sequences and padded graph rows;
• shared collective arena byte accounting.

Production-path distributed integration tests run with:

• 2 GPUs using NCCL;
• 4 GPUs using NCCL.

The distributed prefill and decode outputs are compared against the unsharded reference, including CUDA Graph replay after input mutation.

Test results:

• test_qsa.py: 201 passed, 1 skipped
• test_qsa_kv_pool.py: 5 passed
• Ruff import and error checks: passed
• Ruff format check: passed
• Python compileall: passed
• git diff --check: passed
• Qwen3.8-Flash-Next 256K/512K/1M prefill benchmark: passed

Scope

This PR only shards QSA cache states for Qwen3.8-Flash-Next-style hybrid QSA models.

The following remain unchanged:

• GDN recurrent states
• PLE states
• MoE expert placement

Pipeline parallelism and PD disaggregation are not supported by QSA cache sharding in this PR.

---

CI States

Latest PR Test (Base): ⏳ Run #42155951214
Latest PR Test (Extra): ❌ Run #42155951040
Latest PR Test (AMD ROCm 10): ❌ Run #42155951136

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[BBuf]

[BBuf Review Agent] Could you split it below 3000 changed lines (now 4625)?

[BBuf]

[BBuf Review Agent] Could you split it below 3000 changed lines (now 4146)?

[BBuf]

[BBuf Review Agent] A few things to fix before this can go in.

[BBuf]

[BBuf Review Agent] With shard size 4, production writes still go through the hybrid layer router directly to full_kv_pool, so this localization never runs. Owner attention reads slot // 4 while the writer uses the global slot, producing wrong KV reads and eventually out-of-bounds writes. Could we apply the mapping on the routed write path?

[BBuf]

[BBuf Review Agent] With shard size 4 and 256 compressed blocks per rank, this skips every local score even though the global merge still has to choose 512 of 1024 blocks. _maybe_merge_distributed_topk then ranks values from torch.empty, so token selection is incorrect. Could we disable the scoring shortcut when distributed merging is active?

[BBuf]

[BBuf Review Agent] On B200 with BF16 queries and head dimension 256, decode returns from _try_fused_kv_attention before reaching this owner dispatch. That path uses global slots against the rank-local cache and skips the owner merge, causing incorrect reads or out-of-bounds accesses. Could we disable fused eligibility for sharded pools until it handles this layout?

[BBuf]

[BBuf Review Agent] Could you split it below 3000 changed lines (now 3107)?

[BBuf]

/tag-and-rerun-ci

[BBuf]

/rerun-failed-ci

[BBuf]

/rerun-failed-ci