Fix DeepSeek V4 loading with RunAI Model Streamer.

[Broduker]

Motivation

Fix DeepSeek V4 loading with RunAI Model Streamer.

DeepSeek V4 performs a special FP8 dequantization pass for .wo_a.weight using the corresponding .wo_a.scale. The previous implementation first materialized the entire weights iterable with list(weights) and then dequantized from a dictionary. That is safe for normal local loading, but it breaks the streaming-loading contract used by runai_streamer.

In cluster validation with DeepSeek-V4-Flash on NVIDIA H20, the model could stream all tensors successfully, but the first decode failed with CUTLASS mixed dtype grouped GEMM / TMA descriptor initialization errors. Local loading of the same model did not reproduce the issue.

Fixes #30239

Modifications

• Always use the streaming-safe DeepSeek V4 FP8 wo_a dequant path instead of materializing the full weight iterator.
• Clone only pending RunAI Model Streamer tensors before retaining them across iterator steps.
• Preserve the old missing-scale behavior: if no .wo_a.scale tensors exist, pass .wo_a.weight through unchanged; if any .wo_a.scale exists, unmatched .wo_a.weight still raises.
• Add unit tests covering:
  • scale before weight and weight before scale
  • streaming output order versus legacy output order
  • cloned pending RunAI streamer tensors
  • missing-scale behavior

Accuracy Tests

No model-output accuracy change is expected. The patch only changes how DeepSeek V4 FP8 wo_a tensors are paired and dequantized during weight loading; the resulting tensor dtype remains bfloat16, matching the existing non-streaming path.

The streaming helper can yield a dequantized .wo_a.weight later than the legacy dict-based helper when .wo_a.weight appears before its .wo_a.scale. This does not change the model state because DeepSeek V4 loads tensors by parameter name. The unit tests compare legacy and streaming results by name and assert the tensor values match even when the yield order differs.

Speed Tests and Profiling

No throughput benchmark was run. This change avoids materializing the RunAI streamer weight iterator for DeepSeek V4 wo_a dequantization, so it should preserve the intended streaming-loading behavior and not affect steady-state inference speed.

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. Not applicable: this is a model loader bug fix with no user-facing API change.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ✅ Run #30328605916
Latest PR Test (Extra): ❌ Run #30328605785

[gemini-code-assist]

Code Review

This pull request introduces streaming dequantization for FP8 weights in the DeepSeek-V4 model via the new _dequant_fp8_wo_a_streaming function, replacing the previous non-streaming implementation. It also adds a helper to clone RunAI streamed tensors to prevent issues during streaming. Comprehensive unit tests have been added to verify weight and scale pairing, correctness against the legacy implementation, cloning behavior, and error handling. There are no review comments, and the changes look solid.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[JustinTong0323]

Rebase onto the latest main and cover the DeepSeek V4 DSpark loader before merging.

[JustinTong0323]

Rebase onto the latest main and use this streaming path in DeepseekV4ForCausalLMDSpark.load_weights too. PR #30261 added deepseek_v4_dspark.py with weights = list(weights), and DSPARK draft models inherit --load-format runai_streamer by default, so that path can still consume RunAI's reused buffers and load corrupted weights.

[Broduker]

Thanks, @JustinTong0323, I have updated it.

[JustinTong0323]

/tag-and-rerun-ci

[JustinTong0323]

/rerun-failed-ci

[JustinTong0323]

LGTM