diff --git a/README.md b/README.md index 6dea64b8b1b..ac92a182309 100644 --- a/README.md +++ b/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. | diff --git a/benchmarks/rocm/bench_fa2_prefill.py b/benchmarks/rocm/bench_fa2_prefill.py index 7355367babe..727174ac6c0 100644 --- a/benchmarks/rocm/bench_fa2_prefill.py +++ b/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) diff --git a/csrc/rocm/batch_decode.cu b/csrc/rocm/batch_decode.cu index 3ffe4388160..a3766e0c900 100644 --- a/csrc/rocm/batch_decode.cu +++ b/csrc/rocm/batch_decode.cu @@ -87,6 +87,7 @@ void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, at::Tensor o, std::optional 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(); diff --git a/csrc/rocm/batch_pod.cu b/csrc/rocm/batch_pod.cu index 74fac8d7c2b..201882d1a5f 100644 --- a/csrc/rocm/batch_pod.cu +++ b/csrc/rocm/batch_pod.cu @@ -53,6 +53,8 @@ void batch_pod_with_kv_cache_tensor( std::optional maybe_mask_indptr_d, std::optional 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)); diff --git a/csrc/rocm/batch_prefill.cu b/csrc/rocm/batch_prefill.cu index 3b47ce5a6fc..f7f58f95c14 100644 --- a/csrc/rocm/batch_prefill.cu +++ b/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 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(layout); @@ -207,6 +208,7 @@ void BatchPrefillWithPagedKVCacheRun(at::Tensor float_workspace_buffer, int64_t layout, int64_t window_left ADDITIONAL_FUNC_PARAMS, std::optional maybe_partial_o = std::nullopt, std::optional 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(layout); diff --git a/csrc/rocm/pod.cu b/csrc/rocm/pod.cu index 7a786adf5e5..38ad8a495a1 100644 --- a/csrc/rocm/pod.cu +++ b/csrc/rocm/pod.cu @@ -48,6 +48,8 @@ void pod_with_kv_cache_tensor( std::optional maybe_mask_indptr_d, std::optional 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(layout_p); uint32_t qo_len_p = q_p.size(0); uint32_t num_qo_heads = q_p.size(1); diff --git a/csrc/rocm/pytorch_extension_utils.h b/csrc/rocm/pytorch_extension_utils.h index 3163d29fcef..1e81249024d 100644 --- a/csrc/rocm/pytorch_extension_utils.h +++ b/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; diff --git a/csrc/rocm/single_decode.cu b/csrc/rocm/single_decode.cu index fac87290862..8249d47ce55 100644 --- a/csrc/rocm/single_decode.cu +++ b/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); diff --git a/csrc/rocm/single_prefill.cu b/csrc/rocm/single_prefill.cu index e7b7fc216b4..bb487b816cd 100644 --- a/csrc/rocm/single_prefill.cu +++ b/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 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; diff --git a/docs/rocm/backends.md b/docs/rocm/backends.md index c264e0d5b81..35ddc20033c 100644 --- a/docs/rocm/backends.md +++ b/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 diff --git a/flashinfer/jit/rocm/modules.py b/flashinfer/jit/rocm/modules.py index e7ce05c81ba..161bff01cb5 100644 --- a/flashinfer/jit/rocm/modules.py +++ b/flashinfer/jit/rocm/modules.py @@ -97,6 +97,132 @@ def _generate_additional_params_hip( generate_additional_params = _generate_additional_params_hip +# dtype_map_hip sends both OCP spellings to the *fnuz* C++ types, whose exponent +# bias is one greater (E4M3 7 vs 8, E5M2 15 vs 16), so either is silently 2x off. +_FP8_OCP_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) +_FP8_FNUZ_DTYPES = (torch.float8_e4m3fnuz, torch.float8_e5m2fnuz) + + +_WIDE_DTYPES = (torch.float16, torch.bfloat16) + +# Only batch_prefill_paged_aiter.cu has an fp8 arm, and it takes e4m3 in both +# spellings because aiter.dtypes.fp8 is arch-dependent, writing bf16 out. +# Mirrors prefill.py, which cannot be imported here -- it imports this module. +_AITER_FP8_DTYPES = (torch.float8_e4m3fnuz, torch.float8_e4m3fn) + + +def _check_aiter_dtypes( + what: str, dtype_q, dtype_kv, dtype_o, *, allow_fp8: bool +) -> None: + """Allowlist for an explicit backend="aiter". + + Equality alone let a mapped-but-unserved dtype through: int8 and uint8 are + in dtype_map_hip, so they built a URI and reached ninja with no kernel. + `allow_fp8` is false wherever the launcher is fp16/bf16 only. + """ + if dtype_q != dtype_kv: + raise NotImplementedError( + f"{what}: AITER requires equal query and KV dtypes; got " + f"{dtype_q} and {dtype_kv}." + ) + served = _WIDE_DTYPES + (_AITER_FP8_DTYPES if allow_fp8 else ()) + for role, dtype in (("query", dtype_q), ("KV", dtype_kv)): + if dtype not in served: + raise NotImplementedError( + f"{what}: AITER {role} dtype {dtype} is not supported. Use " + + ( + "float16, bfloat16, or the e4m3 fp8 encoding this GPU uses " + "(aiter.dtypes.fp8)." + if allow_fp8 + else "float16 or bfloat16; AITER serves fp8 from its paged " + "batch-prefill kernel only." + ) + ) + # The launchers pin the output exactly: `o == q_dtype` for a wide query + # (single_prefill_aiter.cu, batch_ragged_prefill_aiter.cu) and bf16 for an + # fp8 one (batch_prefill_paged_aiter.cu). Anything looser builds, then fails. + expected_o = torch.bfloat16 if dtype_q in _AITER_FP8_DTYPES else dtype_q + if dtype_o != expected_o: + raise NotImplementedError( + f"{what}: AITER writes output dtype {expected_o} for query dtype " + f"{dtype_q}; got {dtype_o}." + ) + + +def _check_aiter_decode_dtypes(what: str, dtype_q, dtype_kv, dtype_o) -> None: + """Allowlist for AITER's paged-decode kernel, which is stricter than prefill. + + `batch_decode_aiter.cu` takes fp16/bf16 only -- no fp8 arm at all -- and + requires k, v and o to share the query's dtype. + """ + for role, dtype in (("query", dtype_q), ("KV", dtype_kv), ("output", dtype_o)): + if dtype not in _WIDE_DTYPES: + raise NotImplementedError( + f"{what}: AITER paged decode {role} dtype {dtype} is not " + "supported; the kernel serves float16 and bfloat16 only." + ) + if not (dtype_q == dtype_kv == dtype_o): + raise NotImplementedError( + f"{what}: AITER paged decode needs one dtype for query, KV and " + f"output; got {dtype_q}, {dtype_kv} and {dtype_o}." + ) + + +def _check_fa2_fp8_dtypes( + what: str, dtype_q, dtype_kv, dtype_o, *, kv_must_match_q: bool +) -> None: + """Refuse dtype combinations the in-tree attention kernels cannot serve. + + An allowlist, so a dtype nobody considered gets a sentence instead of a + ninja log. Raised at module-build time rather than from IsInvalid(), which + fires host-side only after a multi-minute JIT. + """ + for role, dtype in (("query", dtype_q), ("output", dtype_o)): + if dtype not in _WIDE_DTYPES: + raise NotImplementedError( + f"{what}: {role} dtype {dtype} is not supported on ROCm. The MFMA " + "path is f16f16f32, so the query and output must be float16 or " + "bfloat16; only the KV cache may be fp8." + ) + if dtype_kv in _FP8_OCP_DTYPES: + raise NotImplementedError( + f"{what}: fp8 KV dtype {dtype_kv} is an OCP encoding. The HIP kernel " + "would read it as the fnuz type of the same shape, whose exponent bias " + "is one greater (E4M3 7 vs 8, E5M2 15 vs 16), and be silently 2x wrong. " + "Re-quantize as torch.float8_e4m3fnuz or torch.float8_e5m2fnuz." + ) + if dtype_kv not in _WIDE_DTYPES + _FP8_FNUZ_DTYPES: + raise NotImplementedError( + f"{what}: KV dtype {dtype_kv} is not supported on ROCm. Use float16, " + "bfloat16, or an fnuz fp8 cache (torch.float8_e4m3fnuz, " + "torch.float8_e5m2fnuz)." + ) + # The MFMA is instantiated from DTypeQ, so a bf16 cache under an fp16 query + # is read as fp16. Only the *plain* decode kernel is exempt: it cast_loads q + # and k into vec_t separately. use_tensor_cores=True decode builds + # the prefill module, so it is bound by this rule like any other prefill. + if kv_must_match_q and dtype_kv in _WIDE_DTYPES and dtype_kv != dtype_q: + raise NotImplementedError( + f"{what}: KV dtype {dtype_kv} differs from query dtype {dtype_q}. A " + "2-byte cache must match the query; only an fp8 cache may differ." + ) + + +def _check_pod_fp8_dtypes(what: str, dtype_q, dtype_kv, dtype_o) -> None: + """Refuse any fp8 on POD, which prefill.cuh's tile arithmetic does not cover. + + pod.cuh carries upstream's CUDA register constant (8, against CDNA's 4) and + sizes the smem budget off the query dtype, so it can instantiate KernelTraits + at a NUM_MMA_KV the prefill dispatchers never produce and nothing has tested. + """ + if dtype_kv in _FP8_FNUZ_DTYPES + _FP8_OCP_DTYPES: + raise NotImplementedError( + f"{what}: an fp8 KV cache is not supported. Use " + "BatchPrefillWithPagedKVCacheWrapper or BatchDecodeWithPagedKVCacheWrapper, " + "which serve fp8 KV, and run the two phases unfused." + ) + _check_fa2_fp8_dtypes(what, dtype_q, dtype_kv, dtype_o, kv_must_match_q=True) + def get_single_decode_uri( dtype_q: torch.dtype, @@ -235,6 +361,11 @@ def gen_single_decode_module( use_sliding_window: bool, use_logits_soft_cap: bool, ): + # Before the URI: it indexes filename_safe_dtype_map, so an unsupported + # dtype would raise KeyError rather than this allowlist's message. + _check_fa2_fp8_dtypes( + "single decode", dtype_q, dtype_kv, dtype_o, kv_must_match_q=False + ) uri = get_single_decode_uri( dtype_q, dtype_kv, @@ -281,6 +412,16 @@ def gen_single_prefill_module( use_logits_soft_cap: bool, use_fp16_qk_reduction: bool, ) -> JitSpec: + # Before the URI: it indexes filename_safe_dtype_map, so an unsupported + # dtype would raise KeyError rather than this allowlist's message. + if backend == "aiter": + _check_aiter_dtypes( + "single prefill", dtype_q, dtype_kv, dtype_o, allow_fp8=False + ) + else: + _check_fa2_fp8_dtypes( + "single prefill", dtype_q, dtype_kv, dtype_o, kv_must_match_q=True + ) uri = get_single_prefill_uri( backend, dtype_q, @@ -366,6 +507,11 @@ def gen_batch_decode_aiter_module( head_dim_qk: int, head_dim_vo: int, ) -> JitSpec: + # First, ahead of both the URI (which indexes filename_safe_dtype_map, so an + # unsupported dtype would raise KeyError) and the aiter import (absent on a + # py3.13 box, which would mask the refusal with ModuleNotFoundError). + _check_aiter_decode_dtypes("batch decode (aiter)", dtype_q, dtype_kv, dtype_o) + import aiter as _aiter_mod # Batch decode resolves its .so as an absolute path from Python via @@ -411,6 +557,11 @@ def gen_batch_decode_module( use_sliding_window: bool, use_logits_soft_cap: bool, ): + # Before the URI: it indexes filename_safe_dtype_map, so an unsupported + # dtype would raise KeyError rather than this allowlist's message. + _check_fa2_fp8_dtypes( + "batch decode", dtype_q, dtype_kv, dtype_o, kv_must_match_q=False + ) uri = get_batch_decode_uri( dtype_q, dtype_kv, @@ -460,6 +611,14 @@ def gen_batch_prefill_module( use_logits_soft_cap: bool, use_fp16_qk_reduction: bool, ) -> JitSpec: + # Before the URI: it indexes filename_safe_dtype_map, so an unsupported + # dtype would raise KeyError rather than this allowlist's message. + if backend == "aiter": + _check_aiter_dtypes("batch prefill", dtype_q, dtype_kv, dtype_o, allow_fp8=True) + else: + _check_fa2_fp8_dtypes( + "batch prefill", dtype_q, dtype_kv, dtype_o, kv_must_match_q=True + ) uri = get_batch_prefill_uri( backend, dtype_q, @@ -563,6 +722,9 @@ def gen_customize_single_decode_module( use_sliding_window: bool = False, use_logits_soft_cap: bool = False, ) -> JitSpec: + _check_fa2_fp8_dtypes( + "single decode", dtype_q, dtype_kv, dtype_o, kv_must_match_q=False + ) gen_directory = FLASHINFER_GEN_SRC_DIR / uri ( additional_params_decl, @@ -648,10 +810,26 @@ def gen_customize_single_prefill_module( use_fp16_qk_reduction: bool = False, fp8_enabled: bool = False, ): - # Upstream routes fp8 through fa3-only sm90 templates; ROCm has no fp8 - # prefill kernel, so refuse rather than silently return a bf16/fp16 one. + # fp8_enabled selects upstream's fa3-only sm90 fp8-*query* templates. An fp8 + # KV cache is served by the fa2 kernel and does not come through here. if fp8_enabled: - raise ValueError("fp8 prefill is not supported on ROCm") + raise ValueError( + "fp8 queries are not supported on ROCm; an fp8 KV cache is served " + "via kv_data_type on the fa2 prefill and decode wrappers" + ) + # AITER compiles its own fp8 dtypes, but its kernels take one dtype for q + # and kv; an explicit backend="aiter" otherwise reaches ninja with a pair + # no kernel exists for. + if backend == "aiter": + _check_aiter_dtypes( + "single prefill", dtype_q, dtype_kv, dtype_o, allow_fp8=False + ) + # Not `== "fa2"`: ROCm logs and ignores an explicit "fa3", routing it to + # the same in-tree kernel, so only AITER compiles its own fp8 dtypes. + if backend != "aiter": + _check_fa2_fp8_dtypes( + "single prefill", dtype_q, dtype_kv, dtype_o, kv_must_match_q=True + ) kwargs = { "variant_decl": variant_decl, @@ -803,6 +981,9 @@ def gen_customize_batch_decode_module( use_sliding_window: bool = False, use_logits_soft_cap: bool = False, ): + _check_fa2_fp8_dtypes( + "batch decode", dtype_q, dtype_kv, dtype_o, kv_must_match_q=False + ) gen_directory = FLASHINFER_GEN_SRC_DIR / uri (additional_params_decl, additional_func_params, additional_params_setter) = ( generate_additional_params( @@ -887,10 +1068,24 @@ def gen_customize_batch_prefill_module( use_fp16_qk_reduction: bool = False, fp8_enabled: bool = False, ): - # Upstream routes fp8 through fa3-only sm90 templates; ROCm has no fp8 - # prefill kernel, so refuse rather than silently return a bf16/fp16 one. + # fp8_enabled selects upstream's fa3-only sm90 fp8-*query* templates. An fp8 + # KV cache is served by the fa2 kernel and does not come through here. if fp8_enabled: - raise ValueError("fp8 prefill is not supported on ROCm") + raise ValueError( + "fp8 queries are not supported on ROCm; an fp8 KV cache is served " + "via kv_data_type on the fa2 prefill and decode wrappers" + ) + # AITER compiles its own fp8 dtypes, but its kernels take one dtype for q + # and kv; an explicit backend="aiter" otherwise reaches ninja with a pair + # no kernel exists for. + if backend == "aiter": + _check_aiter_dtypes("batch prefill", dtype_q, dtype_kv, dtype_o, allow_fp8=True) + # Not `== "fa2"`: ROCm logs and ignores an explicit "fa3", routing it to + # the same in-tree kernel, so only AITER compiles its own fp8 dtypes. + if backend != "aiter": + _check_fa2_fp8_dtypes( + "batch prefill", dtype_q, dtype_kv, dtype_o, kv_must_match_q=True + ) kwargs = { "variant_decl": variant_decl, @@ -1109,6 +1304,7 @@ def gen_pod_module( use_sliding_window_d: bool, use_logits_soft_cap_d: bool, ) -> JitSpec: + _check_pod_fp8_dtypes("POD", dtype_q, dtype_kv, dtype_o) uri = get_pod_uri( dtype_q, dtype_kv, @@ -1158,6 +1354,7 @@ def gen_batch_pod_module( use_sliding_window_d: bool, use_logits_soft_cap_d: bool, ) -> JitSpec: + _check_pod_fp8_dtypes("batch POD", dtype_q, dtype_kv, dtype_o) uri = get_batch_pod_uri( dtype_q, dtype_kv, @@ -1212,6 +1409,9 @@ def _gen_customize_pod_like_module( use_fp16_qk_reduction: bool, ) -> JitSpec: """Shared body for single (prefix="pod") and batch (prefix="batch_pod") POD JIT modules.""" + # Both customize entry points reach the kernel only through here, so this is + # the seam that keeps them from bypassing gen_pod_module's refusal. + _check_pod_fp8_dtypes(f"{prefix} (customize)", dtype_q, dtype_kv, dtype_o) gen_directory = FLASHINFER_GEN_SRC_DIR / uri kwargs = { diff --git a/flashinfer/jit/utils.py b/flashinfer/jit/utils.py index 6306102612f..04f1d051cad 100644 --- a/flashinfer/jit/utils.py +++ b/flashinfer/jit/utils.py @@ -72,6 +72,7 @@ def write_if_different(path: pathlib.Path, content: str) -> None: torch.float8_e4m3fn: "e4m3", torch.float8_e4m3fnuz: "e4m3fnuz", torch.float8_e5m2: "e5m2", + torch.float8_e5m2fnuz: "e5m2fnuz", torch.int8: "i8", torch.uint8: "u8", torch.int32: "i32", @@ -94,9 +95,13 @@ def filename_safe_dtype_map_kv(dtype: torch.dtype) -> str: dtype_map_hip = { torch.float16: "__half", torch.bfloat16: "__hip_bfloat16", + # NOTE: the two OCP entries map onto the *fnuz* C++ types, whose exponent bias + # is one greater (E4M3 7 vs 8, E5M2 15 vs 16). Callers that reach a HIP kernel + # with an OCP tensor must refuse it rather than rely on this (rocm/prefill.py). torch.float8_e4m3fn: "__hip_fp8_e4m3_fnuz", torch.float8_e4m3fnuz: "__hip_fp8_e4m3_fnuz", torch.float8_e5m2: "__hip_fp8_e5m2_fnuz", + torch.float8_e5m2fnuz: "__hip_fp8_e5m2_fnuz", torch.int8: "int8_t", torch.uint8: "uint8_t", torch.int32: "int32_t", diff --git a/flashinfer/rocm/arch_caps.py b/flashinfer/rocm/arch_caps.py index 6ca0415d312..d2818135686 100644 --- a/flashinfer/rocm/arch_caps.py +++ b/flashinfer/rocm/arch_caps.py @@ -473,31 +473,31 @@ def _archs(gfx942: ArchSupport, gfx950: ArchSupport) -> Mapping[str, ArchSupport "single_decode", "hip", _archs(_HIP_942, _HIP_950), - note="MHA / GQA / MQA.", + note="MHA / GQA / MQA; fp8 KV-cache (E4M3FNUZ, E5M2FNUZ). No fp8 module is prebuilt, so first use pays a cold build.", ), Capability( "batch_decode", "hip", _archs(_HIP_942, _HIP_950), - note="MHA / GQA / MQA; fp8 KV-cache (E4M3FNUZ) and CUDA-graph capture.", + note="MHA / GQA / MQA; fp8 KV-cache (E4M3FNUZ, E5M2FNUZ) on both the plain and `use_tensor_cores=True` paths, and CUDA-graph capture.", ), Capability( "single_prefill", "hip", _archs(_HIP_942, _HIP_950), - note="MHA / GQA / MQA, including custom attention masks.", + note="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.", ), Capability( "batch_prefill", "hip", _archs(_HIP_942, _HIP_950), - note="Paged and ragged; MHA / GQA / MQA, including custom attention masks.", + note="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.", ), Capability( "block_sparse", "hip", _archs(_HIP_942, _HIP_950), - note="`BlockSparseAttentionWrapper` and the variable-block variant. Native HIP FA2 only -- `determine_attention_backend` never returns `aiter` here.", + note="`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`.", ), Capability( "cascade", @@ -510,7 +510,7 @@ def _archs(gfx942: ArchSupport, gfx950: ArchSupport) -> Mapping[str, ArchSupport "pod", "hip", _archs(_HIP_942, _HIP_950), - note="`PODWithPagedKVCacheWrapper` and the batch variant. JIT-only, excluded from AOT as upstream.", + note="`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.", ), Capability( "rope", diff --git a/flashinfer/rocm/decode.py b/flashinfer/rocm/decode.py index 2a7ffcc33b0..c2c301cf0a1 100644 --- a/flashinfer/rocm/decode.py +++ b/flashinfer/rocm/decode.py @@ -35,10 +35,12 @@ ) from .aiter_utils import handle_aiter_probe_failure from ..jit.core import logger +from ..jit.rocm.modules import _check_fa2_fp8_dtypes from ..page import get_seq_lens from .prefill import ( _aiter_bootstrap_lock, _auto_select_prefill_backend, + _check_kv_dtypes_match, get_batch_prefill_jit_module, get_batch_prefill_module, get_single_prefill_module, @@ -126,6 +128,9 @@ def _warn_deprecated_plan_positional_args(api_name: str) -> None: @functools.cache def get_single_decode_module(*args): + # Before the URI: it indexes filename_safe_dtype_map, so an unsupported + # dtype would raise KeyError rather than the allowlist's message. + _check_fa2_fp8_dtypes("single decode", *args[:3], kv_must_match_q=False) uri = get_single_decode_uri(*args) module = gen_single_decode_module(*args).build_and_load() run_func = module.run.default @@ -265,6 +270,9 @@ def _fake_run_batch_decode( @functools.cache def get_batch_decode_module(*args): + # Before the URI: it indexes filename_safe_dtype_map, so an unsupported + # dtype would raise KeyError rather than the allowlist's message. + _check_fa2_fp8_dtypes("batch decode", *args[:3], kv_must_match_q=False) uri = get_batch_decode_uri(*args) mod = gen_batch_decode_module(*args).build_and_load() plan_func = mod.plan.default @@ -787,6 +795,8 @@ def single_decode_with_kv_cache( rope_theta = 1e4 num_qo_heads = q.shape[0] + _check_kv_dtypes_match(k, v) + lse = None if return_lse: lse = torch.empty((num_qo_heads,), dtype=torch.float32, device=q.device) diff --git a/flashinfer/rocm/prefill.py b/flashinfer/rocm/prefill.py index af9af4b51a6..8a17642443d 100755 --- a/flashinfer/rocm/prefill.py +++ b/flashinfer/rocm/prefill.py @@ -16,7 +16,6 @@ """ import functools -import logging import math import os import threading @@ -30,6 +29,7 @@ from .arch_caps import capability_reason, require_capability from ..jit.core import logger from ..jit.rocm import aiter_variants as _variants +from ..jit.rocm.modules import _check_fa2_fp8_dtypes from ..jit import ( gen_batch_prefill_module, gen_customize_batch_prefill_module, @@ -206,6 +206,11 @@ def get_customize_batch_prefill_module( @functools.cache def get_single_prefill_module(backend, *args): + # Before the URI: it indexes filename_safe_dtype_map, so an unsupported + # dtype would raise KeyError rather than the allowlist's message. + # AITER compiles its own fp8 dtypes, so it is checked at its own seam. + if backend != "aiter": + _check_fa2_fp8_dtypes("single prefill", *args[:3], kv_must_match_q=True) uri = get_single_prefill_uri(backend, *args) module = gen_single_prefill_module(backend, *args).build_and_load() run_func = module.run.default @@ -522,19 +527,33 @@ def _require_native_fp8_dtype(dtype_q: torch.dtype) -> None: ) +def _check_kv_dtypes_match(k: torch.Tensor, v: torch.Tensor) -> None: + """The kernel is built for one DTypeKV, taken from k, and reads v with it. + + Only 2-byte KV used to reach the kernel, so a mismatch was caught by the + 8-bit static_assert; with an fp8 cache it returns NaN instead. + """ + if k.dtype != v.dtype: + raise ValueError( + f"k has dtype {k.dtype} but v has {v.dtype}; ROCm attention compiles a " + "single KV dtype and would reinterpret v's bytes as k's." + ) + + def _reject_fp8_on_fa2(dtype_q: torch.dtype, backend: str) -> None: - """Raise if an fp8 prefill resolved to fa2, which has no fp8 kernel. + """Raise if an fp8 *query* prefill resolved to fa2, which serves KV only. - Without this the refusal surfaces from ninja: the in-tree kernel rejects - 8-bit types with a static_assert, so the caller gets a compiler log. + Without this the refusal surfaces from ninja: the in-tree kernel rejects an + 8-bit query with a static_assert, so the caller gets a compiler log. """ if backend == "fa2" and dtype_q in FP8_PREFILL_DTYPES: raise NotImplementedError( - f"fp8 prefill (dtype={dtype_q}) has no in-tree fa2 kernel on ROCm -- " - "include/flashinfer/rocm/attention/prefill.cuh rejects 8-bit types at " - "compile time. AITER serves fp8 only on paged batch prefill, via " + f"fp8 prefill query (dtype={dtype_q}) has no in-tree fa2 kernel on " + "ROCm -- the MFMA path is f16f16f32, so prefill.cuh takes a 2-byte " + "query. AITER serves an fp8 query on paged batch prefill only, via " "BatchPrefillWithPagedKVCacheWrapper with per-tensor scale_q/scale_k/" - "scale_v. Otherwise cast q/k/v to bf16 or fp16." + "scale_v. To stay on fa2, cast the query to bf16 or fp16 and keep the " + "KV cache fp8 -- an fnuz fp8 KV cache is served in-tree." ) @@ -1127,6 +1146,11 @@ def _aiter_batch_ragged_available( @functools.cache def get_batch_prefill_module(backend, *args): + # Before the URI: it indexes filename_safe_dtype_map, so an unsupported + # dtype would raise KeyError rather than the allowlist's message. + # AITER compiles its own fp8 dtypes, so it is checked at its own seam. + if backend != "aiter": + _check_fa2_fp8_dtypes("batch prefill", *args[:3], kv_must_match_q=True) if backend == "aiter": module = gen_batch_prefill_module(backend, *args).build_and_load() _c_paged_run = module.paged_run.default @@ -1907,6 +1931,7 @@ def single_prefill_with_kv_cache( _check_pos_encoding_mode(pos_encoding_mode) _check_kv_layout(kv_layout) + _check_kv_dtypes_match(k, v) tmp = _get_cache_buf("single_prefill_with_kv_cache_tmp", 32 * 1024 * 1024, q.device) if logits_soft_cap is None: logits_soft_cap = 0.0 @@ -1937,11 +1962,19 @@ def single_prefill_with_kv_cache( lse = torch.empty((q.size(0), q.size(1)), dtype=torch.float32, device=q.device) if is_float8(q): - # FP8 quant enabled, do sanity check: - # 1. unsupported feature - # 2. dtype check - assert window_left == -1 - assert q.dtype == k.dtype == v.dtype + # Not asserts: these reach here before any backend is chosen, and no + # backend serves either pair -- a sentence beats a bare AssertionError. + if window_left != -1: + raise NotImplementedError( + "single prefill: an fp8 query does not support a sliding window; " + f"got window_left={window_left}." + ) + if not (q.dtype == k.dtype == v.dtype): + raise NotImplementedError( + f"single prefill: an fp8 query needs an fp8 k and v, got q={q.dtype}, " + f"k={k.dtype}, v={v.dtype}. The in-tree kernel refuses an fp8 query " + "outright, and AITER compiles one dtype for all three." + ) assert q.shape[-1] == k.shape[-1] == v.shape[-1] if scale_q is None: scale_q = torch.ones(q.shape[1], dtype=torch.float32, device=q.device) @@ -4292,15 +4325,6 @@ def run( check_shape_dtype_device( out, q.shape[:-1] + v.shape[-1:], q.dtype, q.device, "out" ) - if is_float8(q): - logging.warning( - "Our current prefill kernel implementation needs f16 input, the f8 inputs " - " are casted to f16, which could result in performance degradation." - ) - q = q.to(torch.float16) - k = k.to(torch.float16) - v = v.to(torch.float16) - if self._custom_mask_buf is not None: mask_mode = MaskMode.CUSTOM.value else: diff --git a/include/flashinfer/rocm/attention/batch_pod.cuh b/include/flashinfer/rocm/attention/batch_pod.cuh index 96318b195fd..20b3f5d8f73 100644 --- a/include/flashinfer/rocm/attention/batch_pod.cuh +++ b/include/flashinfer/rocm/attention/batch_pod.cuh @@ -144,6 +144,10 @@ hipError_t BatchPODWithKVCacheTensorDispatched(PrefillParams prefill_params, typename DecodeParams::DTypeO* tmp_v_d, float* tmp_s_d, bool /*enable_pdl*/, hipStream_t stream, int* sm_aware_sched) { + // POD sizes its tiles independently of the prefill dispatchers, so an fp8 + // KV would reach a geometry nothing has tested. Refused in Python too. + static_assert(!is_fp8_kv_v, + "fp8 KV cache is not supported for POD attention"); static_assert(std::is_same::value); static_assert( std::is_same::value); diff --git a/include/flashinfer/rocm/attention/pod.cuh b/include/flashinfer/rocm/attention/pod.cuh index 8556819d37c..97695d16843 100644 --- a/include/flashinfer/rocm/attention/pod.cuh +++ b/include/flashinfer/rocm/attention/pod.cuh @@ -151,6 +151,10 @@ hipError_t PODWithKVCacheTensorDispatched(PrefillParams prefill_params, DecodeParams decode_params, typename DecodeParams::DTypeO* tmp_v, float* tmp_s, bool /*enable_pdl*/, hipStream_t stream) { + // POD sizes its tiles independently of the prefill dispatchers, so an fp8 + // KV would reach a geometry nothing has tested. Refused in Python too. + static_assert(!is_fp8_kv_v, + "fp8 KV cache is not supported for POD attention"); static_assert(std::is_same::value); static_assert( std::is_same::value); diff --git a/include/flashinfer/rocm/attention/prefill.cuh b/include/flashinfer/rocm/attention/prefill.cuh index 8ed41d62c8b..54d96dceaac 100644 --- a/include/flashinfer/rocm/attention/prefill.cuh +++ b/include/flashinfer/rocm/attention/prefill.cuh @@ -15,6 +15,7 @@ #include "flashinfer/rocm/mma.h" #include "flashinfer/rocm/platform.hpp" #include "flashinfer/rocm/utils.cuh" +#include "flashinfer/rocm/vec_dtypes.h" #ifdef FP16_QK_REDUCTION_SUPPORTED #include "../../fp16.h" @@ -38,6 +39,20 @@ using mma::MMAMode; constexpr uint32_t WARP_SIZE = kWarpSize; +// Keyed on the fp8 types, not on sizeof()==1: an int8/uint8 KV cache has no +// float interpretation, and widening it here would send it through +// vec_cast's primary template, which reads 0..255 as a float value. +template +inline constexpr bool is_fp8_kv_v = + std::is_same_v || std::is_same_v; + +// KV arrives from HBM as fp8 but lives in LDS as 2-byte: the MFMA path is +// f16f16f32, so a 2-byte tile keeps swizzle and fragment addressing identical +// to bf16. The dispatchers size the smem budget with this too, so the budget +// and SharedStorage cannot disagree. +template +using kv_smem_type_t = std::conditional_t, DTypeQ, DTypeKV>; + constexpr uint32_t get_num_warps_q(const uint32_t cta_tile_q) { if (cta_tile_q > 16) { return 4; @@ -99,13 +114,18 @@ struct KernelTraits { static constexpr PosEncodingMode POS_ENCODING_MODE = POS_ENCODING_MODE_; using DTypeQ = DTypeQ_; + // The HBM type. Everything past the global->smem produce path uses + // DTypeKVSmem instead, so an fp8 cache costs HBM bandwidth only. using DTypeKV = DTypeKV_; + using DTypeKVSmem = kv_smem_type_t; using DTypeO = DTypeO_; using DTypeQKAccum = DTypeQKAccum_; using IdType = IdType_; using AttentionVariant = AttentionVariant_; - static_assert(sizeof(DTypeKV_) != 1, "8-bit types not supported for CDNA3"); + static_assert(sizeof(DTypeQ_) == 2, "fp8 queries are not supported: the MFMA path is f16f16f32"); + static_assert(sizeof(DTypeKVSmem) == 2, + "KV must be 2-byte in LDS: a non-float 8-bit cache has no dequantization"); using SmemBasePtrTy = uint2; static constexpr uint32_t NUM_THREADS = NUM_WARPS_Q * NUM_WARPS_KV * 64; @@ -139,9 +159,9 @@ struct KernelTraits { static constexpr uint32_t UPCAST_STRIDE_Q = HEAD_DIM_QK / upcast_size(); static constexpr uint32_t UPCAST_STRIDE_K = - HEAD_DIM_QK / upcast_size(); + HEAD_DIM_QK / upcast_size(); static constexpr uint32_t UPCAST_STRIDE_V = - HEAD_DIM_VO / upcast_size(); + HEAD_DIM_VO / upcast_size(); static constexpr uint32_t UPCAST_STRIDE_O = HEAD_DIM_VO / upcast_size(); @@ -149,13 +169,11 @@ struct KernelTraits { return ((NUM_MMA_D_VO < 4) || (NUM_MMA_D_VO == 4 && NUM_MMA_KV % 2 == 1) || (POS_ENCODING_MODE == PosEncodingMode::kRoPELlama && NUM_MMA_D_VO > 4 && NUM_MMA_D_VO % (2 * NUM_WARPS_Q) != 0) || - (NUM_MMA_Q * (8 * NUM_MMA_D_VO + 2 * sizeof(DTypeQKAccum) * NUM_MMA_KV) >= 256) || - (sizeof(DTypeKV) == 1 && NUM_MMA_KV * 2 % NUM_WARPS_Q != 0) || - (sizeof(DTypeKV) == 1 && POS_ENCODING_MODE == PosEncodingMode::kRoPELlama)); + (NUM_MMA_Q * (8 * NUM_MMA_D_VO + 2 * sizeof(DTypeQKAccum) * NUM_MMA_KV) >= 256)); } using SharedStorage = SharedStorageQKVO; + HEAD_DIM_VO, DTypeQ, DTypeKVSmem, DTypeO>; #ifdef FP16_QK_REDUCTION_SUPPORTED template static constexpr DT getNegInf() { @@ -280,41 +298,46 @@ __device__ __forceinline__ void q_frag_apply_llama_rope_with_pos(T* x_first_half } } -template -__device__ __forceinline__ void produce_kv_impl( - uint32_t warp_idx, uint32_t lane_idx, - smem_t smem, uint32_t* smem_offset, - typename KTraits::DTypeKV** gptr, const uint32_t stride_n, const uint32_t kv_idx_base, - const uint32_t kv_len) { +/*! + * \brief Move one lane granule -- 4 KV elements -- from global into the 2-byte smem tile. + * + * Forwards verbatim for a 2-byte cache, so the fp16/bf16 path keeps the + * instructions it had. The fp8 arm dequantizes through float on purpose: + * vec_cast decodes fnuz with the OCP bias and is 2x wrong. + */ +template +__device__ __forceinline__ void produce_kv_unit( + smem_t smem, uint32_t smem_offset, + const typename KTraits::DTypeKV* gptr, bool predicate) { using DTypeKV = typename KTraits::DTypeKV; - constexpr uint32_t KV_THR_LAYOUT_COL = KTraits::KV_THR_LAYOUT_COL; // 16 - constexpr uint32_t NUM_WARPS = KTraits::NUM_WARPS; - constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; - constexpr uint32_t NUM_WARPS_Q = KTraits::NUM_WARPS_Q; - constexpr uint32_t NUM_MMA_D = produce_v ? KTraits::NUM_MMA_D_VO : KTraits::NUM_MMA_D_QK; - constexpr uint32_t UPCAST_STRIDE = - produce_v ? KTraits::UPCAST_STRIDE_V : KTraits::UPCAST_STRIDE_K; - constexpr uint32_t VECTOR_BIT_WIDTH = KTraits::VECTOR_BIT_WIDTH; - - // NOTE: NUM_MMA_KV*4/NUM_WARPS_Q = NUM_WARPS_KV*NUM_MMA_KV*4/num_warps - static_assert(NUM_MMA_KV * 4 % NUM_WARPS_Q == 0); - uint32_t kv_idx = kv_idx_base + warp_idx * 4 + lane_idx / KV_THR_LAYOUT_COL; + using DTypeKVSmem = typename KTraits::DTypeKVSmem; -#pragma unroll - for (uint32_t i = 0; i < NUM_MMA_KV * 4 / NUM_WARPS_Q; ++i) { -#pragma unroll - for (uint32_t j = 0; j < NUM_MMA_D / (8 / sizeof(DTypeKV)); ++j) { - smem.template load_vector_async(*smem_offset, *gptr, kv_idx < kv_len); - *smem_offset = smem.template advance_offset_by_column<16>(*smem_offset, j); - *gptr += 16 * upcast_size(); + if constexpr (std::is_same_v) { + smem.template load_vector_async(smem_offset, gptr, predicate); + } else { + constexpr uint32_t GRANULE = upcast_size(); + b64_t* smem_ptr = smem.base + smem_offset; + static_assert(GRANULE * sizeof(DTypeKV) == sizeof(uint32_t)); + if (predicate) { + // Stage the KV bytes as one dword, or gfx950 issues GRANULE separate + // global_load_ubyte. bit_cast, not a reinterpret_cast of the uint32_t + // (that puns it) and not a DTypeKV array (gfx950 cannot default-construct + // the fnuz types). wide is 8-byte aligned: vec_cast reads it as float2. + struct alignas(sizeof(uint32_t)) Staged { + DTypeKV e[GRANULE]; + }; + uint32_t raw; + __builtin_memcpy(&raw, gptr, sizeof(raw)); + const Staged packed = __builtin_bit_cast(Staged, raw); + alignas(sizeof(b64_t)) float wide[GRANULE]; + alignas(sizeof(b64_t)) DTypeKVSmem narrow[GRANULE]; + vec_cast::template cast(wide, packed.e); + vec_cast::template cast(narrow, wide); + __builtin_memcpy(smem_ptr, narrow, sizeof(b64_t)); + } else if constexpr (fill_mode == SharedMemFillMode::kFillZero) { + *smem_ptr = make_uint2(0, 0); } - kv_idx += NUM_WARPS * 4; - *smem_offset = smem.template advance_offset_by_row(*smem_offset) - - (sizeof(DTypeKV) * NUM_MMA_D * 2); - *gptr += NUM_WARPS * 4 * stride_n - - sizeof(DTypeKV) * NUM_MMA_D * 2 * upcast_size(); } - *smem_offset -= KTraits::CTA_TILE_KV * UPCAST_STRIDE; } /*! @@ -335,7 +358,7 @@ __device__ __forceinline__ void produce_kv( typename KTraits::DTypeKV** gptr, const uint32_t stride_n, const uint32_t kv_idx_base, const uint32_t kv_len, const dim3 tid = threadIdx) { const uint32_t warp_idx = get_warp_idx(tid.y, tid.z), lane_idx = tid.x; - using DTypeKV = typename KTraits::DTypeKV; + using DTypeKVSmem = typename KTraits::DTypeKVSmem; constexpr uint32_t KV_THR_LAYOUT_COL = KTraits::KV_THR_LAYOUT_COL; // 16 constexpr uint32_t NUM_WARPS = KTraits::NUM_WARPS; constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; @@ -352,16 +375,16 @@ __device__ __forceinline__ void produce_kv( #pragma unroll for (uint32_t i = 0; i < NUM_MMA_KV * 4 / NUM_WARPS_Q; ++i) { #pragma unroll - for (uint32_t j = 0; j < NUM_MMA_D / (8 / sizeof(DTypeKV)); ++j) { - smem.template load_vector_async(*smem_offset, *gptr, kv_idx < kv_len); + for (uint32_t j = 0; j < NUM_MMA_D / (8 / sizeof(DTypeKVSmem)); ++j) { + produce_kv_unit(smem, *smem_offset, *gptr, kv_idx < kv_len); *smem_offset = smem.template advance_offset_by_column<16>(*smem_offset, j); - *gptr += 16 * upcast_size(); + *gptr += 16 * upcast_size(); } kv_idx += NUM_WARPS * 4; *smem_offset = smem.template advance_offset_by_row(*smem_offset) - - (sizeof(DTypeKV) * NUM_MMA_D * 2); + (sizeof(DTypeKVSmem) * NUM_MMA_D * 2); *gptr += NUM_WARPS * 4 * stride_n - - sizeof(DTypeKV) * NUM_MMA_D * 2 * upcast_size(); + sizeof(DTypeKVSmem) * NUM_MMA_D * 2 * upcast_size(); } *smem_offset -= KTraits::CTA_TILE_KV * UPCAST_STRIDE; } @@ -373,6 +396,7 @@ __device__ __forceinline__ void page_produce_kv( const uint32_t kv_idx_base, const size_t* thr_local_kv_offset, const uint32_t kv_len, const dim3 tid = threadIdx) { using DTypeKV = typename KTraits::DTypeKV; + using DTypeKVSmem = typename KTraits::DTypeKVSmem; constexpr SharedMemFillMode fill_mode = produce_v ? SharedMemFillMode::kFillZero : SharedMemFillMode::kNoFill; constexpr uint32_t KV_THR_LAYOUT_COL = KTraits::KV_THR_LAYOUT_COL; @@ -395,15 +419,15 @@ __device__ __forceinline__ void page_produce_kv( DTypeKV* gptr = produce_v ? paged_kv.v_data + thr_local_kv_offset[i] : paged_kv.k_data + thr_local_kv_offset[i]; #pragma unroll - for (uint32_t j = 0; j < NUM_MMA_D / (8 / sizeof(DTypeKV)); ++j) { - smem.template load_vector_async(*smem_offset, gptr, kv_idx < kv_len); + for (uint32_t j = 0; j < NUM_MMA_D / (8 / sizeof(DTypeKVSmem)); ++j) { + produce_kv_unit(smem, *smem_offset, gptr, kv_idx < kv_len); *smem_offset = smem.template advance_offset_by_column<16>(*smem_offset, j); - gptr += 16 * upcast_size(); + gptr += 16 * upcast_size(); } kv_idx += NUM_WARPS * 4; *smem_offset = smem.template advance_offset_by_row(*smem_offset) - - (sizeof(DTypeKV) * NUM_MMA_D * 2); + (sizeof(DTypeKVSmem) * NUM_MMA_D * 2); } *smem_offset -= KTraits::CTA_TILE_KV * UPCAST_STRIDE; } @@ -625,7 +649,9 @@ __device__ __forceinline__ void k_smem_inplace_apply_rotary( const uint32_t kv_idx_base, smem_t* k_smem, uint32_t* k_smem_offset_r, float (*rope_freq)[4], const dim3 tid = threadIdx) { - using DTypeKV = typename KTraits::DTypeKV; + // Rotates fragments already resident in K smem, so this is the smem type -- + // with the HBM type an fp8 cache would reinterpret 2-byte smem as fp8. + using DTypeKV = typename KTraits::DTypeKVSmem; static_assert(sizeof(DTypeKV) == 2); constexpr uint32_t UPCAST_STRIDE_K = KTraits::UPCAST_STRIDE_K; constexpr uint32_t THREADS_PER_BMATRIX_ROW_SET = KTraits::THREADS_PER_BMATRIX_ROW_SET; @@ -753,11 +779,7 @@ __device__ __forceinline__ void compute_qk( #pragma unroll for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { - if constexpr (sizeof(typename KTraits::DTypeKV) == 1) { - static_assert(false, "FP8 support not yet implemented for CDNA3"); - } else { - k_smem->load_fragment(*k_smem_offset_r, b_frag); - } + k_smem->load_fragment(*k_smem_offset_r, b_frag); *k_smem_offset_r = k_smem->template advance_offset_by_row<16, UPCAST_STRIDE_K>(*k_smem_offset_r); @@ -777,19 +799,11 @@ __device__ __forceinline__ void compute_qk( } } } - if constexpr (sizeof(typename KTraits::DTypeKV) == 1) { - if (mma_d % 2 == 1) { - *k_smem_offset_r = k_smem->template advance_offset_by_column( - *k_smem_offset_r, mma_d / 2); - } - *k_smem_offset_r -= KTraits::NUM_MMA_KV * 16 * UPCAST_STRIDE_K; - } else { - *k_smem_offset_r = - k_smem->template advance_offset_by_column( - *k_smem_offset_r, mma_d, k_col_idx) - - KTraits::NUM_MMA_KV * 16 * UPCAST_STRIDE_K; - k_col_idx += QK_SMEM_COLUMN_ADVANCE; - } + *k_smem_offset_r = + k_smem->template advance_offset_by_column( + *k_smem_offset_r, mma_d, k_col_idx) - + KTraits::NUM_MMA_KV * 16 * UPCAST_STRIDE_K; + k_col_idx += QK_SMEM_COLUMN_ADVANCE; } *q_smem_offset_r -= KTraits::NUM_MMA_D_QK * QK_SMEM_COLUMN_ADVANCE; *k_smem_offset_r -= KTraits::NUM_MMA_D_QK * (QK_SMEM_COLUMN_ADVANCE); @@ -1009,11 +1023,7 @@ __device__ __forceinline__ void compute_sfm_v( #pragma unroll for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { uint32_t b_frag[INT32_ELEMS_PER_THREAD]; - if constexpr (sizeof(typename KTraits::DTypeKV) == 1) { - static_assert(false, "FP8 V path not implemented for CDNA3 yet"); - } else { - v_smem->load_matrix_m16n16_trans(*v_smem_offset_r, b_frag); - } + v_smem->load_matrix_m16n16_trans(*v_smem_offset_r, b_frag); #pragma unroll for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { if constexpr (std::is_same_v) { @@ -1024,17 +1034,10 @@ __device__ __forceinline__ void compute_sfm_v( o_frag[mma_q][mma_d], (uint32_t*)s_frag[mma_q][mma_kv], b_frag); } } - if constexpr (sizeof(typename KTraits::DTypeKV) == 1) { - if (mma_d % 2 == 1) { - *v_smem_offset_r = v_smem->template advance_offset_by_column( - *v_smem_offset_r, mma_d / 2); - } - } else { - *v_smem_offset_r = - v_smem->template advance_offset_by_column( - *v_smem_offset_r, mma_d, v_col_idx); - v_col_idx += V_SMEM_COLUMN_ADVANCE; - } + *v_smem_offset_r = + v_smem->template advance_offset_by_column( + *v_smem_offset_r, mma_d, v_col_idx); + v_col_idx += V_SMEM_COLUMN_ADVANCE; } *v_smem_offset_r = v_smem->template advance_offset_by_row<16, UPCAST_STRIDE_V>(*v_smem_offset_r) - @@ -1505,12 +1508,14 @@ __device__ __forceinline__ void SinglePrefillWithKVCacheDevice( DTypeKV* k_ptr = k + (chunk_start + warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL) * k_stride_n + kv_head_idx * k_stride_h + - (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(); + (lane_idx % KV_THR_LAYOUT_COL) * + upcast_size(); DTypeKV* v_ptr = v + (chunk_start + warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL) * v_stride_n + kv_head_idx * v_stride_h + - (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(); + (lane_idx % KV_THR_LAYOUT_COL) * + upcast_size(); uint32_t k_smem_offset_r = k_smem.template get_permuted_offset( get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, (lane_idx / 16)); @@ -1678,8 +1683,9 @@ hipError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DT : 0u; constexpr uint32_t min_valid_mma_kv_ = (HEAD_DIM_VO / 16u == 4u) ? 2u : 1u; const uint32_t max_num_mma_kv_smem = std::max( - min_valid_mma_kv_, static_cast(kv_budget_ / ((HEAD_DIM_QK + HEAD_DIM_VO) * 16 * - NUM_WARPS_KV * sizeof(DTypeKV)))); + min_valid_mma_kv_, + static_cast(kv_budget_ / ((HEAD_DIM_QK + HEAD_DIM_VO) * 16 * NUM_WARPS_KV * + sizeof(kv_smem_type_t)))); // control NUM_MMA_KV for maximum warp occupancy DISPATCH_NUM_MMA_KV(min(max_num_mma_kv_smem, max_num_mma_kv_reg), NUM_MMA_KV, { @@ -1940,13 +1946,15 @@ __global__ __launch_bounds__(KTraits::NUM_THREADS) void BatchPrefillWithRaggedKV lane_idx / KV_THR_LAYOUT_COL) * k_stride_n + kv_head_idx * k_stride_h + - (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(); + (lane_idx % KV_THR_LAYOUT_COL) * + upcast_size(); DTypeKV* v_ptr = v + (kv_indptr[request_idx] + chunk_start + warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL) * v_stride_n + kv_head_idx * v_stride_h + - (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(); + (lane_idx % KV_THR_LAYOUT_COL) * + upcast_size(); produce_kv(k_smem, &k_smem_offset_w, &k_ptr, k_stride_n, 0, chunk_size, tid); @@ -2213,7 +2221,9 @@ __device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( page_iter, entry_idx); thr_local_kv_offset[i] = paged_kv.protective_get_kv_offset( page_iter, kv_head_idx, entry_idx, - (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(), last_indptr); + (lane_idx % KV_THR_LAYOUT_COL) * + upcast_size(), + last_indptr); } page_produce_kv(k_smem, &k_smem_offset_w, paged_kv, 0, thr_local_kv_offset, chunk_size, tid); @@ -2255,7 +2265,9 @@ __device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( page_iter, entry_idx); thr_local_kv_offset[i] = paged_kv.protective_get_kv_offset( page_iter, kv_head_idx, entry_idx, - (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(), last_indptr); + (lane_idx % KV_THR_LAYOUT_COL) * + upcast_size(), + last_indptr); } memory::wait_group<1>(); block.sync(); @@ -2461,9 +2473,10 @@ hipError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Param ? static_cast(max_smem_per_threadblock) / 2u - q_smem_bytes_ : 0u; constexpr uint32_t min_valid_mma_kv_ = (HEAD_DIM_VO / 16u == 4u) ? 2u : 1u; - const uint32_t max_num_mma_kv_smem = std::max( - min_valid_mma_kv_, static_cast(kv_budget_ / ((HEAD_DIM_QK + HEAD_DIM_VO) * 16 * - NUM_WARPS_KV * sizeof(DTypeKV)))); + const uint32_t max_num_mma_kv_smem = + std::max(min_valid_mma_kv_, + static_cast(kv_budget_ / ((HEAD_DIM_QK + HEAD_DIM_VO) * 16 * NUM_WARPS_KV * + sizeof(kv_smem_type_t)))); DISPATCH_NUM_MMA_KV(min(max_num_mma_kv_smem, max_num_mma_kv_reg), NUM_MMA_KV, { using KTraits = @@ -2570,9 +2583,10 @@ hipError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params ? static_cast(max_smem_per_threadblock) / 2u - q_smem_bytes_ : 0u; constexpr uint32_t min_valid_mma_kv_ = (HEAD_DIM_VO / 16u == 4u) ? 2u : 1u; - const uint32_t max_num_mma_kv_smem = std::max( - min_valid_mma_kv_, static_cast(kv_budget_ / ((HEAD_DIM_QK + HEAD_DIM_VO) * 16 * - NUM_WARPS_KV * sizeof(DTypeKV)))); + const uint32_t max_num_mma_kv_smem = + std::max(min_valid_mma_kv_, + static_cast(kv_budget_ / ((HEAD_DIM_QK + HEAD_DIM_VO) * 16 * NUM_WARPS_KV * + sizeof(kv_smem_type_t)))); DISPATCH_NUM_MMA_KV(min(max_num_mma_kv_smem, max_num_mma_kv_reg), NUM_MMA_KV, { using KTraits = diff --git a/tests/rocm/test_fp8_kv_prefill.py b/tests/rocm/test_fp8_kv_prefill.py new file mode 100644 index 00000000000..1b9f4f0e103 --- /dev/null +++ b/tests/rocm/test_fp8_kv_prefill.py @@ -0,0 +1,832 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: Apache-2.0 +"""fp8 KV cache on the in-tree (fa2) prefill kernel, and the refusals around it. + +The oracle is that fp8 -> fp16 is lossless, so feeding the kernel an fp8 cache +must equal feeding the trusted 2-byte kernel the same values dequantized by +torch. Both runs select the same tile geometry -- the LDS tile is 2-byte either +way -- so the accumulation order matches and the results are bit-identical. + +pos_encoding_mode is NONE throughout: the fa2 RoPE path returns a different +answer on every call for *any* KV dtype (measured on an unmodified kernel), so +it cannot serve as a reference here. +""" + +import logging +import math + +import pytest +import torch +from attention_reference import naive_attention +from jit_utils import gen_decode_attention_modules, gen_prefill_attention_modules + +import flashinfer +from flashinfer.jit.core import logger + +logger.setLevel(logging.ERROR) + +FNUZ_DTYPES = [torch.float8_e4m3fnuz, torch.float8_e5m2fnuz] +OCP_DTYPES = [torch.float8_e4m3fn, torch.float8_e5m2] +WORKSPACE = 128 * 1024 * 1024 + + +def _skip_without_aiter(): + """Building an AITER spec imports aiter. The wheel is cp312-only, so a + py3.13 image has none and the refusal tests still have to run.""" + pytest.importorskip("aiter") + + +@pytest.fixture(autouse=True, scope="module") +def warmup_jit(): + common = ([0], [False, True], [False], [False]) # posenc, swa, softcap, f16qk + specs = ( + gen_prefill_attention_modules( + [torch.float16], [torch.float16] + FNUZ_DTYPES, [64, 128, 256], *common + ) + + gen_prefill_attention_modules( + [torch.bfloat16], [torch.bfloat16] + FNUZ_DTYPES, [128], *common + ) + # Decode too: single decode has its own kernel, which would otherwise + # compile serially inside the test body. Only the shapes the tests use. + + gen_decode_attention_modules( + [torch.float16], + [torch.float16] + FNUZ_DTYPES, + [64, 128], + [0], + [False], + [False], + ) + ) + # Both calls emit the shared helper specs; ninja rejects duplicate rules. + flashinfer.jit.build_jit_specs( + list({spec.name: spec for spec in specs}.values()), verbose=False + ) + yield + + +def _quantized_pair(shape, fp8_dtype, wide_dtype=torch.float16, device="cuda:0"): + """An fp8 tensor and its exact dequantization into wide_dtype.""" + src = torch.randn(shape, dtype=torch.float16, device=device) + q8 = src.to(fp8_dtype) + return q8, q8.to(wide_dtype) + + +def _assert_matches(got, ref): + # Bit-equality, not a tolerance: both runs use a 2-byte LDS tile and so the + # same geometry and accumulation order. A tolerance here would swallow a + # tile-selection change, which is the one thing this oracle exists to catch. + assert torch.equal(got, ref), ( + f"not bit-identical: max |delta| " + f"{(got.float() - ref.float()).abs().max().item():.3e}" + ) + + +@pytest.mark.parametrize("q_dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +@pytest.mark.parametrize("head_dim", [64, 128, 256]) +@pytest.mark.parametrize("qo_len,kv_len", [(1, 64), (37, 128), (77, 396), (577, 2048)]) +@pytest.mark.parametrize("causal", [False, True]) +def test_single_prefill_fp8_kv_matches_dequantized( + q_dtype, fp8_dtype, head_dim, qo_len, kv_len, causal +): + if q_dtype is torch.bfloat16 and head_dim != 128: + pytest.skip("bf16 queries are warmed up at head_dim 128 only") + torch.manual_seed(0) + num_qo_heads, num_kv_heads = 32, 8 + q = torch.randn(qo_len, num_qo_heads, head_dim, dtype=q_dtype, device="cuda:0") + k8, k16 = _quantized_pair((kv_len, num_kv_heads, head_dim), fp8_dtype, q_dtype) + v8, v16 = _quantized_pair((kv_len, num_kv_heads, head_dim), fp8_dtype, q_dtype) + + kwargs = dict(causal=causal, backend="fa2", pos_encoding_mode="NONE") + got = flashinfer.single_prefill_with_kv_cache(q, k8, v8, **kwargs) + ref = flashinfer.single_prefill_with_kv_cache(q, k16, v16, **kwargs) + _assert_matches(got, ref) + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +@pytest.mark.parametrize("head_dim", [64, 128]) +@pytest.mark.parametrize("page_size", [1, 16]) +def test_paged_prefill_fp8_kv_matches_dequantized(fp8_dtype, head_dim, page_size): + torch.manual_seed(0) + bs, qo_len, num_qo_heads, num_kv_heads = 3, 41, 32, 8 + # kv_len must exceed qo_len or the bottom-right causal mask blanks most + # query rows and both sides trivially agree on zeros. + pages_per_seq = max(5, -(-2 * qo_len // page_size)) + num_pages = bs * pages_per_seq + kv8, kv16 = _quantized_pair( + (num_pages, 2, page_size, num_kv_heads, head_dim), fp8_dtype + ) + dev = "cuda:0" + qo_indptr = torch.arange(0, bs + 1, dtype=torch.int32, device=dev) * qo_len + kv_indptr = torch.arange(0, bs + 1, dtype=torch.int32, device=dev) * pages_per_seq + kv_indices = torch.arange(0, num_pages, dtype=torch.int32, device=dev) + last_page_len = torch.full((bs,), page_size, dtype=torch.int32, device=dev) + q = torch.randn( + bs * qo_len, num_qo_heads, head_dim, dtype=torch.float16, device=dev + ) + + outs = [] + for cache, kv_dtype in ((kv8, fp8_dtype), (kv16, torch.float16)): + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + torch.empty(WORKSPACE, dtype=torch.uint8, device=dev), "NHD", backend="fa2" + ) + wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + causal=True, + q_data_type=torch.float16, + kv_data_type=kv_dtype, + ) + outs.append(wrapper.run(q, cache)) + _assert_matches(*outs) + + +def _decode( + kv_indptr, + kv_indices, + last_page_len, + q, + cache, + kv_dtype, + *, + tc, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, +): + wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + torch.empty(WORKSPACE, dtype=torch.uint8, device=q.device), + "NHD", + use_tensor_cores=tc, + backend="fa2", + ) + wrapper.plan( + kv_indptr, + kv_indices, + last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + q_data_type=torch.float16, + kv_data_type=kv_dtype, + ) + return wrapper.run(q, cache) + + +def _decode_inputs( + head_dim, + fp8_dtype, + bs=4, + num_qo_heads=32, + num_kv_heads=8, + page_size=16, + pages_per_seq=4, +): + torch.manual_seed(0) + dev = "cuda:0" + num_pages = bs * pages_per_seq + cache8, cache16 = _quantized_pair( + (num_pages, 2, page_size, num_kv_heads, head_dim), fp8_dtype + ) + shapes = dict( + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + page_size=page_size, + ) + idx = ( + torch.arange(0, bs + 1, dtype=torch.int32, device=dev) * pages_per_seq, + torch.arange(0, num_pages, dtype=torch.int32, device=dev), + torch.full((bs,), page_size, dtype=torch.int32, device=dev), + ) + q = torch.randn(bs, num_qo_heads, head_dim, dtype=torch.float16, device=dev) + return idx, q, cache8, cache16, shapes + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +@pytest.mark.parametrize("head_dim", [64, 128]) +def test_tensor_core_decode_fp8_kv(fp8_dtype, head_dim): + """The reported defect: use_tensor_cores=True routes decode at the prefill kernel.""" + idx, q, cache8, cache16, shapes = _decode_inputs(head_dim, fp8_dtype) + got = _decode(*idx, q, cache8, fp8_dtype, tc=True, **shapes) + ref = _decode(*idx, q, cache16, torch.float16, tc=True, **shapes) + _assert_matches(got, ref) + + +@pytest.mark.parametrize("kv_dtype", [torch.float16] + FNUZ_DTYPES) +@pytest.mark.parametrize("head_dim", [64, 128]) +def test_tensor_core_decode_matches_plain_decode(kv_dtype, head_dim): + """Cross-check the two decode kernels against each other. + + Nothing covered the tensor-core path with any dtype, which is why an fp8 + cache there reached a user as a ninja log. This also exercises plain decode + with e5m2fnuz, newly reachable now that the dtype maps name it. + """ + fp8 = kv_dtype if kv_dtype in FNUZ_DTYPES else torch.float8_e4m3fnuz + idx, q, cache8, cache16, shapes = _decode_inputs(head_dim, fp8) + cache = cache16 if kv_dtype == torch.float16 else cache8 + tc = _decode(*idx, q, cache, kv_dtype, tc=True, **shapes) + plain = _decode(*idx, q, cache, kv_dtype, tc=False, **shapes) + # Different kernels and reduction orders, so agreement is numerical. + torch.testing.assert_close(tc.float(), plain.float(), rtol=2e-2, atol=2e-2) + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +def test_ragged_prefill_fp8_kv_matches_dequantized(fp8_dtype): + torch.manual_seed(0) + dev = "cuda:0" + bs, qo_len, kv_len, num_qo_heads, num_kv_heads, head_dim = 3, 37, 128, 32, 8, 128 + q = torch.randn( + bs * qo_len, num_qo_heads, head_dim, dtype=torch.float16, device=dev + ) + k8, k16 = _quantized_pair((bs * kv_len, num_kv_heads, head_dim), fp8_dtype) + v8, v16 = _quantized_pair((bs * kv_len, num_kv_heads, head_dim), fp8_dtype) + qo_indptr = torch.arange(0, bs + 1, dtype=torch.int32, device=dev) * qo_len + kv_indptr = torch.arange(0, bs + 1, dtype=torch.int32, device=dev) * kv_len + + outs = [] + for k, v, kv_dtype in ((k8, v8, fp8_dtype), (k16, v16, torch.float16)): + wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper( + torch.empty(WORKSPACE, dtype=torch.uint8, device=dev), "NHD", backend="fa2" + ) + wrapper.plan( + qo_indptr, + kv_indptr, + num_qo_heads, + num_kv_heads, + head_dim, + causal=True, + q_data_type=torch.float16, + kv_data_type=kv_dtype, + ) + outs.append(wrapper.run(q, k, v)) + _assert_matches(*outs) + + +def test_ragged_prefill_fp8_query_is_refused_at_plan(): + """run() used to silently cast q, k AND v to f16 on an fp8 query, which would + now discard a caller's fp8 KV cache. The refusal moved to plan(), so that + branch is unreachable and has been removed.""" + dev = "cuda:0" + bs, qo_len, kv_len, num_qo_heads, num_kv_heads, head_dim = 2, 16, 64, 32, 8, 128 + qo_indptr = torch.arange(0, bs + 1, dtype=torch.int32, device=dev) * qo_len + kv_indptr = torch.arange(0, bs + 1, dtype=torch.int32, device=dev) * kv_len + wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper( + torch.empty(WORKSPACE, dtype=torch.uint8, device=dev), "NHD", backend="fa2" + ) + with pytest.raises((NotImplementedError, ValueError), match="(?i)fp8"): + wrapper.plan( + qo_indptr, + kv_indptr, + num_qo_heads, + num_kv_heads, + head_dim, + causal=True, + q_data_type=torch.float8_e4m3fnuz, + kv_data_type=torch.float8_e4m3fnuz, + ) + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +def test_mismatched_k_v_dtypes_are_refused(fp8_dtype): + """An fp8 k with a 2-byte v used to die in the 8-bit static_assert; with the + kernel serving fp8 it returned NaN instead.""" + dev = "cuda:0" + q = torch.randn(32, 32, 128, dtype=torch.float16, device=dev) + k = torch.randn(64, 8, 128, dtype=torch.float16, device=dev).to(fp8_dtype) + v = torch.randn(64, 8, 128, dtype=torch.float16, device=dev) + with pytest.raises(ValueError, match="single KV dtype"): + flashinfer.single_prefill_with_kv_cache(q, k, v, causal=True, backend="fa2") + + +@pytest.mark.parametrize("int_dtype", [torch.int8, torch.uint8]) +def test_integer_kv_is_refused(int_dtype): + """1-byte integers have no float interpretation; widening them would read + 0..255 as values. The dtype allowlist refuses them before any build.""" + dev = "cuda:0" + q = torch.randn(32, 32, 128, dtype=torch.float16, device=dev) + k = torch.randint(0, 8, (64, 8, 128), dtype=int_dtype, device=dev) + with pytest.raises(NotImplementedError, match="(?i)kv dtype"): + flashinfer.single_prefill_with_kv_cache(q, k, k, causal=True, backend="fa2") + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES + OCP_DTYPES) +def test_fp8_query_is_refused(fp8_dtype): + dev = "cuda:0" + q = torch.randn(32, 8, 128, dtype=torch.float16, device=dev).to(fp8_dtype) + k = torch.randn(64, 8, 128, dtype=torch.float16, device=dev).to(fp8_dtype) + with pytest.raises((NotImplementedError, ValueError), match="(?i)fp8"): + flashinfer.single_prefill_with_kv_cache(q, k, k, causal=True, backend="fa2") + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES + OCP_DTYPES) +def test_fp8_query_with_wide_kv_is_refused(fp8_dtype): + """A typed refusal, never a bare assert. + + The fnuz spellings are already refused earlier; the OCP ones reached the + fp8-query branch before any backend was chosen and failed its + `assert q.dtype == k.dtype == v.dtype` with no message. + """ + dev = "cuda:0" + q = torch.randn(32, 8, 128, dtype=torch.float16, device=dev).to(fp8_dtype) + k = torch.randn(64, 8, 128, dtype=torch.float16, device=dev) + with pytest.raises(NotImplementedError, match="(?i)fp8|dtype"): + flashinfer.single_prefill_with_kv_cache(q, k, k, causal=True) + + +@pytest.mark.parametrize("fp8_dtype", OCP_DTYPES) +def test_fp8_query_with_sliding_window_is_refused(fp8_dtype): + """Also a sentence, not `AssertionError: assert window_left == -1`. + + OCP only: an fnuz query is refused earlier -- e4m3fnuz by the fa2 fp8-query + rejection, e5m2fnuz by the allowlist -- so OCP is the one spelling that + reaches this branch, and it was the last bare assert on the fp8-query path. + """ + dev = "cuda:0" + q = torch.randn(32, 8, 128, dtype=torch.float16, device=dev).to(fp8_dtype) + k = torch.randn(64, 8, 128, dtype=torch.float16, device=dev).to(fp8_dtype) + with pytest.raises(NotImplementedError, match="(?i)sliding window"): + flashinfer.single_prefill_with_kv_cache(q, k, k, causal=True, window_left=16) + + +@pytest.mark.parametrize("ocp_dtype", OCP_DTYPES) +def test_ocp_fp8_kv_is_refused(ocp_dtype): + """dtype_map_hip sends OCP to the fnuz type, whose bias is one greater.""" + dev = "cuda:0" + q = torch.randn(32, 32, 128, dtype=torch.float16, device=dev) + k = torch.randn(64, 8, 128, dtype=torch.float16, device=dev).to(ocp_dtype) + with pytest.raises(NotImplementedError, match="(?i)ocp|bias"): + flashinfer.single_prefill_with_kv_cache(q, k, k, causal=True, backend="fa2") + + +@pytest.mark.parametrize( + "dtype_q,dtype_kv,dtype_o,single_allowed,batch_allowed", + [ + (torch.int8, torch.int8, torch.int8, False, False), + (torch.uint8, torch.uint8, torch.uint8, False, False), + (torch.float32, torch.float32, torch.float32, False, False), + (torch.float8_e5m2fnuz, torch.float8_e5m2fnuz, torch.bfloat16, False, False), + (torch.float16, torch.bfloat16, torch.float16, False, False), + (torch.float16, torch.float16, torch.int8, False, False), + (torch.float16, torch.float16, torch.float16, True, True), + (torch.bfloat16, torch.bfloat16, torch.bfloat16, True, True), + # e4m3 in both arch spellings: served by batch_prefill_paged_aiter.cu, + # which is the only AITER launcher with an fp8 arm. single_prefill and + # batch_ragged are fp16/bf16 only, so the single seam must refuse it. + (torch.float8_e4m3fnuz, torch.float8_e4m3fnuz, torch.bfloat16, False, True), + (torch.float8_e4m3fn, torch.float8_e4m3fn, torch.bfloat16, False, True), + ], +) +def test_aiter_generator_dtype_allowlist( + dtype_q, dtype_kv, dtype_o, single_allowed, batch_allowed +): + """Equality alone let int8/uint8 through: both are in dtype_map_hip, so they + built a URI and reached ninja with no kernel behind them.""" + from flashinfer.jit.rocm.modules import ( + gen_batch_prefill_module, + gen_single_prefill_module, + ) + + cases = ( + ( + single_allowed, + lambda: gen_single_prefill_module( + "aiter", dtype_q, dtype_kv, dtype_o, 128, 128, 0, False, False, False + ), + ), + ( + batch_allowed, + lambda: gen_batch_prefill_module( + "aiter", + dtype_q, + dtype_kv, + dtype_o, + torch.int32, + 128, + 128, + 0, + False, + False, + False, + ), + ), + ) + for allowed, call in cases: + if allowed: + _skip_without_aiter() + call() + else: + with pytest.raises(NotImplementedError, match="(?i)aiter"): + call() + + +def test_every_public_generator_refuses_before_its_uri(): + """Swept rather than spot-checked: each of these builds a URI by indexing + filename_safe_dtype_map, so an unmapped dtype raised a bare KeyError unless + the allowlist ran first. float32 is the reachable unmapped case. + """ + from flashinfer.jit.rocm import modules as m + + dt = torch.float32 + calls = { + "gen_single_prefill_module": lambda: m.gen_single_prefill_module( + "fa2", dt, dt, dt, 128, 128, 0, False, False, False + ), + "gen_batch_prefill_module": lambda: m.gen_batch_prefill_module( + "fa2", dt, dt, dt, torch.int32, 128, 128, 0, False, False, False + ), + "gen_single_decode_module": lambda: m.gen_single_decode_module( + dt, dt, dt, 128, 128, 0, False, False + ), + "gen_batch_decode_module": lambda: m.gen_batch_decode_module( + dt, dt, dt, torch.int32, 128, 128, 0, False, False + ), + "gen_batch_decode_aiter_module": lambda: m.gen_batch_decode_aiter_module( + dt, dt, dt, 128, 128 + ), + "gen_pod_module": lambda: m.gen_pod_module( + dt, dt, dt, 128, 0, False, False, False, torch.int32, 0, False, False + ), + "gen_batch_pod_module": lambda: m.gen_batch_pod_module( + dt, dt, dt, 128, 0, False, False, False, torch.int32, 0, False, False + ), + "gen_customize_single_decode_module": lambda: m.gen_customize_single_decode_module( + "sweep_csd", dt, dt, dt, 128, 128, ["v"], ["float"], [], [], "V", "v" + ), + "gen_customize_single_prefill_module": lambda: m.gen_customize_single_prefill_module( + "fa2", "sweep_csp", dt, dt, dt, 128, 128, ["v"], ["float"], [], [], "V", "v" + ), + "gen_customize_batch_decode_module": lambda: m.gen_customize_batch_decode_module( + "sweep_cbd", + dt, + dt, + dt, + torch.int32, + 128, + 128, + ["v"], + ["float"], + [], + [], + "V", + "v", + ), + "gen_customize_batch_prefill_module": lambda: m.gen_customize_batch_prefill_module( + "fa2", + "sweep_cbp", + dt, + dt, + dt, + torch.int32, + 128, + 128, + ["v"], + ["float"], + [], + [], + "V", + "v", + ), + "gen_customize_pod_module": lambda: m.gen_customize_pod_module( + "sweep_cpod", dt, dt, dt, torch.int32, 128, "V", "V" + ), + "gen_customize_batch_pod_module": lambda: m.gen_customize_batch_pod_module( + "sweep_cbpod", dt, dt, dt, torch.int32, 128, "V", "V" + ), + } + # Derived, not hand-listed: a new generator joins the sweep by existing. + # gen_customize_* take a uri, so they cannot be called generically above, + # but they must still be covered -- assert the enumeration is complete. + discovered = { + n + for n in dir(m) + if n.startswith("gen_") and n.endswith("_module") and callable(getattr(m, n)) + } + assert discovered <= set(calls), ( + f"generators missing from this sweep: {sorted(discovered - set(calls))}" + ) + for name, call in calls.items(): + # Named, so a regression says which builder stopped refusing. + with pytest.raises(NotImplementedError, match="(?i)not supported|fp8"): + try: + call() + except KeyError as exc: # pragma: no cover - the bug being guarded + raise AssertionError( + f"{name} built its URI before validating: KeyError({exc})" + ) from exc + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +def test_pod_refuses_fp8_kv(fp8_dtype): + """POD sizes its tiles with upstream's CUDA register constant, so it can ask + for a NUM_MMA_KV the prefill dispatchers never produce. Guarded at the + generator, which is the single seam both POD wrappers build through.""" + from flashinfer.jit.rocm.modules import gen_batch_pod_module, gen_pod_module + + common = dict( + head_dim=128, + pos_encoding_mode_p=0, + use_sliding_window_p=False, + use_logits_soft_cap_p=False, + use_fp16_qk_reduction=False, + dtype_idx=torch.int32, + pos_encoding_mode_d=0, + use_sliding_window_d=False, + use_logits_soft_cap_d=False, + ) + for gen in (gen_pod_module, gen_batch_pod_module): + with pytest.raises(NotImplementedError, match="(?i)fp8"): + gen(torch.float16, fp8_dtype, torch.float16, **common) + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +def test_single_decode_mismatched_k_v_dtypes_are_refused(fp8_dtype): + """The batch wrapper reaches the batch C++ macro, so it cannot cover this. + + Covers the Python pre-check only, which saves a cold JIT. `_check_kv_dtypes_match` + runs before any module is built, so no argument to this entry point reaches + `single_decode.cu`'s `CHECK_KV_DTYPES_MATCH` -- that one backstops callers + who build the module themselves, which is the point of putting it in C++. + """ + dev = "cuda:0" + nq, nkv, hd, kv_len = 32, 8, 128, 64 + q = torch.randn(nq, hd, dtype=torch.float16, device=dev) + k8 = torch.randn(kv_len, nkv, hd, dtype=torch.float16, device=dev).to(fp8_dtype) + v16 = torch.randn(kv_len, nkv, hd, dtype=torch.float16, device=dev) + with pytest.raises((ValueError, RuntimeError), match="single KV dtype"): + flashinfer.decode.single_decode_with_kv_cache(q, k8, v16) + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +@pytest.mark.parametrize("head_dim", [64, 128]) +def test_single_decode_fp8_kv_matches_dequantized(fp8_dtype, head_dim): + """SingleDecodeWithKVCacheKernel is its own kernel, not the batch one, and + e5m2fnuz only became expressible when this series added it to dtype_map_hip. + + Tolerance, not bit-equality: an fp8 cache selects a different geometry from + the 2-byte reference (vec_size 16 vs 8, tile_size_per_bdx 2 vs 1), so this + cannot pin the fp8 shape the way the prefill oracle pins prefill's. atol is + 5x the worst measured delta (6.1e-5) and 6x under a 1% output shift. + """ + torch.manual_seed(0) + dev = "cuda:0" + nq, nkv, kv_len = 32, 8, 64 + q = torch.randn(nq, head_dim, dtype=torch.float16, device=dev) + k8, k16 = _quantized_pair((kv_len, nkv, head_dim), fp8_dtype) + v8, v16 = _quantized_pair((kv_len, nkv, head_dim), fp8_dtype) + + got = flashinfer.decode.single_decode_with_kv_cache(q, k8, v8) + ref = flashinfer.decode.single_decode_with_kv_cache(q, k16, v16) + torch.testing.assert_close(got.float(), ref.float(), rtol=1e-3, atol=3e-4) + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +def test_decode_mismatched_k_v_dtypes_are_refused(fp8_dtype): + """Decode specializes on k.dtype and casts both pointers to it. Measured + before the guard: an fp8 k with an fp16 v returned NaN and no error.""" + dev = "cuda:0" + bs, nq, nkv, hd, page_size, pages = 2, 32, 8, 128, 16, 4 + k8 = torch.randn( + bs * pages, page_size, nkv, hd, dtype=torch.float16, device=dev + ).to(fp8_dtype) + v16 = torch.randn(bs * pages, page_size, nkv, hd, dtype=torch.float16, device=dev) + kv_indptr = torch.arange(0, bs + 1, dtype=torch.int32, device=dev) * pages + kv_indices = torch.arange(0, bs * pages, dtype=torch.int32, device=dev) + last = torch.full((bs,), page_size, dtype=torch.int32, device=dev) + q = torch.randn(bs, nq, hd, dtype=torch.float16, device=dev) + + wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + torch.empty(WORKSPACE, dtype=torch.uint8, device=dev), "NHD", backend="fa2" + ) + wrapper.plan( + kv_indptr, + kv_indices, + last, + nq, + nkv, + hd, + page_size, + q_data_type=torch.float16, + kv_data_type=fp8_dtype, + ) + # RuntimeError, not ValueError: plan() already built the module, so the C++ + # entry point is what refuses this rather than a Python pre-check. + with pytest.raises(RuntimeError, match="single KV dtype"): + wrapper.run(q, (k8, v16)) + + +def test_mixed_wide_kv_dtype_is_refused(): + """The MFMA is instantiated from DTypeQ, so a bf16 cache under an fp16 query + is read as fp16. aot.py already excludes the pair from the AOT matrix.""" + dev = "cuda:0" + q = torch.randn(32, 32, 128, dtype=torch.float16, device=dev) + k = torch.randn(64, 8, 128, dtype=torch.bfloat16, device=dev) + with pytest.raises(NotImplementedError, match="(?i)differs from query dtype"): + flashinfer.single_prefill_with_kv_cache(q, k, k, causal=True, backend="fa2") + + +@pytest.mark.parametrize("ocp_dtype", OCP_DTYPES) +def test_customize_decode_generators_refuse_ocp(ocp_dtype): + """gen_{single,batch}_decode_module delegate here, and both customize + generators are publicly re-exported, so guarding only the wrappers left the + public JIT API able to compile OCP fp8 as fnuz.""" + from flashinfer.jit.rocm.modules import ( + gen_customize_batch_decode_module, + gen_customize_single_decode_module, + ) + + common = dict( + dtype_q=torch.float16, + dtype_kv=ocp_dtype, + dtype_o=torch.float16, + head_dim_qk=128, + head_dim_vo=128, + additional_tensor_names=[], + additional_tensor_dtypes=[], + additional_scalar_names=[], + additional_scalar_dtypes=[], + variant_name="v", + variant_decl="", + ) + with pytest.raises(NotImplementedError, match="(?i)ocp"): + gen_customize_single_decode_module(uri="u", **common) + with pytest.raises(NotImplementedError, match="(?i)ocp"): + gen_customize_batch_decode_module(uri="u", idtype=torch.int32, **common) + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +@pytest.mark.parametrize("window_left", [0, 63]) +def test_single_prefill_fp8_kv_sliding_window(fp8_dtype, window_left): + """Windowing changes which KV tiles are produced and predicated, which is + the kFillZero/kNoFill split the fp8 produce arm rewrites.""" + torch.manual_seed(0) + qo_len, kv_len, num_qo_heads, num_kv_heads, head_dim = 77, 396, 32, 8, 128 + q = torch.randn( + qo_len, num_qo_heads, head_dim, dtype=torch.float16, device="cuda:0" + ) + k8, k16 = _quantized_pair((kv_len, num_kv_heads, head_dim), fp8_dtype) + v8, v16 = _quantized_pair((kv_len, num_kv_heads, head_dim), fp8_dtype) + + kwargs = dict( + causal=True, backend="fa2", pos_encoding_mode="NONE", window_left=window_left + ) + got = flashinfer.single_prefill_with_kv_cache(q, k8, v8, **kwargs) + ref = flashinfer.single_prefill_with_kv_cache(q, k16, v16, **kwargs) + _assert_matches(got, ref) + + +def test_unmapped_dtype_refuses_before_the_uri(): + """The URI builder indexes filename_safe_dtype_map, so a dtype absent from it + raised KeyError before the allowlist ran. float32 is the reachable case.""" + dev = "cuda:0" + q = torch.randn(32, 32, 128, dtype=torch.float32, device=dev) + with pytest.raises(NotImplementedError, match="(?i)query dtype"): + flashinfer.single_prefill_with_kv_cache(q, q, q, causal=True, backend="fa2") + + +def _block_sparse_out(k, v, kv_dtype, q, n, nh, head_dim, dev): + block = torch.ones(n // 16, n // 16, dtype=torch.bool) + indptr = torch.zeros(block.size(0) + 1, dtype=torch.int32) + indptr[1:] = torch.cumsum(block.sum(dim=1), 0) + indices = block.nonzero()[:, 1].to(torch.int32) + wrapper = flashinfer.BlockSparseAttentionWrapper( + torch.empty(WORKSPACE, dtype=torch.uint8, device=dev) + ) + wrapper.plan( + indptr.to(dev), + indices.to(dev), + n, + n, + 16, + 16, + nh, + nh, + head_dim, + q_data_type=torch.float16, + kv_data_type=kv_dtype, + ) + return wrapper.run(q, k, v) + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +def test_block_sparse_fp8_kv_matches_dequantized(fp8_dtype): + """Block-sparse builds through the batch-prefill module, so deleting the + blanket static_assert made fp8 reachable here too.""" + torch.manual_seed(0) + n, nh, head_dim, dev = 128, 4, 128, "cuda:0" + q = torch.randn(n, nh, head_dim, dtype=torch.float16, device=dev) + k8, k16 = _quantized_pair((n, nh, head_dim), fp8_dtype) + v8, v16 = _quantized_pair((n, nh, head_dim), fp8_dtype) + args = (q, n, nh, head_dim, dev) + got = _block_sparse_out(k8, v8, fp8_dtype, *args) + ref = _block_sparse_out(k16, v16, torch.float16, *args) + _assert_matches(got, ref) + + +def test_block_sparse_mismatched_k_v_dtypes_are_refused(): + """sparse.py is upstream and unshadowed, so the C++ entry point is the only + seam that can refuse this without carrying an edit into every sync.""" + torch.manual_seed(0) + n, nh, head_dim, dev = 128, 4, 128, "cuda:0" + q = torch.randn(n, nh, head_dim, dtype=torch.float16, device=dev) + k16 = torch.randn(n, nh, head_dim, dtype=torch.float16, device=dev) + v_bf = torch.randn(n, nh, head_dim, dtype=torch.bfloat16, device=dev) + with pytest.raises(RuntimeError, match="single KV dtype"): + _block_sparse_out(k16, v_bf, torch.float16, q, n, nh, head_dim, dev) + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +@pytest.mark.parametrize("causal", [False, True]) +def test_fp8_kv_scales_fold_as_documented(fp8_dtype, causal): + """k_scale/v_scale are how a real cache is read, and they are folded outside + the kernel: k_scale into sm_scale, v_scale onto the output. + + Asserting those two identities rather than comparing against a rescaled + cache -- that comparison is dominated by rounding the output to fp16 while + it is still in the un-scaled domain, not by whether the scales are right. + """ + torch.manual_seed(0) + qo_len, kv_len, num_qo_heads, num_kv_heads, head_dim = 77, 396, 32, 8, 128 + dev = "cuda:0" + q = torch.randn(qo_len, num_qo_heads, head_dim, dtype=torch.float16, device=dev) + k8, _ = _quantized_pair((kv_len, num_kv_heads, head_dim), fp8_dtype) + v8, _ = _quantized_pair((kv_len, num_kv_heads, head_dim), fp8_dtype) + base_sm_scale = 1.0 / math.sqrt(head_dim) + common = dict(causal=causal, backend="fa2", pos_encoding_mode="NONE") + + def run(**kw): + return flashinfer.single_prefill_with_kv_cache(q, k8, v8, **common, **kw) + + # k_scale folds into sm_scale, so naming it is the same call. + k_scale = 0.37 + torch.testing.assert_close( + run(sm_scale=base_sm_scale, k_scale=k_scale).float(), + run(sm_scale=base_sm_scale * k_scale).float(), + rtol=0, + atol=0, + ) + + # v_scale is a post-multiply on the output, so it must be exactly linear. + # atol is one fp16 subnormal: v_scale multiplies the fp16 output in place, + # and only results that land in the subnormal range round differently here. + v_scale = 0.25 # a power of two, so nothing else can round + torch.testing.assert_close( + run(sm_scale=base_sm_scale, v_scale=v_scale).float(), + run(sm_scale=base_sm_scale).float() * v_scale, + rtol=0, + atol=6e-8, + ) + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +def test_fp8_kv_quantization_quality(fp8_dtype): + """What the fp8 cache costs in accuracy, against the *unquantized* values. + + The reference must be the pre-quantization tensor: dequantizing the same + fp8 cache the kernel reads makes the quantization error zero on both sides + and measures nothing. Bound is relative L2, which the per-element max is + too outlier-driven to give. + """ + torch.manual_seed(0) + qo_len, kv_len, num_heads, head_dim = 64, 256, 8, 128 + dev = "cuda:0" + q = torch.randn(qo_len, num_heads, head_dim, dtype=torch.float16, device=dev) + k_src = torch.randn(kv_len, num_heads, head_dim, dtype=torch.float16, device=dev) + v_src = torch.randn(kv_len, num_heads, head_dim, dtype=torch.float16, device=dev) + k8, v8 = k_src.to(fp8_dtype), v_src.to(fp8_dtype) + + got = flashinfer.single_prefill_with_kv_cache(q, k8, v8, causal=True, backend="fa2") + ref, _ = naive_attention(q.float(), k_src.float(), v_src.float(), causal=True) + rel_l2 = ((got.float() - ref).norm() / ref.norm()).item() + # e5m2fnuz keeps 2 mantissa bits against e4m3fnuz's 3, so it earns ~3x. + bound = 0.05 if fp8_dtype is torch.float8_e4m3fnuz else 0.16 + assert rel_l2 < bound, f"{fp8_dtype} relative L2 {rel_l2:.4f} exceeds {bound}" + + +@pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) +def test_fp8_kv_matches_its_own_dequantization_in_fp32(fp8_dtype): + """The fp32 twin of the fp16 oracle: same quantized values both sides, so + any gap is the kernel's accumulation rather than the cache's coarseness.""" + torch.manual_seed(0) + qo_len, kv_len, num_heads, head_dim = 64, 256, 8, 128 + dev = "cuda:0" + q = torch.randn(qo_len, num_heads, head_dim, dtype=torch.float16, device=dev) + k8, _ = _quantized_pair((kv_len, num_heads, head_dim), fp8_dtype) + v8, _ = _quantized_pair((kv_len, num_heads, head_dim), fp8_dtype) + + got = flashinfer.single_prefill_with_kv_cache(q, k8, v8, causal=True, backend="fa2") + # naive_attention, not F.scaled_dot_product_attention: torch's is_causal + # anchors the mask top-left, FlashInfer bottom-right, and qo_len != kv_len here. + ref, _ = naive_attention( + q.float(), k8.to(torch.float32), v8.to(torch.float32), causal=True + ) + torch.testing.assert_close(got.float(), ref.float(), rtol=2e-2, atol=2e-2)