Metal グラフと圧縮 KV の実装を読む: compressor、indexer、HC、tail-only RoPE【第3回/全8回】
本シリーズは、DeepSeek V4 Flash / Pro 専用推論エンジン DwarfStar（
ds4）のコードを読み解く連載です。
第3回は、推論エンジンの心臓部である Metal グラフと圧縮 KV キャッシュの実装を追います。主な参照箇所:
ds4.c,ds4_metal.m,metal/dsv4_kv.metal,metal/dsv4_misc.metal,metal/dsv4_hc.metal,metal/flash_attn.metal
連載「DwarfStar(ds4) を読む」全8回
- 第1回 なぜ専用エンジンを書くのか
- 第2回 非対称2bit量子化とimatrix
- 第3回 Metalグラフと圧縮KV（本記事）
- 第4回 ディスクKVキャッシュ
- 第5回 サーバとDSMLツール呼び出し
- 第6回 TCPパイプライン分散推論
- 第7回 ネイティブエージェント
- 第8回 ステアリング・MTP・評価基盤
TL;DR
- DS4 のアテンションは、単純な全 KV キャッシュではない。raw sliding-window KV、圧縮 KV、ratio-4 層用の indexer を組み合わせる。
- Flash の圧縮率は固定ではない。
ds4_expected_layer_compress_ratio()によると、Flash は層 0-1 が非圧縮、以降は偶数層が ratio 4、奇数層が ratio 128。Pro は層 0-1 も ratio 128 で、その後は同じ交互構成。 - compressor は単純平均ではない。
wkvとwgateで KV 候補とスコアを作り、APE を足し、softmax 加重プーリングして圧縮行を作る。 - ratio-4 層では、アテンション用 compressor と indexer 用 compressor の 2 レーンを持つ。indexer は圧縮行全体から top-k を選び、FlashAttention 側は raw SWA と選ばれた圧縮行を同時に attend する。
- DS4 Flash は Hyper-Connection
n_hc=4を持ち、Metal 側では Sinkhorn 正規化済みの 4x4 混合行列を作る専用カーネルがある。 - RoPE はヘッド全体ではなく末尾 64 次元だけにかかる。圧縮層では通常の RoPE base ではなく圧縮用 base を使う。
1. まず「圧縮率」は層ごとに違う
前回までの説明では「4 トークンごとに 1 圧縮行」という ratio-4 を中心に見てきました。しかしコードを読むと、DS4 の圧縮はそれだけではありません。
ds4.c の ds4_expected_layer_compress_ratio() は、モデルの variant ごとの期待圧縮率を明示しています。
static uint32_t ds4_expected_layer_compress_ratio(uint32_t il) {
switch (DS4_MODEL_VARIANT) {
case DS4_VARIANT_FLASH:
if (il < 2) return 0;
return (il & 1u) == 0 ? 4u : 128u;
case DS4_VARIANT_PRO:
if (il < 2) return 128u;
return (il & 1u) == 0 ? 4u : 128u;
}
}
Flash の場合:
| 層 | 圧縮 |
|---|---|
| 0-1 | なし |
| 偶数層 2,4,6,... | ratio 4 |
| 奇数層 3,5,7,... | ratio 128 |
Pro の場合:
| 層 | 圧縮 |
|---|---|
| 0-1 | ratio 128 |
| 偶数層 2,4,6,... | ratio 4 |
| 奇数層 3,5,7,... | ratio 128 |
この値は単にコード側で仮定しているだけではなく、GGUF メタデータの deepseek4.attention.compress_ratios と照合されます。不一致なら起動時に落ちます。
validate_compress_ratio_metadata(m);
DS4 は「この GGUF ならたぶん動く」ではなく、「この形状・このメタデータ・このテンソルレイアウトでなければ動かさない」設計です。専用エンジンらしい割り切りです。
2. KV キャッシュの基本構造
CPU 経路のコメントが、DS4 の KV キャッシュを理解する最短ルートです。
* The CPU path is the correctness reference. It maintains raw SWA KV rows,
* optional compressed KV rows, the indexer mask for ratio-4 layers, and a
* reusable decode scratch arena ...
実体は ds4_layer_cache です。
typedef struct {
float *raw_kv;
uint32_t n_raw;
uint32_t cap_raw;
uint32_t compress_ratio;
uint32_t comp_cap;
uint32_t n_comp;
float *attn_comp_kv;
float *attn_state_kv;
float *attn_state_score;
uint32_t n_index_comp;
float *index_comp_kv;
float *index_state_kv;
float *index_state_score;
} ds4_layer_cache;
ここから、各層が持つ状態は大きく 3 種類だと分かります。
| 状態 | 役割 |
|---|---|
raw_kv |
直近トークン用の sliding-window KV |
attn_comp_kv |
アテンション本体が読む圧縮 KV |
index_comp_kv |
ratio-4 層で圧縮行を選ぶための indexer KV |
raw_kv は全コンテキスト分を持ちません。DS4_N_SWA、つまり Flash の形状では 128 トークン分が上限です。
static uint32_t ds4_default_raw_cap(uint32_t ctx_size) {
uint32_t raw_cap = DS4_N_SWA;
if (raw_cap > ctx_size) raw_cap = ctx_size;
...
}
これにより、直近の局所文脈は高精度に保持し、古い文脈は圧縮行に畳み込む構造になります。
3. prefill は 4096 トークンチャンクが基本
長いプロンプトを一括で Metal グラフに流すと、スクラッチや一時バッファが巨大になります。DS4 は長いプロンプトの prefill をチャンク化します。
ds4_default_prefill_cap_for_prompt() は、環境変数 DS4_METAL_PREFILL_CHUNK がなければ、プロンプトが 4096 トークンを超える場合にチャンク上限を 4096 にします。
const char *env = getenv("DS4_METAL_PREFILL_CHUNK");
...
} else if (prompt_len > 4096) {
cap = 4096u;
}
README の Benchmarking 節も、長いプロンプトは既定で 4096 トークンチャンクで prefill されると説明しています。チャンクサイズを変えると KV チェックポイント / logits の経路が変わるので、比較時には実行設定として明示する必要があります。
このチャンク化は単なるメモリ節約ではありません。compressor/indexer の絶対位置境界を維持しながら、同じ範囲対応の layer-major グラフを再利用するための設計です。
4. compressor は「平均」ではなく学習済み加重プーリング
圧縮 KV と聞くと、単純に 4 トークンを平均して 1 行にしているように想像しがちです。DS4 の compressor はそうではありません。
デコード参照実装 compressor_decode_one() は、1 トークンごとに次の処理を行います。
-
wkvで圧縮 KV 候補を作る -
wgateでスコアを作る - APE、つまり学習済みの位置バイアスをスコアに足す
- 再帰状態に
kvとscoreを格納する - ratio 境界に来たら softmax 加重プーリングで 1 圧縮行を作る
- RMS 正規化をかける
- tail-only RoPE をかける
- アテンション圧縮 KV なら FP8 量子化の往復、indexer 圧縮 KV なら QAT の往復を通す
スコア + APE の格納は Metal カーネルでも専用化されています。
kernel void kernel_dsv4_compressor_store_one(
...,
device const float * kv,
device const float * score,
device const char * ape,
device float * state_kv,
device float * state_score,
uint gid [[thread_position_in_grid]]) {
...
state_kv[dst] = kv[gid];
state_score[dst] = score[gid] + ape_v;
}
プーリングは softmax 加重平均です。compressor_pool_decode_state() は次元ごとにスコアの最大値を取り、exp(score - max) を重みとして KV を平均します。
const float w = expf(state_score[r * width + j] - max_score);
denom += w;
sum += w * state_kv[r * width + j];
out[j] = denom > 0.0f ? sum / denom : 0.0f;
Metal 側にも融合カーネル kernel_dsv4_softmax_pool があります。
for (int64_t ir = 0; ir < args.ne00; ++ir) {
const float s = ... score ...
const float w = exp(s - max_s);
const float v = ... kv ...
sum += w;
acc += v*w;
}
dst = acc/sum;
したがって、compressor は「学習済みスコアによる次元ごとのプーリング」です。これは長文脈を捨てるのではなく、モデルが学習した圧縮表現に変換していると見るべきです。
5. ratio-4 層は 2 レーンの状態を持つ
ratio 4 の層は特別です。compressor_pool_decode_state() のコメントにある通り、ratio-4 層ではアテンション圧縮と indexer 圧縮の 2 レーンを持ちます。
/* Ratio-4 layers keep two lanes: attention compression and indexer compression. */
const uint32_t coff = compress_ratio == 4 ? 2u : 1u;
const uint32_t width = coff * head_dim;
ratio 4 の場合、状態は 4 行ではなく 8 行相当を扱います。前半と後半を持つ回転 frontier です。圧縮行を吐いた後、Metal カーネル kernel_dsv4_ratio4_shift_f32 が後半を前半に移します。
ratio-4 frontier（8 行）の流れを図にすると次の通りです。
frontier (ratio-4, 8 行) ── 4 トークンごとに 1 圧縮行を emit
┌─────────── 前サイクル ───────────┐┌─────────── 現サイクル ───────────┐
│ row0 row1 row2 row3 ││ row4 row5 row6 row7 │
└──────────────────────────────────┘└──────────────────────────────────┘
↑ pos%4 で row4..7 に書き込み
境界 (pos+1)%4==0 で:
(1) 8 行を softmax 加重プーリング → 1 圧縮行
(2) RMS 正規化 → tail RoPE → FP8 往復 → attn_comp_kv へ append
(3) row0..3 ← row4..7 にシフト (kernel_dsv4_ratio4_shift_f32)
// Ratio-4 compression keeps two 4-row halves of recurrent state.
kernel void kernel_dsv4_ratio4_shift_f32(
...,
device float * state_kv,
device float * state_score,
uint gid [[thread_position_in_grid]]) {
const uint n = 4u * args.width;
state_kv[gid] = state_kv[n + gid];
state_score[gid] = state_score[n + gid];
}
この 2 レーン構造は、アテンション本体が読む圧縮 KV と、indexer が読む圧縮 KV を同じ ratio 境界で更新するために必要です。
6. indexer は圧縮行から top-k を選ぶ
Flash の形状では n_indexer_head = 64, n_indexer_head_dim = 128, n_indexer_top_k = 512 です。通常のアテンションヘッドの次元は 512 ですが、indexer は 128 幅の細い表現で圧縮行をスコアリングします。
ds4.c のテンソル検証では、ratio 4 の層だけ indexer 関連テンソルを要求します。
if (compress_ratio == 4) {
l->indexer_attn_q_b = required_tensorf(...);
l->indexer_proj = required_tensorf(...);
l->indexer_compressor_ape = required_tensorf(...);
l->indexer_compressor_kv = required_tensorf(...);
l->indexer_compressor_gate = required_tensorf(...);
l->indexer_compressor_norm = required_tensorf(...);
}
Metal 側にはデコード専用の kernel_dsv4_indexer_score_one_direct があります。1 圧縮行をスレッドグループに置き、64 個の indexer ヘッドを 4 ヘッドのグループで処理してスコアを作る、とコメントされています。
indexer の出力は「どの圧縮行をアテンション本体に見せるか」です。全圧縮行に密にアテンションするのではなく、選ばれた top-k 行だけを見ることで、1M コンテキストでもアテンションのコストを抑えます。
7. raw SWA + 選ばれた圧縮行を同時に attend する
ratio-4 経路のアテンションは、kernel_dsv4_indexed_mixed_attention_heads8 が分かりやすいです。コメントに「DS4 ratio-4 indexed mixed attention」とあり、密な top-k マスク経路を置き換えるカーネルです。
処理は大きく 2 段階です。
- 直近の raw sliding-window 行を attend する
- top-k で選ばれた圧縮行を attend する
raw 側はウィンドウを見て、qpos から見える範囲だけを処理します。
const uint window_first = (args.window != 0u && qpos + 1u > args.window) ?
qpos + 1u - args.window : 0u;
uint first = max(first_raw_pos, window_first);
uint last = min(qpos, raw_last_pos);
圧縮側は indexer が作った topk を読みます。
for (uint i = 0; i < args.top_k; i++) {
const int32_t idx = row_topk[i];
if (idx < 0) continue;
if ((uint)idx >= visible) break;
...
attend comp row idx
}
このカーネルは online softmax の形で M と S を更新しながら、raw 行と圧縮行を同じアテンション出力に統合します。直近は raw、古い文脈は選ばれた圧縮行、という設計がカーネルの中に直接現れています。
8. tail-only RoPE と圧縮層用の base
DeepSeek V4 の RoPE はヘッド全体ではなく、ヘッドの末尾 n_rot = 64 次元にだけかかります。Flash のヘッド次元は 512 なので、先頭 448 次元は非 RoPE、末尾 64 次元だけが回転します。
rope_tail_ext_inplace() はその名の通り末尾（tail）を指します。
const uint32_t n_nope = head_dim - n_rot;
float *tail = x + h * head_dim + n_nope;
さらに DS4 は、圧縮層では通常の DS4_ROPE_FREQ_BASE ではなく DS4_COMPRESS_ROPE_FREQ_BASE を使います。
static float layer_rope_freq_base(uint32_t il) {
return ds4_layer_compress_ratio(il) != 0 &&
DS4_COMPRESS_ROPE_FREQ_BASE > 0.0f
? DS4_COMPRESS_ROPE_FREQ_BASE
: DS4_ROPE_FREQ_BASE;
}
圧縮層で YaRN 拡張を使う条件もここに入っています。長文脈を成立させるための位置処理が、dense 層と圧縮層で分岐している点が重要です。
9. 圧縮 KV は FP8 往復、indexer は FP4/QAT 往復
compressor_decode_one() は圧縮行を作った後、ヘッド次元に応じて別の往復を通します。
if (head_dim == DS4_N_HEAD_DIM) {
dsv4_fp8_kv_quantize_row_inplace_cpu(out_comp, head_dim, DS4_N_ROT);
} else if (head_dim == DS4_N_INDEXER_HEAD_DIM) {
dsv4_indexer_qat_row_inplace_cpu(out_comp, head_dim);
}
通常のアテンション圧縮 KV は 512 幅で、非 RoPE 部分に E4M3 形式の FP8 往復をかけます。indexer 圧縮 KV は 128 幅で、Hadamard 128 と FP4 の活性量子化を通します。
これは「保存形式だけを軽くする」ためではなく、公式 DeepSeek V4 グラフと一致させるための往復です。ds4.c のコメントにも、indexer Q と indexer compressor KV に QAT 往復が必要で、これがないと top-k 圧縮行選択がモデルのグラフではなくなる、と書かれています。
10. Hyper-Connection: 4 本のストリームを Sinkhorn で混ぜる
Flash の形状では n_hc = 4, n_hc_sinkhorn_iter = 20 です。DS4 は隠れ状態を単一ストリームとして持つのではなく、Hyper-Connection の 4 ストリームを扱います。
Metal カーネル kernel_dsv4_hc_split_sinkhorn は、ミキサ行を以下に分解します。
- pre weights（事前重み）
- post gates（事後ゲート）
- HC 間の結合行列
HC=4 の高速経路では、4x4 行列を作ったあと Sinkhorn 正規化を繰り返します。
for (int iter = 1; iter < args.sinkhorn_iters; ++iter) {
r0 *= 1.0f / (row_sum0 + epsv);
...
col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv);
r0 *= col_inv;
...
}
Sinkhorn 正規化は、行方向・列方向の正規化を交互に行い、行列を二重確率行列に近づける手法です。ここでは 4 本の残差ストリームをどのように混ぜるかを、学習済みミキサから作っています。
DS4 が「普通の Transformer ランナー」ではないことは、このあたりのモデル固有グラフによく表れています。
11. Metal ランタイムは専用パイプラインの集合
ds4_metal.m の冒頭には、専用パイプライン状態が大量に並んでいます。
static id<MTLComputePipelineState> g_hc_split_sinkhorn_pipeline;
static id<MTLComputePipelineState> g_dsv4_fp8_kv_quantize_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexer_qat_pipeline;
static id<MTLComputePipelineState> g_dsv4_ratio4_shift_pipeline;
static id<MTLComputePipelineState> g_dsv4_softmax_pool_pipeline;
static id<MTLComputePipelineState> g_dsv4_topk_mask_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexer_score_one_direct_pipeline;
static id<MTLComputePipelineState> g_dsv4_compressor_store_one_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads8_pipeline;
この一覧だけでも、DwarfStar が「汎用テンソルグラフ実行器」ではなく、DeepSeek V4 のホットパスを直接カーネル化していることが分かります。
例えば compressor store は汎用グラフなら APE コピー・スコア加算・set_rows を複数ディスパッチするところを、kernel_dsv4_compressor_store_one で 1 ディスパッチにしています。ratio-4 shift も汎用コピーではなく専用カーネルです。
この種の特化は、保守対象モデルを広げるほど重荷になります。だから DS4 は第1回で見た通り、「一度に1モデル」という狭い賭けを選んでいます。
12. この記事の要点
第1回では DS4 の圧縮アテンションを概念として見ました。コードまで読むと、重要なのは次の点です。
- 圧縮率は 4 だけではなく、Flash では 0/4/128 が層ごとに混在する
- raw KV は全コンテキストではなく SWA 128 行
- compressor は学習済みスコア + APE + softmax プーリング
- ratio-4 層だけ indexer 経路を持ち、top-k 圧縮行を選ぶ
- アテンションカーネルは raw 行と選ばれた圧縮行を同じ online softmax で統合する
- RoPE、FP8/QAT 往復、HC Sinkhorn まで DeepSeek V4 専用グラフとして実装されている
次回は、この KV 状態をどう保存し、どう再利用するかを扱います。DwarfStar の README が言う「KV cache is a first-class disk citizen」は、単なる思想ではなく、DSV4 ペイロードと KVC ストアという具体的なファイル形式に落ちています。
本記事は Quick Iterate のローカル LLM 研究の一環として、公開リポジトリ antirez/ds4 のコードを読み解いたものです。行番号・定数・ベンチ値は閲覧コミット ba00a8a（2026-05-30）／README 取得日 2026-06-01 時点のものです。ds4-agent は alpha、エンジン本体は beta 品質で活発に変化するため、引用箇所は各自で最新の README / ソースに当たって再確認してください。
クイックイタレート株式会社 ― IoT / 電力監視 / AI / 衛星・無線通信 / システムインテグレーション
ローカル LLM・エージェント基盤に関するお問い合わせはお気軽にどうぞ。