mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-11-05 01:07:28 +00:00
[GH-ISSUE #2226] [Bug] TP2 FP8 layerwise prefill corrupts expert weights by reusing staging slots before all GPU readers finish #1307
Labels
No milestone
No project
No assignees
1 participant
Notifications
Due date
No due date set.
Dependencies
No dependencies set.
Reference
kvcache-ai/ktransformers#1307
Loading…
Add table
Add a link
Reference in a new issue
No description provided.
Delete branch "%!s()"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
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_fp8in KTransformers' SGLang fork reuses twoCPU 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
a099dee70ccfcd8d5dda56aaa0b60cb8ecadabc9.approachingai/ktransformers, digestsha256:e66dbc25bbeb90c2b5e72d0f5b5a682d5087b4ead4a4f178e823f6407cbe5dfa.kt_ep_wrapper.pySHA-256:9eac600149223dc2958a7eb3f0a03cad673e045f509ffc5ddf4152f815f1bc02.scheduling and CUDA graphs disabled; GPU-prefill threshold 3.
SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE=2048in 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.
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-slotbuffer. 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.patchuses the same orderingbut 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:
For a lower-level audit, wrap
_prepare_weight_fp8, call the original method,then copy
self.gpu_layer's fourWEIGHT_NAMES_FP8tensors to CPU and hash theirbytes 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.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.
@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:
output IDs and selected-token log probabilities; maximum delta 0.0.
maximum delta 0.0.
maximum delta 0.0.
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:
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_fp8method. 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.
@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.