9月3日、Lihuan Zhangら複数の著者が「Enabling DeepSeek-V4-Flash Training on AMD Instinct MI355X GPUs with Primus」と題した記事を公開した。AMD Instinct MI355X GPU上でDeepSeek-V4-Flashの事前学習をエンドツーエンドで動作させるまでの取り組みを詳述した内容で、284B MoEモデルに3種類のアテンション機構を混在させた異色のアーキテクチャを、288GB（利用可能268.2GiB）のHBM3Eに約4%の余裕でギリギリ収める設計の解説が中心となっている。アーキテクチャ詳細・設定方法・カーネル最適化について詳しく紹介されている。
DeepSeek-V4-Flash とは
DeepSeek-AIが2025年4月24日に公開したDeepSeek-V4-Flashは、総パラメータ数284B（アクティベート数13B）、コンテキスト長100万トークンに対応するMixture-of-Experts（MoE）モデルだ。ライセンスはMITで、商用利用も可能。
同モデルの最大の特徴は、43層のTransformerに3種類のアテンション機構を混在させている点にある。また、従来の残差接続（Residual Connection）の代わりに「マニフォールド制約ハイパーコネクション（mHC、manifold-constrained hyperconnection）」を採用するなど、標準的なMegatron-LMのコードが前提とする設計を複数の箇所で逸脱している。これが「動かすだけ」なら難しくないが、「速く動かす」のに大きなエンジニアリングコストがかかる理由だ。
Primus とは
記事中で学習フレームワークとして使われているPrimusは、AMD ROCmエコシステム向けに開発されたオープンソースの大規模モデル学習ライブラリだ。Megatron-LMをベースにAMD GPU向けの最適化を重ねており、YAMLチェーンによるモデル定義・カーネルバックエンドの柔軟な切り替え・メモリ使用量の事前試算（Primus Projection）といった機能を備える。リポジトリはAMD公式GitHubで公開されており、本記事で紹介する設定ファイルや実験コードもそこで参照できる。
アーキテクチャの核心：3種類のアテンション
V4-Flashで最もエンジニアの関心を引くのが、レイヤーごとに異なるアテンション方式を使う設計だ。
各レイヤーにはcompress_ratioと呼ばれる値が設定されており、その値によって以下の3パスのいずれかに振り分けられる。
- Dense Attention（スライディングウィンドウ）：最初の2層のみ。128トークンのウィンドウで通常のアテンションを行うウォームアップ的な存在。
- CSA（Compressed Sparse Attention）：4トークンを1つのKVエントリに圧縮後、"lightning indexer"でスコアリングしてトップ512件を残す。21層に適用。
- HCA（Heavily Compressed Attention）：128トークンを1エントリに圧縮し、選別なしで全エントリに対してアテンションを取る。20層に適用。
なお、CSAで使われるlightning indexerとは、圧縮済みKVエントリに対して軽量なスコアリングを行い、上位k件だけを次のアテンション計算に渡す選別機構だ。フルアテンションを計算する前にエントリ数を絞り込むことで、長文脈時のメモリ・計算コストを大幅に削減する役割を担う。
この設計の効果は数字に如実に現れる。コンテキスト長100万トークン時、HCAレイヤーのクエリが参照するKVエントリは約7,900件、CSAは640件──dense attentionの100万件と比べると桁違いに少ない。これにより、DeepSeek-V4-FlashはDeepSeek-V3.2と比較して**推論FLOPsを約10%、KVキャッシュを約7%**に削減できると報告されている。
なお、CSAはグループをオーバーラップさせる設計のため（前4トークンとの重複）、KV側の射影がHCAの2つに対して4つ必要になる。この非対称性が後述のカーネル実装に影響する。
Primus での設定方法
Primusでは、モデルをYAMLファイルのチェーンとして記述する。V4-Flashの場合、3段階の継承構造になっている。
primus/configs/models/megatron/llama_base.yaml 汎用デコーダーのデフォルト
└─ deepseek_v4_base.yaml V4ファミリー共通設定
└─ deepseek_v4_flash.yaml Flash固有のパラメータ
V4固有の設定フィールドの主なものを以下に示す。
| フィールド | Flash値 | 役割 |
|---|---|---|
compress_ratios |
[0, 0, 4, 128, …] |
レイヤーごとのアテンション種別 |
index_topk |
512 | lightning indexerが保持するエントリ数 |
hc_mult / hc_sinkhorn_iters |
4 / 20 | mHCの残差ストリーム数とSinkhorn-Knoppの反復回数 |
o_groups / o_lora_rank |
8 / 1024 | グループ化低ランク出力射影 |
num_hash_layers |
3 | ハッシュルーティングを使う先頭MoE層数 |
ここで重要なのが**model_type: deepseek_v4**というフィールドだ。この1行がなければ、Primusはその他のパラメータに正しい値を入れてもV3形状のモデルを組み立ててしまう。このフィールドがあってはじめて、deepseek_v4_builders.pyに処理が移り、compress_ratiosから正しいハイブリッドアテンションスタックが構築される。
実験設定ファイル（examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml）ではアテンションカーネルをパスごとに分けて指定できる点も特徴だ。
use_v4_attention_backend: ${PRIMUS_USE_V4_ATTENTION_BACKEND:turbo} # dense/HCA用
use_v4_csa_attention_backend: ${PRIMUS_USE_V4_CSA_ATTENTION_BACKEND:turbo} # CSA用
CSAはindexerとtop-k選択を含むため、dense/HCAとは異なるカーネル問題になる。そのため別々にバックエンドを切り替えられる設計になっている。すべてのパラメータが環境変数経由で渡せるため、1つの最適化を単独でスイープすることも容易だ。
メモリ収支：MI355Xにギリギリ収まる設計
学習クラスタを予約する前に、Primus Projectionでメモリ使用量を事前試算できる。V4-Flashの場合、PP=4・EP=8の構成で以下の結果になる。
- HBM3E搭載量：MI355Xは288GB（利用可能268.2GiB）
- ランク0の使用量：257.62GiB（約4%の余裕）
- 最大バケット：FP32のオプティマイザ状態（グラジェントバッファ＋メインパラメータコピー＋2つのモーメント）で142.67GiB（全体の55.4%）
- BF16ウェイト：20.38GiB（FP32状態の約1/7）
- アクティベーション：94.56GiB
モデルカードに記載の284Bに対して、実際のパラメータ総数は290.80Bだ。差分の6.61BはMTPモジュール（Multi-Token Prediction）で、公表値には含まれていない。
また、MoEレイヤーが全パラメータの**95.7%を占めており、43層のアテンションモジュール合計は1.70%（4.94B）に過ぎない。単一MoEレイヤー内では、256個のルーティングエキスパートだけで99.59%**を占める。
実装上の注意点：ドキュメントの誤りに注意
記事ではmHCの実装について、以下の点を明示的に指摘している。
RMSNormは「崩壊（collapse）の後」に位置するのであって、前ではない。V4のいくつかの公開ダイアグラムはこの順序を逆に描いている。
ポーティング時にはこの点に注意が必要だ。
精度オプション・FP8/FP4について
コード・設定ファイル・ベンチマーク結果はすべてPrimusのオープンソースリポジトリで公開されており、同一環境があれば再現可能とのことだ。本記事で紹介した設定はBF16学習を対象としており、FP8・FP4の混合精度学習への対応については元記事末尾のエンドノートに詳細が記載されている。精度オプションの選択を検討する場合は合わせて参照されたい。
詳細はEnabling DeepSeek-V4-Flash Training on AMD Instinct MI355X GPUs with Primusを参照していただきたい。
