kvcache-ai/sglang

Fork 0

mirror of https://github.com/kvcache-ai/sglang.git synced 2026-11-05 01:07:39 +00:00

[PR #102] [Fix] Fence all TP readers before FP8 staging-slot reuse #103

New issue

Open

opened 2026-10-29 21:07:14 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 21:07:14 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/sglang/pull/102
Author: @f-wunderlich
Created: 10/10/2026
Status: 🔄 Open

Base: main ← Head: fix-fp8-staging-reader-fence

📝 Commits (1)

db97622 Fence all TP readers before reusing FP8 staging slots

📊 Changes

2 files changed (+251 additions, -3 deletions)

View changed files

📝 python/sglang/srt/layers/moe/kt_ep_wrapper.py (+14 -3)
➕ test/unit/test_fp8_staging_readers.py (+237 -0)

📄 Description

Motivation

Fixes https://github.com/kvcache-ai/ktransformers/issues/2226.

SharedFullContext._prepare_weight_fp8 double-buffers host expert weights,
but rank zero waits only for its own GPU reader before the native writer
overwrites every TP rank's slot. A slower peer can still be copying expert E
when the writer replaces that slot with E+2. The existing publication barrier
after the write cannot protect those earlier reads.

The same issue remains on this fork's current main
18c0de72e59359baba793f7a5675f983aa656fc9. The fix is submitted here because
this code belongs to KTransformers' SGLang fork, not the parent SGLang main.

Modifications

Drain prior CUDA readers on every TP rank, then collectively fence before
the first write of a new layer load.
Before slot reuse, have every rank synchronize its local copy event and
collectively fence before rank zero submits the next native write.
Preserve the existing publication barriers, TP1's collective-free path,
quantization and kernel arithmetic. Skip the reader wait when no subsequent
write will reuse the slot.
Add a stdlib-only CPU schedule test executing the actual method in isolation
with simulated asynchronous CUDA/native boundaries. Queued copies retain
references to the host slots; the writer rejects overwriting a slot while
any rank still has a pending read. This checks a state invariant, not merely
the presence of synchronization calls in the source.

This intentionally changes only the reproduced generic FP8 method. Formats
delegating to it inherit the fences; separate FP8-channel/BF16 transports and
the specialized GLM-5-Next path are not modified or qualified here.

Accuracy Tests

New CPU fixture:

python test/unit/test_fp8_staging_readers.py

Three test methods / seven schedules pass: TP2 reuse with 3/4/5 experts,
two successive loads with 1/2/5 experts, and TP1 without collectives.
Running the same fixture against the original main source fails all six
TP2 schedules with premature slot reuse; TP1 passes.
Python compilation, isort 7.0.0, black 26.1.0, focused Ruff and whitespace
checks pass for the new test and changed code as applicable.

The production diff is the same fence ordering already deployed and GPU-validated
in our pinned Qwen integration; only upstream line offsets differ. Historical
runtime evidence from #2226, not a new full GPU run of this PR branch:

Qwen3.5-122B-A10B-FP8, TP2, AVX2 CPU offload, two RTX PRO 4500 SM120 GPUs.
The first differing trace entry was TP1's loaded layer-zero expert weights,
before expert arithmetic. One unfenced load had 253/256 down-projection
scale tensors matching expert E+2 rather than the intended expert.
With reader fences, all 30,720 inspected weight/scale tensors matched
their checkpoint TP shards across three requests, both ranks and five chunks.
Removing the fences restored corruption on TP1 and token-zero output changes.
The packaged image without diagnostic hooks passed fresh/replayed prefixes,
two tool-evidence extensions and a structured tool-call/result round trip.

The simulated test does not replace real GPU validation. These results prove
the reproduced transport race, not bit-equivalence of CPU/GPU arithmetic or
universal prefix-cache/agent correctness.

Benchmarking and Profiling

The fences have a real prefill cost; the faster unfenced timings are not a
correct-execution baseline. Historical deployed 2K-chunk measurements:

Input tokens	Fresh TTFT	Replay TTFT	Reused tokens
16,384	75.15 s	9.47 s	14,336
32,768	151.99 s	9.65 s	30,720

Each request generated 16 tokens. Decode remained about 14.5 tok/s in that
configuration. These are pinned-integration observations, not benchmarks of
the entire latest fork or a full-window qualification.

Checklist

Focused CPU regression, negative control, syntax, formatting and whitespace checks completed.
Existing hardware validation and its scope are linked above.
No new GPU serving rollout or pristine-current-main end-to-end run claimed.
The KTransformers SGLang submodule can be advanced by maintainers after this fix is merged.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/sglang/pull/102 **Author:** [@f-wunderlich](https://github.com/f-wunderlich) **Created:** 10/10/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix-fp8-staging-reader-fence` --- ### 📝 Commits (1) - [`db97622`](https://github.com/kvcache-ai/sglang/commit/db97622427ed84ef8776aeac7911c1cd785ca85c) Fence all TP readers before reusing FP8 staging slots ### 📊 Changes **2 files changed** (+251 additions, -3 deletions) <details> <summary>View changed files</summary> 📝 `python/sglang/srt/layers/moe/kt_ep_wrapper.py` (+14 -3) ➕ `test/unit/test_fp8_staging_readers.py` (+237 -0) </details> ### 📄 Description ## Motivation Fixes https://github.com/kvcache-ai/ktransformers/issues/2226. `SharedFullContext._prepare_weight_fp8` double-buffers host expert weights, but rank zero waits only for its own GPU reader before the native writer overwrites every TP rank's slot. A slower peer can still be copying expert E when the writer replaces that slot with E+2. The existing publication barrier after the write cannot protect those earlier reads. The same issue remains on this fork's current main `18c0de72e59359baba793f7a5675f983aa656fc9`. The fix is submitted here because this code belongs to KTransformers' SGLang fork, not the parent SGLang main. ## Modifications - Drain prior CUDA readers on every TP rank, then collectively fence before the first write of a new layer load. - Before slot reuse, have every rank synchronize its local copy event and collectively fence before rank zero submits the next native write. - Preserve the existing publication barriers, TP1's collective-free path, quantization and kernel arithmetic. Skip the reader wait when no subsequent write will reuse the slot. - Add a stdlib-only CPU schedule test executing the actual method in isolation with simulated asynchronous CUDA/native boundaries. Queued copies retain references to the host slots; the writer rejects overwriting a slot while any rank still has a pending read. This checks a state invariant, not merely the presence of synchronization calls in the source. This intentionally changes only the reproduced generic FP8 method. Formats delegating to it inherit the fences; separate FP8-channel/BF16 transports and the specialized GLM-5-Next path are not modified or qualified here. ## Accuracy Tests New CPU fixture: ```sh python test/unit/test_fp8_staging_readers.py ``` - Three test methods / seven schedules pass: TP2 reuse with 3/4/5 experts, two successive loads with 1/2/5 experts, and TP1 without collectives. - Running the same fixture against the original main source fails all six TP2 schedules with premature slot reuse; TP1 passes. - Python compilation, isort 7.0.0, black 26.1.0, focused Ruff and whitespace checks pass for the new test and changed code as applicable. The production diff is the same fence ordering already deployed and GPU-validated in our pinned Qwen integration; only upstream line offsets differ. Historical runtime evidence from #2226, **not a new full GPU run of this PR branch**: - Qwen3.5-122B-A10B-FP8, TP2, AVX2 CPU offload, two RTX PRO 4500 SM120 GPUs. - The first differing trace entry was TP1's loaded layer-zero expert weights, before expert arithmetic. One unfenced load had 253/256 down-projection scale tensors matching expert E+2 rather than the intended expert. - With reader fences, all **30,720** inspected weight/scale tensors matched their checkpoint TP shards across three requests, both ranks and five chunks. - Removing the fences restored corruption on TP1 and token-zero output changes. - The packaged image without diagnostic hooks passed fresh/replayed prefixes, two tool-evidence extensions and a structured tool-call/result round trip. The simulated test does not replace real GPU validation. These results prove the reproduced transport race, not bit-equivalence of CPU/GPU arithmetic or universal prefix-cache/agent correctness. ## Benchmarking and Profiling The fences have a real prefill cost; the faster unfenced timings are not a correct-execution baseline. Historical deployed 2K-chunk measurements: | Input tokens | Fresh TTFT | Replay TTFT | Reused tokens | | ---: | ---: | ---: | ---: | | 16,384 | 75.15 s | 9.47 s | 14,336 | | 32,768 | 151.99 s | 9.65 s | 30,720 | Each request generated 16 tokens. Decode remained about 14.5 tok/s in that configuration. These are pinned-integration observations, not benchmarks of the entire latest fork or a full-window qualification. ## Checklist - Focused CPU regression, negative control, syntax, formatting and whitespace checks completed. - Existing hardware validation and its scope are linked above. - No new GPU serving rollout or pristine-current-main end-to-end run claimed. - The KTransformers SGLang submodule can be advanced by maintainers after this fix is merged. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-29 21:07:14 +00:00

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

kvcache-ai/sglang#103

No description provided.

Rows
Columns