Kernel and Memory Optimization（カーネル・メモリ最適化）

「カタログ FLOPS」と「実 throughput / step time」を分けるのは、framework (tech-73) でも HW (tech-70 / tech-71) でもなく、その間の runtime kernel + memory 層である。HBM 帯域・SMEM tile・mixed precision・activation 再計算・fusion・FlashAttention・GEMM 分解・allocator・CUDA Graph を 10 節 + 8 ピットフォールで俯瞰。サービング (tech-69) / 高速化 (tech-46) / 分散 (tech-48) / optimizer (tech-51) / storage (tech-75) はこの層を前提に動く。

GPU メモリ階層と Roofline

主要 GPU の HBM BW / 容量 / L2 / SMEM/SM / SM / FP16 TFLOPS: A100 SXM4 2.0 TB/s / 80 GB / 40 MB / 164 KB / 108 / 312、H100 SXM5 3.35 TB/s / 80 GB / 50 MB / 228 KB / 132 / 989、H200 SXM5 4.8 TB/s / 141 GB / 50 MB / 228 KB / 132 / 989、B100/B200 (未確認) ~8 TB/s / 192 GB / ~200 MB / ~4.5 PFLOPS FP8、MI300X ~5.3 TB/s / 192 GB / 304 CU (L2/MALL ~256 MB 未確認)。Register file A100/H100 とも 256 KB/SM (65,536×32-bit)。Roofline: AI=FLOPs/B が ridge (peak FLOP/s ÷ peak BW) 未満で memory-bound、H100 BF16 ridge ≈ 295 FLOP/B。decode (batch=1) AI ≈ 1 で BW-bound、prefill は大 batch × seq で compute-bound。HBM 2×/世代対し compute 4–8×/世代で memory wall 拡大、本稿最適化の根拠。

メモリレイアウト — coalescing / bank / contiguity

Warp 32 thread が連続 32×dtype B 読むと L2 line 128 B + HBM sector 32 B (coalescing)、stride 128 B では 32 別 transaction → 32× 浪費。cuBLAS column-major、PyTorch row-major で torch.mm(A,B) 内部 (B^T)(A^T)。.t()/.permute() 後の非 contiguous は .contiguous() 暗黙 copy。cuDNN conv NHWC (memory_format=torch.channels_last) で A100+ TensorCore 20–30% 速 (未確認)。SMEM 32 banks × 4 B/cycle、warp 全 thread 同 bank → 32-way bank conflict で 32× 遅、回避は padding と CUTLASS swizzle。FA は SMEM swizzle で bank conflict 消去。

混合精度 — FP16 / BF16 / FP8 / FP4 / INT4

TensorCore: V100 (FP16 HMMA) → A100 (TF32/FP16/BF16/INT8/INT4, mma.sync 16×8×16) → H100 (FP8 E4M3/E5M2 + 非同期 WGMMA 64×N×K) → Blackwell (FP4 MX, E2M1, 未確認)。BF16 exp 8 bit で FP32 同レンジ→loss scaling 不要、LLM 訓練 default。FP16 exp 5 bit は GradScaler 動的調整。FP8 Transformer Engine: E4M3 (max 448) forward 重み・活性、E5M2 (max 57344) backward 勾配。per-tensor より per-block (128×128) scale が精度高く DeepSeek-V3 (arXiv:2412.19437) が H800 native FP8 訓練成立。OCP MXFP8/MXFP4 8/16 要素ブロック共有指数で標準化、FP4 は FP8 比 2× (未確認)。PTQ は GPTQ (arXiv:2210.17323) / AWQ (arXiv:2306.00978) / SmoothQuant (arXiv:2211.10438, activation→weight W8A8) で tech-46、optimizer FP8 は tech-51。

Activation メモリと再計算

Transformer 1 層 activation ≈ seq × batch × hidden × (34 + 5×heads×seq/hidden) B (Korthikanti 2022, arXiv:2205.05198)、Llama-3 70B (hidden=8192, L=80, seq=4096, batch=1) 素 ~520 GB。Chen 2016 (arXiv:1604.06174) の full checkpoint は O(√N) memory + 再計算 +33% FLOPs。Selective recomputation は softmax / dropout / QK^T のみ捨て activation 30–40% 保持 + 再計算 +5–10%。FA は Q,K,V,O,L のみ保持で N×N 不要、backward は L から tile 再計算。PyTorch checkpoint(fn, x, use_reentrant=False) (FSDP2 / torch.compile 互換)。memory 予算は tech-48、本稿は機構面担当。

カーネル融合 — TorchInductor / XLA / Triton

launch overhead ~5–10 μs/call、decode 80 層 × 20 ops × 7 μs ≈ 11 ms 対し計算 0.5 ms → 22× オーバヘッド。fusion は中間 tensor を HBM 書戻さず SMEM/register で済ます事。GEMM→ReLU 分離 3N MB HBM、融合 N MB。TorchDynamo + TorchInductor graph 捕獲 + Triton emit、max_autotune で algo 探索 (pointwise / horizontal / vertical 自動)。XLA instruction + output fusion (TPU で特に効く)。LayerNorm Apex FusedLayerNorm で 3→2-pass、RMSNorm 単 pass、Softmax は Milakov 2018 (arXiv:1805.02867) online softmax で 1-pass (FA の核)。cuBLASLt / CUTLASS は GEMM + bias + ReLU/GELU + residual を epilogue 融合。vLLM fused MoE routing + grouped GEMM + gather 1 Triton kernel 2–4× (未確認)、Mamba selective scan は並列 prefix scan で PyTorch 比 ~10× (未確認)。

FlashAttention 系列と PagedAttention

FA-1 (Dao 2022, arXiv:2205.14135) は online softmax を 2D tile 拡張、N×N softmax 不要で O(N) memory + SMEM tiling。FA-2 (Dao 2023, arXiv:2307.08691) Q outer loop で SM 占有率向上、非 matmul FLOP 削減 + GQA で A100 ~73% peak。FA-3 (Shah 2024, arXiv:2407.08608, H100/H200 のみ) は TMA (HBM→SMEM 非同期 DMA、N-D stride) + WGMMA (4 warp × 128 thread で 64×N×K) + 2-stage 二重 buffer の perfect pipelining、H100 BF16 ~75% peak + FP8 path。PagedAttention (Kwon 2023, arXiv:2309.06180) は serving KV を 16-token block + logical table で OS 仮想記憶風管理、無駄 60–80%→<4%、continuous batching + prefix cache 可能化。Flash-Decoding KV 長軸並列で長文 decode 8–50× (未確認)、PyTorch 2.5+ FlexAttention は lambda → Triton codegen で ALiBi / SW / 任意 mask。Ring Attention (arXiv:2310.01889) は tech-48、serving は tech-46。

GEMM 最適化 — CUTLASS / WGMMA / TMA

LLM の QKV / output / MLP / lm_head 全て GEMM で prefill FLOPs の ~90%。cuBLAS 自動選択、cuBLASLt algo 列挙 + epilogue 融合制御。CUTLASS は CTA tile (128×128 SMEM 常駐) → Warp tile (mma.sync / WGMMA) → Thread tile (register) 3 階層 + SMEM 二重 buffer で「計算しつつ次 tile load」、3.x CuTe Layout<Shape, Stride> が任意 N-D layout を型表現。Hopper WGMMA 4 warp 協調で mma.sync (16×8×16) より granularity 大、TMA 任意 stride N-D copy が async で FA-3 overlap 元。SplitK GEMM は M×N 小・K 大 (decode) で K 分割 + 部分 reduce、batch=1 で SM 使い切り。GroupedGEMM (CUTLASS GemmGrouped) は MoE 不均衡 M を 1 launch (2–5× 未確認)。Megatron-LM TP column/row-parallel の AllReduce 2→1 は tech-48。

メモリアロケータと expandable_segments

PyTorch は cudaMalloc/Free 毎回呼ばず caching allocator で free block を pool 保持。サイズ多様化で 断片化「10 GB 空いて OOM」、兆候 memory_reserved() >> memory_allocated()、確認 memory_stats() の fragmentation.*。緩和は PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 と PyTorch 2.1+ expandable_segments:True (未確認)。後者は CUDA virtual memory API (cuMemAddressReserve/cuMemCreate/cuMemMap, CUDA 10.2+) で仮想連続+物理分割 segment 伸縮、seq_len 変動 / 動的 batch serving の OOM 激減。cudaMallocAsync stream-ordered で CPU-GPU 同期不要 (PyTorch 部分採用, 未確認)。empty_cache() は pool を driver に返すが単一 forward+backward peak は減らず (call 間のみ) は典型誤解。

CUDA Graph と launch オーバヘッド削減

cudaLaunchKernel CPU 側 ~5–10 μs、GEMM 大物で誤差だが decode 小物連発で支配的。CUDA Graph は launch/memcpy 列を record → instantiate → replay で 1 μs 未満に圧縮、torch.cuda.CUDAGraph() capture + graph.replay()。torch.compile(mode="reduce-overhead") 自動 capture、max-autotune GEMM algo 探索 + cudnn.benchmark 追加。制約: shape static、Python 制御フロー不可、memory address replay 間固定 → vLLM/TRT-LLM/SGLang は batch=1,2,4,8,…,max ごと graph pool 保持。warmup 2–3 step + cold start 課題。CUDA 12.4+ conditional graph node が限定 branching 許可で speculative decoding 1 graph 化 (未確認 API)。serving は tech-69。

ピットフォールとフロンティア

P1 長文 OOM: FA で attention O(N) でも MLP/QKV は O(L × seq × hidden)、seq=128K + 70B 素 ~100 GB → FA + selective recompute 必須。P2 TF32 silent drift: PyTorch 1.7+ CUDA matmul TF32 (mantissa 10 bit) default、CPU FP32 と数値差、deterministic は cuda.matmul.allow_tf32=False + cudnn.allow_tf32=False。P3 CUDA Graph pool OOM: graph 捕獲時 peak を pool 予約 → batch 多数で additive、pool 制限 + max_split_size_mb。P4 retain_graph リーク: backward(retain_graph=True) 残置で graph 無限伸長、memory_allocated() 単調増加で検知。P5 backward peak forward 2×: activation + grad 2–3×、checkpoint で compute +33%、backward もプロファイル。P6 ZeRO-Offload PCIe 律速: PCIe 4.0 32 GB/s 対 HBM 3.35 TB/s で 100× 差、70B (~140 GB) 毎 step 転送 4.4 秒/step 実用外、7–13B が上限。ZeRO-Infinity (NVMe ~7 GB/s) は load 専用。P7 paged 8-bit optimizer: bitsandbytes.PagedAdamW8bit (m,v) per-block INT8 + OOM 時 CPU ページング、70B Adam states 56→14 GB、+5% 時間。P8 torch.compile graph break: dynamic shape / print / data-dep if で graph 分裂 → 融合機会喪失、TORCH_COMPILE_DEBUG=1 + torch.compiler.disable(fn)。フロンティア: B200 native FP4、DeepSeek-V3 per-block FP8、OCP MXFP4/MXFP6、SMEM 拡大 FA-N、CUDA 12.4+ conditional graph、cutoff 後再確認。

Closes iAmTheSuperHiro/magi-system#372.