[GH-ISSUE #2226] [Bug] TP2 FP8 layerwise prefill corrupts expert weights by reusing staging slots before all GPU readers finish #1307

Open
opened 2026-10-29 14:51:05 +00:00 by bsakaguchi · 2 comments
Owner

Originally created by @f-wunderlich on GitHub (Oct 7, 2026).
Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2226

TP2 FP8 expert staging-buffer corruption

Scope and cause

SharedFullContext._prepare_weight_fp8 in KTransformers' SGLang fork reuses two
CPU staging slots while copying a full MoE layer to each TP rank's GPU. Rank
zero's native writer populates every rank's CPU buffers. Before reusing a slot,
the original code waits only for rank zero's CUDA copy event. The other GPU
can still be reading its buffer when the producer overwrites it for expert
E+2. The subsequent publication barrier happens too late to protect that read.

This is weight-byte corruption, not merely FP8 rounding. It occurs without
radix caching and before the first GPU expert calculation. Similar-looking
FP8-channel/BF16 loaders have not been qualified by this experiment; the patch
is intentionally limited to the reproduced block-FP8 path.

The same rank-zero-only wait is present in upstream fork commit
18c0de72e59359baba793f7a5675f983aa656fc9. The fork disables its issue tracker,
so the report belongs in the KTransformers tracker.

Reproduced environment

  • Qwen/Qwen3.5-122B-A10B-FP8, revision
    a099dee70ccfcd8d5dda56aaa0b60cb8ecadabc9.
  • Two EPYC 7443 CPUs, 1 TiB RAM, two RTX PRO 4500 Blackwell 32 GiB GPUs.
  • Public base image approachingai/ktransformers, digest
    sha256:e66dbc25bbeb90c2b5e72d0f5b5a682d5087b4ead4a4f178e823f6407cbe5dfa.
  • PyTorch 2.9.1+cu128, Triton 3.5.1, KTransformers 0.7.0.post1.
  • Original installed kt_ep_wrapper.py SHA-256:
    9eac600149223dc2958a7eb3f0a03cad673e045f509ffc5ddf4152f815f1bc02.
  • TP2, 48 CPU threads/two NUMA pools, AVX2 native FP8, zero resident GPU experts.
  • Triton attention, linear attention, MoE and FP8 GEMM; 2,048-token chunks.
  • Isolation controls: radix disabled, deterministic inference enabled, overlap
    scheduling and CUDA graphs disabled; GPU-prefill threshold 3.
  • Deterministic 2K chunks require
    SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE=2048 in this pinned runtime.

Causal validation (October 7)

Identical, uncached 9,763-token requests were run serially. The first differing
trace entry was the loaded layer-zero expert weights on TP1. Embeddings,
attention/projections, router outputs and expert inputs before that entry
matched. TP1's expert output differed next; the MLP all-reduce propagated the
difference to TP0.

Loaded GPU bytes were hashed per expert and compared with the checkpoint's TP
shards: gate/up rows and down-projection columns, with BF16 checkpoint scales
converted to the FP32 representation exported by the native loader.

Control Observed result
Original loader One load had 252/256 incorrect TP1 down-projection tensors; TP0 matched.
Reader fences enabled All 30,720 checked layer-zero weight/scale tensors matched across three requests, both ranks and five chunks.
Reader fences enabled, 16 output tokens Two fresh requests had identical token IDs, selected log probabilities and top-five probabilities at every position; maximum log-probability delta 0.0.
Fences removed TP1 corruption returned across all five chunks and the response differed at token zero.

In one unfenced load, 253/256 down-projection scale tensors were exactly the
checkpoint scales for expert E+2, identifying reuse of the same two-slot
buffer. Many weight tensors contained mixed bytes rather than a complete other
expert. This supports the producer/reader race rather than a matrix-arithmetic
or recurrent-state-restoration explanation.

The diagnostic fence drained previous readers at load entry and waited for
each rank's local copy event before a collective reader-completion barrier.
The packaged fp8-staging-reader-fence.patch uses the same ordering
but avoids an unnecessary reader wait when there is no subsequent slot reuse.
It also avoids introducing multi-rank barriers for TP1. Existing publication
barriers remain; quantization and kernels are unchanged.

Request reproducer

Launch the text-only server with the configuration above. The deployment and
image launcher preserve the rest of the tested configuration; the isolation
flags are temporary diagnostics, not production defaults. Run the following
against the native SGLang API, not a proxy:

import json
import urllib.request

text = "Incident evidence follows.\n" + (
  "The pod is Ready; an intermittent 503 is not a confirmed root cause.\n" * 512
) + "\nSummarize what is known without inventing causes.\nSummary:"
text += "\nNew tool evidence: the backend returned HTTP 200.\n"
payload = {
  "text": text, "return_logprob": True, "top_logprobs_num": 5,
  "sampling_params": {"temperature": 0, "max_new_tokens": 16, "ignore_eos": True},
}
results = []
for repeat in range(3):
  request = urllib.request.Request(
    "http://127.0.0.1:8000/generate", json.dumps(payload).encode(),
    {"Content-Type": "application/json"},
  )
  with urllib.request.urlopen(request, timeout=1200) as response:
    result = json.loads(response.read())
  meta = result["meta_info"]
  assert meta["prompt_tokens"] == 9763 and meta["cached_tokens"] == 0
  results.append(meta["output_token_logprobs"])
  print(repeat, result["text"], results[-1])
print("identical selected IDs/probabilities:", all(r == results[0] for r in results[1:]))

For a lower-level audit, wrap _prepare_weight_fp8, call the original method,
then copy self.gpu_layer's four WEIGHT_NAMES_FP8 tensors to CPU and hash their
bytes per expert before GPU quantization post-processing. Audit layer zero on
both ranks over multiple chunks. This synchronization perturbs timing, but the
race remained reproducible with the audit enabled; the uninstrumented requests
also varied. Weight hashes, response metadata and diagnostic scripts were
retained on the evaluation node, not added as bulky repository logs.

Fresh-run equivalence does not by itself qualify prefix restoration or remove
CPU W8A16 versus GPU W8A8 arithmetic differences. Those are separate serving
qualification checks after deploying the packaged fix.

Proposed fix

Apply in the SGLang fork with patch --strip=1 --directory=python/sglang <fp8-staging-reader-fence.patch.

--- a/srt/layers/moe/kt_ep_wrapper.py
+++ b/srt/layers/moe/kt_ep_wrapper.py
@@ -2098,6 +2098,12 @@
         tp_world_size = get_tensor_model_parallel_world_size()
         do_write = tp_rank == 0 and wrapper is not None
 
+        # Rank 0 writes every rank's CPU staging slots. Drain all previous
+        # readers before starting a new load, including reuse across layers.
+        if dist.is_initialized() and tp_world_size > 1:
+            torch.cuda.synchronize(device)
+            dist.barrier(group=get_tp_group().device_group)
+
         if do_write:
             # Calculate per-expert byte sizes (buffer is double-buffered: [2, ...])
             w13_weight_buf = self.cpu_buffers["w13_weight"]
@@ -2156,6 +2162,14 @@
         for idx, e in enumerate(cpu_expert_ids):
             slot = idx % 2  # Double buffering based on iteration index
 
+            # The next write reuses the slot read by expert idx-1. Every
+            # rank must finish that H2D copy before rank 0 can overwrite it;
+            # the publication barrier below alone does not protect readers.
+            if idx > 0 and idx + 1 < len(cpu_expert_ids):
+                events[idx - 1].synchronize()
+                if dist.is_initialized() and tp_world_size > 1:
+                    dist.barrier(group=get_tp_group().device_group)
+
             # Sync write for expert e, submit write for next CPU expert
             if do_write:
                 wrapper.sync_write_weight_scale_to_buffer()
@@ -2163,8 +2177,5 @@
                 if idx + 1 < len(cpu_expert_ids):
                     next_slot = (idx + 1) % 2
-                    # Before writing to next_slot, ensure copy from that slot is complete.
-                    if idx > 0:
-                        events[idx - 1].synchronize()
                     submit_write_expert(cpu_expert_ids[idx + 1], next_slot)
 
             # Barrier to ensure all ranks see the written data

The instrumented reader-fence control above is validated. The packaged patch removes a redundant final-slot wait; a rebuilt image is undergoing additional end-to-end prefix-cache qualification. I will add those results separately rather than claim cache-path correctness from this fresh-prefill control.

Originally created by @f-wunderlich on GitHub (Oct 7, 2026). Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2226 # TP2 FP8 expert staging-buffer corruption ## Scope and cause `SharedFullContext._prepare_weight_fp8` in KTransformers' SGLang fork reuses two CPU staging slots while copying a full MoE layer to each TP rank's GPU. Rank zero's native writer populates every rank's CPU buffers. Before reusing a slot, the original code waits only for rank zero's CUDA copy event. The other GPU can still be reading its buffer when the producer overwrites it for expert `E+2`. The subsequent publication barrier happens too late to protect that read. This is weight-byte corruption, not merely FP8 rounding. It occurs without radix caching and before the first GPU expert calculation. Similar-looking FP8-channel/BF16 loaders have not been qualified by this experiment; the patch is intentionally limited to the reproduced block-FP8 path. The same rank-zero-only wait is present in upstream fork commit `18c0de72e59359baba793f7a5675f983aa656fc9`. The fork disables its issue tracker, so the report belongs in the KTransformers tracker. ## Reproduced environment - Qwen/Qwen3.5-122B-A10B-FP8, revision `a099dee70ccfcd8d5dda56aaa0b60cb8ecadabc9`. - Two EPYC 7443 CPUs, 1 TiB RAM, two RTX PRO 4500 Blackwell 32 GiB GPUs. - Public base image `approachingai/ktransformers`, digest `sha256:e66dbc25bbeb90c2b5e72d0f5b5a682d5087b4ead4a4f178e823f6407cbe5dfa`. - PyTorch 2.9.1+cu128, Triton 3.5.1, KTransformers 0.7.0.post1. - Original installed `kt_ep_wrapper.py` SHA-256: `9eac600149223dc2958a7eb3f0a03cad673e045f509ffc5ddf4152f815f1bc02`. - TP2, 48 CPU threads/two NUMA pools, AVX2 native FP8, zero resident GPU experts. - Triton attention, linear attention, MoE and FP8 GEMM; 2,048-token chunks. - Isolation controls: radix disabled, deterministic inference enabled, overlap scheduling and CUDA graphs disabled; GPU-prefill threshold 3. - Deterministic 2K chunks require `SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE=2048` in this pinned runtime. ## Causal validation (October 7) Identical, uncached 9,763-token requests were run serially. The first differing trace entry was the loaded layer-zero expert weights on TP1. Embeddings, attention/projections, router outputs and expert inputs before that entry matched. TP1's expert output differed next; the MLP all-reduce propagated the difference to TP0. Loaded GPU bytes were hashed per expert and compared with the checkpoint's TP shards: gate/up rows and down-projection columns, with BF16 checkpoint scales converted to the FP32 representation exported by the native loader. | Control | Observed result | | --- | --- | | Original loader | One load had 252/256 incorrect TP1 down-projection tensors; TP0 matched. | | Reader fences enabled | All 30,720 checked layer-zero weight/scale tensors matched across three requests, both ranks and five chunks. | | Reader fences enabled, 16 output tokens | Two fresh requests had identical token IDs, selected log probabilities and top-five probabilities at every position; maximum log-probability delta 0.0. | | Fences removed | TP1 corruption returned across all five chunks and the response differed at token zero. | In one unfenced load, 253/256 down-projection scale tensors were exactly the checkpoint scales for expert `E+2`, identifying reuse of the same two-slot buffer. Many weight tensors contained mixed bytes rather than a complete other expert. This supports the producer/reader race rather than a matrix-arithmetic or recurrent-state-restoration explanation. The diagnostic fence drained previous readers at load entry and waited for each rank's local copy event before a collective reader-completion barrier. The packaged `fp8-staging-reader-fence.patch` uses the same ordering but avoids an unnecessary reader wait when there is no subsequent slot reuse. It also avoids introducing multi-rank barriers for TP1. Existing publication barriers remain; quantization and kernels are unchanged. ## Request reproducer Launch the text-only server with the configuration above. The deployment and image launcher preserve the rest of the tested configuration; the isolation flags are temporary diagnostics, not production defaults. Run the following against the native SGLang API, not a proxy: ```python import json import urllib.request text = "Incident evidence follows.\n" + ( "The pod is Ready; an intermittent 503 is not a confirmed root cause.\n" * 512 ) + "\nSummarize what is known without inventing causes.\nSummary:" text += "\nNew tool evidence: the backend returned HTTP 200.\n" payload = { "text": text, "return_logprob": True, "top_logprobs_num": 5, "sampling_params": {"temperature": 0, "max_new_tokens": 16, "ignore_eos": True}, } results = [] for repeat in range(3): request = urllib.request.Request( "http://127.0.0.1:8000/generate", json.dumps(payload).encode(), {"Content-Type": "application/json"}, ) with urllib.request.urlopen(request, timeout=1200) as response: result = json.loads(response.read()) meta = result["meta_info"] assert meta["prompt_tokens"] == 9763 and meta["cached_tokens"] == 0 results.append(meta["output_token_logprobs"]) print(repeat, result["text"], results[-1]) print("identical selected IDs/probabilities:", all(r == results[0] for r in results[1:])) ``` For a lower-level audit, wrap `_prepare_weight_fp8`, call the original method, then copy `self.gpu_layer`'s four `WEIGHT_NAMES_FP8` tensors to CPU and hash their bytes per expert before GPU quantization post-processing. Audit layer zero on both ranks over multiple chunks. This synchronization perturbs timing, but the race remained reproducible with the audit enabled; the uninstrumented requests also varied. Weight hashes, response metadata and diagnostic scripts were retained on the evaluation node, not added as bulky repository logs. Fresh-run equivalence does not by itself qualify prefix restoration or remove CPU W8A16 versus GPU W8A8 arithmetic differences. Those are separate serving qualification checks after deploying the packaged fix. ## Proposed fix Apply in the SGLang fork with `patch --strip=1 --directory=python/sglang <fp8-staging-reader-fence.patch`. ```diff --- a/srt/layers/moe/kt_ep_wrapper.py +++ b/srt/layers/moe/kt_ep_wrapper.py @@ -2098,6 +2098,12 @@ tp_world_size = get_tensor_model_parallel_world_size() do_write = tp_rank == 0 and wrapper is not None + # Rank 0 writes every rank's CPU staging slots. Drain all previous + # readers before starting a new load, including reuse across layers. + if dist.is_initialized() and tp_world_size > 1: + torch.cuda.synchronize(device) + dist.barrier(group=get_tp_group().device_group) + if do_write: # Calculate per-expert byte sizes (buffer is double-buffered: [2, ...]) w13_weight_buf = self.cpu_buffers["w13_weight"] @@ -2156,6 +2162,14 @@ for idx, e in enumerate(cpu_expert_ids): slot = idx % 2 # Double buffering based on iteration index + # The next write reuses the slot read by expert idx-1. Every + # rank must finish that H2D copy before rank 0 can overwrite it; + # the publication barrier below alone does not protect readers. + if idx > 0 and idx + 1 < len(cpu_expert_ids): + events[idx - 1].synchronize() + if dist.is_initialized() and tp_world_size > 1: + dist.barrier(group=get_tp_group().device_group) + # Sync write for expert e, submit write for next CPU expert if do_write: wrapper.sync_write_weight_scale_to_buffer() @@ -2163,8 +2177,5 @@ if idx + 1 < len(cpu_expert_ids): next_slot = (idx + 1) % 2 - # Before writing to next_slot, ensure copy from that slot is complete. - if idx > 0: - events[idx - 1].synchronize() submit_write_expert(cpu_expert_ids[idx + 1], next_slot) # Barrier to ensure all ranks see the written data ``` The instrumented reader-fence control above is validated. The packaged patch removes a redundant final-slot wait; a rebuilt image is undergoing additional end-to-end prefix-cache qualification. I will add those results separately rather than claim cache-path correctness from this fresh-prefill control.
Author
Owner

@f-wunderlich commented on GitHub (Oct 7, 2026):

The proposed patch has now been built into a clean image and deployed without
diagnostic instrumentation. The installed patched source SHA-256 is
e009a5811c685825c35b433c8f6b6509a7342b94adfdb78ae62f8d59c4c33a74.

With the original serving settings (TP2, radix enabled, CUDA graphs enabled,
2K chunks, GPU-prefill threshold 256), the previously failing 9.7K-prefix suite
passes without changing its numerical gates:

  • Two independently cache-flushed 9,748-token calculations: identical 16 greedy
    output IDs and selected-token log probabilities; maximum delta 0.0.
  • Cached replay, reusing 8,192 tokens: identical IDs and selected probabilities;
    maximum delta 0.0.
  • Replay after an unrelated request: again identical IDs/probabilities,
    maximum delta 0.0.
  • Two tool-evidence extensions, reusing all 9,748 prefix tokens: identical
    16 greedy IDs against fresh references, maximum selected-logprob deltas
    0.0742 and 0.0161 (unchanged tolerance 0.1).

The short extensions use CPU W8A16 experts whereas the fresh references use
GPU W8A8 prefill, so these last checks do not assert bit-equivalence between
the arithmetic paths. Direct chat and a structured mock tool-call/result round
trip also pass.

The first cold request took 56.46 s including 16 output tokens; the independent
fresh repeat took 47.51 s. Cached replay took 10.40 s and the first incremental
round took 2.01 s. The extra fences have a real prefill cost; earlier unfenced
timings are not a correct-execution baseline. Six fixed 256-token single-request
decode samples had a median of 14.50 tokens/s; three paired samples had median
aggregate throughput of 14.71 tokens/s, essentially unchanged from the original
CPU expert decode configuration.

Isolated prefill timings with 2K chunks, after layerwise initialization:

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. These are serving timings, not numerical
equivalence checks at 16K/32K or full-window production qualification.

Scope clarification: this patch targets the generic _prepare_weight_fp8
method. Other formats delegating to it inherit the fence, but only Qwen's
block-FP8 path was validated here. Separate FP8-channel/BF16 loaders were not
modified. Passing this fixture is not a full 128K-window or all-model cache
correctness claim.

<!-- gh-comment-id:6031402423 --> @f-wunderlich commented on GitHub (Oct 7, 2026): The proposed patch has now been built into a clean image and deployed without diagnostic instrumentation. The installed patched source SHA-256 is `e009a5811c685825c35b433c8f6b6509a7342b94adfdb78ae62f8d59c4c33a74`. With the original serving settings (TP2, radix enabled, CUDA graphs enabled, 2K chunks, GPU-prefill threshold 256), the previously failing 9.7K-prefix suite passes without changing its numerical gates: - Two independently cache-flushed 9,748-token calculations: identical 16 greedy output IDs and selected-token log probabilities; maximum delta **0.0**. - Cached replay, reusing 8,192 tokens: identical IDs and selected probabilities; maximum delta **0.0**. - Replay after an unrelated request: again identical IDs/probabilities, maximum delta **0.0**. - Two tool-evidence extensions, reusing all 9,748 prefix tokens: identical 16 greedy IDs against fresh references, maximum selected-logprob deltas **0.0742** and **0.0161** (unchanged tolerance 0.1). The short extensions use CPU W8A16 experts whereas the fresh references use GPU W8A8 prefill, so these last checks do not assert bit-equivalence between the arithmetic paths. Direct chat and a structured mock tool-call/result round trip also pass. The first cold request took 56.46 s including 16 output tokens; the independent fresh repeat took 47.51 s. Cached replay took 10.40 s and the first incremental round took 2.01 s. The extra fences have a real prefill cost; earlier unfenced timings are not a correct-execution baseline. Six fixed 256-token single-request decode samples had a median of 14.50 tokens/s; three paired samples had median aggregate throughput of 14.71 tokens/s, essentially unchanged from the original CPU expert decode configuration. Isolated prefill timings with 2K chunks, after layerwise initialization: | 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. These are serving timings, not numerical equivalence checks at 16K/32K or full-window production qualification. Scope clarification: this patch targets the generic `_prepare_weight_fp8` method. Other formats delegating to it inherit the fence, but only Qwen's block-FP8 path was validated here. Separate FP8-channel/BF16 loaders were not modified. Passing this fixture is not a full 128K-window or all-model cache correctness claim.
Author
Owner

@f-wunderlich commented on GitHub (Oct 10, 2026):

The fix is now submitted to the SGLang fork that owns this code:
https://github.com/kvcache-ai/sglang/pull/102.

It preserves the deployed reader-fence ordering and adds a CPU-only regression
executing the actual transport method with delayed simulated DMA readers.
All seven schedules pass with the fix; the original source fails all six TP2
schedules, while TP1 passes. This covers slot reuse and successive layer loads.

The PR clearly separates that new schedule test from our already reported
hardware validation; no new full GPU run of the upstream branch is claimed.
Once merged, the KTransformers SGLang submodule can be advanced to include it.

<!-- gh-comment-id:6094447866 --> @f-wunderlich commented on GitHub (Oct 10, 2026): The fix is now submitted to the SGLang fork that owns this code: https://github.com/kvcache-ai/sglang/pull/102. It preserves the deployed reader-fence ordering and adds a CPU-only regression executing the actual transport method with delayed simulated DMA readers. All seven schedules pass with the fix; the original source fails all six TP2 schedules, while TP1 passes. This covers slot reuse and successive layer loads. The PR clearly separates that new schedule test from our already reported hardware validation; no new full GPU run of the upstream branch is claimed. Once merged, the KTransformers SGLang submodule can be advanced to include it.
Sign in to join this conversation.
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/ktransformers#1307
No description provided.