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固有の設定フィールドの主なものを以下に示す。

ここで重要なのが**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を参照していただきたい。