[DeepSeek-V4] Offload fp8 wo_a dequant to GPU to speed up weight loading - #2
Open
AMD-yanfeiwang wants to merge 2 commits into
Open
[DeepSeek-V4] Offload fp8 wo_a dequant to GPU to speed up weight loading#2AMD-yanfeiwang wants to merge 2 commits into
AMD-yanfeiwang wants to merge 2 commits into
Conversation
When SGLANG_OPT_FP8_WO_A_GEMM is disabled, the fp8 `wo_a` weights are dequantized to bf16 at load time via `_dequant_fp8`. This ran on the host because the checkpoint tensors are loaded on CPU, making weight loading CPU-bound and very slow for large models like DeepSeek-V4 (the dequant is a single-threaded elementwise multiply over large tensors while the GPUs sit idle). Move the elementwise multiply to the current CUDA device when available and move the result back to the input device, so downstream weight-loading behavior is unchanged while the heavy math is GPU-accelerated. Using `torch.cuda.current_device()` keeps this correct under tensor parallelism where each rank binds to its own device.
…/...) Address review: the previous condition gated on torch.cuda.is_available(), so only CUDA/HIP benefited and NPU (and XPU/HPU) fell back to the slow CPU dequant. Offload any host (CPU) tensor to the current accelerator via the device-agnostic get_device() helper (already used across the repo, and consistent with the torch.get_device_module() usage in this file), then move the result back to the input device. This covers CUDA/HIP/NPU/XPU/... and resolves to the current rank's device under TP; on a CPU-only host get_device() returns "cpu" so the move is a no-op.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
When
SGLANG_OPT_FP8_WO_A_GEMMis disabled, the fp8wo_aweights of DeepSeek-V4 are dequantized to bf16 at load time in_dequant_fp8(python/sglang/srt/models/deepseek_v4.py). Because the checkpoint tensors are loaded on the host, this dequant ran on the CPU, making weight loading CPU-bound and very slow for large models.Observed on an 8x MI355 node serving
DeepSeek-V4-Pro(tp=8): after the shards finish loading, startup stalls for several minutes onExecute dequant fp8 wo_awith each TP scheduler pinned at ~30% CPU and the GPUs sitting at 0% utilization.Fix
torch.cuda.current_device()so it stays correct under tensor parallelism, where each rank is bound to its own GPU (avoids piling every rank ontocuda:0).The dequant is a trivially parallel elementwise operation, so offloading it to the GPU turns a multi-minute, single-threaded CPU step into a sub-second one.
Test plan
SGLANG_OPT_FP8_WO_A_GEMM=falseand confirm weight loading no longer stalls onExecute dequant fp8 wo_a.SGLANG_OPT_FP8_WO_A_GEMM=truepath is unaffected.CI States
Latest PR Test (Base): ❌ Run #27998902302
Latest PR Test (Extra): ❌ Run #27998902194