feat(deepseek-v4): CP adaption and gfx942 (MI308X/CDNA3) adaptation — 128k CP + full-model multi-node SFT - #994
Open
aaab8b wants to merge 1 commit into
Open
Conversation
…model multi-node SFT Make DeepSeek-V4-Flash trainable on gfx942/CDNA3 (the upstream sparse-MLA kernels target gfx950/CDNA4), covering three areas: gfx942 kernel fixes: - DSA sparse-MLA backward overflowed the 64 KB LDS budget at compile time (512-latent Q tile + Triton's stage double-buffering -> 73728 B). Fixed with num_stages=1 (PRIMUS_DSA_BWD_NUM_STAGES), an amd_pingpong_disabled() compile scope that turns off primus_turbo's global ping-pong/async-copy knobs (which also double-buffer LDS and are a measured pessimization here), and tunable tile sizes as a fallback. The forward is pinned to the non-ping-pong schedule too, and the backward skips provably-zero rope MMAs. Long-context / context parallelism: - CSA/HCA CP: left-boundary KV exchange (P2P) + global compressed-pool all_gather in deepseek_v4_cp.py; the indexer scoring kernel gains Q_OFFSET so visibility is judged on the query's global position under CP; streaming chunked top-K (PRIMUS_INDEXER_TOPK_CHUNK) avoids materializing the full [B,S,P] score row; a full-Triton indexer path avoids an eager einsum that would materialize [B,S,H,P]. Multi-node full-model SFT: - NNODES=2/3 branches with PP layouts; optimizer CPU offload with a tunable fraction; socket-network notes (NCCL_ALGO=Ring fixes a Tree-over-TCP cross-node all_reduce deadlock on fabrics without RDMA). One-click launchers and docs: gfx942/run_128k_dense_hca_csa.sh, gfx942/run_full_4k_multinode.sh, plus READMEs. Verified: 128k 4-layer (CP=8/TP=1, 42.30 GB peak); full model 4k on 3 nodes, 10-step and 100-step runs with loss 11.90 -> 9.19 monotone, grad norm 20 -> 2.8, 0 nan iterations, exit 0. All changes are gfx942-gated or env-var-defaulted-off and do not change gfx950 numerics or behavior. Co-Authored-By: Claude <noreply@anthropic.com>
aaab8b
requested review from
Xiaoming-AMD,
limou102 and
wenxie-amd
as code owners
August 18, 2026 07:29
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
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.
What
Make DeepSeek-V4-Flash trainable on gfx942 / CDNA3 (MI308X). The upstream
sparse-MLA kernels target gfx950/CDNA4; this PR adapts them plus the long-context
and multi-node paths so the full model trains on CDNA3.
gfx942 kernel fixes
plus Triton's stage double-buffering asks for 73728 B > the 64 KB CDNA3 LDS
budget. Fixed with
num_stages=1(PRIMUS_DSA_BWD_NUM_STAGES), anamd_pingpong_disabled()compile scope that disables primus_turbo's globalping-pong / async-copy knobs (which double-buffer LDS and are a measured
pessimization for these kernels), and tunable tile sizes as a fallback.
rope MMAs.
Long-context / context parallelism
all_gather (
deepseek_v4_cp.py).Q_OFFSETso visibility is judged on the query'sglobal position under CP (0 without CP → identical to before).
PRIMUS_INDEXER_TOPK_CHUNK) avoids materializing thefull
[B,S,P]score row; full-Triton indexer path avoids an eager einsum thatwould materialize
[B,S,H,P].Multi-node full-model SFT
NNODES=2/3branches with PP layouts; optimizer CPU offload with a tunablefraction; socket-network note:
NCCL_ALGO=Ringfixes a Tree-over-TCPcross-node all_reduce deadlock on fabrics without RDMA.
Scripts & docs
One-click launchers
examples/deepseek-v4/gfx942/run_128k_dense_hca_csa.shandrun_full_4k_multinode.sh, plus READMEs.Verification
monotone, grad norm 20 → 2.8, 0 nan iterations, exit 0.
All changes are gfx942-gated or env-var-defaulted-off and do not change
gfx950 numerics or behavior.
Test plan
pre-commit run --all-filesgreen on the changed files