From 518806385ffed94452c80211ac7b069cd05b3650 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sat, 12 Sep 2026 10:27:48 -0400 Subject: [PATCH 01/18] feat(rocm): serve an fp8 KV cache in the in-tree prefill kernel fp16 query with an fp8 KV cache failed at JIT build on every fa2 prefill entry point, with a ninja log rather than an error. Measured on gfx942 at 6dc0b0fc0: 54 compiler error lines from BatchDecode(use_tensor_cores=True) and from BatchPrefillWithPagedKVCacheWrapper, 27 from single_prefill_with_kv_cache. Every one traces to the blanket static_assert at prefill.cuh:108, which kills SharedStorage; the two static_assert(false) arms in compute_qk/compute_sfm_v never fire because the kernel body is never instantiated. KV now arrives from HBM as fp8 and is dequantized into a 2-byte LDS tile (DTypeKVSmem) in the global->smem produce path. This keeps swizzle, fragment addressing and the smem budget bit-identical to bf16, so permuted_smem.cuh is untouched and the fp16/bf16 path emits the instructions it did before -- produce_kv_unit forwards verbatim when the cache is already 2-byte. Measured on gfx942, fp16 single prefill before and after, llama3-8b shapes: 0.074/0.240/0.744/2.263/8.701 ms against 0.074/0.239/0.744/2.277/8.708. The fp8 win is HBM bandwidth and cache capacity, both properties of the HBM type. kv_smem_type_t keys on the fp8 types, not on sizeof()==1: an int8/uint8 cache has no float interpretation, and widening it would send it through vec_cast's primary template, which reads 0..255 as a value. Keying on size would have turned the old blanket static_assert from a refusal into silent garbage. Dequant goes through vec_cast, not vec_cast: the latter's convert_e4m3x2_to_f16x2 decodes fnuz with the OCP bias (-7 where the format's bias is 8) and is 2x wrong. It stays unreachable. It is also built from raw arrays rather than vec_t, because on gfx950 HIP_FP8_TYPE_FNUZ=0 leaves vec_t with a host-only default constructor -- identical source compiled clean on gfx942 and failed on gfx950 with `no matching constructor` at vec_dtypes.h:853. The granule is staged through one 4-byte register load because vec_cast is a per-element loop and converting straight off gptr can emit four global_load_ubyte where the 2-byte path emits one global_load_dwordx2. This measured no difference on gfx942 (tensor-core decode identical to three decimal places), so it is kept for not depending on the compiler coalescing, not for speed. The dead sizeof(DTypeKV)==1 arms are deleted rather than repaired: they carry upstream's 8-half CUDA fragment geometry where a CDNA fragment is 4 halves, and they call advance_offset_by_column without the k_col_idx that k128B_16Row requires. The two IsInvalid() fp8 clauses go with them -- at CTA_TILE_Q=128/head_dim=128 the budget gives NUM_MMA_KV=1, so 1*2 % 4 != 0 would have failed a working config at runtime. Not chosen: an fp8 smem tile with in-loop fragment dequant. NUM_MMA_KV is register-capped (ELEMS_PER_FRAGMENT = 16*16/64 = 4), so halving the KV smem type widens exactly one cell of the tile table, and there is no async copy to forfeit -- commit_group/wait_group are no-ops and load_64b is a plain assignment, so every KV byte already passes through VGPRs. The seam back is the one conditional_t in kv_smem_type_t. Co-Authored-By: Claude --- .../flashinfer/rocm/attention/batch_pod.cuh | 4 + include/flashinfer/rocm/attention/pod.cuh | 4 + include/flashinfer/rocm/attention/prefill.cuh | 204 +++++++++--------- 3 files changed, 115 insertions(+), 97 deletions(-) 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..597f2d49519 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,42 @@ __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) { + // One register load, then convert: vec_cast is a per-element loop. Raw + // arrays because gfx950 leaves vec_t without a device ctor, and + // memcpy rather than a cast because -O3 keeps strict aliasing on. + // wide is 8-byte aligned because vec_cast reads it through float2. + alignas(sizeof(uint32_t)) DTypeKV packed[GRANULE]; + __builtin_memcpy(packed, gptr, sizeof(uint32_t)); + alignas(sizeof(b64_t)) float wide[GRANULE]; + alignas(sizeof(b64_t)) DTypeKVSmem narrow[GRANULE]; + vec_cast::template cast(wide, packed); + 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 +354,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 +371,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 +392,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 +415,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 +645,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 +775,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 +795,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 +1019,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 +1030,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 +1504,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 +1679,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 +1942,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 +2217,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 +2261,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 +2469,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 +2579,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 = From 719560c2a02f326d2f861f89996ff175b9d6e54c Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sat, 12 Sep 2026 13:24:23 -0400 Subject: [PATCH 02/18] fix(rocm): refuse the fp8 shapes the prefill kernel cannot serve Guards at module-build time rather than in IsInvalid(), which fires host-side only after a multi-minute JIT. fp8 query and output: the MFMA path is f16f16f32, so both must be 2-byte. KernelTraits asserts sizeof(DTypeQ)==2 to keep the C++ side honest. The OCP encodings: dtype_map_hip sends torch.float8_e4m3fn and torch.float8_e5m2 onto the *fnuz* C++ types, whose exponent bias is 8 against OCP's 7, so a caller passing one would get a silently 2x-wrong result. Refused for prefill and for decode, which shares that map -- without the decode half, the claim in backends.md and arch_caps that the OCP spellings are refused would have been false for plain decode. The guard sits at gen_customize_{single,batch}_prefill_module, not at the outer generators: both delegate there, and so does the jit_args custom variant path (rocm/prefill.py:2110, 3397), which the outer guards missed entirely. It is keyed on backend != "aiter" rather than == "fa2" because ROCm logs and ignores an explicit "fa3" and routes it to the same in-tree kernel -- keying on "fa2" let backend="fa3" past. AITER keeps its own fp8 dtypes: it compiles kernels against aiter.dtypes.fp8, which is OCP e4m3fn on gfx950. fp8 KV on POD: pod.cuh and batch_pod.cuh carry upstream's CUDA register constant (8, against CDNA's ELEMS_PER_FRAGMENT/NUM_MMA_Q = 4) and size the smem budget off the query dtype, so they can instantiate KernelTraits at a NUM_MMA_KV the prefill dispatchers never produce. Before this series the blanket static_assert caught that; deleting it would have let POD compile fp8 at a geometry nothing has tested. The POD-wide refusal is raised before the dtype one, or the user is told to re-quantize into something POD also rejects. Mismatched k/v dtypes were never checked: the module is built from k.dtype alone and v is read with it. Only 2-byte KV used to reach the kernel so the 8-bit static_assert caught it; with an fp8 cache it returned NaN silently. Measured on gfx942, k=e4m3fnuz with v=fp16: "NO ERROR -- out dtype=torch.float16 finite=False". torch.float8_e5m2fnuz was absent from both filename_safe_dtype_map and dtype_map_hip, so it raised KeyError from the URI builder rather than anything legible -- hit while running the oracle. The fp8_enabled raise in the customize generators keeps its behaviour but loses its message: it gates upstream's fa3-only sm90 fp8-*query* templates, and read as a blanket refusal of what this series implements. BatchPrefillWithRaggedKVCacheWrapper.run() silently cast q, k AND v to f16 on an fp8 query, which would now discard a caller's fp8 KV cache. That branch is deleted rather than converted to a raise: plan() refuses an fp8 q_data_type, and _check_cached_qkv_data_type rejects a runtime q that disagrees with the plan, so it is unreachable. Four gaps the PR review found, all reproduced before accepting: The allowlist admitted mixed 2-byte operands. The MFMA is instantiated from DTypeQ, so a bf16 cache under an fp16 query is read as fp16; aot.py already skips that pair and DISPATCH_PYTORCH_QKV_DTYPE_TO_CTYPE admits only kv == q or wide-q + fp8-kv. A 2-byte KV must now equal the query. The decode guard sat on gen_{single,batch}_decode_module, but both delegate to the customize generators, which flashinfer/jit/__init__.py and jit/rocm/api.py re-export -- so the public JIT API could still compile OCP fp8 as fnuz. Moved down, as the prefill guards already were. Mismatched k/v reached both decode entry points too: decode specializes on k.dtype and casts both pointers to it. Measured on gfx942, paged decode with k=e4m3fnuz and v=fp16: "NO ERROR -- finite=False". The prefill helper is now shared through the import edge decode.py already has on prefill.py, rather than a new module. Two more from later rounds. The allowlist ran after the URI was built, and the URI indexes filename_safe_dtype_map -- so a dtype absent from that map (float32 is the reachable one) still raised a bare KeyError instead of the refusal the docs advertise. The check moved ahead of get_*_uri in all four prefill/decode loaders, keyed on backend != "aiter" for prefill so The k/v check was added to flashinfer/sparse.py, which has no prior AMD commit -- "git log origin/amd-integration -- flashinfer/sparse.py" returns upstream commits only -- so it would have carried a conflict hunk into every upstream sync. It is now CHECK_KV_DTYPES_MATCH at the six C++ entry points instead, and sparse.py is untouched. That covers strictly more: both block-sparse wrappers call paged_run, and POD has no shadow module under flashinfer/rocm/, so neither could be reached from Python without editing an upstream file. The Python check stays as well, for a different job: it refuses before the URI and before any AITER bootstrap, so a mismatch costs no JIT build. A/B, with the sparse.py check deleted: block-sparse with k=fp16, v=bf16 raises "RuntimeError: paged_k_cache has dtype Half but paged_v_cache has BFloat16" -- only the C++ can produce that. The same probe found block-sparse fp8 KV is bitwise identical to the dequantized reference for both fnuz dtypes. Co-Authored-By: Claude --- csrc/rocm/batch_decode.cu | 1 + csrc/rocm/batch_pod.cu | 2 + csrc/rocm/batch_prefill.cu | 2 + csrc/rocm/pod.cu | 2 + csrc/rocm/pytorch_extension_utils.h | 8 ++ csrc/rocm/single_decode.cu | 2 +- csrc/rocm/single_prefill.cu | 1 + flashinfer/jit/rocm/modules.py | 157 ++++++++++++++++++++++++++-- flashinfer/jit/utils.py | 5 + flashinfer/rocm/decode.py | 10 ++ flashinfer/rocm/prefill.py | 59 +++++++---- 11 files changed, 224 insertions(+), 25 deletions(-) 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/flashinfer/jit/rocm/modules.py b/flashinfer/jit/rocm/modules.py index e7ce05c81ba..b30a9d84fcb 100644 --- a/flashinfer/jit/rocm/modules.py +++ b/flashinfer/jit/rocm/modules.py @@ -97,6 +97,99 @@ 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) + +# AITER prefill serves e4m3 only, and both spellings because aiter.dtypes.fp8 is +# arch-dependent; it writes bf16 for an fp8 query. 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) -> 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. + """ + if dtype_q != dtype_kv: + raise NotImplementedError( + f"{what}: AITER requires equal query and KV dtypes; got " + f"{dtype_q} and {dtype_kv}." + ) + for role, dtype in (("query", dtype_q), ("KV", dtype_kv)): + if dtype not in _WIDE_DTYPES + _AITER_FP8_DTYPES: + 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 dtype_o not in _WIDE_DTYPES: + raise NotImplementedError( + f"{what}: AITER output dtype {dtype_o} is not supported; it writes " + "float16 or bfloat16, and bfloat16 for an fp8 query." + ) + + +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)." + ) + # Prefill only: the MFMA is instantiated from DTypeQ, so a bf16 cache under + # an fp16 query is read as fp16. Decode cast_loads q and k into vec_t + # separately and accumulates in float, so mixed 2-byte dtypes are fine there. + 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, @@ -281,6 +374,14 @@ 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) + 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, @@ -460,6 +561,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) + 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 +672,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 +760,24 @@ 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) + # 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 +929,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 +1016,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) + # 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 +1252,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 +1302,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, 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/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 844b45caaf0..b84cf3f3795 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, @@ -201,6 +201,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 @@ -518,19 +523,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." ) @@ -1071,6 +1090,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 @@ -1851,6 +1875,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 @@ -1885,7 +1910,14 @@ def single_prefill_with_kv_cache( # 1. unsupported feature # 2. dtype check assert window_left == -1 - assert q.dtype == k.dtype == v.dtype + # Not an assert: a wide KV cache under an fp8 query reaches here before + # any backend is chosen, and no backend serves that pair. + 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) @@ -4217,15 +4249,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: From 8e761617d1d878354311f85696fb5e37309d6c01 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sat, 12 Sep 2026 13:24:36 -0400 Subject: [PATCH 03/18] test(rocm): cover fp8 KV on the prefill kernel, and the refusals 56 cases on gfx942, all green. The equivalence oracle is that fp8 -> fp16 is lossless, so an fp8 cache must give what the trusted 2-byte kernel gives when torch does the dequant. Both runs pick the same tile geometry -- the LDS tile is 2-byte either way -- so the bar is bit-equality, and every case meets it rather than merely landing inside 1e-3. test_tensor_core_decode_fp8_kv covers the reported defect. Its float16 parametrization is not padding: no test anywhere planned a tensor-core decode on the fa2 path with any dtype, which is why an fp8 cache there reached a user as a ninja log. 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 prefill.cuh at 6dc0b0fc0, spread 0.2-1.4 over repeated identical calls, 6 runs giving 6 distinct results, at kv_len=16 where there is a single KV tile. It is pre-existing and orthogonal to fp8, but it means RoPE cannot serve as a reference here. test_fp8_kv_quantization_quality is deliberately separate from the oracle: it measures quantization error against fp32, so a failure there means the cache is too coarse rather than that the kernel disagrees with its fp16 twin. It uses naive_attention rather than F.scaled_dot_product_attention -- torch anchors is_causal top-left and FlashInfer bottom-right, which at qo_len=64/kv_len=256 mismatches 93% of elements and looks exactly like a kernel bug. Co-Authored-By: Claude --- tests/rocm/test_fp8_kv_prefill.py | 648 ++++++++++++++++++++++++++++++ 1 file changed, 648 insertions(+) create mode 100644 tests/rocm/test_fp8_kv_prefill.py diff --git a/tests/rocm/test_fp8_kv_prefill.py b/tests/rocm/test_fp8_kv_prefill.py new file mode 100644 index 00000000000..98068cb2e32 --- /dev/null +++ b/tests/rocm/test_fp8_kv_prefill.py @@ -0,0 +1,648 @@ +# 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_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 + + +@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], *common + ) + gen_prefill_attention_modules( + [torch.bfloat16], [torch.bfloat16] + FNUZ_DTYPES, [128], *common + ) + # 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]) +@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("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,allowed", + [ + (torch.int8, torch.int8, torch.int8, False), + (torch.uint8, torch.uint8, torch.uint8, False), + (torch.float32, torch.float32, torch.float32, False), + (torch.float8_e5m2fnuz, torch.float8_e5m2fnuz, torch.bfloat16, False), + (torch.float16, torch.bfloat16, torch.float16, False), + (torch.float16, torch.float16, torch.int8, False), + (torch.float16, torch.float16, torch.float16, True), + (torch.bfloat16, torch.bfloat16, torch.bfloat16, True), + # AITER's own fp8 prefill, both arch spellings of e4m3: must stay open. + (torch.float8_e4m3fnuz, torch.float8_e4m3fnuz, torch.bfloat16, True), + (torch.float8_e4m3fn, torch.float8_e4m3fn, torch.bfloat16, True), + ], +) +def test_aiter_generator_dtype_allowlist(dtype_q, dtype_kv, dtype_o, 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, + ) + + calls = ( + lambda: gen_single_prefill_module( + "aiter", dtype_q, dtype_kv, dtype_o, 128, 128, 0, False, False, False + ), + lambda: gen_batch_prefill_module( + "aiter", + dtype_q, + dtype_kv, + dtype_o, + torch.int32, + 128, + 128, + 0, + False, + False, + False, + ), + ) + for call in calls: + if allowed: + call() + else: + with pytest.raises(NotImplementedError, match="(?i)aiter"): + call() + + +@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_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) From 2835a84a9acf65ceb92b8516cc8b66d4232f66f8 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sat, 12 Sep 2026 13:26:51 -0400 Subject: [PATCH 04/18] docs(rocm): record fp8 KV on the HIP path, and what it refuses The batch_decode hip row already advertised "fp8 KV-cache (E4M3FNUZ)", which was true of the plain path and false of use_tensor_cores=True. It now covers both, plus E5M2FNUZ. single_prefill and batch_prefill gain the capability; pod records that it refuses fp8 KV and why. backends.md's "fp8 on the HIP path" section was two bullets; it now states the four constraints a caller actually hits -- 2-byte query and output, fnuz spellings only, no POD, JIT-only. Its AITER section claimed the in-tree kernel "rejects 8-bit types in a static_assert", which this series makes false for KV. README regenerated with scripts/gen_arch_support_matrix.py. Co-Authored-By: Claude --- README.md | 10 ++++---- docs/rocm/backends.md | 40 +++++++++++++++++++++++++++---- flashinfer/rocm/arch_caps.py | 10 ++++---- tests/rocm/test_fp8_kv_prefill.py | 4 +++- 4 files changed, 49 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index 37aeb4ba9cd..90f7470f76b 100644 --- a/README.md +++ b/README.md @@ -170,12 +170,12 @@ library's actual routing. Do not edit it by hand; run | `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. | +| `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/docs/rocm/backends.md b/docs/rocm/backends.md index ef315bb66f2..88532fada16 100644 --- a/docs/rocm/backends.md +++ b/docs/rocm/backends.md @@ -366,6 +366,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. @@ -705,9 +712,32 @@ 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), batch decode, and decode with +`use_tensor_cores=True` all accept an fp8 KV-cache in `float8_e4m3fnuz` or +`float8_e5m2fnuz`. The cache is dequantized into a 2-byte LDS tile on the way in, +so the saving is HBM bandwidth and cache capacity, not math. + +* **The query and output stay 2-byte.** The MFMA path is f16f16f32, so an fp8 + query is refused at module build. +* **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. @@ -735,7 +765,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/rocm/arch_caps.py b/flashinfer/rocm/arch_caps.py index e7bb228a2f6..16ef915e7b0 100644 --- a/flashinfer/rocm/arch_caps.py +++ b/flashinfer/rocm/arch_caps.py @@ -461,25 +461,25 @@ def _archs(gfx942: ArchSupport, gfx950: ArchSupport) -> Mapping[str, ArchSupport "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", @@ -492,7 +492,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/tests/rocm/test_fp8_kv_prefill.py b/tests/rocm/test_fp8_kv_prefill.py index 98068cb2e32..7c109b507c6 100644 --- a/tests/rocm/test_fp8_kv_prefill.py +++ b/tests/rocm/test_fp8_kv_prefill.py @@ -488,7 +488,9 @@ def test_single_prefill_fp8_kv_sliding_window(fp8_dtype, window_left): 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") + 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) From 4fd9e14c8251e6e3b3616d2f732425527b33c8cd Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sat, 12 Sep 2026 13:26:51 -0400 Subject: [PATCH 05/18] bench(rocm): give bench_fa2_prefill a KV-dtype axis _bytes() hardcoded 2 bytes per element for q, o, k and v alike, so every fp8 achieved-bandwidth number would have read ~2x high. The q read and o write stay 2-byte; k and v now follow --kv-dtype. Both batch wrappers also had to declare q_data_type/kv_data_type on plan(), which otherwise default the cache to the query dtype. Co-Authored-By: Claude --- benchmarks/rocm/bench_fa2_prefill.py | 76 +++++++++++++++++----------- 1 file changed, 47 insertions(+), 29 deletions(-) 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) From 7a88247495fba30e77012b721b214ec8ef5bca1a Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 11:53:46 -0400 Subject: [PATCH 06/18] fix(rocm): pin the AITER output dtype to what its launchers accept The allowlist admitted any 2-byte output, but the generated launchers pin it exactly: single_prefill_aiter.cu:85 and batch_ragged_prefill_aiter.cu:56 require o == q_dtype, and batch_prefill_paged_aiter.cu:67 requires bf16 for an fp8 query. fp16 q/kv with a bf16 output therefore passed the check, paid a cold JIT build, and only then failed in TORCH_CHECK -- the outcome the allowlist exists to replace. Matches how the callers already derive it (prefill.py:2663, :3955): FP8_PREFILL_OUT_DTYPE for an fp8 query, else q_data_type. Co-Authored-By: Claude --- flashinfer/jit/rocm/modules.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/flashinfer/jit/rocm/modules.py b/flashinfer/jit/rocm/modules.py index b30a9d84fcb..bd55427b52a 100644 --- a/flashinfer/jit/rocm/modules.py +++ b/flashinfer/jit/rocm/modules.py @@ -129,10 +129,14 @@ def _check_aiter_dtypes(what: str, dtype_q, dtype_kv, dtype_o) -> None: "float16, bfloat16, or the e4m3 fp8 encoding this GPU uses " "(aiter.dtypes.fp8)." ) - if dtype_o not in _WIDE_DTYPES: + # 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 output dtype {dtype_o} is not supported; it writes " - "float16 or bfloat16, and bfloat16 for an fp8 query." + f"{what}: AITER writes output dtype {expected_o} for query dtype " + f"{dtype_q}; got {dtype_o}." ) From f50c834679d309bb4c827b2a0c70b5da2dc92df1 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 11:53:46 -0400 Subject: [PATCH 07/18] test(rocm): cover the single-decode k/v mismatch directly test_decode_mismatched_k_v_dtypes_are_refused drives the batch wrapper, which reaches BatchDecodeWithPagedKVCacheRun -- so neither single-decode seam was covered: the Python pre-check that saves a cold JIT, nor single_decode.cu's CHECK_KV_DTYPES_MATCH. Removing either would not have failed a test. Co-Authored-By: Claude --- tests/rocm/test_fp8_kv_prefill.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/rocm/test_fp8_kv_prefill.py b/tests/rocm/test_fp8_kv_prefill.py index 7c109b507c6..3b10a62243e 100644 --- a/tests/rocm/test_fp8_kv_prefill.py +++ b/tests/rocm/test_fp8_kv_prefill.py @@ -407,6 +407,22 @@ def test_pod_refuses_fp8_kv(fp8_dtype): 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. + + Guards both single-decode seams at once: the Python pre-check that saves a + cold JIT, and `single_decode.cu`'s own `CHECK_KV_DTYPES_MATCH`. + """ + 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) def test_decode_mismatched_k_v_dtypes_are_refused(fp8_dtype): """Decode specializes on k.dtype and casts both pointers to it. Measured From 34262802fc836df571d8256795519a485f86adc6 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 13:21:28 -0400 Subject: [PATCH 08/18] fix(rocm): refuse before the URI at every public generator, not three of them An audit of every gen_* in modules.py that takes dtype_q/dtype_kv found the guard placed after URI construction on the two decode builders, and absent entirely from gen_batch_decode_aiter_module -- all three publicly re-exported through flashinfer/jit/__init__.py. Measured before the fix: gen_single_decode_module(float32, ...) -> KeyError: torch.float32 gen_batch_decode_module(float32, ...) -> KeyError: torch.float32 gen_batch_decode_aiter_module(int8, ...) -> NO RAISE, spec built gen_batch_decode_aiter_module(e4m3fnuz, ...) -> NO RAISE, spec built The AITER decode pair is the worse half: batch_decode_aiter.cu:109 serves fp16/bf16 only, so an fp8 or int8 module built and could only fail at run time. _check_aiter_decode_dtypes is separate from the prefill allowlist because the decode kernel has no fp8 arm at all. The new test sweeps all six builders rather than spot-checking, so the next entry point that indexes filename_safe_dtype_map before validating fails a test instead of a review round. Co-Authored-By: Claude --- flashinfer/jit/rocm/modules.py | 32 ++++++++++++++++++++++++++++++ tests/rocm/test_fp8_kv_prefill.py | 33 +++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/flashinfer/jit/rocm/modules.py b/flashinfer/jit/rocm/modules.py index bd55427b52a..5a1a1b7b834 100644 --- a/flashinfer/jit/rocm/modules.py +++ b/flashinfer/jit/rocm/modules.py @@ -140,6 +140,25 @@ def _check_aiter_dtypes(what: str, dtype_q, dtype_kv, dtype_o) -> None: ) +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: @@ -332,6 +351,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, @@ -478,6 +502,9 @@ def gen_batch_decode_aiter_module( # needs no variant store. aiter_jit_dir = os.path.join(os.path.dirname(_aiter_mod.__file__), "jit") + # Before the URI: it indexes filename_safe_dtype_map, so an unsupported + # dtype would raise KeyError rather than this allowlist's message. + _check_aiter_decode_dtypes("batch decode (aiter)", dtype_q, dtype_kv, dtype_o) uri = get_batch_decode_aiter_uri( dtype_q, dtype_kv, dtype_o, head_dim_qk, head_dim_vo ) @@ -516,6 +543,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, diff --git a/tests/rocm/test_fp8_kv_prefill.py b/tests/rocm/test_fp8_kv_prefill.py index 3b10a62243e..60686604a07 100644 --- a/tests/rocm/test_fp8_kv_prefill.py +++ b/tests/rocm/test_fp8_kv_prefill.py @@ -384,6 +384,39 @@ def test_aiter_generator_dtype_allowlist(dtype_q, dtype_kv, dtype_o, allowed): 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 + ), + } + for name, call in calls.items(): + with pytest.raises(NotImplementedError): + call() + + @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 From 54cdcaecf9b46075d3606591b0b2f9658c9ae046 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 13:59:56 -0400 Subject: [PATCH 09/18] fix(rocm): keep the fp8 KV load off an fp8 array, which gfx950 cannot construct The strict-aliasing fix in b3e3c0ea9 staged the global load in a `DTypeKV packed[GRANULE]`. gfx950 builds with HIP_FP8_TYPE_FNUZ=0, which leaves __hip_fp8_e4m3_fnuz without a default constructor, so an array of it does not compile: prefill.cuh:326:41: error: no matching constructor for initialization of 'DTypeKV[4]' (aka '__hip_fp8_e4m3_fnuz[4]') gfx942 has the constructor and compiled it, so no gfx942 run could catch this; every commit from b3e3c0ea9 on was broken on CDNA4, and neither conformance nor pre-commit compiles kernels. Staging in a uint32_t keeps both properties: no fp8 array, so gfx950 compiles, and the load is still a __builtin_memcpy, so it stays aliasing-clean under -O3. Only the vec_cast read-back casts, as it did before b3e3c0ea9. Verified on gfx950 (MI350X, job 67920205, rocm 7.15): builds, and the fp8 kernel tests pass. gfx942 (MI300X) re-run green. Co-Authored-By: Claude --- include/flashinfer/rocm/attention/prefill.cuh | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/include/flashinfer/rocm/attention/prefill.cuh b/include/flashinfer/rocm/attention/prefill.cuh index 597f2d49519..62a9a7baf26 100644 --- a/include/flashinfer/rocm/attention/prefill.cuh +++ b/include/flashinfer/rocm/attention/prefill.cuh @@ -319,15 +319,15 @@ __device__ __forceinline__ void produce_kv_unit( b64_t* smem_ptr = smem.base + smem_offset; static_assert(GRANULE * sizeof(DTypeKV) == sizeof(uint32_t)); if (predicate) { - // One register load, then convert: vec_cast is a per-element loop. Raw - // arrays because gfx950 leaves vec_t without a device ctor, and - // memcpy rather than a cast because -O3 keeps strict aliasing on. - // wide is 8-byte aligned because vec_cast reads it through float2. - alignas(sizeof(uint32_t)) DTypeKV packed[GRANULE]; - __builtin_memcpy(packed, gptr, sizeof(uint32_t)); + // One register load, then convert: vec_cast is a per-element loop. uint32_t + // rather than DTypeKV[GRANULE], which gfx950 cannot default-construct; + // memcpy rather than a cast, because -O3 keeps strict aliasing on. + uint32_t packed; + __builtin_memcpy(&packed, gptr, sizeof(packed)); alignas(sizeof(b64_t)) float wide[GRANULE]; alignas(sizeof(b64_t)) DTypeKVSmem narrow[GRANULE]; - vec_cast::template cast(wide, packed); + vec_cast::template cast(wide, + reinterpret_cast(&packed)); vec_cast::template cast(narrow, wide); __builtin_memcpy(smem_ptr, narrow, sizeof(b64_t)); } else if constexpr (fill_mode == SharedMemFillMode::kFillZero) { From c63616a7286f56e595d269a36be0b70e40d1d192 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 14:00:17 -0400 Subject: [PATCH 10/18] fix(rocm): match the AITER guards to the launchers, and run them before the import Two defects in the guard layer, both surfaced by the first gfx950 run. `_check_aiter_dtypes` admitted e4m3 fp8 at every AITER prefill seam, but only batch_prefill_paged_aiter.cu has an fp8 arm. single_prefill_aiter.cu and batch_ragged_prefill_aiter.cu are fp16/bf16 only, so a single-prefill fp8 module built and could never execute: single_prefill aiter fp8: NO RAISE -> spec built (before) single_prefill aiter fp8: NotImplementedError: ... (after) batch_prefill aiter fp8: NO RAISE -> spec built (both; paged serves it) single_prefill aiter fp16: NO RAISE -> spec built (both; no over-refusal) `gen_batch_decode_aiter_module` imported aiter as its first statement, ahead of the dtype check the previous commit had just moved up. The refusal was therefore reachable only where aiter is installed; the gfx950 image is py3.13 and the wheel is cp312-only, so it raised ModuleNotFoundError instead. With the import blocked: batch_decode_aiter f32 (no aiter): NotImplementedError: ... (after) The tests carried the same assumption: the allowed AITER cases call a generator that imports aiter unconditionally, so they failed rather than skipped on a box without it. Guarded with importorskip; the refusal cases need no guard, which is the point of checking first. gfx950 (MI350X, job 67920205) and gfx942 (MI300X) both green. Co-Authored-By: Claude --- flashinfer/jit/rocm/modules.py | 43 +++++++++++----- tests/rocm/test_fp8_kv_prefill.py | 85 ++++++++++++++++++++----------- 2 files changed, 83 insertions(+), 45 deletions(-) diff --git a/flashinfer/jit/rocm/modules.py b/flashinfer/jit/rocm/modules.py index 5a1a1b7b834..d09d40242ae 100644 --- a/flashinfer/jit/rocm/modules.py +++ b/flashinfer/jit/rocm/modules.py @@ -105,29 +105,38 @@ def _generate_additional_params_hip( _WIDE_DTYPES = (torch.float16, torch.bfloat16) -# AITER prefill serves e4m3 only, and both spellings because aiter.dtypes.fp8 is -# arch-dependent; it writes bf16 for an fp8 query. Mirrors prefill.py, which -# cannot be imported here -- it imports this module. +# 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) -> None: +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 _WIDE_DTYPES + _AITER_FP8_DTYPES: + 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)." + + ( + "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 @@ -405,7 +414,9 @@ def gen_single_prefill_module( # 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) + _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 @@ -495,6 +506,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 @@ -502,9 +518,6 @@ def gen_batch_decode_aiter_module( # needs no variant store. aiter_jit_dir = os.path.join(os.path.dirname(_aiter_mod.__file__), "jit") - # Before the URI: it indexes filename_safe_dtype_map, so an unsupported - # dtype would raise KeyError rather than this allowlist's message. - _check_aiter_decode_dtypes("batch decode (aiter)", dtype_q, dtype_kv, dtype_o) uri = get_batch_decode_aiter_uri( dtype_q, dtype_kv, dtype_o, head_dim_qk, head_dim_vo ) @@ -600,7 +613,7 @@ def gen_batch_prefill_module( # 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) + _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 @@ -807,7 +820,9 @@ def gen_customize_single_prefill_module( # 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) + _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": @@ -1063,7 +1078,7 @@ def gen_customize_batch_prefill_module( # 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) + _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": diff --git a/tests/rocm/test_fp8_kv_prefill.py b/tests/rocm/test_fp8_kv_prefill.py index 60686604a07..ba29045774a 100644 --- a/tests/rocm/test_fp8_kv_prefill.py +++ b/tests/rocm/test_fp8_kv_prefill.py @@ -31,6 +31,12 @@ 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 @@ -335,22 +341,26 @@ def test_ocp_fp8_kv_is_refused(ocp_dtype): @pytest.mark.parametrize( - "dtype_q,dtype_kv,dtype_o,allowed", + "dtype_q,dtype_kv,dtype_o,single_allowed,batch_allowed", [ - (torch.int8, torch.int8, torch.int8, False), - (torch.uint8, torch.uint8, torch.uint8, False), - (torch.float32, torch.float32, torch.float32, False), - (torch.float8_e5m2fnuz, torch.float8_e5m2fnuz, torch.bfloat16, False), - (torch.float16, torch.bfloat16, torch.float16, False), - (torch.float16, torch.float16, torch.int8, False), - (torch.float16, torch.float16, torch.float16, True), - (torch.bfloat16, torch.bfloat16, torch.bfloat16, True), - # AITER's own fp8 prefill, both arch spellings of e4m3: must stay open. - (torch.float8_e4m3fnuz, torch.float8_e4m3fnuz, torch.bfloat16, True), - (torch.float8_e4m3fn, torch.float8_e4m3fn, torch.bfloat16, True), + (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, allowed): +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 ( @@ -358,26 +368,33 @@ def test_aiter_generator_dtype_allowlist(dtype_q, dtype_kv, dtype_o, allowed): gen_single_prefill_module, ) - calls = ( - lambda: gen_single_prefill_module( - "aiter", dtype_q, dtype_kv, dtype_o, 128, 128, 0, False, False, False + cases = ( + ( + single_allowed, + lambda: gen_single_prefill_module( + "aiter", dtype_q, dtype_kv, dtype_o, 128, 128, 0, False, False, False + ), ), - lambda: gen_batch_prefill_module( - "aiter", - dtype_q, - dtype_kv, - dtype_o, - torch.int32, - 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 call in calls: + for allowed, call in cases: if allowed: + _skip_without_aiter() call() else: with pytest.raises(NotImplementedError, match="(?i)aiter"): @@ -413,8 +430,14 @@ def test_every_public_generator_refuses_before_its_uri(): ), } for name, call in calls.items(): - with pytest.raises(NotImplementedError): - call() + # 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) From 9be6d0960f43b03c65e8437f5a49f2722b269575 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 16:10:22 -0400 Subject: [PATCH 11/18] fix(rocm): close the customize-POD bypass of the fp8 refusal gen_pod_module and gen_batch_pod_module refuse an fp8 KV cache, but the two customize entry points reached the kernel without passing through either. They built a spec the POD kernel's static_assert then rejected as a ninja log -- the failure _check_pod_fp8_dtypes exists to replace -- and an unmapped dtype still raised a bare KeyError: gen_pod_module fp8kv: NotImplementedError (before and after) gen_customize_pod_module fp8kv: NO RAISE -> spec built -> NotImplementedError gen_customize_batch_pod_module fp8kv: NO RAISE -> spec built -> NotImplementedError gen_customize_pod_module f32: KeyError: torch.float32 -> NotImplementedError Both delegate to _gen_customize_pod_like_module, so one guard there covers them, matching how CHECK_KV_DTYPES_MATCH is a single C++ seam rather than a check per launcher. The sweep test grew the four customize prefill/decode generators plus gen_batch_pod_module and both customize POD ones, and now asserts its own completeness against dir(modules) -- the gap is what let this through, since the sweep was billed as covering every public generator and listed 6 of 13. gfx942: 119 passed, 16 skipped, 0 failed. Co-Authored-By: Claude --- flashinfer/jit/rocm/modules.py | 3 ++ tests/rocm/test_fp8_kv_prefill.py | 57 +++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/flashinfer/jit/rocm/modules.py b/flashinfer/jit/rocm/modules.py index d09d40242ae..bb2ae55093a 100644 --- a/flashinfer/jit/rocm/modules.py +++ b/flashinfer/jit/rocm/modules.py @@ -1408,6 +1408,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/tests/rocm/test_fp8_kv_prefill.py b/tests/rocm/test_fp8_kv_prefill.py index ba29045774a..5de8174d370 100644 --- a/tests/rocm/test_fp8_kv_prefill.py +++ b/tests/rocm/test_fp8_kv_prefill.py @@ -428,7 +428,64 @@ def test_every_public_generator_refuses_before_its_uri(): "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"): From 81f747d2063815a22b242c9f98588dfd9fd0dbfa Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 23:26:01 -0400 Subject: [PATCH 12/18] fix(rocm): stage the fp8 KV load with bit_cast, not a punned uint32_t 54cdcaecf fixed the gfx950 build by staging the global load in a uint32_t and reading it back through DTypeKV*. That is a strict-aliasing violation: vec_cast reads src[i].__x, a struct member, out of a live uint32_t object, and -O3 keeps strict aliasing on. Deleting the staging and reading straight from gptr is alias-clean but costs CDNA4 four loads for one. Measured on the real vec_cast header, per kernel, hipcc -O3: form gfx942 gfx950 direct from gptr 1x load_dword 4x load_ubyte uint32_t + reinterpret_cast 1x load_dword 1x load_dword __builtin_bit_cast into DTypeKV[] 1x load_dword 1x load_dword bit_cast materializes real DTypeKV objects from the bytes, so the read is alias-clean, and unlike an array declaration it needs no default constructor -- which is what gfx950 withholds when HIP_FP8_TYPE_FNUZ=0. The fix is source-level only: compiling the real produce_kv_unit both ways and normalising the per-TU build-id hash gives byte-identical assembly on both arches, so this cannot have changed behaviour where the punned form already worked. gfx942 md5 ee85057f91aafff7f2a4e1c88d2d12cd (both forms) gfx950 md5 bdce7a2da18b4ad53573a1301e0839d6 (both forms) gfx950 (MI350X, asrock-1w300-e2-3, aiter 0.1.21.post2): 119 passed, 16 skipped, 0 failed. Co-Authored-By: Claude --- include/flashinfer/rocm/attention/prefill.cuh | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/include/flashinfer/rocm/attention/prefill.cuh b/include/flashinfer/rocm/attention/prefill.cuh index 62a9a7baf26..54d96dceaac 100644 --- a/include/flashinfer/rocm/attention/prefill.cuh +++ b/include/flashinfer/rocm/attention/prefill.cuh @@ -319,15 +319,19 @@ __device__ __forceinline__ void produce_kv_unit( b64_t* smem_ptr = smem.base + smem_offset; static_assert(GRANULE * sizeof(DTypeKV) == sizeof(uint32_t)); if (predicate) { - // One register load, then convert: vec_cast is a per-element loop. uint32_t - // rather than DTypeKV[GRANULE], which gfx950 cannot default-construct; - // memcpy rather than a cast, because -O3 keeps strict aliasing on. - uint32_t packed; - __builtin_memcpy(&packed, gptr, sizeof(packed)); + // 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, - reinterpret_cast(&packed)); + 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) { From 1ed1569bc4d722eb75e7d731140acba0ae640525 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 00:31:20 -0400 Subject: [PATCH 13/18] test(rocm): cover single-decode fp8 numerics, and say which decode is exempt Single decode has its own kernel, and this series newly made e5m2fnuz expressible there by adding it to dtype_map_hip, but the only coverage was the k/v-mismatch refusal -- nothing checked the output. Compares an fp8 cache against its exact dequantization, as the prefill tests do. The oracle is sensitive rather than vacuous; substituting an unrelated cache misses by ~650x the tolerance against an output RMS of 0.2: e4m3fnuz real 1.5e-05 wrong-cache 1.30 tol 2e-3 e5m2fnuz real 1.9e-06 wrong-cache 1.27 tol 2e-3 Tolerance rather than the prefill tests' bit-equality because decode accumulates in float with no pinned reduction order. The kv_must_match_q comment said "decode" where it meant the plain decode kernel. use_tensor_cores=True builds the prefill module, so it is bound by the rule; measured, the wrapper refuses fp16-q with bf16-KV at use_tensor_cores=True and accepts it at False. gfx950 (MI350X): 6 passed, 0 failed. Co-Authored-By: Claude --- flashinfer/jit/rocm/modules.py | 7 ++++--- tests/rocm/test_fp8_kv_prefill.py | 20 ++++++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/flashinfer/jit/rocm/modules.py b/flashinfer/jit/rocm/modules.py index bb2ae55093a..161bff01cb5 100644 --- a/flashinfer/jit/rocm/modules.py +++ b/flashinfer/jit/rocm/modules.py @@ -197,9 +197,10 @@ def _check_fa2_fp8_dtypes( "bfloat16, or an fnuz fp8 cache (torch.float8_e4m3fnuz, " "torch.float8_e5m2fnuz)." ) - # Prefill only: the MFMA is instantiated from DTypeQ, so a bf16 cache under - # an fp16 query is read as fp16. Decode cast_loads q and k into vec_t - # separately and accumulates in float, so mixed 2-byte dtypes are fine there. + # 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 " diff --git a/tests/rocm/test_fp8_kv_prefill.py b/tests/rocm/test_fp8_kv_prefill.py index 5de8174d370..980b9e5d1c0 100644 --- a/tests/rocm/test_fp8_kv_prefill.py +++ b/tests/rocm/test_fp8_kv_prefill.py @@ -536,6 +536,26 @@ def test_single_decode_mismatched_k_v_dtypes_are_refused(fp8_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 rather than bit-equality: decode accumulates in float and its + reduction order is not pinned the way the prefill LDS tile pins prefill's. + """ + 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=2e-3, atol=2e-3) + + @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 From fb9ab1b14db7356c63e2d14f2646e8622c3722f4 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 01:10:59 -0400 Subject: [PATCH 14/18] test(rocm): tighten the single-decode tolerance, and document the path it covers atol=2e-3 was 33x the worst real delta, so a ~1% systematic shift -- the size a wrong exponent bias or mis-scaled dequant produces on a bounded input -- would have passed the test written to catch it. Measured across all four parametrizations on gfx950: dtype hd max|delta| wrong-cache rms(ref) e4m3fnuz 64 0.0 1.07 0.188 e4m3fnuz 128 1.5e-05 1.07 0.195 e5m2fnuz 64 1.2e-07 1.07 0.188 e5m2fnuz 128 6.1e-05 1.16 0.198 atol=3e-4 keeps 5x headroom over the worst case and sits 6x under a 1% shift. The docstring said the looseness was an unpinned reduction order. It is not: an fp8 cache selects a different geometry from the 2-byte reference (BatchDecodeVecSize gives 16 vs 8, tile_size_per_bdx 2 vs 1), so the test compares two kernel shapes and cannot pin the fp8 one. Said plainly, since the prefill oracle's bit-equality does make that guarantee and the contrast is the thing a reader needs. Single decode accepts fp8 KV via kv_must_match_q=False and is now tested, but every doc surface still said otherwise -- arch_caps, the generated README row, and the backends.md enumeration all omitted it. Warmup builds the decode specs too; the four single-decode modules were compiling serially inside the test body on a cold cache. gfx950 (MI350X): 6 passed, 0 failed. Co-Authored-By: Claude --- README.md | 2 +- docs/rocm/backends.md | 6 +++--- flashinfer/rocm/arch_caps.py | 2 +- tests/rocm/test_fp8_kv_prefill.py | 32 +++++++++++++++++++++++-------- 4 files changed, 29 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 90f7470f76b..d850ec30746 100644 --- a/README.md +++ b/README.md @@ -169,7 +169,7 @@ 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. | +| `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. | diff --git a/docs/rocm/backends.md b/docs/rocm/backends.md index df375b4a3d1..504436d51c6 100644 --- a/docs/rocm/backends.md +++ b/docs/rocm/backends.md @@ -758,9 +758,9 @@ partial case — its own kernels are HIP, but what it calls is not: ### fp8 on the HIP path: KV cache only -Single prefill, batch prefill (paged and ragged), batch decode, and decode with -`use_tensor_cores=True` all accept an fp8 KV-cache in `float8_e4m3fnuz` or -`float8_e5m2fnuz`. The cache is dequantized into a 2-byte LDS tile on the way in, +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 cache is dequantized into a 2-byte LDS tile on the way in, so the saving is HBM bandwidth and cache capacity, not math. * **The query and output stay 2-byte.** The MFMA path is f16f16f32, so an fp8 diff --git a/flashinfer/rocm/arch_caps.py b/flashinfer/rocm/arch_caps.py index 21fbb702da6..97b4aa902ed 100644 --- a/flashinfer/rocm/arch_caps.py +++ b/flashinfer/rocm/arch_caps.py @@ -472,7 +472,7 @@ 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", diff --git a/tests/rocm/test_fp8_kv_prefill.py b/tests/rocm/test_fp8_kv_prefill.py index 980b9e5d1c0..8eacb1ae47d 100644 --- a/tests/rocm/test_fp8_kv_prefill.py +++ b/tests/rocm/test_fp8_kv_prefill.py @@ -19,7 +19,7 @@ import pytest import torch from attention_reference import naive_attention -from jit_utils import gen_prefill_attention_modules +from jit_utils import gen_decode_attention_modules, gen_prefill_attention_modules import flashinfer from flashinfer.jit.core import logger @@ -40,10 +40,23 @@ def _skip_without_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], *common - ) + gen_prefill_attention_modules( - [torch.bfloat16], [torch.bfloat16] + FNUZ_DTYPES, [128], *common + specs = ( + gen_prefill_attention_modules( + [torch.float16], [torch.float16] + FNUZ_DTYPES, [64, 128], *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( @@ -541,8 +554,11 @@ def test_single_decode_mismatched_k_v_dtypes_are_refused(fp8_dtype): 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 rather than bit-equality: decode accumulates in float and its - reduction order is not pinned the way the prefill LDS tile pins prefill's. + + 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" @@ -553,7 +569,7 @@ def test_single_decode_fp8_kv_matches_dequantized(fp8_dtype, head_dim): 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=2e-3, atol=2e-3) + torch.testing.assert_close(got.float(), ref.float(), rtol=1e-3, atol=3e-4) @pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) From 8e60d07f9e53aa0b727aa365043eab3dc2034000 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 09:49:04 -0400 Subject: [PATCH 15/18] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Debasis Mandal --- docs/rocm/backends.md | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/rocm/backends.md b/docs/rocm/backends.md index abb5884b902..35ddc20033c 100644 --- a/docs/rocm/backends.md +++ b/docs/rocm/backends.md @@ -780,11 +780,14 @@ partial case — its own kernels are HIP, but what it calls is not: 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 cache is dequantized into a 2-byte LDS tile on the way in, -so the saving is HBM bandwidth and cache capacity, not math. - -* **The query and output stay 2-byte.** The MFMA path is f16f16f32, so an fp8 - query is refused at module build. +`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 From d86993cf4b96c974d92dba8a91075482205eeae2 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 10:57:12 -0400 Subject: [PATCH 16/18] test(rocm): cover fp8 KV at head_dim 256 The allowlist does not restrict head_dim and the in-tree kernel's own matrix runs 256 (test_single_prefill_kernels.py), so fp8 KV at 256 was reachable and unverified -- and the tile count and LDS traversal both change there. Measured before adding, since gfx950 has pre-existing failures at large head dims and a test that fails for someone else's reason is worse than no test. It is exact: hd=256 e4m3fnuz bit_equal=True max|d|=0.0 finite=True hd=256 e5m2fnuz bit_equal=True max|d|=0.0 finite=True Only the fp16 warmup arm grows: the test already skips bf16 away from 128, so the bf16 x 256 cells skip rather than pulling in specs nothing runs. gfx950 (MI350X, aiter 0.1.21.post2): 139 passed, 32 skipped, 0 failed. Co-Authored-By: Claude --- tests/rocm/test_fp8_kv_prefill.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/rocm/test_fp8_kv_prefill.py b/tests/rocm/test_fp8_kv_prefill.py index 8eacb1ae47d..5fde638245c 100644 --- a/tests/rocm/test_fp8_kv_prefill.py +++ b/tests/rocm/test_fp8_kv_prefill.py @@ -42,7 +42,7 @@ 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], *common + [torch.float16], [torch.float16] + FNUZ_DTYPES, [64, 128, 256], *common ) + gen_prefill_attention_modules( [torch.bfloat16], [torch.bfloat16] + FNUZ_DTYPES, [128], *common @@ -84,7 +84,7 @@ def _assert_matches(got, ref): @pytest.mark.parametrize("q_dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("fp8_dtype", FNUZ_DTYPES) -@pytest.mark.parametrize("head_dim", [64, 128]) +@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( From cccee3b68b1eb58a15f484d7bfcad3316d7f5b80 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 11:23:43 -0400 Subject: [PATCH 17/18] test(rocm): stop claiming coverage of a C++ guard this test cannot reach The docstring said it guarded "both single-decode seams at once", the Python pre-check and single_decode.cu's CHECK_KV_DTYPES_MATCH. It reaches only the first: _check_kv_dtypes_match runs before any module is built, so Python always wins and the C++ line never executes. ValueError: k has dtype torch.float8_e4m3fnuz but v has torch.float16 ... The batch test reaches its C++ macro by passing a (k, v) tuple to run(), which Python does not re-validate; single decode has no equivalent path. The macro still earns its place -- it backstops callers who build the module themselves -- but a test that cannot execute a line should not say it covers it, which is how a removed guard would have gone unnoticed. Co-Authored-By: Claude --- tests/rocm/test_fp8_kv_prefill.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/rocm/test_fp8_kv_prefill.py b/tests/rocm/test_fp8_kv_prefill.py index 5fde638245c..d055317fa41 100644 --- a/tests/rocm/test_fp8_kv_prefill.py +++ b/tests/rocm/test_fp8_kv_prefill.py @@ -537,8 +537,10 @@ def test_pod_refuses_fp8_kv(fp8_dtype): 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. - Guards both single-decode seams at once: the Python pre-check that saves a - cold JIT, and `single_decode.cu`'s own `CHECK_KV_DTYPES_MATCH`. + 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 From 4ce31b2f82c295625d60b7e0023ec5269164e9ab Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 12:13:29 -0400 Subject: [PATCH 18/18] fix(rocm): replace the last bare assert on the fp8-query path with a sentence `assert window_left == -1` (pre-existing, 7ccb007ea) sat ahead of the typed fp8-query refusal this series added, so a sliding window produced a bare AssertionError -- against a PR whose point is that every unservable combination fails with a sentence. The reachable set is narrower than it looks, which the test now reflects. Measured with window_left=16: e4m3fnuz NotImplementedError: fp8 prefill query ... has no in-tree ... e5m2fnuz NotImplementedError: single prefill: query dtype ... not supported e4m3fn NotImplementedError: an fp8 query does not support a sliding window Both fnuz spellings are refused earlier by a different check, so OCP is the only one that reaches this branch. Parametrizing the test over all four and matching on "sliding window" would have failed on the two fnuz cases. gfx950 (MI350X, aiter 0.1.21.post2): 141 passed, 32 skipped, 0 failed. Co-Authored-By: Claude --- flashinfer/rocm/prefill.py | 13 +++++++------ tests/rocm/test_fp8_kv_prefill.py | 15 +++++++++++++++ 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/flashinfer/rocm/prefill.py b/flashinfer/rocm/prefill.py index d373d5f6033..8a17642443d 100755 --- a/flashinfer/rocm/prefill.py +++ b/flashinfer/rocm/prefill.py @@ -1962,12 +1962,13 @@ 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 - # Not an assert: a wide KV cache under an fp8 query reaches here before - # any backend is chosen, and no backend serves that pair. + # 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}, " diff --git a/tests/rocm/test_fp8_kv_prefill.py b/tests/rocm/test_fp8_kv_prefill.py index d055317fa41..1b9f4f0e103 100644 --- a/tests/rocm/test_fp8_kv_prefill.py +++ b/tests/rocm/test_fp8_kv_prefill.py @@ -343,6 +343,21 @@ def test_fp8_query_with_wide_kv_is_refused(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."""