RFC: Port DeepSeek TileKernels FP8 quant + MoE routing to ROCm/HIP for llama.cpp

[kmbandy]

Overview

DeepSeek TileKernels is a recently released library of production-quality GPU kernels used in DeepSeek's internal training and inference. Written in TileLang (a Python DSL that compiles to GPU code), it includes highly optimized implementations of:

• FP8/FP4 quantization (per-token, per-block, per-channel) with fused SwiGLU
• MoE routing (fused top-k gating, token-to-expert mapping, fused expand/reduce)
• Batched transpose

The library is currently NVIDIA SM90/SM100 only (H100/H200, Blackwell). However, the algorithms themselves are hardware-agnostic — the SM90 requirement comes from Hopper-specific pipeline features that are already disabled in the default configs. This makes the kernels strong candidates for ROCm/HIP ports targeting our gfx1030 (RX 6900 XT) and gfx1201 (R9700) hardware.

---

Why This Matters for llama.cpp

1. FP8 KV Cache (gfx1201 / R9700 priority)

The tile_kernels/quant/ module provides production-quality per-token and per-block FP8 cast kernels. The algorithm:

1. Load input tile into registers
2. Find per-token or per-block absmax via warp-level reduction
3. Compute scale factor (sf) and its inverse
4. Multiply + cast to FP8, store quantized output + scale factors

This is exactly what's needed for FP8 KV cache in llama.cpp. The R9700 (gfx1201 / RDNA4) has native FP8 hardware support via ROCm — a HIP port of these kernels would give us production-quality FP8 KV cache scaling that matches what's running on H100s. The 6900 XT (gfx1030 / RDNA2) has no hardware FP8, so FP8 quant there would be software-emulated and lower priority.

Relevant files:

• tile_kernels/quant/per_token_cast_kernel.py — per-token (row-wise) FP8 scaling
• tile_kernels/quant/per_block_cast_kernel.py — per-block FP8 scaling (block sizes 32x32, 128x128)
• tile_kernels/quant/swiglu_forward_and_per_token_cast_kernel.py — fused SwiGLU + FP8 cast

2. MoE Routing (both GPUs — high priority)

This is the more immediately impactful work for our stack. We're running Qwen3.5 122B-A10B and 97B-A10B MoE models in llama.cpp. The router fires on every token at every layer — llama.cpp's current ROCm MoE path does this with generic unfused kernels.

TileKernels' tile_kernels/moe/ provides:

• top2_sum_gate_kernel.py — fused top-k selection + score normalization in a single warp-level pass. Supports sigmoid, sqrtsoftplus, softmax scoring. Uses warp shuffle reductions (T.shfl_xor, T.shfl_sync → direct HIP equivalents: __shfl_xor, __shfl_sync).
• get_fused_mapping_kernel.py — fused token-to-expert mapping
• expand_to_fused_kernel.py / reduce_fused_kernel.py — fused expert expansion and weighted reduction
• normalize_weight_kernel.py — expert weight normalization
• topk_gate_kernel.py — general top-k gating

Currently llama.cpp handles these as separate dispatch steps. Fusing them would reduce kernel launch overhead and memory round-trips that matter at the batch sizes we run.

---

Portability Analysis

What makes these portable

1. TL_DISABLE_WARP_SPECIALIZED: True is set on all kernels — this disables Hopper's warp specialization / CTA pipeline, the most hardware-specific part. The compiled output is closer to "optimized CUDA" than "Hopper tensor core code."

2. TMA (Tensor Memory Accelerator) is already bypassed — disable_tma=True is set on all memory copies in the quant kernels. TMA is H100's DMA engine; these kernels don't rely on it.

3. Warp-level primitives map directly to HIP — T.shfl_xor, T.shfl_sync, T.warp_reduce_max all have 1:1 HIP equivalents.

4. TileLang has a HIP backend — it's worth testing whether @tilelang.jit with a ROCm target can compile these directly before resorting to manual HIP rewrites.

Key gotcha: wavefront size

RDNA default wavefront is 64 lanes. The MoE gate kernel assumes warp_size = 32 and uses 5-step butterfly reductions:

for i in T.unroll(0, 5):   # 5 steps covers 32 lanes
    x += T.shfl_xor(x, 1 << (4 - i))

On wave64 this would need 6 steps and the lane indexing throughout the top-k loop would need adjustment. Solution: compile with -mattr=+wavefrontsize32 (wave32 mode) on RDNA2/3/4 — this is already a common pattern in ROCm HIP kernel work and maps the warp logic 1:1.

Hardware targets

Kernel	gfx1030 (6900 XT, RDNA2)	gfx1201 (R9700, RDNA4)
FP8 per-token cast	Emulated (no hw FP8)	✅ Native FP8
FP8 per-block cast	Emulated	✅ Native FP8
FP8 KV cache	Low value	High value
MoE top-k gate	✅ wave32 mode	✅ wave32 mode
Fused expand/reduce	✅	✅

---

Proposed Work

Phase 1: TileLang HIP target probe (low effort)

• Install TileLang, attempt to compile the quant and MoE kernels targeting ROCm/HIP
• If successful, thin Python wrappers + llama.cpp ggml-hip integration
• If TileLang HIP backend is too immature, move to Phase 2

Phase 2: Manual HIP ports

Starting with highest-value kernels:

1. per_token_cast → HIP — straightforward register-level FP8 cast, ~150 lines of HIP. Target gfx1201 first (native FP8), add gfx1030 software path.
2. top2_sum_gate → HIP — warp-level top-k, replace 5-step butterfly with wave32 compile flag. Wire into llama.cpp's MoE dispatch path.
3. Fused expand/reduce → HIP — replace separate MoE dispatch steps.

Phase 3: Upstream candidates

Clean, well-tested ROCm MoE routing improvements are a strong upstream PR candidate for mainline llama.cpp. The FP8 KV cache kernels are gfx1201-specific but the MoE work benefits any ROCm user.

---

References

• TileKernels repo: https://github.com/deepseek-ai/TileKernels
• TileLang repo: https://github.com/tile-ai/tilelang
• Relevant hardware: gfx1030 (RX 6900 XT), gfx1201 (R9700 / RDNA4)
• Related: GGML_HIP_NO_VMM=ON currently set in our build — HIP VMM dynamic KV allocation is a parallel workstream

[github-actions]

This issue was closed because it has been inactive for 14 days since being marked as stale.