fp8 KV cache in the in-tree prefill kernel

README.md

@@ -169,13 +169,13 @@ library's actual routing. Do not edit it by hand; run
 | `silu_and_mul` | `aiter` -- opt-in | ✅ | ✅ | `aiter::silu_and_mul`, linked at the C++ level. Opt-in; matches native in fp16, lower in bf16. |
 | `fused_moe` | `aiter` -- only backend | ✅ | ✅ | `aiter_fused_moe`; bf16/fp16. Weights must be pre-shuffled with `shuffle_moe_weight` or results are silently wrong. |
 | `fused_moe_fp8` | `aiter` -- only backend | ✅ | ✅ | `aiter_fused_moe` with fp8 weights in `moe_fp8_dtype()` plus both scales; activations are quantized per token in the shim. |
-| `single_decode` | `hip` -- only backend | ✅ | ✅ | MHA / GQA / MQA. |
-| `batch_decode` | `hip` -- fallback; auto tries `aiter` first | ✅ | ✅ | MHA / GQA / MQA; fp8 KV-cache (E4M3FNUZ) and CUDA-graph capture. |
-| `single_prefill` | `hip` -- fallback; auto tries `aiter` first | ✅ | ✅ | MHA / GQA / MQA, including custom attention masks. |
-| `batch_prefill` | `hip` -- fallback; auto tries `aiter` first | ✅ | ✅ | Paged and ragged; MHA / GQA / MQA, including custom attention masks. |
-| `block_sparse` | `hip` -- only backend | ✅ | ✅ | `BlockSparseAttentionWrapper` and the variable-block variant. Native HIP FA2 only -- `determine_attention_backend` never returns `aiter` here. |
+| `single_decode` | `hip` -- only backend | ✅ | ✅ | MHA / GQA / MQA; fp8 KV-cache (E4M3FNUZ, E5M2FNUZ). No fp8 module is prebuilt, so first use pays a cold build. |
+| `batch_decode` | `hip` -- fallback; auto tries `aiter` first | ✅ | ✅ | MHA / GQA / MQA; fp8 KV-cache (E4M3FNUZ, E5M2FNUZ) on both the plain and `use_tensor_cores=True` paths, and CUDA-graph capture. |
+| `single_prefill` | `hip` -- fallback; auto tries `aiter` first | ✅ | ✅ | MHA / GQA / MQA, including custom attention masks. fp8 KV-cache (E4M3FNUZ, E5M2FNUZ) with a 2-byte query, JIT-only; the OCP spellings are refused, since they would be read under the fnuz exponent bias. |
+| `batch_prefill` | `hip` -- fallback; auto tries `aiter` first | ✅ | ✅ | Paged and ragged; MHA / GQA / MQA, including custom attention masks. fp8 KV-cache (E4M3FNUZ, E5M2FNUZ) with a 2-byte query, JIT-only -- no fp8 prefill module is prebuilt, so first use pays a cold build. |
+| `block_sparse` | `hip` -- only backend | ✅ | ✅ | `BlockSparseAttentionWrapper` and the variable-block variant. Native HIP FA2 only -- `determine_attention_backend` never returns `aiter` here. Both wrappers build through the batch-prefill module, so the fp8 KV-cache dtypes and refusals are the same as `batch_prefill`. |
 | `cascade` | `hip` -- merge only; levels are auto-routed and can be `aiter` | ✅ | ✅ | Two-level shared-prefix attention. The `hip` backend is the merge kernels only: each level runs through the ordinary prefill/decode entry points at `backend="auto"`, so it can reach AITER, and no cascade wrapper exposes `backend=` to override that. `FLASHINFER_HIP_FUSED_CASCADE=1` threads partial state through the levels of `MultiLevelCascadeAttentionWrapper` only; AITER levels and both shared-prefix wrappers still merge post-hoc. |
-| `pod` | `hip` -- only backend | ✅ | ✅ | `PODWithPagedKVCacheWrapper` and the batch variant. JIT-only, excluded from AOT as upstream. |
+| `pod` | `hip` -- only backend | ✅ | ✅ | `PODWithPagedKVCacheWrapper` and the batch variant. JIT-only, excluded from AOT as upstream. An fp8 KV-cache is refused: POD sizes its tiles independently of the prefill dispatchers and no geometry there has been tested. |
 | `rope` | `hip` -- auto picks this | ✅ | ✅ | LLaMA and LLaMA 3.1 scaling; fused RoPE + fp8 quant + paged-KV append (E4M3FNUZ, E5M2FNUZ). |
 | `append_paged_kv_cache` | `hip` -- auto picks this | ✅ | ✅ | fp8 KV-cache supported. Sustains 3.62 TB/s against AITER's 2.86 on gfx942, so `auto` picks this. |
 | `rmsnorm` | `hip` -- auto picks this | ✅ | ✅ | What `auto` always picks: level with AITER on speed and more accurate. |

benchmarks/rocm/bench_fa2_prefill.py

@@ -144,8 +144,25 @@
         f"Default: {_DEFAULT_BATCH_Q_LEN} (chunked-prefill burst, matches AITER bench)."
     ),
 )
+_bench_parser.add_argument(
+    "--kv-dtype",
+    choices=["fp16", "fp8_e4m3fnuz", "fp8_e5m2fnuz"],
+    default="fp16",
+    help=(
+        "KV cache dtype. The query and output stay fp16 -- the MFMA path is "
+        "f16f16f32 -- so fp8 buys HBM bandwidth, not math. Default: fp16."
+    ),
+)
 _bench_args, _ = _bench_parser.parse_known_args()
 
+_KV_DTYPES = {
+    "fp16": torch.float16,
+    "fp8_e4m3fnuz": torch.float8_e4m3fnuz,
+    "fp8_e5m2fnuz": torch.float8_e5m2fnuz,
+}
+_kv_dtype: torch.dtype = _KV_DTYPES[_bench_args.kv_dtype]
+_kv_itemsize: int = _kv_dtype.itemsize
+
 _counters = _bench_args.counters
 _model: str = _bench_args.model
 _q_len_asym: int = _bench_args.q_len
@@ -160,6 +177,9 @@
     if _bench_args.label is not None
     else ("fa2" if _counters == "roofline" else f"fa2_{_counters}")
 )
+# Without this an fp8 run overwrites the fp16 CSV it is meant to be compared to.
+if _bench_args.label is None and _kv_dtype is not torch.float16:
+    _label = f"{_label}_{_bench_args.kv_dtype}"
 
 # ---------------------------------------------------------------------------
 # Shape presets (seq_len, num_qo_heads, num_kv_heads, head_dim, causal)
@@ -204,10 +224,19 @@ def _flops(
     return attended * num_qo_heads * head_dim * 4
 
 
+def _kv_randn(*shape: int) -> torch.Tensor:
+    """KV tensor in the selected cache dtype; randn has no fp8 overload."""
+    return torch.randn(*shape, dtype=torch.half, device="cuda").to(_kv_dtype)
+
+
 def _bytes(
     q_len: int, kv_len: int, num_qo_heads: int, num_kv_heads: int, head_dim: int
 ) -> int:
-    return 2 * head_dim * (2 * q_len * num_qo_heads + 2 * kv_len * num_kv_heads)
+    # q read + o write are always 2-byte; k and v follow the cache dtype, so an
+    # fp8 cache halves the KV term. Hardcoding 2 here reports fp8 bandwidth ~2x.
+    qo = 2 * 2 * q_len * num_qo_heads * head_dim
+    kv = 2 * kv_len * num_kv_heads * head_dim * _kv_itemsize
+    return qo + kv
 
 
 # ---------------------------------------------------------------------------
@@ -223,12 +252,8 @@ def _make_configs() -> list[KernelConfig]:
         q = torch.randn(
             seq_len, num_qo_heads, head_dim, dtype=torch.half, device="cuda"
         )
-        k = torch.randn(
-            seq_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda"
-        )
-        v = torch.randn(
-            seq_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda"
-        )
+        k = _kv_randn(seq_len, num_kv_heads, head_dim)
+        v = _kv_randn(seq_len, num_kv_heads, head_dim)
 
         flops = _flops(seq_len, seq_len, num_qo_heads, head_dim, causal)
         theo_bytes = _bytes(seq_len, seq_len, num_qo_heads, num_kv_heads, head_dim)
@@ -258,12 +283,8 @@ def _make_configs() -> list[KernelConfig]:
             q = torch.randn(
                 _q_len_asym, num_qo_heads, head_dim, dtype=torch.half, device="cuda"
             )
-            k = torch.randn(
-                kv_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda"
-            )
-            v = torch.randn(
-                kv_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda"
-            )
+            k = _kv_randn(kv_len, num_kv_heads, head_dim)
+            v = _kv_randn(kv_len, num_kv_heads, head_dim)
 
             flops = _flops(_q_len_asym, kv_len, num_qo_heads, head_dim, causal=True)
             theo_bytes = _bytes(
@@ -300,12 +321,8 @@ def _make_batch_configs(batch_size: int, q_len: int) -> list[KernelConfig]:
         q = torch.randn(
             batch_size * q_len, num_qo_heads, head_dim, dtype=torch.half, device="cuda"
         )
-        k = torch.randn(
-            batch_size * kv_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda"
-        )
-        v = torch.randn(
-            batch_size * kv_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda"
-        )
+        k = _kv_randn(batch_size * kv_len, num_kv_heads, head_dim)
+        v = _kv_randn(batch_size * kv_len, num_kv_heads, head_dim)
         qo_indptr = torch.arange(
             0, (batch_size + 1) * q_len, q_len, dtype=torch.int32, device="cuda"
         )
@@ -317,7 +334,14 @@ def _make_batch_configs(batch_size: int, q_len: int) -> list[KernelConfig]:
             workspace, "NHD", backend="fa2"
         )
         wrapper.plan(
-            qo_indptr, kv_indptr, num_qo_heads, num_kv_heads, head_dim, causal=True
+            qo_indptr,
+            kv_indptr,
+            num_qo_heads,
+            num_kv_heads,
+            head_dim,
+            causal=True,
+            q_data_type=torch.float16,
+            kv_data_type=_kv_dtype,
         )
 
         flops = batch_size * _flops(q_len, kv_len, num_qo_heads, head_dim, causal=True)
@@ -357,15 +381,7 @@ def _make_paged_configs(
         q = torch.randn(
             batch_size * q_len, num_qo_heads, head_dim, dtype=torch.half, device="cuda"
         )
-        kv_cache = torch.randn(
-            total_pages,
-            2,
-            page_size,
-            num_kv_heads,
-            head_dim,
-            dtype=torch.half,
-            device="cuda",
-        )
+        kv_cache = _kv_randn(total_pages, 2, page_size, num_kv_heads, head_dim)
 
         qo_indptr = torch.arange(
             0, (batch_size + 1) * q_len, q_len, dtype=torch.int32, device="cuda"
@@ -398,6 +414,8 @@ def _make_paged_configs(
             head_dim,
             page_size,
             causal=True,
+            q_data_type=torch.float16,
+            kv_data_type=_kv_dtype,
         )
 
         flops = batch_size * _flops(q_len, kv_len, num_qo_heads, head_dim, causal=True)

csrc/rocm/batch_decode.cu

@@ -87,6 +87,7 @@ void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
                                     at::Tensor o, std::optional<at::Tensor> maybe_lse,
                                     int64_t kv_layout_code,
                                     int64_t window_left ADDITIONAL_FUNC_PARAMS) {
+  CHECK_KV_DTYPES_MATCH(paged_k_cache, paged_v_cache);
   DecodePlanInfo plan_info;
   const size_t size = plan_info_vec.numel();
   const int64_t* first = plan_info_vec.data_ptr<int64_t>();

csrc/rocm/batch_pod.cu

@@ -53,6 +53,8 @@ void batch_pod_with_kv_cache_tensor(
     std::optional<at::Tensor> maybe_mask_indptr_d, std::optional<at::Tensor> maybe_alibi_slopes_d,
     double logits_soft_cap_d, double sm_scale_d, double rope_rcp_scale_d, double rope_rcp_theta_d,
     bool enable_pdl, at::Tensor sm_aware_sched) {
+  CHECK_KV_DTYPES_MATCH(paged_k_cache_p, paged_v_cache_p);
+  CHECK_KV_DTYPES_MATCH(paged_k_cache_d, paged_v_cache_d);
   // Prefill setup
   PrefillPlanInfo plan_info_p;
   plan_info_p.FromVector(tensor_to_vec(plan_info_vec_p));

csrc/rocm/batch_prefill.cu

@@ -77,6 +77,7 @@ void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer,
                                       at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o,
                                       std::optional<at::Tensor> maybe_lse, int64_t mask_mode_code,
                                       int64_t layout, int64_t window_left ADDITIONAL_FUNC_PARAMS) {
+  CHECK_KV_DTYPES_MATCH(k, v);
   PrefillPlanInfo plan_info;
   plan_info.FromVector(tensor_to_vec(plan_info_vec));
   QKVLayout kv_layout = static_cast<QKVLayout>(layout);
@@ -207,6 +208,7 @@ void BatchPrefillWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
                                      int64_t layout, int64_t window_left ADDITIONAL_FUNC_PARAMS,
                                      std::optional<at::Tensor> maybe_partial_o = std::nullopt,
                                      std::optional<at::Tensor> maybe_partial_lse = std::nullopt) {
+  CHECK_KV_DTYPES_MATCH(paged_k_cache, paged_v_cache);
   PrefillPlanInfo plan_info;
   plan_info.FromVector(tensor_to_vec(plan_info_vec));
   QKVLayout kv_layout = static_cast<QKVLayout>(layout);

csrc/rocm/pod.cu

@@ -48,6 +48,8 @@ void pod_with_kv_cache_tensor(
     std::optional<at::Tensor> maybe_mask_indptr_d, std::optional<at::Tensor> maybe_alibi_slopes_d,
     double logits_soft_cap_d, double sm_scale_d, double rope_rcp_scale_d, double rope_rcp_theta_d,
     bool enable_pdl) {
+  CHECK_KV_DTYPES_MATCH(k_p, v_p);
+  CHECK_KV_DTYPES_MATCH(paged_k_cache_d, paged_v_cache_d);
   QKVLayout kv_layout_p = static_cast<QKVLayout>(layout_p);
   uint32_t qo_len_p = q_p.size(0);
   uint32_t num_qo_heads = q_p.size(1);

csrc/rocm/pytorch_extension_utils.h

@@ -270,6 +270,14 @@ inline constexpr uint32_t pack_u16(uint16_t a, uint16_t b) {
 
 #define CHECK_GE(a, b) TORCH_CHECK((a) >= (b), "CHECK_GE(" #a ", " #b ") failed. ", a, " vs ", b)
 
+// The kernel compiles one KV dtype, taken from k, and reads v with it, so a
+// mismatch reinterprets v's bytes instead of converting them.
+#define CHECK_KV_DTYPES_MATCH(k, v)                                                        \
+  TORCH_CHECK((k).scalar_type() == (v).scalar_type(), #k " has dtype ", (k).scalar_type(), \
+              " but " #v " has ", (v).scalar_type(),                                       \
+              "; ROCm attention compiles a single KV dtype and would "                     \
+              "reinterpret " #v "'s bytes as " #k "'s")
+
 inline bool is_float8_tensor(const at::Tensor& tensor) {
   return tensor.scalar_type() == at::ScalarType::Float8_e4m3fnuz ||
          tensor.scalar_type() == at::ScalarType::Float8_e5m2fnuz;

csrc/rocm/single_decode.cu

@@ -47,7 +47,7 @@ void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::T
   CHECK_DIM(3, v);
   CHECK_SHAPE(k, v);
   CHECK_EQ(q.size(1), k.size(2));
-  CHECK_EQ(v.scalar_type(), k.scalar_type());
+  CHECK_KV_DTYPES_MATCH(k, v);
   unsigned int num_qo_heads = q.size(0);
   unsigned int head_dim_qk = q.size(1);
   unsigned int head_dim_vo = v.size(2);

csrc/rocm/single_prefill.cu

@@ -39,6 +39,7 @@ void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::
                                   at::Tensor o, std::optional<at::Tensor> maybe_lse,
                                   int64_t mask_mode_code, int64_t layout,
                                   int64_t window_left ADDITIONAL_FUNC_PARAMS) {
+  CHECK_KV_DTYPES_MATCH(k, v);
   auto device = q.device();
   unsigned int head_dim_qk = q.size(2);
   unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads;

docs/rocm/backends.md

@@ -386,6 +386,13 @@ runs it in CI.
 
 ## Known limitations
 
+**`pos_encoding_mode="ROPE_LLAMA"` on the fa2 prefill kernel is
+non-deterministic.** It returns a different result on every call from identical
+inputs — measured on gfx942, spread 0.2-1.4 over repeats, at every KV dtype
+including plain fp16, and already present at a single KV tile. `"NONE"` is
+bitwise reproducible. No prefill test covers RoPE, which is how it survived;
+treat in-kernel RoPE prefill as unusable and rotate q/k beforehand instead.
+
 AITER constraints fall into two groups. The first errors out under
 `backend="aiter"` and triggers fallback under `backend="auto"`. The second
 is worse: the call runs, but the flag is silently dropped.
@@ -769,9 +776,35 @@ partial case — its own kernels are HIP, but what it calls is not:
   OnlineSoftmax / SamplingFromLogits), the **logits processor** pipeline,
   and **quantization** (`packbits`, `segment_packbits`).
 
-### fp8 on the HIP path
-
-* Batch decode accepts an fp8 KV-cache (`float8_e4m3fnuz`).
+### fp8 on the HIP path: KV cache only
+
+Single prefill, batch prefill (paged and ragged), single decode, batch decode,
+and decode with `use_tensor_cores=True` all accept an fp8 KV-cache in
+`float8_e4m3fnuz` or `float8_e5m2fnuz`. The prefill-backed paths dequantize
+the cache into a 2-byte LDS tile on the way in; plain decode keeps an fp8 LDS
+tile and converts values to float when consuming them. Either way, the saving
+is HBM traffic and cache capacity, not math.
+
+* **The query and output stay 2-byte.** The prefill/tensor-core MFMA path is
+  f16f16f32, so an fp8 query is refused at module build; plain decode is limited
+  to the same public query and output dtype allowlist.
+* **Only the fnuz spellings.** `torch.float8_e4m3fn` and `torch.float8_e5m2` are
+  OCP encodings, and the HIP types they map to use an exponent bias one greater
+  (E4M3 7 vs 8, E5M2 15 vs 16), so passing one would be silently 2x wrong. Both
+  are refused.
+* **POD refuses fp8 KV.** `pod.cuh` sizes its tiles independently of the prefill
+  dispatchers, so it can ask for a geometry nothing has tested. Run the two
+  phases unfused.
+* **JIT-only.** No fp8 prefill or decode module is prebuilt, so first use pays a
+  cold `hipcc` build.
+* **The dequantize is hardware on gfx942, software on gfx950.** CDNA4 sets
+  `HIP_FP8_TYPE_FNUZ=0`, so the fnuz->float leg is a ~15-op routine rather than
+  a `cvt` instruction. Correctness is verified on both.
+* **Measured as capacity, not speed.** Prefill came out flat and only
+  `use_tensor_cores=True` decode gained. Note the produce path still issues one
+  4-byte load per lane where a 2-byte cache issues 8, so the instruction count
+  does not fall with the byte count -- the flat result is not yet known to be
+  inherent.
 * RoPE has a fused RoPE + fp8-quantize + paged-KV-append path covering
   `float8_e4m3fnuz` and `float8_e5m2fnuz`, alongside LLaMA and LLaMA 3.1
   scaling.
@@ -799,7 +832,9 @@ causal). The constraints, all of them AITER's:
   flat-gather path raise `NotImplementedError` naming fp8 on either backend;
   single prefill raises `NotImplementedError` on fa2 and `RuntimeError` from
   AITER's own dtype check on `backend="aiter"`. What none of them do is reach
-  the fa2 kernel's `static_assert`, which surfaces as a compiler log.
+  a kernel `static_assert`, which surfaces as a compiler log. All of this is
+  about an fp8 *query*; an fp8 KV cache with a 2-byte query is served in tree,
+  on fa2 — see above.
 
 ## Tests