[Perf] KVShard: 0SM KV Cache Allgather for KV Cache Sharding by Shunkangz · Pull Request #44029 · sgl-project/sglang

Shunkangz · 2026-10-13T09:12:08Z

Motivation

Zero-SM KV-Shard Gather

The sharded pool's per-layer prefix allgather previously used NCCL ring/LL kernels. When overlapped with attention and GEMMs, these kernels competed for SMs and HBM bandwidth.

SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER enables zero-SM gathering and is on by default. The gather communicator is created with NCCL_CTA_POLICY_ZERO, allowing NCCL to run the in-place allgather on copy engines instead of SM kernels.

Memory Allocation

Both scratch slots are carved from a single ncclMemAlloc buffer registered as a symmetric window. Each slot uses a 4 KiB-aligned offset that is identical across all ranks.

_scratch_tensor_specs returns a mapping of name -> (shape, dtype).
The mixin handles allocation.
The buffer is deregistered and freed when its last tensor view is destroyed.

Activation and Fallbacks

The engine defaults NCCL_CUMEM_ENABLE=1 when zero-SM gathering applies to a single-node deployment.

The pool retains Torch allocation when any of the following applies:

SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER is disabled.
A memory saver or custom memory pool is active.
The NCCL version is older than 2.28.
The shard group spans multiple nodes. NCCL would use its hierarchical inter-node RMA path in this case, which remains unverified.

If window registration returns a NULL window because NCCL cuMem is disabled, the pool logs one warning and gathering remains on SM kernels.

NCCL 2.30 Race Workaround

With NVLS enabled, NCCL 2.30 lazily builds the multicast team during the first copy-engine collective's synchronization, without a barrier. This race can cause the first gather to hang.

The pool works around the race by building the multicast team during registration and issuing a barrier afterward.

GPU Test Coverage

The GPU test profiles a batch of gathers and verifies that:

No NCCL kernels are launched.
Peer copies occur.

A gather communicator that loses NCCL_CTA_POLICY_ZERO therefore fails the test.

Modifications

Accuracy Tests

Speed Tests and Profiling

GPUs	MiB/rank	Gather SM	Gather zero-SM	Speedup	Overlap SM	Overlap zero-SM	Speedup
2	0.0625	17.07	24.22	0.71×	95.09	93.29	1.02×
2	1	20.23	25.54	0.79×	93.83	92.40	1.02×
2	16	63.05	48.32	1.30×	111.52	99.62	1.12×
2	64	173.72	113.68	1.53×	204.30	116.76	1.75×
4	0.0625	23.74	35.76	0.66×	96.04	92.61	1.04×
4	1	35.51	38.99	0.91×	106.54	97.96	1.09×
4	16	120.49	117.89	1.02×	174.36	114.66	1.52×
4	64	407.31	312.13	1.30×	424.19	314.97	1.35×
8	0.0625	32.16	53.24	0.60×	97.67	93.78	1.04×
8	1	58.67	55.91	1.05×	117.02	94.57	1.24×
8	16	222.52	220.26	1.01×	270.86	217.99	1.24×
8	64	834.24	736.59	1.13×	853.80	732.58	1.17×

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review and Merge Process

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ⏳ Run #40061210028
Latest PR Test (Extra): ❌ Run #40061209948
Latest PR Test (AMD ROCm 10): ⏳ Run #40061210004

The sharded pool's per-layer prefix allgather ran on NCCL ring/LL kernels that compete with attention and GEMMs for SMs and HBM while overlapped. With SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER (default on) the gather communicator is created with NCCL_CTA_POLICY_ZERO and both scratch slots are carved, at 4 KiB-aligned offsets identical on every rank, out of one ncclMemAlloc buffer registered as a symmetric window, so NCCL runs the in-place allgather on copy engines. - _scratch_tensor_specs returns name -> (shape, dtype); the mixin allocates. - Torch allocation is kept when the env is off, under the memory saver or a custom mem pool, with NCCL < 2.28, or when the shard group spans nodes (NCCL would take its hierarchical inter-node RMA path there, unverified). A NULL window (NCCL cuMem off) logs one warning and the gather stays on SM kernels. - The buffer is deregistered and freed when the last tensor view dies. - The engine defaults NCCL_CUMEM_ENABLE=1 when the zero-SM gather applies on a single-node deployment. - Works around an NCCL 2.30 race: with NVLS, the multicast team is built lazily inside the first CE collective's sync without a barrier, so the first gather could hang; the pool builds it at registration and barriers. - The GPU test profiles a gather batch and asserts no NCCL kernel launches and peer copies do, so a comm that loses CTA_POLICY_ZERO fails it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Shunkangz · 2026-10-14T07:27:59Z

/tag-run-ci-label

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Preserve the zero-SM symmetric buffer types alongside upstream NCCL registration diagnostics and updated DSA kernel APIs.

Fridge003 · 2026-10-22T06:55:18Z

Please post accuracy results (for verifying the correctness of 0sm kv cache all gather), and the performance improvement results (how much speed up it can get)

Shunkangz · 2026-10-22T06:59:03Z

Please post accuracy results (for verifying the correctness of 0sm kv cache all gather), and the performance improvement results (how much speed up it can get)

For both of accuracy and performance speedup, we might need to wait for the #44447 merged. After that, we can directly reuse the E2E test for Qwen3.8 Flash Next to verify accuracy and performance.

mintlify · 2026-10-23T02:49:44Z

Preview approval required

This pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit 967e583.

Approve preview

Require NCCL 2.31.2 for patched copy-engine synchronization and release symmetric allocations on initialization failures. Compare zero-SM gathers byte-for-byte with normal NCCL on two and four GPUs, including graph replay and profiler checks.

Preserve upstream DSA test relocation alongside zero-SM gather correctness and lifetime coverage. Check the actual pool device before enabling symmetric scratch so CPU pools on CUDA hosts retain ordinary allocation.

Preserve the upstream page-interleave module imports and MTP coverage. Move zero-SM gather and symmetric-buffer lifetime tests into the dedicated multi-GPU suite, retaining two- and four-GPU CI registrations.

Shunkangz requested review from CatherineSue, JustinTong0323, Ying1123, alphabetc1, ch-wan, hanming-lu, hnyls2002, huangtingwei9988, hzh0425, ispobock, merrymercy, slin1237, xiezhq-hermann and yizhang2077 as code owners October 13, 2026 09:12

Shunkangz added 2 commits October 13, 2026 02:47

[kv-shard] Share NCCL symmetric buffer allocation

68bb77e

[kv-shard] Use NCCL environment policy for copy-engine gathers

c1f09cb

Shunkangz requested review from sogalin, wisclmy0611 and zijiexia as code owners October 13, 2026 10:03

github-actions Bot added the documentation Improvements or additions to documentation label Oct 13, 2026

[kv-shard] Scope zero-SM policy to the KV gather communicator

71d4672

github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Oct 14, 2026

Shunkangz and others added 2 commits October 14, 2026 00:40

[kv-shard] Require NCCL 2.29.2 for the symmetric gather scratch

ac5aedf

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Merge branch 'main' into kvshard/zero-sm-gather

e26413a

fivechenxi mentioned this pull request Oct 19, 2026

GLM-5.3 6000D：7K 多前缀稳态优化补丁与复现记录 fivechenxi/sglang#28

Open

5 of 6 tasks

Shunkangz changed the title ~~0SM KV Cache Allgather for KV Cache Sharding~~ [Perf] 0SM KV Cache Allgather for KV Cache Sharding Oct 19, 2026

This comment has been minimized.

Sign in to view

Shunkangz changed the title ~~[Perf] 0SM KV Cache Allgather for KV Cache Sharding~~ [Perf] Comm: 0SM KV Cache Allgather for KV Cache Sharding Oct 19, 2026

Shunkangz changed the title ~~[Perf] Comm: 0SM KV Cache Allgather for KV Cache Sharding~~ [Perf] KVShard: 0SM KV Cache Allgather for KV Cache Sharding Oct 19, 2026

hzh0425 mentioned this pull request Oct 19, 2026

Agentic Cache Scheduling System Roadmap (Q4) #43098

Open

Merge latest main into kvshard/zero-sm-gather

4ef519d

Preserve the zero-SM symmetric buffer types alongside upstream NCCL registration diagnostics and updated DSA kernel APIs.

Shunkangz mentioned this pull request Oct 21, 2026

[Perf] CP: QSA Context Parallel Optimization #43792

Merged

5 tasks done

HZY-Wade mentioned this pull request Oct 21, 2026

[Feat] KVShard: Integrate GLM DSA/MTP with PD and CPU HiCache #42973

Merged

6 of 8 tasks

Fridge003 reviewed Oct 22, 2026

View reviewed changes

Comment thread python/sglang/srt/distributed/device_communicators/pynccl_allocator.py Resolved

Comment thread test/registered/unit/mem_cache/test_page_interleave_shard.py Resolved

Comment thread python/sglang/srt/distributed/device_communicators/pynccl_allocator.py Resolved

Shunkangz force-pushed the kvshard/zero-sm-gather branch from a1ce245 to 5e3feac Compare October 23, 2026 02:56

Merge latest main into kvshard/zero-sm-gather

4e34c20

Preserve upstream DSA test relocation alongside zero-SM gather correctness and lifetime coverage. Check the actual pool device before enabling symmetric scratch so CPU pools on CUDA hosts retain ordinary allocation.

Fridge003 approved these changes Oct 24, 2026

View reviewed changes

Merge latest main into kvshard/zero-sm-gather

967e583

Preserve the upstream page-interleave module imports and MTP coverage. Move zero-SM gather and symmetric-buffer lifetime tests into the dedicated multi-GPU suite, retaining two- and four-GPU CI registrations.

Fridge003 merged commit 3ea7e21 into sgl-project:main Oct 24, 2026
161 of 189 checks passed

AMD-yanfeiwang mentioned this pull request Oct 28, 2026

feat(cco): fused SDMA wait+put post, and an engine-wait in-place all-gather example ROCm/mori#866

Merged

6 of 7 tasks

Navigation Menu

Repository navigation

[Perf] KVShard: 0SM KV Cache Allgather for KV Cache Sharding - #44029

[Perf] KVShard: 0SM KV Cache Allgather for KV Cache Sharding#44029
Fridge003 merged 10 commits into
sgl-project:mainfrom
Shunkangz:kvshard/zero-sm-gather

Shunkangz commented Oct 13, 2026 •

edited by github-actions Bot

Loading

Uh oh!

Shunkangz commented Oct 14, 2026

Uh oh!

This comment has been minimized.

Fridge003 commented Oct 22, 2026

Uh oh!

Shunkangz commented Oct 22, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mintlify Bot commented Oct 23, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Shunkangz commented Oct 13, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Zero-SM KV-Shard Gather

Memory Allocation

Activation and Fallbacks

NCCL 2.30 Race Workaround

GPU Test Coverage

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

Review and Merge Process

CI States

Uh oh!

Shunkangz commented Oct 14, 2026

Uh oh!

This comment has been minimized.

Fridge003 commented Oct 22, 2026

Uh oh!

Shunkangz commented Oct 22, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mintlify Bot commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Preview approval required

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Shunkangz commented Oct 13, 2026 •

edited by github-actions Bot

Loading

mintlify Bot commented Oct 23, 2026 •

edited

Loading