#1685 [RFC/PoC] Native Chunked Online Softmax Prefill for Dense Models (Qwen3.6) on HIP (Tested on gfx1031)

本文

## Summary & Context

Building on top of @aflin's excellent **`qwen36-35b-support`** branch (commit `20f4d4b`), which brought foundational Qwen3.6-35B-A3B (`qwen35moe`) capabilities to Strata, I tested the implementation on a physical AMD GPU environment: **AMD Radeon RX 6750 GRE 10GB (`gfx1031`, RDNA2, Wave32, Windows 10, ROCm 7.x via TheRock wheels)**.

In commit `20f4d4b`, Aaron noted:
> *"strata serve on a card without the batched prompt path for the model (qwen35moe below sm_75, or HIP) ... reads every prompt through the decode windows, as one-shot generate already did. ... Not run on a real Pascal or AMD card."*

On my physical `gfx1031` card, I verified this fallback behavior as well as an `IQ2_S` startup issue. As an amateur enthusiast, I put together a working, lightweight proof-of-concept (PoC) that unlocks the native batched prefill path for dense models on HIP without relying on Tensor Cores or WMMA.

For an 8,726-token prompt, cold prefill time dropped from **78.51s down to 27.98s ~ 29.40s** (a **2.8x speedup**, saving over 64% wall-clock time), while maintaining full generation speed (~91–100 tok/s with 85.7% MTP draft acceptance).

I would like to share the root cause analysis, kernel design, and empirical benchmarks in case they provide helpful reference for maintainers adapting native prompt prefill across AMD architectures.

---

## 1. Bottleneck Analysis

### Issue A: Missing `IQ2_S` in Down-Expert Formats
* **Symptom**: Serving models with `IQ2_S` MoE layers failed at startup with:  
  `layer 10's experts are IQ2_S/IQ2_S (ggml types 22/22), which this engine has no GPU kernels for`
* **Root Cause**: `STRATA_MMVQ_FMTS` in `src/kernels/cuda/iq_kernels.cu` already had `X(22)` (`GGML_TYPE_IQ2_S`), but `STRATA_D_FMTS` (the down-projection expert macro) was missing it.
* **Fix**: Adding `X(22)` into `STRATA_D_FMTS` instantiates the down kernel and resolves startup cleanly.

### Issue B: No Batched Dense Prompt Attention on AMD HIP
* **Symptom**: During model startup, the engine logged:  
  `this GPU has no batched prompt path for a qwen35moe model (NVIDIA sm_75+): prompts are read through the decode windows (slower)`  
  An 8,726-token prompt took **78.51 seconds** because prompts were digested 1–6 tokens at a time via decode windows.
* **Root Cause**: Models without a sparse indexer (`!g.has_indexer`, such as Qwen3.6) rely on dense attention. In `src/prefill/prefill.cpp`, `Prefill::supports` explicitly returned `false` on HIP when `!g.has_indexer` because upstream prompt attention was implemented using NVIDIA Tensor Core PTX (`mma.sync`).

---

## 2. Kernel Design: Non-Fused Chunked Online Softmax

### Why Avoid Fused WMMA on RDNA2 (`gfx1030` / `gfx1031`)?
Consumer RDNA2 cards have no native matrix cores (`v_wmma`). Emulating synthetic matrix fragments via Clang often leads to high register pressure, spilling, or driver instability.

### The 128-Thread Online Softmax Design
Instead of matrix cores, I implemented a 128-thread chunked Online Softmax kernel with warp shuffles (`hip_dense_prompt_attn_kernel`) in `src/kernels/cuda/qsa_prompt_attn.cu`:

1. **Exact 1:1 Dimension-to-Thread Mapping**:
   - `blockDim.x = 128` threads (4 warps under Wave32).
   - `const int d = threadIdx.x;` exactly maps to the 128 dimensions of Head Dim (`HD = 128`).
   - Each thread is responsible for exactly 1 feature dimension of $Q, K, V$ and writes 1 element to output $O$. This avoids out-of-bounds access and minimizes register pressure.
2. **Online Softmax Accumulation**:
   - Maintains running maximum $m$ and normalization sum $l$ entirely in registers across the prompt width $W$.
   - Eliminates the need for any intermediate $O(W \times W)$ attention matrix in global memory or scratch buffers.
3. **Wave32 Shuffle Reduction**:
   - In-warp dot product reduction via `__shfl_xor_sync(0xffffffffu, pdot[h], offset, 32)` across offsets 16, 8, 4, 2, 1.
   - Inter-warp reduction across the 4 warps through a small $12 \times 4$ float `__shared__` buffer.
4. **On-the-Fly KV Dequantization**:
   - Directly decodes `KV_MODE 0` (FP16), `1` (INT8 `q8_0`), `3` (Hybrid), and `4` (`q4_0`) from Strata's paged memory pools (`k_pool`, `k_q`, `k_q4`).

---

## 3. Benchmark Telemetry (Verified on `gfx1031`)

* **Environment**: Windows 10, ROCm 7.x (via TheRock wheels), AMD Radeon RX 6750 GRE 10GB (`gfx1031`, Wave32)
* **Model**: `Qwen3.6-35B-A3B` (APEX-MINI, INT8 KV cache, MTP Speculative Draft = 5)
* **Prompt Length**: 8,726 tokens (financial quantitative pipeline full context)

| Metric | Upstream Decode Fallback (Commit 20f4d4b) | With Native HIP Prefill PoC | Improvement |
| :--- | :--- | :--- | :--- |
| **Prefill Wall-Clock Time** | **78.51 s** (78,507 ms) | **27.98 s ~ 28.49 s** | **⚡ -50.5 s (-64.3%)** |
| **Prefill Throughput** | 111.1 tokens/s | **311.8 tokens/s** | **🚀 2.8x speedup** |
| **Decode Token Speed** | 89.4 tokens/s | **90.6 ~ 100.6 tokens/s** | Zero regression |
| **MTP Acceptance Rate** | 85.6% | **85.7% (2626 / 3063)** | High precision |
| **Numerical Integrity** | Verified | **100% Passed Audit** | Monotonicity & rules valid |

---

## 4. Generalization Considerations for Maintainers

If upstream intends to adapt this to all AMD architectures, here are a few key points to consider:

1. **Wavefront Size (Wave32 vs. Wave64)**:
   - RDNA2 / RDNA3 (`gfx103x`, `gfx110x`) default to Wave32, where 128 threads = 4 warps.
   - CDNA / GCN (`gfx906`, MI200) use Wave64, where 128 threads = 2 warps.
   - The warp shuffle loop can be generalized using `warpSize` or `#if defined(__AMDGCN_WAVEFRONT_SIZE__)`:
     ```cpp
     for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
         pdot[h] += __shfl_xor_sync(0xffffffffu, pdot[h], offset, warpSize);
     }
     ```
2. **Head Dimensions beyond 128**:
   - If future models feature `HD != 128`, threads can loop over dimensions: `for (int d = tid; d < HD; d += blockDim.x)`.

---

## 5. Reference Implementation

The working patch tested on `gfx1031` includes:
- `src/kernels/cuda/iq_kernels.cu`: adding `X(22)` to `STRATA_D_FMTS`
- `src/prefill/prefill.cpp`: unlocking `Prefill::supports` for HIP dense attention
- `src/kernels/cuda/qsa_prompt_attn.cu`: `hip_dense_prompt_attn_kernel<KV_MODE>` and dispatch

Huge thanks to @Niko1221 for creating Strata and to @aflin for the initial Qwen3.6 integration! Hope these empirical measurements and PoC implementation provide helpful insights for official HIP prefill support across AMD GPUs!
(Note: The benchmark measurements and kernel were physically tested on my local GPU; this summary writeup was drafted with AI assistance and is provided for reference only.)
（注：基准测试与内核已在我的本地显卡上通过真机实测；本文由 AI 辅助整理起草，仅供参考交流。）