Cross-posting from internal kernel-team analysis by @carlushuang for broader visibility. The source thread tracks per-kernel ownership and ETAs internally; this issue is the public-facing technical inventory only.

Executive Summary

This issue inventories every kernel/operation needed to run DeepSeek-V4 (Flash + Pro) on AMD MI3xx/MI350X/MI355X, cross-referenced against:

• HF reference (deepseek-ai/DeepSeek-V4-Flash) — TileLang kernels
• vLLM implementation — FlashMLA/FlashInfer + custom fused kernels + multi-stream
• Current AMD stack — ATOM (model), aiter (kernels), CK (GEMM)

Bottom line: 5 kernels are functionally missing on AMD (sparse attention, compressor, Hadamard, Sinkhorn, FP4 upcast-at-load). 3 fusion opportunities are proven by vLLM but not yet implemented in aiter. FP4 expert GEMM works via FusedMoE but needs a native/upcast path for peak throughput.

Model Config Summary

Notation: In all kernel tables below:
- T = number of input tokens (sequence length)
- H = number of attention heads (H=64 for Flash, H=128 for Pro)
- G = number of O-projection groups (G=8 for Flash, G=16 for Pro)
- SK = number of gathered KV tokens per query (see sparse_attn row for case-by-case formulas)
- max_B = maximum batch size (number of concurrent sequences the pre-allocated buffers can hold)
- max_S = maximum sequence length (up to 1,048,576 for 1M context; determines KV cache and compressor buffer sizes at init)
Where Flash and Pro differ, both shapes are shown. Where only one shape is shown, both variants use the same dimensions.
Note: max_B and max_S appear in statically pre-allocated buffers (HF/ATOM eager mode). vLLM/SGLang use paged KV cache with dynamic block allocation instead.

For per-query KV token counts, see ATOM#664 Section 7.

1. MLA / Attention Block

1.1 sparse_attn Deep Dive (P0 — Biggest Gap)

What it does: Index-gather attention — NOT standard FlashAttention.

# For each (batch, seq_pos):
#   1. Gather KV[topk_idxs] by index (topk = 640 Flash CSA, 160 Flash HCA at 4K)
#   2. Q @ KV_gathered^T → scores
#   3. Online softmax (running max + sum exp)
#   4. Add attn_sink bias AFTER normalization
#   5. Softmax scores @ KV_gathered → output

Why generic FA won't work: Dense FA mask is O(S²) memory — prohibitive for 1M context. The gather indices are sparse and query-dependent.

vLLM approach: Integrates FlashMLA/FlashInfer backends but still requires custom sparse-gather logic.

AMD need: Custom aiter kernel with:

• Index gather by topk_idxs (can be -1 for invalid)
• Online softmax over gathered KV
• attn_sink as pseudo-token added to denominator
• Support both Flash (topk=512) and Pro (topk=1024) for CSA; HCA uses topk=8192 (effectively all entries for ≤1M)

Perf target: Memory-bound on KV gather. Optimize gather coalescing above compute.

2. Indexer Block

The Indexer runs on a separate CUDA stream in vLLM, overlapped with main attention prep.

2.1 Indexer Memory at 1M Context

Important: At full 1M prefill, index_score = [1, 1M, 250K] at FP32 ≈ 1 TB — infeasible in a single pass. Chunked prefill is mandatory. The table below shows per-chunk shapes at a typical chunk size of T=4K (matching ATOM#664 examples), plus the full-context column for reference.

index_score is the largest intermediate. FP4 precision for Q/KV is essential, and chunked prefill caps the per-chunk index_score to manageable sizes.

3. Compressor Block

There are two Compressor instances per CSA layer (Attention + Indexer) and one per HCA layer.

Three compressor variants share the same pipeline (project → ape → reshape → pool → norm → RoPE → quant → cache write) but differ in head_dim, compress_ratio, overlap, rotation, and output quantization. All shapes are shown in one table below.

3.1 KV Cache Shapes

3.2 Persistent Buffer (kv_state / score_state) — Decode Only

vLLM treats compressor state as sliding-window KV with sliding_window = coff * compress_ratio, registered under the same hybrid KV cache manager. Shift operation after each compression: kv_state[:, :ratio] = kv_state[:, ratio:] (current becomes previous).

4. MoE Block

4.1 FP4 GEMM Deep Dive (P0)

HF reference:

# C[M,N] = A_fp8[M,K] @ B_fp4[N,K]^T
# Act scale:  per-1×128 on K-dim, E8M0/FP32
# Weight scale: per-1×32 on K-dim, E8M0
# B stored as [N, K//2] — 2 FP4 values packed per byte along K

vLLM: Ships with native FP4 MoE weights. Uses custom FP4 backend (likely upcast-at-load).

AMD current: aiter FusedMoE works via:

• gfx942 (MI300X): ASM/CK path
• gfx950 (MI355X): Triton path with CDNA4MXScaleLayout
• swiglu_limit applied in triton post-kernel (not in fused SwiGLU — 9× amplitude loss bug)

AMD gap: No native FP4 tensor core on MI3xx confirmed. Options:

1. Upcast-at-load (like TileLang): Load FP4 sub-blocks → upcast to FP8 in shared mem → run CK GemmMXFP8. Recommended for bring-up.
2. Pre-upcast weights to FP8: ~2× weight memory overhead.
3. Native FP4: Verify MI350+ sub-byte support first.

Weight shapes (V4-Flash):

w13_weight:       (257, 4096, 2048)   torch.float4_e2m1fn_x2   # 256 routed + 1 shared
w13_weight_scale: (257, 4096, 128)    torch.float8_e8m0fnu     # per-1×32 along K
w2_weight:        (257, 4096, 1024)   torch.float4_e2m1fn_x2
w2_weight_scale:  (257, 4096, 64)     torch.float8_e8m0fnu

Weight shapes (V4-Pro):

w13_weight:       (385, 6144, 3584)   torch.float4_e2m1fn_x2   # 384 routed + 1 shared
w13_weight_scale: (385, 6144, 224)    torch.float8_e8m0fnu     # per-1×32 along K
w2_weight:        (385, 7168, 1536)   torch.float4_e2m1fn_x2
w2_weight_scale:  (385, 7168, 96)     torch.float8_e8m0fnu

5. mHC (Manifold-Constrained Hyper-Connections) Block

5.1 Sinkhorn Bottleneck

Cost per forward: 20 iterations × 43 layers (Flash) = 860 loops.

• PyTorch fallback: ~100–500µs per call → 86–430ms total
• aiter mhc_pre kernel: ~1–5µs per call → 0.9–4.3ms total

aiter's mhc_pre is already kernelized. The main blocker was device-mismatch bug (#2916, now fixed). No new kernel needed here — just the bugfix merge.

5.2 Fused Block Kernel (Future Optimization)

vLLM has not yet fused the full block, but it's the next logical step:

__global__ void block_forward(...) {
    // 1. hc_pre: Sinkhorn → 1 copy
    // 2. RMSNorm
    // 3. Attention or MoE
    // 4. hc_post: combine with 4-copy residual
}

This would eliminate 4 kernel launches per layer.

6. Quantization & Utility Kernels

7. vLLM-Proven Fusions (Not Yet in aiter)

vLLM has deployed these fusions with measured speedups. They are targets for aiter/CK on AMD:

8. AMD Gap Summary & Priority Matrix

Missing Kernels (Need New Implementation)

Working but Need Hardening

9. Kernel Implementation Inventory

Implementation approach legend:

• ASM — Pure GPU assembly (hand-optimized microkernels) - contact: @niels-zhang
• flyDSL — AMD DSL for declarative patterns (elementwise, reduction, transforms) - contact: @coderfeli
• Triton/Gluon — Triton-based kernels (aiter’s fused-op path) - contact: @vgokhale
• HIP/C++ — Raw HIP kernel, Opus C++ kernel, or Composable Kernel (CK) - contact: @valarLip @carlushuang

(1) Attention

• sparse_attn — ❌ Torch fallback (P0 gap)
  • Index-gather + online softmax + attn_sink bias over variable topk gathered KV entries
  • NOT FlashAttention — sparse, query-dependent gather indices
  • Must handle -1 invalid indices, attn_sink as pseudo-token in softmax denominator
  • See §1.1 for full spec
  • ☑️ ASM    ☑️ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

(2) Compression

• Indexer/Compressor kernel

  • CSA Flash: Q[T,64,512] × KV[640,512], CSA Pro: Q[T,128,512] × KV[1152,512]
  • HCA: Q[T,H,512] × KV[128+T/128, 512] (grows with context)
  • ⬜ ASM    ☑️ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

• Indexer/Compressor misc (linear, elem-wise, fusions) — ❌ Torch fallback

  • CSA (ratio=4, overlap=True): wkv/wgate FP32 GEMM → +ape broadcast → overlap transform [T/4,4,1024]→[T/4,8,512] → softmax gated pooling → RMSNorm → RoPE → act_quant FP8 → cache write (§3.1)
  • HCA (ratio=128, overlap=False): same but no overlap transform, pool over 128 tokens (§3.3)
  • Indexer compressor (head_dim=128, rotate=True): same pipeline but Hadamard → FP4 quant instead of split FP8 (§3.2)
  • vLLM fuses entire pipeline into one kernel (~1.4–3× speedup)
  • ⬜ ASM    ☑️ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

(3) KV Cache Codesign

DeepSeek-V4 uses a hybrid KV cache with three storage paths:

Buffer mechanics:

• kv_state / score_state double-buffer for CSA: [max_B, 8, 256] or [max_B, 8, 1024]
• HCA rolling window: [max_B, 128, 512]
• vLLM integrates both under a single hybrid KV cache manager with sliding_window = coff * compress_ratio
• ⬜ ASM    ⬜ flyDSL    ⬜ Triton/Gluon    ☑️ HIP/C++

Source: §3.3–3.4 for detailed shape traces and vLLM integration

(4) MoE

• A8W4 (FP8 act × FP4 weight) expert GEMM — includes SwiGLU + clamp — ⚠️ aiter FusedMoE

  • Expert w1/w3/w2 GEMM + fused SwiGLU + clamp (§4 rows 8-10)
  • Flash: [6,4096]@[4096,2048] / [6,2048]@[2048,4096] per expert
  • Pro: [6,7168]@[7168,3072] / [6,3072]@[3072,7168] per expert
  • ⬜ ASM    ☑️ flyDSL    ☑️ Triton/Gluon    ⬜ HIP/C++

• sqrtsoftplus scoring — ✅ aiter

  • sqrt(softplus(x)) on [B,T,256/384] (§4 row 2)
  • ⬜ ASM    ⬜ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

• Hash routing (layers 0–2) — ✅ PyTorch gather

  • tid2eid[input_ids] lookup (§4 row 5)
  • ⬜ ASM    ⬜ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

• MegaMoE (Expert-Parallel All-to-All + FP4 GEMM) — ❌ Not available on AMD (mid-term target)

  • vLLM's DeepseekV4MegaMoEExperts: fused all-to-all dispatch + fp8_fp4_mega_moe() via DeepGEMM, symmetric NCCL buffer management
  • Combines inter-GPU expert routing (NCCL-based) with per-expert FP8×FP4 GEMM in a single fused op
  • NV: DeepGEMM fp8_fp4_mega_moe() + get_symm_buffer_for_mega_moe() (requires SM100/Blackwell)
  • AMD dependency: mori's NCCL-GIN equivalent for the symmetric communication buffer / all-to-all overlap
  • AMD GEMM dependency: aiter/CK A8W4 expert GEMM (see above)
  • ⬜ ASM    ☑️ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

(5) GEMM

• A8W8 (FP8 act × FP8 weight) — ✅ CK GemmMXFP8

  • Attention: wq_a, wq_b, wkv, wo_b (§1 rows 1,3,6,14)
  • Indexer: wq_b (§2 row 2)
  • MoE: gate scoring projection (§4 row 1)
  • ⬜ ASM    ☑️ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

• Grouped matmul (BF16) — ⚠️ PyTorch einsum

  • Attention: wo_a grouped LoRA — einsum("bsgd,grd->bsgr") (§1 row 13)
  • [T,G,4096] × [G,1024,4096] → [T,G,1024], G=8 (Flash) / G=16 (Pro)
  • ⬜ ASM    ☑️ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

• Small BF16 matmul — ⚠️ PyTorch

  • Indexer: weights_proj [T,dim]@[dim,64] (§2 row 6) — tiny output dim
  • mHC: hc_fn [T,hc×dim]@[hc×dim,24] (§5 row 2) — tiny output dim
  • ⬜ ASM    ☑️ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

(6) Norm

• RMSNorm — ✅ aiter

  • Attention: Q norm [T,1024/1536], KV norm [T,512], input/output norms [T,dim] (§1 rows 2,7,15)
  • Compressor: post-pool norm [T/4,512] (§3.1 row 6)
  • ⬜ ASM    ⬜ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

• Per-head Q Norm (inline) — ⚠️ PyTorch fallback

  • Attention: q *= rsqrt(mean(q², dim=-1) + eps) on [T,H,512] (§1 row 4)
  • vLLM fuses into Q/KV pre-attn kernel (~10–20× speedup)
  • ⬜ ASM    ⬜ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

• mHC inline norm (FP32) — ⚠️ PyTorch

  • rsqrt(mean(x²) + eps) on flattened [T, hc×dim] = [T,16384/28672] (§5 row 1)
  • ⬜ ASM    ⬜ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

(7) mHC

• hc_split_sinkhorn — ✅ aiter mhc_pre (needs #2916 bugfix merge)

  • 20-iter Sinkhorn on 4×4 matrix, 860 calls per forward (20 iters × 43 layers)
  • ⬜ ASM    ⬜ flyDSL    ⬜ Triton/Gluon    ☑️ HIP/C++

• hc_post combine — ✅ aiter mhc_post

  • post * x + comb * residual with 4 HC copies
  • ⬜ ASM    ⬜ flyDSL    ⬜ Triton/Gluon    ☑️ HIP/C++

(8) Quantization & Transform

• act_quant (BF16→FP8 e4m3, per-1×128) — ✅ aiter

  • Attention: KV non-RoPE dims [T,448] (§1 row 9)
  • Compressor: compressed KV non-RoPE dims [T/4,448] (§3.1 row 8)
  • ⬜ ASM    ⬜ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

• fp4_act_quant (BF16→FP4 e2m1, per-1×32) — ❌ Torch fallback

  • Indexer: Q after Hadamard [T,64,128] (§2 row 5)
  • Indexer compressor: compressed KV [T/4,128] (§3.2)
  • ⬜ ASM    ⬜ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

• rotate_activation (Hadamard transform) — ❌ Torch fallback

  • Indexer: Q before FP4 quant [T,64,128] (§2 row 4)
  • Indexer compressor: KV before FP4 quant [T/4,128] (§3.2)
  • Uses fast_hadamard_transform library, scale = dim^{-0.5}
  • ⬜ ASM    ⬜ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

• Inverse RoPE + FP8 quant (fused) — ⚠️ Separate ops on AMD

  • Attention: de-rotate o[..., -64:] then FP8-quant before wo_a (§1 row 12)
  • vLLM fuses these (~2–3× speedup)
  • ⬜ ASM    ⬜ flyDSL    ☑️ Triton/Gluon    ☑️ HIP/C++

10. Reference Links

Compiled from: HF reference (inference/kernel.py, inference/model.py), vLLM blog (Apr 24 2026), ATOM PR#650, AITER PR#2916, and AMD kernel team analysis.