Skip to content

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
AMD-AGI:mainfrom
aaab8b:feat/deepseek-v4/gfx942-adaptation
Open

feat(deepseek-v4): CP adaption and gfx942 (MI308X/CDNA3) adaptation — 128k CP + full-model multi-node SFT #994
aaab8b wants to merge 1 commit into
AMD-AGI:mainfrom
aaab8b:feat/deepseek-v4/gfx942-adaptation

Conversation

@aaab8b

@aaab8b aaab8b commented Aug 18, 2026

Copy link
Copy Markdown

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

  • DSA sparse-MLA backward LDS overflow (compile-time): the 512-latent Q tile
    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), an
    amd_pingpong_disabled() compile scope that disables primus_turbo's global
    ping-pong / async-copy knobs (which double-buffer LDS and are a measured
    pessimization for these kernels), and tunable tile sizes as a fallback.
  • Forward pinned to the non-ping-pong schedule; backward skips provably-zero
    rope MMAs.

Long-context / context parallelism

  • CSA/HCA CP: left-boundary KV exchange (P2P) + global compressed-pool
    all_gather (deepseek_v4_cp.py).
  • Indexer scoring gains Q_OFFSET so visibility is judged on the query's
    global position under CP (0 without CP → identical to before).
  • Streaming chunked top-K (PRIMUS_INDEXER_TOPK_CHUNK) avoids materializing the
    full [B,S,P] score row; 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 note: NCCL_ALGO=Ring fixes a Tree-over-TCP
    cross-node all_reduce deadlock on fabrics without RDMA.

Scripts & docs

One-click launchers examples/deepseek-v4/gfx942/run_128k_dense_hca_csa.sh and
run_full_4k_multinode.sh, plus READMEs.

Verification

  • 128k, 4-layer (CP=8/TP=1): 42.30 GB peak.
  • Full model, 4k, 3 nodes: 10-step and 100-step runs, 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.

Test plan

  • pre-commit run --all-files green on the changed files
  • 128k 4-layer smoke on 1×8 MI308X
  • full-model 4k 10-step & 100-step on 3×8 MI308X (no NaN)

…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant