Skip to content

[RL] DeepSeek V4 weight-update fixes: UE8M0 packed-scale shards, attn_sink cache refresh, per-process DeepGEMM cache - #41043

Open
yueming-yuan wants to merge 1 commit into
mainfrom
yueming/rl-model-side-weight-update-fixes
Open

[RL] DeepSeek V4 weight-update fixes: UE8M0 packed-scale shards, attn_sink cache refresh, per-process DeepGEMM cache#41043
yueming-yuan wants to merge 1 commit into
mainfrom
yueming/rl-model-side-weight-update-fixes

Conversation

@yueming-yuan

@yueming-yuan yueming-yuan commented Sep 24, 2026 •
edited by github-actions Bot
Loading

Copy link
Copy Markdown
Collaborator

Motivation

RL weight updates on DeepSeek V4 hit three model-side problems on main: a UE8M0 packed scale that cannot be TP-sharded when the shard boundary falls inside a packed word, an attn_sink per-rank cache that keeps serving pre-update values, and colocated engines racing on one DeepGEMM JIT cache dir.

Modifications

  • layers/linear.py (RowParallelLinear): when a BlockQuantScaleParameter with format_ue8m0 receives a full-size int32 packed scale, unpack to fp32, narrow in scale-column units to the TP shard, repack (_load_ue8m0_packed_scale).
  • models/deepseek_v4*.py: refresh_attn_sink_cache() re-slices the cached per-rank attn_sink; post_load_weights calls it for every attention module (target and nextn). DSpark pads the sink to the same width as q.
  • SGLANG_DG_CACHE_DIR_PER_PROCESS (new Envs entry, default off): append pid_<pid> to SGLANG_DG_CACHE_DIR so colocated RL engines on a host do not share one JIT cache.
  • deepseek_v4.py: the unloaded-weights warning now reports a count instead of the full list, since partial RL updates trigger it by design.

Squash of sglang-miles #27131, #27603, #27604, #27728, #30421. Supersedes #29916. Not included: #30742 (FP8 tuple k-size belongs to the on-policy port) and #33478 (obsolete: main builds the unified-loader mappings inside deepseek_weight_loader).

Testing

pre-commit clean. File runs via /rerun-test: the tests referencing linear.py block-quant loading, deep_gemm_wrapper, and the DSv4 e2e set.

CI States

Latest PR Test (Base): ❌ Run #35957646125
Latest PR Test (Extra): ❌ Run #35957646181
Latest PR Test (AMD ROCm 10): ❌ Run #35957646268

@yueming-yuan
[RL] DeepSeek V4 weight-update fixes: UE8M0 packed-scale shards, attn…
0b8982c 
…_sink cache refresh, per-process DeepGEMM cache

Model-side fixes RL weight updates on DeepSeek V4 need on main:

- RowParallelLinear: a DeepGEMM UE8M0 packed scale (int32, 4 scale columns
  per word) cannot be narrowed to the TP shard directly when the shard
  boundary sits inside a word; unpack to fp32, narrow in scale columns,
  repack.
- DeepSeek V4 / DSpark: the per-rank attn_sink slice is cached on first use;
  a weight update rewrites attn_sink in place, so refresh the cache in
  post_load_weights or decode keeps serving the pre-update sink.
- SGLANG_DG_CACHE_DIR_PER_PROCESS: colocated RL engines on one host get a
  per-process DeepGEMM JIT cache dir instead of racing on shared files.
- Shorten the unloaded-weights warning, which fires on every partial RL
  update by design.

Squash of sglang-miles #27131, #27603, #27604, #27728, #30421. Supersedes
#29916. #30742 (FP8 tuple k-size) stays with the on-policy port and #33478
is obsolete on main, where the unified-loader mappings live in
deepseek_weight_loader.
@yueming-yuan

yueming-yuan commented Sep 24, 2026

Copy link
Copy Markdown
Collaborator Author

/tag-run-ci-label

@yueming-yuan

yueming-yuan commented Sep 24, 2026

Copy link
Copy Markdown
Collaborator Author

/rerun-test test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py test/registered/quant/test_fp8_utils.py test/registered/unit/layers/test_fp8_bpreshuffle_scale.py test/registered/gemm/test_linear_bf16_fp32_hpc.py test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py test/registered/attention/unittests/dsv4/test_deepseek_v4.py test/registered/unit/models/test_deepseek_v4_rope_policy.py test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py test/registered/e2e/dsv4/test_dsv4_fp8_cast.py test/registered/e2e/dsv4/test_dsv4_unified_fp8_decode.py test/registered/e2e/dsv4/test_dsv4_unified_fp8_prefill.py test/registered/e2e/dsv4/test_dsv4_unified_fp8_backend_prefill.py test/registered/e2e/dsv4/test_dsv4_unified_fp8_qk_norm_rope.py test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py test/registered/e2e/dsv4/test_dsv4_unified_fp8_compress_store.py test/registered/spec/dspark/test_dspark_draft_path_default.py test/registered/spec/dspark/test_dspark_kernel_parity.py test/registered/kernels/ops/attention/test_dsa_sink_pad_cache.py test/registered/unit/model_executor/test_derived_weight_cache.py test/registered/unit/model_executor/test_compensated_mhc_update_guard.py test/registered/rl/test_update_weights_from_distributed.py test/registered/rl/test_update_weights_from_tensor.py

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Sep 24, 2026
@github-actions

github-actions Bot commented Sep 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Results for /rerun-test test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py test/registered/quant/test_fp8_utils.py test/registered/unit/layers/test_fp8_bpreshuffle_scale.py test/registered/gemm/test_linear_bf16_fp32_hpc.py test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py test/registered/attention/unittests/dsv4/test_deepseek_v4.py test/registered/unit/models/test_deepseek_v4_rope_policy.py test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py test/registered/e2e/dsv4/test_dsv4_fp8_cast.py test/registered/e2e/dsv4/test_dsv4_unified_fp8_decode.py test/registered/e2e/dsv4/test_dsv4_unified_fp8_prefill.py test/registered/e2e/dsv4/test_dsv4_unified_fp8_backend_prefill.py test/registered/e2e/dsv4/test_dsv4_unified_fp8_qk_norm_rope.py test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py test/registered/e2e/dsv4/test_dsv4_unified_fp8_compress_store.py test/registered/spec/dspark/test_dspark_draft_path_default.py test/registered/spec/dspark/test_dspark_kernel_parity.py test/registered/kernels/ops/attention/test_dsa_sink_pad_cache.py test/registered/unit/model_executor/test_derived_weight_cache.py test/registered/unit/model_executor/test_compensated_mhc_update_guard.py test/registered/rl/test_update_weights_from_distributed.py test/registered/rl/test_update_weights_from_tensor.py:

🚀 4-gpu-b200 (2 tests): ✅ View workflow run 

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/attention/unittests/dsv4/test_deepseek_v4.py

🚀 1-gpu-5090 (3 tests): ✅ View workflow run 

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/spec/dspark/test_dspark_kernel_parity.py
cd test/ && python3 registered/rl/test_update_weights_from_tensor.py

🚀 1-gpu-h100 (5 tests): ✅ View workflow run 

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/quant/test_fp8_utils.py
cd test/ && python3 registered/gemm/test_linear_bf16_fp32_hpc.py
cd test/ && python3 registered/attention/unittests/dsv4/test_deepseek_v4.py
cd test/ && python3 registered/kernels/ops/attention/test_dsa_sink_pad_cache.py

🚀 ubuntu-latest (9 tests): ✅ View workflow run 

cd test/ && python3 registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py
cd test/ && python3 registered/unit/layers/test_fp8_bpreshuffle_scale.py
cd test/ && python3 registered/unit/layers/moe/test_mega_moe_deepgemm_api.py
cd test/ && python3 registered/unit/models/test_deepseek_v4_rope_policy.py
cd test/ && python3 registered/unit/models/test_deepseek_v4_shared_expert_fusion.py
cd test/ && python3 registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py
cd test/ && python3 registered/spec/dspark/test_dspark_draft_path_default.py
cd test/ && python3 registered/unit/model_executor/test_derived_weight_cache.py
cd test/ && python3 registered/unit/model_executor/test_compensated_mhc_update_guard.py

🚀 2-gpu-h100 (1 test): ✅ View workflow run 

cd test/ && python3 registered/rl/test_update_weights_from_distributed.py

⛔ test/registered/e2e/dsv4/test_dsv4_fp8_cast.py: test/registered/e2e/dsv4/test_dsv4_fp8_cast.py is registered for AMD (suite stage-b-test-1-gpu-small-amd-mi35x), not for CUDA or CPU; rerun-test.yml has no AMD job. Rerun it with /rerun-failed-ci, or dispatch the AMD workflow manually.

⛔ test/registered/e2e/dsv4/test_dsv4_unified_fp8_decode.py: test/registered/e2e/dsv4/test_dsv4_unified_fp8_decode.py is registered for AMD (suite stage-b-test-1-gpu-small-amd-mi35x), not for CUDA or CPU; rerun-test.yml has no AMD job. Rerun it with /rerun-failed-ci, or dispatch the AMD workflow manually.

⛔ test/registered/e2e/dsv4/test_dsv4_unified_fp8_prefill.py: test/registered/e2e/dsv4/test_dsv4_unified_fp8_prefill.py is registered for AMD (suite stage-b-test-1-gpu-small-amd-mi35x), not for CUDA or CPU; rerun-test.yml has no AMD job. Rerun it with /rerun-failed-ci, or dispatch the AMD workflow manually.

⛔ test/registered/e2e/dsv4/test_dsv4_unified_fp8_backend_prefill.py: test/registered/e2e/dsv4/test_dsv4_unified_fp8_backend_prefill.py is registered for AMD (suite stage-b-test-1-gpu-small-amd-mi35x), not for CUDA or CPU; rerun-test.yml has no AMD job. Rerun it with /rerun-failed-ci, or dispatch the AMD workflow manually.

⛔ test/registered/e2e/dsv4/test_dsv4_unified_fp8_qk_norm_rope.py: test/registered/e2e/dsv4/test_dsv4_unified_fp8_qk_norm_rope.py is registered for AMD (suite stage-b-test-1-gpu-small-amd-mi35x), not for CUDA or CPU; rerun-test.yml has no AMD job. Rerun it with /rerun-failed-ci, or dispatch the AMD workflow manually.

⛔ test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py: test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py is registered for AMD (suite stage-b-test-1-gpu-small-amd-mi35x), not for CUDA or CPU; rerun-test.yml has no AMD job. Rerun it with /rerun-failed-ci, or dispatch the AMD workflow manually.

⛔ test/registered/e2e/dsv4/test_dsv4_unified_fp8_compress_store.py: test/registered/e2e/dsv4/test_dsv4_unified_fp8_compress_store.py is registered for AMD (suite stage-b-test-1-gpu-small-amd-mi35x), not for CUDA or CPU; rerun-test.yml has no AMD job. Rerun it with /rerun-failed-ci, or dispatch the AMD workflow manually.

@yueming-yuan yueming-yuan mentioned this pull request Sep 24, 2026
Open
8 of 33 tasks

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner
@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner
@Fridge003 Fridge003 Awaiting requested review from Fridge003 Fridge003 is a code owner
@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner
@HaiShaw HaiShaw Awaiting requested review from HaiShaw HaiShaw is a code owner
@ch-wan ch-wan Awaiting requested review from ch-wan ch-wan is a code owner
@BBuf BBuf Awaiting requested review from BBuf BBuf is a code owner
@Edwardf0t1 Edwardf0t1 Awaiting requested review from Edwardf0t1 Edwardf0t1 is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

deepseek run-ci CI: run the baseline test suite on this PR

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@yueming-yuan