From 617ebb306c6aae06bb0ff403d57a00dc684c52f5 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Wed, 17 Jun 2026 09:04:24 +0900 Subject: [PATCH 01/23] feat(attention): NVFP4 KV cache read path for FA2 paged attention (SM120) Implements the e2m1 (NVFP4) KV-cache read path for the FA2 paged-prefill and decode kernels on consumer Blackwell (SM120): - FP4 E2M1 decode vector casts and KV scale-stride plumbing through the kernels - pass JIT scalar args for NVFP4 paged prefill; disambiguate FP4 KV JIT module names so distinct dtype configs don't collide in the JIT cache - size the attention output from V rather than Q (asymmetric QK vs VO head dims) - allow GQA group size 6 in decode dispatch - guard the FP4 paged-prefill JIT flags and alias the CUTLASS DSL OperandMajorMode for the FP4 path Signed-off-by: Jetha Chan (cherry picked from commit 77758eeb2a761fee8ea6e998b0a62bb3d79985a4) --- csrc/batch_decode.cu | 9 +- csrc/batch_prefill.cu | 12 +-- csrc/batch_prefill_customize_config.jinja | 8 ++ flashinfer/cute_dsl/utils.py | 6 ++ flashinfer/decode.py | 8 +- flashinfer/jit/attention/modules.py | 48 +++++++--- flashinfer/jit/utils.py | 9 ++ flashinfer/prefill.py | 74 +++++++++++---- include/flashinfer/attention/persistent.cuh | 23 +++-- include/flashinfer/attention/prefill.cuh | 100 +++++++++++++------- include/flashinfer/page.cuh | 66 ++++++++++++- include/flashinfer/utils.cuh | 3 + include/flashinfer/vec_dtypes.cuh | 99 +++++++++++++++++++ tests/jit/test_attention_utils.py | 73 ++++++++++++++ 14 files changed, 452 insertions(+), 86 deletions(-) create mode 100644 tests/jit/test_attention_utils.py diff --git a/csrc/batch_decode.cu b/csrc/batch_decode.cu index ddfa5584d20..3e4d3f6533f 100644 --- a/csrc/batch_decode.cu +++ b/csrc/batch_decode.cu @@ -172,15 +172,10 @@ void BatchDecodeWithPagedKVCacheRun(TensorView float_workspace_buffer, const auto q_stride_n = q.stride(0); const auto q_stride_h = q.stride(1); - // get kv_cache_strides - const int64_t* kv_cache_strides = nullptr; + // get kv-cache strides auto k_strides = paged_k_cache.strides(); auto v_strides = paged_v_cache.strides(); TVM_FFI_ICHECK_EQ(k_strides.size(), v_strides.size()); - for (int i = 0; i < k_strides.size(); ++i) { - TVM_FFI_ICHECK_EQ(k_strides[i], v_strides[i]); - } - kv_cache_strides = k_strides.data(); ffi::CUDADeviceGuard device_guard(q.device().device_id); const cudaStream_t stream = get_stream(q.device()); @@ -191,7 +186,7 @@ void BatchDecodeWithPagedKVCacheRun(TensorView float_workspace_buffer, paged_kv_t paged_kv( num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, static_cast(paged_k_cache.data_ptr()), - static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_v_cache.data_ptr()), k_strides.data(), v_strides.data(), static_cast(paged_kv_indices.data_ptr()), static_cast(paged_kv_indptr.data_ptr()), static_cast(paged_kv_last_page_len.data_ptr())); diff --git a/csrc/batch_prefill.cu b/csrc/batch_prefill.cu index e610a3e5838..41a8dab571a 100644 --- a/csrc/batch_prefill.cu +++ b/csrc/batch_prefill.cu @@ -275,13 +275,10 @@ void BatchPrefillWithPagedKVCacheRun(TensorView float_workspace_buffer, const auto q_stride_n = q.stride(0); const auto q_stride_h = q.stride(1); - // get kv_cache_strides - const int64_t* kv_cache_strides = paged_k_cache.strides().data(); + // get kv-cache strides + auto k_cache_strides = paged_k_cache.strides(); + auto v_cache_strides = paged_v_cache.strides(); TVM_FFI_ICHECK_EQ(paged_k_cache.ndim(), paged_v_cache.ndim()); - for (int i = 0; i < paged_k_cache.ndim(); ++i) { - TVM_FFI_ICHECK_EQ(paged_k_cache.stride(i), paged_v_cache.stride(i)) - << "k/v strides differs at " << i; - } ffi::CUDADeviceGuard device_guard(float_workspace_buffer.device().device_id); const cudaStream_t stream = get_stream(float_workspace_buffer.device()); @@ -296,7 +293,8 @@ void BatchPrefillWithPagedKVCacheRun(TensorView float_workspace_buffer, paged_kv_t paged_kv( num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, static_cast(paged_k_cache.data_ptr()), - static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_v_cache.data_ptr()), k_cache_strides.data(), + v_cache_strides.data(), static_cast(paged_kv_indices.data_ptr()), static_cast(paged_kv_indptr.data_ptr()), static_cast(paged_kv_last_page_len.data_ptr())); diff --git a/csrc/batch_prefill_customize_config.jinja b/csrc/batch_prefill_customize_config.jinja index e4f60b1143a..6f5f1d9aa42 100644 --- a/csrc/batch_prefill_customize_config.jinja +++ b/csrc/batch_prefill_customize_config.jinja @@ -6,6 +6,7 @@ #include #include #include +#include #define ADDITIONAL_FUNC_PARAMS {{ additional_func_params }} #define ADDITIONAL_PARAMS_SETTER {{ additional_params_setter }} @@ -23,6 +24,13 @@ using DTypeQ = {{ dtype_q }}; using DTypeKV = {{ dtype_kv }}; using DTypeO = {{ dtype_o }}; using IdType = {{ idtype }}; +{% if require_fp4_kv_cache %} +#ifndef FLASHINFER_ENABLE_FP4_E2M1 +#error "NVFP4 KV paged prefill compiled without FLASHINFER_ENABLE_FP4_E2M1" +#endif +static_assert(std::is_same_v, + "NVFP4 KV paged prefill must build with the packed FP4 KV container type"); +{% endif %} constexpr int HEAD_DIM_QK = {{ head_dim_qk }}; constexpr int HEAD_DIM_VO = {{ head_dim_vo }}; constexpr bool USE_FP16_QK_REDUCTION = {{ use_fp16_qk_reduction }}; diff --git a/flashinfer/cute_dsl/utils.py b/flashinfer/cute_dsl/utils.py index a5f576ed579..b6b6bd85c69 100644 --- a/flashinfer/cute_dsl/utils.py +++ b/flashinfer/cute_dsl/utils.py @@ -28,6 +28,12 @@ from cutlass.cutlass_dsl import dsl_user_op from cutlass.cute.typing import AddressSpace, Numeric, Pointer, Type +if not hasattr(cute.nvgpu, "OperandMajorMode"): + try: + cute.nvgpu.OperandMajorMode = cute.nvgpu.tcgen05.OperandMajorMode + except AttributeError: + pass + def ceil_div(a: int, b: int) -> int: """Ceiling division.""" diff --git a/flashinfer/decode.py b/flashinfer/decode.py index c8d0507f123..427c50f743e 100644 --- a/flashinfer/decode.py +++ b/flashinfer/decode.py @@ -2029,7 +2029,13 @@ def run( out_dtype = getattr(self, "_cached_o_data_type", None) or q.dtype # For NVFP4 KV (uint8 packed), v_cache last dim is head_dim//2; # use q's head_dim for output instead - out_head_dim = q.shape[-1] if kv_cache_sf is not None else v_cache.shape[-1] + # NVFP4 packed: unpacked VO width is packed bytes * 2 (supports + # asymmetric QK/VO plans; q.shape[-1] assumed QK == VO). + out_head_dim = ( + v_cache.shape[-1] * 2 + if kv_cache_sf is not None + else v_cache.shape[-1] + ) out = torch.empty( q.shape[:-1] + (out_head_dim,), dtype=out_dtype, device=q.device ) diff --git a/flashinfer/jit/attention/modules.py b/flashinfer/jit/attention/modules.py index dd6a6eb9707..d2444bfbe69 100755 --- a/flashinfer/jit/attention/modules.py +++ b/flashinfer/jit/attention/modules.py @@ -23,6 +23,7 @@ from .. import env as jit_env from ..core import ( JitSpec, + common_nvcc_flags, gen_jit_spec, logger, sm90a_nvcc_flags, @@ -33,6 +34,7 @@ dtype_map, dtype_map_kv, filename_safe_dtype_map, + filename_safe_dtype_map_kv, mask_mode_literal, pos_encoding_mode_literal, write_if_different, @@ -54,7 +56,7 @@ def get_single_decode_uri( ) -> str: return ( f"single_decode_with_kv_cache_dtype_q_{filename_safe_dtype_map[dtype_q]}_" - f"dtype_kv_{filename_safe_dtype_map[dtype_kv]}_" + f"dtype_kv_{filename_safe_dtype_map_kv(dtype_kv)}_" f"dtype_o_{filename_safe_dtype_map[dtype_o]}_" f"head_dim_qk_{head_dim_qk}_" f"head_dim_vo_{head_dim_vo}_" @@ -77,7 +79,7 @@ def get_batch_decode_uri( ) -> str: return ( f"batch_decode_with_kv_cache_dtype_q_{filename_safe_dtype_map[dtype_q]}_" - f"dtype_kv_{filename_safe_dtype_map[dtype_kv]}_" + f"dtype_kv_{filename_safe_dtype_map_kv(dtype_kv)}_" f"dtype_o_{filename_safe_dtype_map[dtype_o]}_" f"dtype_idx_{filename_safe_dtype_map[dtype_idx]}_" f"head_dim_qk_{head_dim_qk}_" @@ -100,7 +102,7 @@ def get_batch_mla_uri( ) -> str: return ( f"batch_mla_attention_dtype_q_{filename_safe_dtype_map[dtype_q]}_" - f"dtype_kv_{filename_safe_dtype_map[dtype_kv]}_" + f"dtype_kv_{filename_safe_dtype_map_kv(dtype_kv)}_" f"dtype_o_{filename_safe_dtype_map[dtype_o]}_" f"dtype_idx_{filename_safe_dtype_map[dtype_idx]}_" f"head_dim_ckv_{head_dim_ckv}_" @@ -217,7 +219,7 @@ def get_batch_decode_mla_uri( ) -> str: return ( f"batch_decode_mla_with_kv_cache_dtype_q_{filename_safe_dtype_map[dtype_q]}_" - f"dtype_kv_{filename_safe_dtype_map[dtype_kv]}_" + f"dtype_kv_{filename_safe_dtype_map_kv(dtype_kv)}_" f"dtype_o_{filename_safe_dtype_map[dtype_o]}_" f"dtype_idx_{filename_safe_dtype_map[dtype_idx]}_" f"head_dim_ckv{head_dim_ckv}_" @@ -329,7 +331,7 @@ def get_single_prefill_uri( ) -> str: return ( f"single_prefill_with_kv_cache_dtype_q_{filename_safe_dtype_map[dtype_q]}_" - f"dtype_kv_{filename_safe_dtype_map[dtype_kv]}_" + f"dtype_kv_{filename_safe_dtype_map_kv(dtype_kv)}_" f"dtype_o_{filename_safe_dtype_map[dtype_o]}_" f"head_dim_qk_{head_dim_qk}_" f"head_dim_vo_{head_dim_vo}_" @@ -356,7 +358,7 @@ def get_pod_uri( ) -> str: return ( f"pod_with_kv_cache_dtype_q_{filename_safe_dtype_map[dtype_q]}_" - f"dtype_kv_{filename_safe_dtype_map[dtype_kv]}_" + f"dtype_kv_{filename_safe_dtype_map_kv(dtype_kv)}_" f"dtype_o_{filename_safe_dtype_map[dtype_o]}_" f"head_dim_{head_dim}_" f"posenc_p_{pos_encoding_mode_p}_" @@ -385,7 +387,7 @@ def get_batch_prefill_uri( ) -> str: return ( f"batch_prefill_with_kv_cache_dtype_q_{filename_safe_dtype_map[dtype_q]}_" - f"dtype_kv_{filename_safe_dtype_map[dtype_kv]}_" + f"dtype_kv_{filename_safe_dtype_map_kv(dtype_kv)}_" f"dtype_o_{filename_safe_dtype_map[dtype_o]}_" f"dtype_idx_{filename_safe_dtype_map[dtype_idx]}_" f"head_dim_qk_{head_dim_qk}_" @@ -410,7 +412,7 @@ def get_batch_prefill_attention_sink_uri( ) -> str: return ( f"batch_prefill_with_attention_sink_kv_cache_dtype_q_{filename_safe_dtype_map[dtype_q]}_" - f"dtype_kv_{filename_safe_dtype_map[dtype_kv]}_" + f"dtype_kv_{filename_safe_dtype_map_kv(dtype_kv)}_" f"dtype_o_{filename_safe_dtype_map[dtype_o]}_" f"dtype_idx_{filename_safe_dtype_map[dtype_idx]}_" f"head_dim_qk_{head_dim_qk}_" @@ -432,7 +434,7 @@ def get_batch_attention_uri( ) -> str: return ( f"batch_attention_with_kv_cache_dtype_q_{filename_safe_dtype_map[dtype_q]}_" - f"dtype_kv_{filename_safe_dtype_map[dtype_kv]}_" + f"dtype_kv_{filename_safe_dtype_map_kv(dtype_kv)}_" f"dtype_o_{filename_safe_dtype_map[dtype_o]}_" f"dtype_idx_{filename_safe_dtype_map[dtype_idx]}_" f"head_dim_qk_{head_dim_qk}_" @@ -1617,6 +1619,20 @@ def gen_customize_batch_prefill_module( use_fp16_qk_reduction: bool = False, fp8_enabled: bool = False, ) -> JitSpec: + require_fp4_kv_cache = dtype_map_kv[dtype_kv] == "__nv_fp4x2_e2m1" + if require_fp4_kv_cache: + missing_sf_tensors = [ + name + for name in ("maybe_k_cache_sf", "maybe_v_cache_sf") + if name not in additional_tensor_names + ] + if missing_sf_tensors: + raise ValueError( + "NVFP4 KV paged prefill JIT modules require scale-factor tensors " + f"{missing_sf_tensors}; pass maybe_k_cache_sf and maybe_v_cache_sf " + "as additional tensors." + ) + kwargs = { "variant_decl": variant_decl, "variant_name": variant_name, @@ -1624,6 +1640,7 @@ def gen_customize_batch_prefill_module( "dtype_kv": dtype_map_kv[dtype_kv], "dtype_o": dtype_map[dtype_o], "idtype": dtype_map[idtype], + "require_fp4_kv_cache": require_fp4_kv_cache, "head_dim_qk": head_dim_qk, "head_dim_vo": head_dim_vo, "pos_encoding_mode": pos_encoding_mode_literal[pos_encoding_mode], @@ -1708,12 +1725,17 @@ def gen_customize_batch_prefill_module( generated_config_path = gen_directory / "batch_prefill_config.inc" write_if_different(generated_config_path, generated_inc_str) + extra_cuda_cflags = _fa2_prefill_head_dim_nvcc_flags( + head_dim_qk, head_dim_vo, dtype_kv + ) + if kwargs["require_fp4_kv_cache"]: + # NVFP4 KV kernels need FLASHINFER_ENABLE_FP4_E2M1 (common flags) even + # when the head_dim helper returns no arch-specific flags. + extra_cuda_cflags = (extra_cuda_cflags or []) + common_nvcc_flags return gen_jit_spec( uri, source_paths, - extra_cuda_cflags=_fa2_prefill_head_dim_nvcc_flags( - head_dim_qk, head_dim_vo, dtype_kv - ), + extra_cuda_cflags=extra_cuda_cflags, ) elif backend == "fa3": gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri @@ -1810,7 +1832,7 @@ def get_fmha_cutlass_sm100a_uri( return "fmha_cutlass_sm100a" # return ( # f"fmha_cutlass_sm100a_dtype_q_{filename_safe_dtype_map[dtype_q]}_" - # f"dtype_kv_{filename_safe_dtype_map[dtype_kv]}_" + # f"dtype_kv_{filename_safe_dtype_map_kv(dtype_kv)}_" # f"dtype_o_{filename_safe_dtype_map[dtype_o]}_" # f"dtype_idx_{filename_safe_dtype_map[dtype_idx]}_" # f"head_dim_qk_{head_dim_qk}_" diff --git a/flashinfer/jit/utils.py b/flashinfer/jit/utils.py index befd37d8673..70ce91ab10f 100644 --- a/flashinfer/jit/utils.py +++ b/flashinfer/jit/utils.py @@ -81,6 +81,15 @@ def write_if_different(path: pathlib.Path, content: str) -> None: if hasattr(torch, "float4_e2m1fn_x2"): filename_safe_dtype_map[torch.float4_e2m1fn_x2] = "fp4_e2m1" + +def filename_safe_dtype_map_kv(dtype: torch.dtype) -> str: + """Return the cache-key dtype name for KV cache kernels.""" + + if dtype_map_kv[dtype] == "__nv_fp4x2_e2m1": + return "fp4x2_e2m1" + return filename_safe_dtype_map[dtype] + + pos_encoding_mode_literal = { 0: "PosEncodingMode::kNone", 1: "PosEncodingMode::kRoPELlama", diff --git a/flashinfer/prefill.py b/flashinfer/prefill.py index 6c9f533f424..feed056a9ed 100755 --- a/flashinfer/prefill.py +++ b/flashinfer/prefill.py @@ -1370,8 +1370,10 @@ def single_prefill_with_kv_cache( if scale_v is None: scale_v = torch.ones(v.shape[1], dtype=torch.float32, device=q.device) - # For NVFP4 KV (uint8 packed), last dim is head_dim//2; output uses q head_dim. - out_head_dim = q.shape[-1] if kv_cache_sf is not None else v.shape[-1] + # For NVFP4 KV (uint8 packed), last dim is head_dim//2 packed bytes: the + # unpacked VO width is v.shape[-1] * 2, which equals head_dim_vo even for + # asymmetric (QK, VO) plans; q.shape[-1] assumed QK == VO. + out_head_dim = v.shape[-1] * 2 if kv_cache_sf is not None else v.shape[-1] if backend == "auto": backend = determine_attention_backend( @@ -1691,9 +1693,11 @@ def __init__( ) # jit_args[7] is additional_tensor_names from gen_customize_batch_prefill_module self._jit_additional_tensor_names = list(jit_args[7]) + self._jit_additional_scalar_names = list(jit_args[9]) else: self._jit_module = None self._jit_additional_tensor_names = [] + self._jit_additional_scalar_names = [] self._kv_layout = kv_layout if backend == "cudnn": @@ -2736,7 +2740,15 @@ def run( # For NVFP4 KV (uint8 packed), v_cache last dim is head_dim//2; # use q's head_dim for output instead - out_head_dim = q.shape[-1] if kv_cache_sf is not None else v_cache.shape[-1] + # For NVFP4 KV (uint8 packed), v_cache last dim is packed bytes + # (2 values per byte): the unpacked VO width is v_cache.shape[-1]*2, + # which equals head_dim_vo even for asymmetric (QK, VO) plans. + # Using q.shape[-1] here assumed head_dim_vo == head_dim_qk and made + # the kernel (which writes head_dim_vo-wide rows) garble a too-wide + # output buffer whenever VO < QK. + out_head_dim = ( + v_cache.shape[-1] * 2 if kv_cache_sf is not None else v_cache.shape[-1] + ) if out is None: # Use cached output data type if available (for FP8 attention with FP16 output) out_dtype = getattr(self, "_cached_o_data_type", None) or q.dtype @@ -2826,19 +2838,46 @@ def run( enable_pdl, ] if self._jit_module is not None: - run_args.extend( - prepare_jit_additional_args( - self._jit_additional_tensor_names, - { - "maybe_custom_mask": self._custom_mask_buf, - "maybe_mask_indptr": self._mask_indptr_buf, - "maybe_alibi_slopes": lambda: _get_cache_alibi_slopes_buf( - q.shape[1], q.device - ), - }, - args, - ) + additional_args = prepare_jit_additional_args( + self._jit_additional_tensor_names, + { + "maybe_custom_mask": self._custom_mask_buf, + "maybe_mask_indptr": self._mask_indptr_buf, + "maybe_alibi_slopes": lambda: _get_cache_alibi_slopes_buf( + q.shape[1], q.device + ), + "maybe_prefix_len_ptr": self._prefix_len_ptr, + "maybe_token_pos_in_items_ptr": self._token_pos_in_items_ptr, + "maybe_max_item_len_ptr": self._max_item_len_ptr, + "maybe_k_cache_sf": key_block_scales, + "maybe_v_cache_sf": value_block_scales, + }, + args, ) + expected_additional_arg_count = len( + self._jit_additional_tensor_names + ) + len(self._jit_additional_scalar_names) + if len(additional_args) < expected_additional_arg_count: + _, scale_q_scalar = _split_scale_param(q_scale) + _, scale_k_scalar = _split_scale_param(k_scale) + _, scale_v_scalar = _split_scale_param(v_scale) + jit_scalar_values = { + "logits_soft_cap": logits_soft_cap, + "sm_scale": sm_scale, + "rope_rcp_scale": 1.0 / rope_scale, + "rope_rcp_theta": 1.0 / rope_theta, + "scale_q_scalar": scale_q_scalar, + "scale_k_scalar": scale_k_scalar, + "scale_v_scalar": scale_v_scalar, + "token_pos_in_items_len": self._token_pos_in_items_len, + } + scalar_start = max( + 0, + len(additional_args) - len(self._jit_additional_tensor_names), + ) + for name in self._jit_additional_scalar_names[scalar_start:]: + additional_args.append(jit_scalar_values[name]) + run_args.extend(additional_args) else: # Extract FP8 scale tensors from *args if q is FP8 fp8_scale_q = None @@ -3854,8 +3893,9 @@ def run( else: k_sf, v_sf = kv_cache_sf.unbind(dim=1) - # For NVFP4 KV (uint8 packed), v last dim is head_dim//2; use q head_dim for output - out_head_dim = q.shape[-1] if kv_cache_sf is not None else v.shape[-1] + # NVFP4 packed: unpacked VO width is packed bytes * 2 (supports + # asymmetric QK/VO; q.shape[-1] assumed QK == VO). + out_head_dim = v.shape[-1] * 2 if kv_cache_sf is not None else v.shape[-1] if out is None: # when input dtype is fp8, we need to use bf16 output out_dtype = torch.bfloat16 if q.dtype.itemsize == 1 else q.dtype diff --git a/include/flashinfer/attention/persistent.cuh b/include/flashinfer/attention/persistent.cuh index 5000fb8c50e..0c21634a6e2 100644 --- a/include/flashinfer/attention/persistent.cuh +++ b/include/flashinfer/attention/persistent.cuh @@ -265,7 +265,10 @@ struct BlockBatchPagedAttentionPersistent { v_smem_offset_w = get_permuted_offset( warp_idx * KTraits::KV_THR_LAYOUT_ROW + lane_idx / KTraits::KV_THR_LAYOUT_COL, lane_idx % KTraits::KV_THR_LAYOUT_COL); - size_t thr_local_kv_offset[NUM_MMA_KV * KTraits::KV_THR_LAYOUT_COL / 2 / KTraits::NUM_WARPS_Q]; + size_t thr_local_kv_offset_k[NUM_MMA_KV * KTraits::KV_THR_LAYOUT_COL / 2 / + KTraits::NUM_WARPS_Q]; + size_t thr_local_kv_offset_v[NUM_MMA_KV * KTraits::KV_THR_LAYOUT_COL / 2 / + KTraits::NUM_WARPS_Q]; #pragma unroll 1 for (IdType work_idx = work_indptr[blockIdx.y]; work_idx < work_indptr[blockIdx.y + 1]; @@ -322,9 +325,12 @@ struct BlockBatchPagedAttentionPersistent { prefetch_offest(block_iter_base + kv_tile_idx * CTA_TILE_KV, packed_kv_bound, kv_head_idx, k_stride_page, k_stride_h, k_stride_n, block_size, - kv_indices, thr_local_kv_offset); + kv_indices, thr_local_kv_offset_k); + prefetch_offest(block_iter_base + kv_tile_idx * CTA_TILE_KV, packed_kv_bound, + kv_head_idx, v_stride_page, v_stride_h, v_stride_n, block_size, + kv_indices, thr_local_kv_offset_v); page_produce_kv(smem_storage, &k_smem_offset_w, k, - kv_start + kv_tile_idx * CTA_TILE_KV, thr_local_kv_offset, + kv_start + kv_tile_idx * CTA_TILE_KV, thr_local_kv_offset_k, kv_end, warp_idx, lane_idx); page_produce_kv_sf( smem_storage, maybe_k_cache_sf, block_iter_base + kv_tile_idx * CTA_TILE_KV, @@ -333,7 +339,7 @@ struct BlockBatchPagedAttentionPersistent { kv_end, warp_idx, lane_idx); cp_async::commit_group(); page_produce_kv(smem_storage, &v_smem_offset_w, v, - kv_start + kv_tile_idx * CTA_TILE_KV, thr_local_kv_offset, + kv_start + kv_tile_idx * CTA_TILE_KV, thr_local_kv_offset_v, kv_end, warp_idx, lane_idx); page_produce_kv_sf( smem_storage, maybe_v_cache_sf, block_iter_base + kv_tile_idx * CTA_TILE_KV, @@ -348,7 +354,10 @@ struct BlockBatchPagedAttentionPersistent { kv_tile_idx + 1 > NUM_STAGES, { prefetch_offest(block_iter_base + (kv_tile_idx - 1) * CTA_TILE_KV, packed_kv_bound, kv_head_idx, k_stride_page, k_stride_h, - k_stride_n, block_size, kv_indices, thr_local_kv_offset); + k_stride_n, block_size, kv_indices, thr_local_kv_offset_k); + prefetch_offest(block_iter_base + (kv_tile_idx - 1) * CTA_TILE_KV, + packed_kv_bound, kv_head_idx, v_stride_page, v_stride_h, + v_stride_n, block_size, kv_indices, thr_local_kv_offset_v); cp_async::wait_group<1>(); __syncthreads(); @@ -376,7 +385,7 @@ struct BlockBatchPagedAttentionPersistent { __syncthreads(); page_produce_kv(smem_storage, &k_smem_offset_w, k, kv_start + (kv_tile_idx - 1) * CTA_TILE_KV, - thr_local_kv_offset, kv_end, warp_idx, lane_idx); + thr_local_kv_offset_k, kv_end, warp_idx, lane_idx); page_produce_kv_sf( smem_storage, maybe_k_cache_sf, block_iter_base + (kv_tile_idx - 1) * CTA_TILE_KV, packed_kv_bound, kv_head_idx, params.k_sf_stride_page, params.k_sf_stride_h, @@ -395,7 +404,7 @@ struct BlockBatchPagedAttentionPersistent { page_produce_kv(smem_storage, &v_smem_offset_w, v, kv_start + (kv_tile_idx - 1) * CTA_TILE_KV, - thr_local_kv_offset, kv_end, warp_idx, lane_idx); + thr_local_kv_offset_v, kv_end, warp_idx, lane_idx); page_produce_kv_sf( smem_storage, maybe_v_cache_sf, block_iter_base + (kv_tile_idx - 1) * CTA_TILE_KV, packed_kv_bound, kv_head_idx, params.v_sf_stride_page, params.v_sf_stride_h, diff --git a/include/flashinfer/attention/prefill.cuh b/include/flashinfer/attention/prefill.cuh index 858e4bf2e2f..0a7a677a73b 100644 --- a/include/flashinfer/attention/prefill.cuh +++ b/include/flashinfer/attention/prefill.cuh @@ -683,8 +683,8 @@ __device__ __forceinline__ void page_produce_kv_on_the_fly( * SF bytes per iteration, advancing by NUM_WARPS * 128 bytes across iterations. * The SF smem layout is a plain flat byte array — no swizzle. * - * SF strides are KV byte strides divided by SF_CONTAINERS (= NVFP4_SF_VEC_SIZE/2 = 8), - * which is exact because all NVFP4-compatible head_dims are divisible by 16. + * SF strides are passed explicitly by the caller instead of being derived from KV data strides. + * This lets runtimes pass scale-factor tensors from interleaved or separately allocated KV pools. * No-op when KTraits::DTypeKV is not FP4. * * \tparam produce_v true → fill v_sf_smem, false → fill k_sf_smem. @@ -705,6 +705,10 @@ __device__ __forceinline__ void page_produce_kv_on_the_fly( * \param warp_idx Global warp index within the CTA. * \param lane_idx Lane index within the warp. */ +#ifndef FLASHINFER_PAGED_V_SF_DESWIZZLE +#define FLASHINFER_PAGED_V_SF_DESWIZZLE 0 +#endif + template __device__ __forceinline__ void page_produce_kv_sf( SmemStorage* smem_storage, uint8_t* sf_ptr, const uint32_t packed_page_iter_base, @@ -743,20 +747,42 @@ __device__ __forceinline__ void page_produce_kv_sf( uint32_t page_iter, entry_idx; const uint32_t packed_block_iter = packed_page_iter_base + sf_smem_row; page_size.divmod(packed_block_iter, page_iter, entry_idx); - const size_t sf_gmem_offset = + const size_t page_head_base = static_cast(packed_block_iter < packed_kv_bound ? indices[page_iter] : 0) * sf_stride_page + - kv_head_idx * sf_stride_h + entry_idx * sf_stride_n + sf_smem_col; - - // V SF must zero-fill out-of-bounds entries: compute_sfm_v reads SF for all CTA_TILE_KV rows - // including padding, and 0 (softmax weight) * NaN (uninitialized SF) = NaN (IEEE 754). - // K SF can use kNoFill since NaN K scores are replaced by -inf via logits_mask before - // update_mdo_states, so they never reach the accumulator. - constexpr auto fill_mode = - produce_v ? cp_async::SharedMemFillMode::kFillZero : cp_async::SharedMemFillMode::kNoFill; - cp_async::pred_load_32b(reinterpret_cast(sf_smem + flat_byte), - reinterpret_cast(sf_ptr + sf_gmem_offset), - in_bounds); + kv_head_idx * sf_stride_h; + + if constexpr (produce_v && FLASHINFER_PAGED_V_SF_DESWIZZLE) { + static_assert(SF_COLS % 4 == 0, + "Paged V-SF de-swizzle requires HEAD_DIM_VO divisible by 64"); + uint32_t packed = 0; + if (in_bounds) { + constexpr uint32_t SF_GROUPS = SF_COLS / 4; + const uint32_t a4 = entry_idx & ~3u; + const uint32_t e = entry_idx & 3u; + uint8_t* packed_bytes = reinterpret_cast(&packed); +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + const uint32_t dcol = sf_smem_col + j; + const uint32_t swz_entry = a4 + dcol / SF_GROUPS; + const uint32_t swz_sd = (dcol % SF_GROUPS) * 4 + e; + packed_bytes[j] = + sf_ptr[page_head_base + static_cast(swz_entry) * sf_stride_n + swz_sd]; + } + } + *reinterpret_cast(sf_smem + flat_byte) = packed; + } else { + const size_t sf_gmem_offset = page_head_base + entry_idx * sf_stride_n + sf_smem_col; + // V SF must zero-fill out-of-bounds entries: compute_sfm_v reads SF for all CTA_TILE_KV + // rows including padding, and 0 (softmax weight) * NaN (uninitialized SF) = NaN (IEEE 754). + // K SF can use kNoFill since NaN K scores are replaced by -inf via logits_mask before + // update_mdo_states, so they never reach the accumulator. + constexpr auto fill_mode = + produce_v ? cp_async::SharedMemFillMode::kFillZero : cp_async::SharedMemFillMode::kNoFill; + cp_async::pred_load_32b( + reinterpret_cast(sf_smem + flat_byte), + reinterpret_cast(sf_ptr + sf_gmem_offset), in_bounds); + } } } } @@ -776,8 +802,8 @@ __device__ __forceinline__ void page_produce_kv_sf( * \param sf_ptr Base pointer to the flat uint8_t SF array (K or V). * \param kv_abs_base Absolute first token index for this CTA tile. * \param kv_head_idx KV head index. - * \param kv_stride_n Byte stride per token in the KV tensor. - * \param kv_stride_h Byte stride per head in the KV tensor. + * \param sf_stride_n Byte stride per token in the SF tensor. + * \param sf_stride_h Byte stride per head in the SF tensor. * \param kv_idx_base First KV row index for this tile within the chunk. * \param kv_len Chunk size; rows at or beyond this are not loaded. * \param warp_idx Global warp index within the CTA. @@ -2943,7 +2969,8 @@ __global__ __launch_bounds__(KTraits::NUM_THREADS) void BatchPrefillWithRaggedKV produce_kv(k_smem, &k_smem_offset_w, &k_ptr, k_stride_n, 0, chunk_size, tid); produce_kv_sf(&smem_storage, maybe_k_cache_sf, kv_abs_base, kv_head_idx, - k_stride_n, k_stride_h, 0, chunk_size, warp_idx, lane_idx); + k_stride_n, k_stride_h, 0, + chunk_size, warp_idx, lane_idx); cp_async::commit_group(); if constexpr (!KTraits::USE_KV_SHARED_SMEM) { // Shared K/V: don't preload V(0) (it would clobber K(0)); V(0) is loaded @@ -3658,8 +3685,13 @@ __device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( smem_t k_smem(smem_storage.k_smem), v_smem(KTraits::USE_KV_SHARED_SMEM ? smem_storage.k_smem : smem_storage.v_smem); constexpr uint32_t NUM_PAGED_KV_OFFSETS = NUM_MMA_KV * KV_THR_LAYOUT_COL / 2 / NUM_WARPS_Q; + // Separate K/V page offsets: the K and V pools may carry different strides + // (e.g. NVFP4 VO-split caches where head_dim_qk != head_dim_vo). The K/V-shared + // smem path computes offsets on the fly, so the arrays collapse to [1] stubs there. [[maybe_unused]] size_t - thr_local_kv_offset[KTraits::USE_KV_SHARED_SMEM ? 1 : NUM_PAGED_KV_OFFSETS]; + thr_local_kv_offset_k[KTraits::USE_KV_SHARED_SMEM ? 1 : NUM_PAGED_KV_OFFSETS]; + [[maybe_unused]] size_t + thr_local_kv_offset_v[KTraits::USE_KV_SHARED_SMEM ? 1 : NUM_PAGED_KV_OFFSETS]; uint32_t k_smem_offset_r = k_smem.template get_permuted_offset( get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + 8 * (lane_idx / 16) + @@ -3711,13 +3743,15 @@ __device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( page_iter, entry_idx); // FP4: GMEM is packed (2 FP4/byte), so the column byte offset is halved relative to fp8 constexpr uint32_t fp4_pack_factor = is_fp4_type_v ? 2 : 1; - 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() / fp4_pack_factor, - last_indptr); + const uint32_t feat_idx = + (lane_idx % KV_THR_LAYOUT_COL) * upcast_size() / fp4_pack_factor; + thr_local_kv_offset_k[i] = paged_kv.protective_get_k_offset( + page_iter, kv_head_idx, entry_idx, feat_idx, last_indptr); + thr_local_kv_offset_v[i] = paged_kv.protective_get_v_offset( + page_iter, kv_head_idx, entry_idx, feat_idx, last_indptr); } page_produce_kv(&smem_storage, &k_smem_offset_w, paged_kv.k_data, 0, - thr_local_kv_offset, chunk_size, warp_idx, lane_idx); + thr_local_kv_offset_k, chunk_size, warp_idx, lane_idx); } page_produce_kv_sf( &smem_storage, maybe_k_cache_sf, packed_page_iter_base, @@ -3727,7 +3761,7 @@ __device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( // Shared K/V loads V(0) inside iter 0 after Q.K^T; preloading it would clobber K(0). if constexpr (!KTraits::USE_KV_SHARED_SMEM) { page_produce_kv(&smem_storage, &v_smem_offset_w, paged_kv.v_data, 0, - thr_local_kv_offset, chunk_size, warp_idx, lane_idx); + thr_local_kv_offset_v, chunk_size, warp_idx, lane_idx); page_produce_kv_sf(&smem_storage, maybe_v_cache_sf, packed_page_iter_base, last_indptr * (uint32_t)paged_kv.page_size, kv_head_idx, v_sf_stride_page, v_sf_stride_h, v_sf_stride_n, @@ -3816,10 +3850,12 @@ __device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( page_iter, entry_idx); // FP4: GMEM is packed (2 FP4/byte), so the column byte offset is halved relative to fp8 constexpr uint32_t fp4_pack_factor = is_fp4_type_v ? 2 : 1; - 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() / fp4_pack_factor, - last_indptr); + const uint32_t feat_idx = + (lane_idx % KV_THR_LAYOUT_COL) * upcast_size() / fp4_pack_factor; + thr_local_kv_offset_k[i] = paged_kv.protective_get_k_offset( + page_iter, kv_head_idx, entry_idx, feat_idx, last_indptr); + thr_local_kv_offset_v[i] = paged_kv.protective_get_v_offset( + page_iter, kv_head_idx, entry_idx, feat_idx, last_indptr); } } // Shared K/V serializes loads (no K/V prefetch overlap) -> drain fully. @@ -3926,8 +3962,8 @@ __device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( cp_async::wait_group<0>(); } else { page_produce_kv(&smem_storage, &k_smem_offset_w, paged_kv.k_data, - (iter + 1) * CTA_TILE_KV, thr_local_kv_offset, chunk_size, - warp_idx, lane_idx); + (iter + 1) * CTA_TILE_KV, thr_local_kv_offset_k, + chunk_size, warp_idx, lane_idx); page_produce_kv_sf( &smem_storage, maybe_k_cache_sf, packed_page_iter_base, last_indptr * (uint32_t)paged_kv.page_size, kv_head_idx, k_sf_stride_page, @@ -3973,8 +4009,8 @@ __device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( packed_page_iter_base = next_packed_page_iter_base; } else { page_produce_kv(&smem_storage, &v_smem_offset_w, paged_kv.v_data, - (iter + 1) * CTA_TILE_KV, thr_local_kv_offset, chunk_size, - warp_idx, lane_idx); + (iter + 1) * CTA_TILE_KV, thr_local_kv_offset_v, + chunk_size, warp_idx, lane_idx); page_produce_kv_sf( &smem_storage, maybe_v_cache_sf, packed_page_iter_base, last_indptr * (uint32_t)paged_kv.page_size, kv_head_idx, v_sf_stride_page, diff --git a/include/flashinfer/page.cuh b/include/flashinfer/page.cuh index a1ee13497d5..c4470176011 100644 --- a/include/flashinfer/page.cuh +++ b/include/flashinfer/page.cuh @@ -43,6 +43,9 @@ struct paged_kv_t { uint32_t stride_page; uint32_t stride_n; uint32_t stride_h; + uint32_t v_stride_page; + uint32_t v_stride_n; + uint32_t v_stride_h; // Internal layout: // [max_num_pages, num_heads, page_size, head_dim] if layout == HND @@ -69,6 +72,9 @@ struct paged_kv_t { stride_page(0), stride_n(0), stride_h(0), + v_stride_page(0), + v_stride_n(0), + v_stride_h(0), k_data(nullptr), v_data(nullptr), indices(nullptr), @@ -103,10 +109,13 @@ struct paged_kv_t { last_page_len(last_page_len), rope_pos_offset(rope_pos_offset) { stride_page = num_heads * page_size * head_dim; + v_stride_page = stride_page; this->k_data = k_data; this->v_data = v_data; stride_n = layout == QKVLayout::kHND ? head_dim : num_heads * head_dim; stride_h = layout == QKVLayout::kHND ? page_size * head_dim : head_dim; + v_stride_n = stride_n; + v_stride_h = stride_h; } /*! @@ -138,10 +147,39 @@ struct paged_kv_t { last_page_len(last_page_len), rope_pos_offset(rope_pos_offset) { stride_page = kv_strides[0]; + v_stride_page = stride_page; this->k_data = k_data; this->v_data = v_data; stride_n = layout == QKVLayout::kHND ? kv_strides[2] : kv_strides[1]; stride_h = layout == QKVLayout::kHND ? kv_strides[1] : kv_strides[2]; + v_stride_n = stride_n; + v_stride_h = stride_h; + } + + /*! + * \brief Construct a paged key-value cache with independent K/V strides. + */ + __host__ __forceinline__ paged_kv_t(uint32_t num_heads, uint32_t page_size, uint32_t head_dim, + uint32_t batch_size, QKVLayout layout, DType* k_data, + DType* v_data, const int64_t* k_strides, + const int64_t* v_strides, IdType* indices, IdType* indptr, + IdType* last_page_len, IdType* rope_pos_offset = nullptr) + : num_heads(num_heads), + page_size(page_size), + head_dim(head_dim), + batch_size(batch_size), + indices(indices), + indptr(indptr), + last_page_len(last_page_len), + rope_pos_offset(rope_pos_offset) { + stride_page = k_strides[0]; + v_stride_page = v_strides[0]; + this->k_data = k_data; + this->v_data = v_data; + stride_n = layout == QKVLayout::kHND ? k_strides[2] : k_strides[1]; + stride_h = layout == QKVLayout::kHND ? k_strides[1] : k_strides[2]; + v_stride_n = layout == QKVLayout::kHND ? v_strides[2] : v_strides[1]; + v_stride_h = layout == QKVLayout::kHND ? v_strides[1] : v_strides[2]; } __host__ __device__ __forceinline__ uint32_t get_length(uint32_t batch_idx) const { @@ -164,6 +202,12 @@ struct paged_kv_t { return page_idx * stride_page + head_idx * stride_h + entry_idx * stride_n + feat_idx; } + __host__ __device__ __forceinline__ size_t get_v_elem_offset(size_t page_idx, size_t head_idx, + size_t entry_idx, + size_t feat_idx) const { + return page_idx * v_stride_page + head_idx * v_stride_h + entry_idx * v_stride_n + feat_idx; + } + /*! * \brief Compute the offset of element inside the page. * \param head_idx The head index @@ -199,13 +243,31 @@ struct paged_kv_t { __device__ __forceinline__ DType* get_v_ptr(IdType page_iter, uint32_t head_idx, uint32_t entry_idx, uint32_t feat_idx) const { - return v_data + get_elem_offset(__ldg(indices + page_iter), head_idx, entry_idx, feat_idx); + return v_data + get_v_elem_offset(__ldg(indices + page_iter), head_idx, entry_idx, feat_idx); + } + + __device__ __forceinline__ size_t protective_get_k_offset(IdType page_iter, uint32_t head_idx, + uint32_t entry_idx, + uint32_t feat_idx, + IdType last_indptr) const { + return protective_get_kv_offset(page_iter, head_idx, entry_idx, feat_idx, last_indptr); + } + + __device__ __forceinline__ size_t protective_get_v_offset(IdType page_iter, uint32_t head_idx, + uint32_t entry_idx, + uint32_t feat_idx, + IdType last_indptr) const { + if (page_iter < last_indptr) { + return get_v_elem_offset(__ldg(indices + page_iter), head_idx, entry_idx, feat_idx); + } else { + return 0; + } } __device__ __forceinline__ DType* protective_get_v_ptr(IdType page_iter, uint32_t head_idx, uint32_t entry_idx, uint32_t feat_idx, IdType last_indptr) const { - return v_data + protective_get_kv_offset(page_iter, head_idx, entry_idx, feat_idx, last_indptr); + return v_data + protective_get_v_offset(page_iter, head_idx, entry_idx, feat_idx, last_indptr); } }; diff --git a/include/flashinfer/utils.cuh b/include/flashinfer/utils.cuh index 15214049c15..fbd1c3e6f11 100644 --- a/include/flashinfer/utils.cuh +++ b/include/flashinfer/utils.cuh @@ -152,6 +152,9 @@ } else if (group_size == 4) { \ constexpr size_t GROUP_SIZE = 4; \ __VA_ARGS__ \ + } else if (group_size == 6) { \ + constexpr size_t GROUP_SIZE = 6; \ + __VA_ARGS__ \ } else if (group_size == 8) { \ constexpr size_t GROUP_SIZE = 8; \ __VA_ARGS__ \ diff --git a/include/flashinfer/vec_dtypes.cuh b/include/flashinfer/vec_dtypes.cuh index 25c3b6fc60d..c64175fcdf5 100644 --- a/include/flashinfer/vec_dtypes.cuh +++ b/include/flashinfer/vec_dtypes.cuh @@ -472,6 +472,44 @@ struct vec_cast { #endif } }; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(float* dst, const __nv_fp4x2_e2m1* src) { + static_assert(vec_size % 2 == 0, "vec_size must be even for fp4x2 dequantization"); +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint32_t fp16x2; + // Valid fp4x2 bytes are at even positions (stride 2); odd positions are padding. + uint32_t b = reinterpret_cast(src)[i * 2]; + asm volatile( + "{\n" + ".reg .b8 fp4_byte;\n" + "mov.b32 {fp4_byte, _, _, _}, %1;\n" + "cvt.rn.f16x2.e2m1x2 %0, fp4_byte;\n" + "}" + : "=r"(fp16x2) + : "r"(b)); + __half2 h2 = *reinterpret_cast<__half2*>(&fp16x2); + reinterpret_cast(dst)[i] = __half22float2(h2); + } +#else + constexpr float lut[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f, + }; +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint8_t b = reinterpret_cast(src)[i * 2]; + dst[i * 2 + 0] = lut[b & 0x0F]; + dst[i * 2 + 1] = lut[(b >> 4) & 0x0F]; + } +#endif + } +}; + template <> struct vec_cast { template @@ -1172,6 +1210,67 @@ struct vec_t<__nv_fp8_e5m2, vec_size> { }; #if defined(FLASHINFER_ENABLE_FP4_E2M1) && CUDA_VERSION >= 12080 +/******************* vec_t<__nv_fp4x2_e2m1> *******************/ + +// __nv_fp4x2_e2m1 is used as a one-byte packed container for FP4 KV-cache data. +// The current FA2 NVFP4 path stores valid packed bytes at even positions and padding at odd +// positions, so vec_size still counts logical elements in the surrounding kernels. +template +struct vec_t<__nv_fp4x2_e2m1, vec_size> { + static_assert(vec_size % 16 == 0, "Invalid vector size"); + int4 data[vec_size / 16]; + + FLASHINFER_INLINE __nv_fp4x2_e2m1& operator[](size_t i) { + return ((__nv_fp4x2_e2m1*)data)[i]; + } + FLASHINFER_INLINE const __nv_fp4x2_e2m1& operator[](size_t i) const { + return ((const __nv_fp4x2_e2m1*)data)[i]; + } + FLASHINFER_INLINE __nv_fp4x2_e2m1* ptr() { return reinterpret_cast<__nv_fp4x2_e2m1*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp4x2_e2m1 val) { + uint32_t val8 = __nv_fp4x2_storage_t(val.__x); + uint32_t val16 = (val8 << 8) | val8; + uint32_t val32 = (val16 << 16) | val16; +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i].x = val32; + data[i].y = val32; + data[i].z = val32; + data[i].w = val32; + } + } + FLASHINFER_INLINE void load(const __nv_fp4x2_e2m1* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ((int4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(__nv_fp4x2_e2m1* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)ptr)[i] = data[i]; + } + } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(__nv_fp4x2_e2m1* dst, const __nv_fp4x2_e2m1* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)dst)[i] = ((int4*)src)[i]; + } + } +}; + /******************* vec_t<__nv_fp4_e2m1> *******************/ // __nv_fp4_e2m1 x 2 diff --git a/tests/jit/test_attention_utils.py b/tests/jit/test_attention_utils.py new file mode 100644 index 00000000000..09f0935a46a --- /dev/null +++ b/tests/jit/test_attention_utils.py @@ -0,0 +1,73 @@ +from pathlib import Path + +import torch + +from flashinfer.jit import env as jit_env +from flashinfer.jit.attention.modules import gen_customize_batch_prefill_module + + +def test_batch_prefill_nvfp4_swa_paged_params_declares_sf_strides( + tmp_path, monkeypatch +): + repo_root = Path(__file__).resolve().parents[2] + monkeypatch.setattr(jit_env, "FLASHINFER_GEN_SRC_DIR", tmp_path / "generated") + monkeypatch.setattr(jit_env, "FLASHINFER_CSRC_DIR", repo_root / "csrc") + + uri = "test_batch_prefill_nvfp4_swa" + gen_customize_batch_prefill_module( + "fa2", + uri, + torch.bfloat16, + torch.uint8, + torch.bfloat16, + torch.int32, + 128, + 128, + ["maybe_k_cache_sf", "maybe_v_cache_sf"], + ["uint8_t", "uint8_t"], + [], + [], + "DefaultAttention", + "struct DefaultAttention {};", + use_sliding_window=True, + ) + + generated = (tmp_path / "generated" / uri / "batch_prefill_config.inc").read_text() + assert "constexpr bool REQUIRE_FP4_KV_CACHE = true;" in generated + assert "constexpr auto USE_SLIDING_WINDOW = true;" in generated + for field in ( + "maybe_k_cache_sf", + "maybe_v_cache_sf", + ): + assert f"uint8_t* {field};" in generated + # SF strides ride on the upstream static param fields (set from the actual + # SF tensors via GetFP4ScaleStrides in the generated params setter). + for field in ("k_sf", "v_sf"): + assert f"uint32_t {field}_stride_page;" in generated + assert f"uint32_t {field}_stride_h;" in generated + assert f"uint32_t {field}_stride_n;" in generated + + +def test_batch_prefill_nvfp4_requires_sf_tensors(): + try: + gen_customize_batch_prefill_module( + "fa2", + "test_batch_prefill_nvfp4_missing_sf", + torch.bfloat16, + torch.uint8, + torch.bfloat16, + torch.int32, + 128, + 128, + [], + [], + [], + [], + "DefaultAttention", + "struct DefaultAttention {};", + ) + except ValueError as exc: + assert "maybe_k_cache_sf" in str(exc) + assert "maybe_v_cache_sf" in str(exc) + else: + raise AssertionError("expected NVFP4 KV prefill without SF tensors to fail") From 0dd5f7224d7df29dea4f69001255c554065d53c3 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Wed, 17 Jun 2026 09:04:25 +0900 Subject: [PATCH 02/23] fix(prefill): shared-memory-aware CtaTileQ; reject infeasible KV tiles FA2DetermineCtaTileQ picked a 1x4 warp layout without accounting for the real shared-memory budget, yielding max_mma_kv=0 at HEAD_DIM_QK=512 under real GQA with 1-byte (fp8/NVFP4) KV. Make the warp-layout choice smem-aware, and cleanly reject KV tiles that do not fit real shared memory instead of failing later. Signed-off-by: Jetha Chan (cherry picked from commit 23ce85af195717fe0646943b76250a20dabc9de6) --- include/flashinfer/attention/prefill.cuh | 2 +- include/flashinfer/attention/scheduler.cuh | 11 ++++----- include/flashinfer/utils.cuh | 28 +++++++++++++++++++--- 3 files changed, 31 insertions(+), 10 deletions(-) diff --git a/include/flashinfer/attention/prefill.cuh b/include/flashinfer/attention/prefill.cuh index 0a7a677a73b..23f1b6bb3ee 100644 --- a/include/flashinfer/attention/prefill.cuh +++ b/include/flashinfer/attention/prefill.cuh @@ -2505,7 +2505,7 @@ cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::D constexpr uint32_t NUM_MMA_D_QK = HEAD_DIM_QK / 16; constexpr uint32_t NUM_MMA_D_VO = HEAD_DIM_VO / 16; int64_t packed_qo_len = qo_len * group_size; - uint32_t cta_tile_q = FA2DetermineCtaTileQ(packed_qo_len, HEAD_DIM_VO); + uint32_t cta_tile_q = FA2DetermineCtaTileQ(packed_qo_len, HEAD_DIM_VO, HEAD_DIM_QK); DISPATCH_CTA_TILE_Q(cta_tile_q, CTA_TILE_Q, { // hd512 uses the 2-Q x 2-KV-warp layout at CTA_TILE_Q=32. FP8 must take it diff --git a/include/flashinfer/attention/scheduler.cuh b/include/flashinfer/attention/scheduler.cuh index fe683555f43..2a8d3075876 100644 --- a/include/flashinfer/attention/scheduler.cuh +++ b/include/flashinfer/attention/scheduler.cuh @@ -548,7 +548,7 @@ inline auto PrefillSplitQOKVIndptr(IdType* qo_indptr_h, IdType* kv_indptr_h, uint32_t page_size, uint32_t max_batch_size_if_split, bool enable_cuda_graph, int32_t window_left, int32_t fixed_split_size, bool disable_split_kv, - int64_t uniform_q_len) { + int64_t uniform_q_len, uint32_t head_dim_qk = 0) { std::vector request_indices, qo_tile_indices, kv_tile_indices, merge_indptr, o_indptr; merge_indptr.push_back(0); o_indptr.push_back(0); @@ -593,7 +593,7 @@ inline auto PrefillSplitQOKVIndptr(IdType* qo_indptr_h, IdType* kv_indptr_h, FLASHINFER_ERROR(err_msg.str()); } } - cta_tile_q = FA2DetermineCtaTileQ(packed_uniform_len, head_dim); + cta_tile_q = FA2DetermineCtaTileQ(packed_uniform_len, head_dim, head_dim_qk); total_num_tiles_q = batch_size * ceil_div(packed_uniform_len, cta_tile_q); } else { // When CUDA graphs are enabled, the lengths of sequences determined by @@ -601,7 +601,7 @@ inline auto PrefillSplitQOKVIndptr(IdType* qo_indptr_h, IdType* kv_indptr_h, // the CUDA graph is created fixes the maximum number of tokens. const uint64_t max_seq_len = total_num_rows - batch_size + 1; uint64_t max_qo_len = uint64_t(max_seq_len) * gqa_group_size; - cta_tile_q = FA2DetermineCtaTileQ(max_qo_len, head_dim); + cta_tile_q = FA2DetermineCtaTileQ(max_qo_len, head_dim, head_dim_qk); // Find an upper bound for the number of tiles, derived from the total // number of rows and the batch size. The sum of qo lengths rounded @@ -615,7 +615,7 @@ inline auto PrefillSplitQOKVIndptr(IdType* qo_indptr_h, IdType* kv_indptr_h, sum_packed_qo_len += packed_qo_len_arr[i]; } const int64_t avg_packed_qo_len = sum_packed_qo_len / batch_size; - cta_tile_q = FA2DetermineCtaTileQ(avg_packed_qo_len, head_dim); + cta_tile_q = FA2DetermineCtaTileQ(avg_packed_qo_len, head_dim, head_dim_qk); total_num_tiles_q = 0; for (uint32_t i = 0; i < batch_size; ++i) { @@ -771,7 +771,6 @@ inline cudaError_t PrefillPlanImpl( int64_t num_colocated_ctas, // for POD attention, limit prefill // splits by #colocated decode CTAs int64_t uniform_q_len, cudaStream_t stream) { - (void)head_dim_qk; (void)sizeof_dtype_o; if (num_qo_heads % num_kv_heads != 0) { std::ostringstream err_msg; @@ -796,7 +795,7 @@ inline cudaError_t PrefillPlanImpl( PrefillSplitQOKVIndptr(qo_indptr_h, kv_indptr_h, total_num_rows, batch_size, num_qo_heads, num_kv_heads, head_dim_vo, page_size, max_batch_size_if_split, enable_cuda_graph, window_left, fixed_split_size, disable_split_kv, - uniform_q_len); + uniform_q_len, head_dim_qk); plan_info.cta_tile_q = cta_tile_q; plan_info.total_num_rows = total_num_rows; diff --git a/include/flashinfer/utils.cuh b/include/flashinfer/utils.cuh index fbd1c3e6f11..d3484111291 100644 --- a/include/flashinfer/utils.cuh +++ b/include/flashinfer/utils.cuh @@ -389,8 +389,13 @@ inline void DebugPrintCUDAArray(T* device_ptr, size_t size, std::string prefix = std::cout << std::endl; } -inline uint32_t FA2DetermineCtaTileQ(int64_t avg_packed_qo_len, uint32_t head_dim) { - if (head_dim >= 512) { +inline uint32_t FA2DetermineCtaTileQ(int64_t avg_packed_qo_len, uint32_t head_dim, + uint32_t head_dim_qk = 0) { + // head_dim is the VO dim at the batch-prefill call sites; head_dim_qk (when + // nonzero) lets asymmetric (QK != VO) configurations report the dim that + // actually drives shared-memory cost. + const uint32_t qk = head_dim_qk ? head_dim_qk : head_dim; + if (qk >= 512) { if (avg_packed_qo_len <= 32) { return 16; // decode / short-q (incl. speculative decode): lean CTA16 } @@ -406,7 +411,24 @@ inline uint32_t FA2DetermineCtaTileQ(int64_t avg_packed_qo_len, uint32_t head_di // avg_packed_qo_len <= 64 return 64; } else { - // avg_packed_qo_len <= 16 + // avg_packed_qo_len <= 16: prefer the 1x4 warp layout (cta_tile_q=16), + // but ONLY if one NUM_MMA_KV step fits shared memory. With 4 kv-warps a + // step costs (qk+vo)*16*4*sizeof(dtype); at (512,256) bf16 that is 96KB + // + a 16KB Q tile > the ~100KB/SM budget of CC 12.x consumer Blackwell + // (and would yield the unrecoverable "Unsupported max_mma_kv: 0" + // dispatch failure). Fall back to cta_tile_q=64 (4x1 layout, 4x + // cheaper KV step) exactly like the Turing branch below. sizeof + // assumed 2 (worst case): only configurations that could not run at + // all are moved. + int dev_id = 0, max_smem_per_sm = 0; + cudaGetDevice(&dev_id); + cudaDeviceGetAttribute(&max_smem_per_sm, cudaDevAttrMaxSharedMemoryPerMultiprocessor, + dev_id); + const uint32_t q_tile_smem = 16 * qk * 2; + const uint32_t kv_step_smem_1x4 = (qk + head_dim) * 16 * 4 * 2; + if (q_tile_smem + kv_step_smem_1x4 > (uint32_t)max_smem_per_sm) { + return 64; + } return 16; } } else { From 9ec7b934bd62fb89985cc253e44762f980fdcf62 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Wed, 17 Jun 2026 09:04:25 +0900 Subject: [PATCH 03/23] fix(prefill): move mask_indptr to custom mask device before segment_packbits Fixes a device mismatch when building packed custom masks for bidirectional (multimodal-prefix) attention. Signed-off-by: Jetha Chan (cherry picked from commit ee9301f500cf39bb982da9cdd78bcb1e6666ad57) --- flashinfer/prefill.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/flashinfer/prefill.py b/flashinfer/prefill.py index feed056a9ed..75cc118f08f 100755 --- a/flashinfer/prefill.py +++ b/flashinfer/prefill.py @@ -2261,9 +2261,13 @@ def plan( ) if packed_custom_mask is None and custom_mask is not None: # create packed custom mask from custom mask + # NOTE(spark-hijinks): the segment_packbits kernel requires the + # indptr on the mask's device (CHECK_DEVICE in quantization.cu), + # but mask_indptr inherits qo_indptr's device, which callers + # (e.g. vLLM) routinely keep on CPU while the mask is on GPU. packed_custom_mask, mask_indptr = segment_packbits( custom_mask.contiguous().view(-1), - mask_indptr, + mask_indptr.to(custom_mask.device), bitorder="little", ) From 3bffad45042ace3f395eb77ce1cb3fc46091306b Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Wed, 17 Jun 2026 09:04:25 +0900 Subject: [PATCH 04/23] feat(sm121): enable DGX Spark (GB10) FP4 dispatch + heuristic test Signed-off-by: Jetha Chan (cherry picked from commit 131bd2a6f8e7922b7466bdbc226da9c4f40c1cc9) --- .github/workflows/nightly-release.yml | 2 +- .github/workflows/release.yml | 2 +- docs/installation.rst | 7 +++++++ flashinfer/gemm/gemm_base.py | 20 ++++++++++++-------- flashinfer/mla/_core.py | 4 ++-- flashinfer/xqa.py | 2 +- tests/gemm/test_mm_fp4.py | 24 ++++++++++++++++++++++++ 7 files changed, 48 insertions(+), 13 deletions(-) diff --git a/.github/workflows/nightly-release.yml b/.github/workflows/nightly-release.yml index 5bd875c5899..091c2fbada0 100644 --- a/.github/workflows/nightly-release.yml +++ b/.github/workflows/nightly-release.yml @@ -153,7 +153,7 @@ jobs: - name: Build wheel in container env: DOCKER_IMAGE: ${{ matrix.arch == 'aarch64' && format('pytorch/manylinuxaarch64-builder:cuda{0}', matrix.cuda) || format('pytorch/manylinux2_28-builder:cuda{0}', matrix.cuda) }} - FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda < '12.9' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || (matrix.cuda < '13.0' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f' || (matrix.arch == 'aarch64' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f')) }} + FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda < '12.9' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || (matrix.cuda < '13.0' && (matrix.arch == 'aarch64' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f 12.1a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f') || (matrix.arch == 'aarch64' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f 12.1a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f')) }} FLASHINFER_DEV_RELEASE_SUFFIX: ${{ needs.setup.outputs.dev_suffix }} run: | # Extract CUDA major and minor versions diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 38e6e0705ff..0f5586a515c 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -182,7 +182,7 @@ jobs: - name: Build wheel in container env: DOCKER_IMAGE: ${{ matrix.arch == 'aarch64' && format('pytorch/manylinuxaarch64-builder:cuda{0}', matrix.cuda) || format('pytorch/manylinux2_28-builder:cuda{0}', matrix.cuda) }} - FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda < '12.9' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || (matrix.cuda < '13.0' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f' || (matrix.arch == 'aarch64' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f')) }} + FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda < '12.9' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || (matrix.cuda < '13.0' && (matrix.arch == 'aarch64' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f 12.1a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f') || (matrix.arch == 'aarch64' && '7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f 12.1a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f')) }} run: | # Extract CUDA major and minor versions CUDA_MAJOR=$(echo "${{ matrix.cuda }}" | cut -d'.' -f1) diff --git a/docs/installation.rst b/docs/installation.rst index d9a2098d705..3ea58dbd5a1 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -110,6 +110,13 @@ You can follow the steps below to install FlashInfer from source code: .. code-block:: bash export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f" + + For DGX Spark / GB10, add the arch-specific SM121 target so JIT-cache + wheels include native ``121a`` artifacts for NVFP4/MXFP4 MMA: + + .. code-block:: bash + + export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f 12.1a" cd flashinfer-jit-cache python -m build --no-isolation --wheel python -m pip install dist/*.whl diff --git a/flashinfer/gemm/gemm_base.py b/flashinfer/gemm/gemm_base.py index 3d9abd7b791..8525395d34a 100755 --- a/flashinfer/gemm/gemm_base.py +++ b/flashinfer/gemm/gemm_base.py @@ -6314,12 +6314,14 @@ def _heuristic_func_mm_fp4( enable_pdl: bool = True, # unused ): r""" - Heuristic function for mm_fp4 backend selection. Routes to either cudnn or cutlass. + Heuristic function for mm_fp4 backend selection. Note: trtllm is not considered in the backend selection because it requires a specific input quantization (swizzling/shuffling) that differs from the preparation used for cudnn and cutlass backends. Logic for which comes first: + - If CUDA version is 13+ and device is SM12x with NVFP4 - use b12x, + then cutlass, then cudnn. - If cuda version is 12 - use cutlass. - If cuda version is 13 and cudnn version is less than 9.15 - use cutlass. - If cuda version is 13 and cudnn version is 9.15 or greater: @@ -6331,12 +6333,14 @@ def _heuristic_func_mm_fp4( # Get compute capability to distinguish between SM100 (10.0) and SM103 (10.3) major, minor = get_compute_capability(a.device) is_sm103 = major == 10 and minor == 3 - is_sm120 = major == 12 and minor == 0 - - # SM120 + CUDA 13: prefer b12x. SM121 (GB10) is intentionally excluded -- b12x - # is supported there as an explicit backend, but cutlass/cudnn are faster in - # most cases, so `auto` keeps using them. - if is_sm120 and use_nvfp4 and cuda_major >= 13: + is_sm12x = major == 12 + + # SM12x + CUDA 13: prefer b12x (warp-level MMA, underfill tile selection). + # We enable BOTH SM120 and SM121 (GB10): upstream defaults SM121 to + # cutlass/cudnn for perf, but the campaign's VO-split / NVFP4 KV path needs + # b12x dispatch on GB10 (sm_121). SM120 and SM121 share the 12.x execution + # constraints; JIT/AOT arch selection is handled by the compilation context. + if is_sm12x and use_nvfp4 and cuda_major >= 13: return [c for c in ("b12x", "cutlass", "cudnn") if c in suitable_backends] # If cuda version is 13 or greater and cudnn version is 9.15 or greater: @@ -6531,7 +6535,7 @@ def mm_fp4( Whether to use 8x4 scale factor layout or 128x4 scale factor layout, defaults to False. backend: Literal["cudnn", "trtllm", "cutlass", "cute-dsl", "b12x", "auto"] - Backend to use, defaults to ``"auto"``. On SM120, ``"auto"`` prefers + Backend to use, defaults to ``"auto"``. On SM12x, ``"auto"`` prefers ``"b12x"`` (NVFP4 only), then ``"cutlass"``, then ``"cudnn"``. On other architectures, ``"auto"`` selects between ``"cudnn"`` and ``"cutlass"`` based on the current CUDA and cuDNN versions. The ``"trtllm"`` and diff --git a/flashinfer/mla/_core.py b/flashinfer/mla/_core.py index bf422ec6e01..4a9fcf03820 100644 --- a/flashinfer/mla/_core.py +++ b/flashinfer/mla/_core.py @@ -2825,7 +2825,7 @@ def trtllm_batch_decode_with_kv_cache_mla( raise ValueError("XQA MLA does not support cum_seq_lens_q / max_q_len") if not is_sm12x_supported(query.device): raise ValueError( - "XQA MLA requires SM120a (CUDA >= 12.8) or SM121a (CUDA >= 13.0)" + "XQA MLA requires SM120a (CUDA >= 12.8) or SM121a (CUDA >= 12.9)" ) fp8_ok = ( query.dtype == torch.float8_e4m3fn and kv_cache.dtype == torch.float8_e4m3fn @@ -3339,7 +3339,7 @@ def xqa_batch_decode_with_kv_cache_mla( ) if not is_sm12x_supported(query.device): raise ValueError( - "XQA MLA requires SM120a (CUDA >= 12.8) or SM121a (CUDA >= 13.0)" + "XQA MLA requires SM120a (CUDA >= 12.8) or SM121a (CUDA >= 12.9)" ) fp8_ok = ( query.dtype == torch.float8_e4m3fn and kv_cache.dtype == torch.float8_e4m3fn diff --git a/flashinfer/xqa.py b/flashinfer/xqa.py index 0fe67cbd351..c8870f70562 100755 --- a/flashinfer/xqa.py +++ b/flashinfer/xqa.py @@ -309,7 +309,7 @@ def xqa( if k_cache.dtype == torch.uint8: assert get_compute_capability(torch.device(device="cuda"))[0] in [12], ( - "XQA NVFP4 KV is only supported on SM120 GPUs" + "XQA NVFP4 KV is only supported on SM12x GPUs" ) assert k_sf_cache is not None, "K SF cache is required when NVFP4 KV is used" assert v_sf_cache is not None, "V SF cache is required when NVFP4 KV is used" diff --git a/tests/gemm/test_mm_fp4.py b/tests/gemm/test_mm_fp4.py index 24e533bd2db..a56ee8edc0c 100644 --- a/tests/gemm/test_mm_fp4.py +++ b/tests/gemm/test_mm_fp4.py @@ -1,3 +1,5 @@ +from types import SimpleNamespace + import pytest import torch import torch.nn.functional as F @@ -203,6 +205,28 @@ def test_mm_fp4_b12x_misaligned_k_raises(): skip_check=False, ) +def test_mm_fp4_auto_prefers_b12x_for_sm121_nvfp4(monkeypatch): + from flashinfer.gemm import gemm_base + + class MockTensor: + device = "cuda" + + monkeypatch.setattr( + gemm_base, "get_cuda_version", lambda: SimpleNamespace(major=13) + ) + monkeypatch.setattr(gemm_base, "get_compute_capability", lambda device: (12, 1)) + + backends = gemm_base._heuristic_func_mm_fp4( + ["cudnn", "cutlass", "b12x"], + MockTensor(), + MockTensor(), + MockTensor(), + MockTensor(), + use_nvfp4=True, + ) + + assert backends == ["b12x", "cutlass", "cudnn"] + def test_mm_fp4_cute_dsl_misaligned_n_raises(): device = torch.device("cuda") From e496d242fa73ce45a51c0ff9c19392a2a78635d3 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Sun, 21 Jun 2026 07:36:15 +0000 Subject: [PATCH 05/23] fix(prefill): key CTA_TILE_Q dispatch on head_dim_qk for asymmetric heads The VO-split rederive left three dispatch gates keyed on head_dim_vo while FA2DetermineCtaTileQ selects CTA_TILE_Q from the QK dim. For asymmetric Gemma-4 heads (qk=512, vo=256) VO-split does not engage (NUM_MMA_D_VO==16), so the full-D path applies and only CTA_TILE_Q=16 fits the 256-register o_frag wall. Align all four sites: - utils.cuh FA2DetermineCtaTileQ: VO>=512 -> {16,32}; QK>=512 & VO<512 -> 16 - prefill.cuh KernelTraits::IsInvalid first clause: key on HEAD_DIM_QK - batch_prefill_{paged,ragged}_kernel_inst.jinja: instantiate {16,32} when head_dim_qk>=512 Validated NVFP4 paged prefill (sm120, RTX): max_abs_err qk256/vo256=0.0048, qk128=0.0046, qk512/vo256=0.0047 (seq128) / 0.0044 (seq512) / 0.0042 (seq1024) / 0.0166 causal. No NaN; output sized by VO. Signed-off-by: Jetha Chan --- csrc/batch_prefill_paged_kernel_inst.jinja | 4 ++-- csrc/batch_prefill_ragged_kernel_inst.jinja | 4 ++-- include/flashinfer/attention/prefill.cuh | 4 ++-- include/flashinfer/utils.cuh | 10 +++++++++- 4 files changed, 15 insertions(+), 7 deletions(-) diff --git a/csrc/batch_prefill_paged_kernel_inst.jinja b/csrc/batch_prefill_paged_kernel_inst.jinja index 2e29edc0ff8..c5d4723a74a 100644 --- a/csrc/batch_prefill_paged_kernel_inst.jinja +++ b/csrc/batch_prefill_paged_kernel_inst.jinja @@ -5,9 +5,9 @@ namespace flashinfer { constexpr auto use_custom_mask = {{ mask_mode }} == MaskMode::kCustom; -{# FA2DetermineCtaTileQ only selects {16, 32} for head_dim_vo >= 512 and +{# FA2DetermineCtaTileQ only selects {16, 32} for head_dim_qk >= 512 and {16, 64, 128} otherwise; don't instantiate unreachable variants. #} -{% for cta_tile_q in ([16, 32] if head_dim_vo | int >= 512 else [16, 64, 128]) %} +{% for cta_tile_q in ([16, 32] if head_dim_qk | int >= 512 else [16, 64, 128]) %} template cudaError_t BatchPrefillWithPagedKVCacheDispatched< /*CTA_TILE_Q=*/{{cta_tile_q}}, {{head_dim_qk}}, {{head_dim_vo}}, {{pos_encoding_mode}}, {{use_fp16_qk_reduction}}, {{mask_mode}}, {{ variant_name }}, PagedParams>(PagedParams params, {{ dtype_o }}* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); diff --git a/csrc/batch_prefill_ragged_kernel_inst.jinja b/csrc/batch_prefill_ragged_kernel_inst.jinja index 22fb4c995e2..6e5eec59e8b 100644 --- a/csrc/batch_prefill_ragged_kernel_inst.jinja +++ b/csrc/batch_prefill_ragged_kernel_inst.jinja @@ -5,9 +5,9 @@ namespace flashinfer { constexpr auto use_custom_mask = {{ mask_mode }} == MaskMode::kCustom; -{# FA2DetermineCtaTileQ only selects {16, 32} for head_dim_vo >= 512 and +{# FA2DetermineCtaTileQ only selects {16, 32} for head_dim_qk >= 512 and {16, 64, 128} otherwise; don't instantiate unreachable variants. #} -{% for cta_tile_q in ([16, 32] if head_dim_vo | int >= 512 else [16, 64, 128]) %} +{% for cta_tile_q in ([16, 32] if head_dim_qk | int >= 512 else [16, 64, 128]) %} template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< /*CTA_TILE_Q=*/{{cta_tile_q}}, {{head_dim_qk}}, {{head_dim_vo}}, {{pos_encoding_mode}}, {{use_fp16_qk_reduction}}, {{mask_mode}}, {{ variant_name }}, RaggedParams>(RaggedParams params, {{ dtype_o }}* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); diff --git a/include/flashinfer/attention/prefill.cuh b/include/flashinfer/attention/prefill.cuh index 23f1b6bb3ee..fdb17d65bc8 100644 --- a/include/flashinfer/attention/prefill.cuh +++ b/include/flashinfer/attention/prefill.cuh @@ -300,8 +300,8 @@ struct KernelTraits { static constexpr bool IsInvalid() { // The first clause prunes (CTA_TILE_Q, head_dim) pairs FA2DetermineCtaTileQ - // never selects: {16, 32} for head_dim_vo >= 512, {16, 64, 128} otherwise. - return ((HEAD_DIM_VO >= 512 ? (CTA_TILE_Q > 32) : (CTA_TILE_Q == 32)) || (NUM_MMA_D_VO < 4) || + // never selects: {16, 32} for head_dim_qk >= 512, {16, 64, 128} otherwise. + return ((HEAD_DIM_QK >= 512 ? (CTA_TILE_Q > 32) : (CTA_TILE_Q == 32)) || (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) || diff --git a/include/flashinfer/utils.cuh b/include/flashinfer/utils.cuh index d3484111291..995cd3c2353 100644 --- a/include/flashinfer/utils.cuh +++ b/include/flashinfer/utils.cuh @@ -395,12 +395,20 @@ inline uint32_t FA2DetermineCtaTileQ(int64_t avg_packed_qo_len, uint32_t head_di // nonzero) lets asymmetric (QK != VO) configurations report the dim that // actually drives shared-memory cost. const uint32_t qk = head_dim_qk ? head_dim_qk : head_dim; - if (qk >= 512) { + if (head_dim >= 512) { + // True VO-split (VO >= 512): the split halves o_frag register pressure, so + // CTA_TILE_Q=32 is feasible for long-q; CTA16 for decode / short-q. if (avg_packed_qo_len <= 32) { return 16; // decode / short-q (incl. speculative decode): lean CTA16 } return 32; // Long-q prefill use CTA_TILE_Q=32 } + if (qk >= 512) { + // Asymmetric large-QK but VO <= 256: VO-split does NOT engage (NUM_MMA_D_VO + // == 16), so o_frag is sized by the full NUM_MMA_D_VO_TILE=16. CTA_TILE_Q>16 + // (NUM_MMA_Q>=2) overflows the 256-register o_frag wall -> only CTA16 is valid. + return 16; + } if (avg_packed_qo_len > 64 && head_dim < 256) { return 128; } else { From 910d409b99c1c22ecd9fce19bf5348a5395274b4 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Tue, 23 Jun 2026 12:00:38 +0900 Subject: [PATCH 06/23] fix(decode): derive NVFP4 output width from KV dtype, not sf presence Address review on #3684: the paged-decode output allocation doubled out_head_dim whenever kv_cache_sf was non-None. Gate it on the KV cache actually being uint8-packed (the NVFP4 layout that stores VO at half width), so a stray scale-factor tensor on a non-uint8 cache fails the shape check instead of silently allocating a mis-sized output. Signed-off-by: Jetha Chan --- flashinfer/decode.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/flashinfer/decode.py b/flashinfer/decode.py index 427c50f743e..9d621eca191 100644 --- a/flashinfer/decode.py +++ b/flashinfer/decode.py @@ -2031,9 +2031,13 @@ def run( # use q's head_dim for output instead # NVFP4 packed: unpacked VO width is packed bytes * 2 (supports # asymmetric QK/VO plans; q.shape[-1] assumed QK == VO). + # Only the NVFP4 packed path (uint8 KV) stores VO at half width; + # derive the doubled output width from the KV dtype, not merely from + # kv_cache_sf being present, so a stray scale-factor tensor on a + # non-uint8 cache can't silently miscompute the output shape. out_head_dim = ( v_cache.shape[-1] * 2 - if kv_cache_sf is not None + if kv_cache_sf is not None and v_cache.dtype == torch.uint8 else v_cache.shape[-1] ) out = torch.empty( From 308a4d8930d8546cefba47e29515117df434091c Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Thu, 25 Jun 2026 10:44:49 +0900 Subject: [PATCH 07/23] fix(attention): disable split-KV for NVFP4 paged prefill/decode NVFP4 paged KV is packed uint8 with a per-16-element FP8 block scale. Split-KV (flash-decoding) chunks the KV range by kv_chunk_size, which is not aligned to the 16-element scale blocks; a chunk boundary landing mid-block -- plus the small per-split chunk tripping the 1-byte-KV NUM_MMA_KV tile floor -- corrupts the dequantized reads. It only surfaces when a short query attends a long KV (qo_len << kv_len), i.e. decode and prefix-cache extend, so dense full-prefill tests miss it while prefix caching / long-context decode break. Gate split-KV off when kv_data_type is NVFP4 (uint8 / float4_e2m1fn_x2) in the paged and ragged prefill plan(). FP8 / 16-bit KV are unaffected. Verified on Gemma-4 E2B/E4B: radix-on retrieval cliffs at ~600 tokens of reused-prefix context without the gate, holds to 1448 (== bf16) with it -- on both sm120 (RTX PRO 6000) and sm121 (GB10 / DGX Spark). Follow-up: make the FP4 split path scale-block-aware (16-token-aligned chunk boundaries + NUM_MMA_KV floor) to restore flash-decoding parallelism for long-context NVFP4 decode. Signed-off-by: Jetha Chan --- flashinfer/prefill.py | 32 +++++++++++++++++++ tests/attention/test_nvfp4_attention_sm120.py | 27 ++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/flashinfer/prefill.py b/flashinfer/prefill.py index 75cc118f08f..0745c76082e 100755 --- a/flashinfer/prefill.py +++ b/flashinfer/prefill.py @@ -1491,6 +1491,26 @@ def _compute_page_mask_indptr( return mask_indptr +def _nvfp4_kv_requires_disabled_split_kv(kv_data_type: torch.dtype) -> bool: + """Whether split-KV must be disabled because the KV cache is NVFP4. + + NVFP4 paged KV is stored as packed ``uint8`` (two E2M1 values per byte) with a + per-16-element FP8 block scale. Split-KV (flash-decoding) partitions the KV + range into chunks sized by ``kv_chunk_size``, which is *not* aligned to the + 16-element scale blocks: a chunk boundary that lands mid-block — together with + the small per-split chunk tripping the 1-byte-KV ``NUM_MMA_KV`` tile floor — + corrupts the dequantized reads. The corruption only surfaces when a short query + attends a long KV (``qo_len << kv_len``), i.e. decode and prefix-cache *extend*, + so dense full-prefill tests miss it while prefix caching breaks. Until the FP4 + split path is made scale-block-aware, disable split-KV for NVFP4 KV. FP8 KV has + no block scales and is unaffected. + """ + if kv_data_type == torch.uint8: # packed NVFP4 (the run path's convention) + return True + native_fp4 = getattr(torch, "float4_e2m1fn_x2", None) + return native_fp4 is not None and kv_data_type == native_fp4 + + class BatchPrefillWithPagedKVCacheWrapper: r"""Wrapper class for prefill/append attention with paged kv-cache for batch of requests. @@ -2478,6 +2498,12 @@ def plan( ] if self._backend == "fa2": args.append(fixed_split_size or -1) # fixed_split_size + if not disable_split_kv and _nvfp4_kv_requires_disabled_split_kv( + kv_data_type + ): + # NVFP4 split-KV corrupts prefix-cached / decode reads; force + # it off. See _nvfp4_kv_requires_disabled_split_kv for why. + disable_split_kv = True args.append(disable_split_kv) # disable_split_kv args.append(0) # num_colocated_ctas args.append(0) # uniform_q_len @@ -3701,6 +3727,12 @@ def plan( ] if self._backend == "fa2": args.append(fixed_split_size or -1) # fixed_split_size + if not disable_split_kv and _nvfp4_kv_requires_disabled_split_kv( + kv_data_type + ): + # NVFP4 split-KV corrupts prefix-cached / decode reads; force + # it off. See _nvfp4_kv_requires_disabled_split_kv for why. + disable_split_kv = True args.append(disable_split_kv) # disable_split_kv args.append(0) # num_colocated_ctas args.append(0) # uniform_q_len diff --git a/tests/attention/test_nvfp4_attention_sm120.py b/tests/attention/test_nvfp4_attention_sm120.py index a0d801ef9ef..c0afa87b753 100644 --- a/tests/attention/test_nvfp4_attention_sm120.py +++ b/tests/attention/test_nvfp4_attention_sm120.py @@ -371,3 +371,30 @@ def test_nvfp4_attention_sm120_causal_mask_column_order(): assert suffix_max <= 1e-5 assert cos_sim >= 0.98 + + +def test_nvfp4_split_kv_gate_dtype_logic(): + """The split-KV gate must fire for NVFP4 KV and only for NVFP4 KV. + + NVFP4 split-KV corrupts prefix-cached / decode reads (a split boundary can land + mid 16-element scale block; the small per-split chunk also trips the 1-byte-KV + NUM_MMA_KV tile floor), so plan() force-disables split-KV when the KV cache is + NVFP4. FP8 / 16-bit KV have no block scales and must keep split-KV. This is a + pure dtype-classification check (no GPU required) guarding that contract. + """ + from flashinfer.prefill import _nvfp4_kv_requires_disabled_split_kv + + # packed NVFP4 (uint8 is the run-path convention) and native fp4 -> gated + assert _nvfp4_kv_requires_disabled_split_kv(torch.uint8) + native_fp4 = getattr(torch, "float4_e2m1fn_x2", None) + if native_fp4 is not None: + assert _nvfp4_kv_requires_disabled_split_kv(native_fp4) + + # 16-bit and FP8 KV split-KV is fine -> not gated + for dtype in ( + torch.float16, + torch.bfloat16, + torch.float8_e4m3fn, + torch.float8_e5m2, + ): + assert not _nvfp4_kv_requires_disabled_split_kv(dtype) From a69f6f936a64f13246ae586c0338ae793f5adae9 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Fri, 26 Jun 2026 00:24:56 +0900 Subject: [PATCH 08/23] fix(attention): harden NVFP4 gating, rope_freq bound, smem SF budget Address CodeRabbit review on #3684 (5 findings): prefill.py: - Gate NVFP4 out_head_dim doubling on `v.dtype == torch.uint8` (the packed storage), not just `kv_cache_sf is not None`, at all three prefill sites. A stray scale-factor tensor on a non-uint8 cache no longer silently doubles the output width. Mirrors the decode-side guard. - Ragged custom-mask plan(): move mask_indptr to custom_mask.device before segment_packbits, mirroring the paged-path fix. mask_indptr inherits qo_indptr's device (often CPU) while custom_mask is on GPU. prefill.cuh: - init_rope_freq: bound the fill loop on NUM_MMA_D_QK/2, not NUM_MMA_D_VO/2. rope_freq is sized [NUM_MMA_D_QK/2][4] and the rotary appliers index it up to NUM_MMA_D_QK/2; the asymmetric VO-split dispatch (qk=512, vo=256) made NUM_MMA_D_QK != NUM_MMA_D_VO reachable, leaving the upper half uninitialized for in-kernel-RoPE callers. - page_produce_kv_sf deswizzle branch: guard the unconditional smem store on flat_byte < SF_TOTAL_BYTES (the rounded-up NUM_SF_ITERS leaves over-range lanes). Live cp_async path was already predicated. - NUM_MMA_KV occupancy budget (all three dispatchers): add the FP4 K/V scale-factor smem term, which scales with CTA_TILE_KV. Prevents the budget from over-selecting a tile whose SharedStorage then exceeds the smem limit. Signed-off-by: Jetha Chan --- flashinfer/prefill.py | 30 +++++++++++++++++++----- include/flashinfer/attention/prefill.cuh | 13 ++++++++-- 2 files changed, 35 insertions(+), 8 deletions(-) diff --git a/flashinfer/prefill.py b/flashinfer/prefill.py index 0745c76082e..4f47b8d2aec 100755 --- a/flashinfer/prefill.py +++ b/flashinfer/prefill.py @@ -1372,8 +1372,14 @@ def single_prefill_with_kv_cache( # For NVFP4 KV (uint8 packed), last dim is head_dim//2 packed bytes: the # unpacked VO width is v.shape[-1] * 2, which equals head_dim_vo even for - # asymmetric (QK, VO) plans; q.shape[-1] assumed QK == VO. - out_head_dim = v.shape[-1] * 2 if kv_cache_sf is not None else v.shape[-1] + # asymmetric (QK, VO) plans; q.shape[-1] assumed QK == VO. Gate on the packed + # (uint8) storage, not just kv_cache_sf, so a stray scale-factor tensor on a + # non-uint8 cache cannot silently double the output width. + out_head_dim = ( + v.shape[-1] * 2 + if kv_cache_sf is not None and v.dtype == torch.uint8 + else v.shape[-1] + ) if backend == "auto": backend = determine_attention_backend( @@ -2777,7 +2783,9 @@ def run( # the kernel (which writes head_dim_vo-wide rows) garble a too-wide # output buffer whenever VO < QK. out_head_dim = ( - v_cache.shape[-1] * 2 if kv_cache_sf is not None else v_cache.shape[-1] + v_cache.shape[-1] * 2 + if kv_cache_sf is not None and v_cache.dtype == torch.uint8 + else v_cache.shape[-1] ) if out is None: # Use cached output data type if available (for FP8 attention with FP16 output) @@ -3439,9 +3447,13 @@ def plan( mask_indptr = _compute_mask_indptr(qo_indptr, kv_indptr) if packed_custom_mask is None and custom_mask is not None: # create packed custom mask from custom mask + # NOTE(spark-hijinks): segment_packbits requires mask_indptr on the + # same device as custom_mask, but mask_indptr inherits qo_indptr's + # device (often CPU) while custom_mask is on GPU. Mirror the paged + # path's .to(device) so the ragged custom-mask flow doesn't crash. packed_custom_mask, mask_indptr = segment_packbits( custom_mask.contiguous().view(-1), - mask_indptr, + mask_indptr.to(custom_mask.device), bitorder="little", ) @@ -3930,8 +3942,14 @@ def run( k_sf, v_sf = kv_cache_sf.unbind(dim=1) # NVFP4 packed: unpacked VO width is packed bytes * 2 (supports - # asymmetric QK/VO; q.shape[-1] assumed QK == VO). - out_head_dim = v.shape[-1] * 2 if kv_cache_sf is not None else v.shape[-1] + # asymmetric QK/VO; q.shape[-1] assumed QK == VO). Gate on the packed + # (uint8) storage so a stray scale-factor tensor on a non-uint8 cache + # can't silently double the output width (matches the decode-side guard). + out_head_dim = ( + v.shape[-1] * 2 + if kv_cache_sf is not None and v.dtype == torch.uint8 + else v.shape[-1] + ) if out is None: # when input dtype is fp8, we need to use bf16 output out_dtype = torch.bfloat16 if q.dtype.itemsize == 1 else q.dtype diff --git a/include/flashinfer/attention/prefill.cuh b/include/flashinfer/attention/prefill.cuh index fdb17d65bc8..faabac2c885 100644 --- a/include/flashinfer/attention/prefill.cuh +++ b/include/flashinfer/attention/prefill.cuh @@ -770,7 +770,13 @@ __device__ __forceinline__ void page_produce_kv_sf( sf_ptr[page_head_base + static_cast(swz_entry) * sf_stride_n + swz_sd]; } } - *reinterpret_cast(sf_smem + flat_byte) = packed; + // NUM_SF_ITERS is rounded up, so the last iter has lanes with + // flat_byte >= SF_TOTAL_BYTES. Unlike the predicated cp_async store below, + // this write is unconditional, so guard it on the buffer bound (not in_bounds, + // which would also skip the intended zero-fill of in-buffer padding rows). + if (flat_byte < SF_TOTAL_BYTES) { + *reinterpret_cast(sf_smem + flat_byte) = packed; + } } else { const size_t sf_gmem_offset = page_head_base + entry_idx * sf_stride_n + sf_smem_col; // V SF must zero-fill out-of-bounds entries: compute_sfm_v reads SF for all CTA_TILE_KV @@ -876,7 +882,10 @@ __device__ __forceinline__ void init_rope_freq(float (*rope_freq)[4], const floa const uint32_t tid_x = threadIdx.x) { const uint32_t lane_idx = tid_x; #pragma unroll - for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO / 2; ++mma_d) { + // rope_freq is sized [NUM_MMA_D_QK/2][4] and the rotary appliers index it up to + // NUM_MMA_D_QK/2; for asymmetric QK/VO (e.g. qk=512, vo=256) NUM_MMA_D_VO/2 would + // leave the upper half of the table uninitialized, so bound on QK. + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_QK / 2; ++mma_d) { #pragma unroll for (uint32_t j = 0; j < 4; ++j) { rope_freq[mma_d][j] = From 6cd9acb07dacd1a98e79d1b583e9dd0557209193 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Fri, 26 Jun 2026 01:45:26 +0900 Subject: [PATCH 09/23] fix(nvfp4-sm120): zero lse before fwd kernel to avoid uninitialized read MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fwd entry point passes the caller's uninitialized lse (torch.empty) straight to run_mha_fwd. The kernel writes `out` for every query row but does not guarantee writing every (batch, head, seq) entry of lse, so lse can read back whatever the allocator handed out — observed as flaky NaNs that pass on a clean allocation and fail when a prior test dirtied the pages. Zero lse before launch, mirroring the existing seq_len==0 branch. Found smoke-testing #3684 on sm120 (RTX PRO 6000): the nvfp4 sm120 accuracy test's lse NaN-check flaked only when run after the paged-prefill nvfp4 test. Signed-off-by: Jetha Chan --- .../nvfp4_attention_sm120_binding.cu | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu b/csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu index b135e7c44a6..28fbd83becc 100644 --- a/csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu +++ b/csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu @@ -302,6 +302,16 @@ void fwd(TensorView q_fp4, TensorView k_fp4, TensorView v_fp4_t, TensorView q_sc return; } + // lse is allocated uninitialized by the caller (torch.empty). The fwd kernel + // writes `out` for every query row but does not guarantee writing every + // (batch, head, seq) entry of lse, so reading it back can surface whatever + // garbage the allocator handed out (observed as flaky NaNs in lse under a + // dirty allocator, depending on test/run ordering). Zero it first, mirroring + // the seq_len==0 branch above, so unwritten entries are a defined 0. + status = cudaMemsetAsync(lse.data_ptr(), 0, numel(lse) * get_element_size(lse), stream); + TVM_FFI_ICHECK(status == cudaSuccess) + << "cudaMemsetAsync(lse) failed: " << cudaGetErrorString(status); + Flash_fwd_params params; set_params_fprop(params, q_fp4, k_fp4, v_fp4_t, q_scale, k_scale, v_scale_t, qk_correction, out, lse, static_cast(sm_scale), causal, per_block_mean); From 9909bb4c41cde66a8e2fb04cdcab21ea03c4c0ee Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Sat, 11 Jul 2026 15:01:00 +0900 Subject: [PATCH 10/23] test(jit): assert the emitted FP4 KV gate string The old assertion expected a literal "constexpr bool REQUIRE_FP4_KV_CACHE = true;" that no code path ever emitted; check the #error guard and static_assert the config template actually renders for FP4 KV. Signed-off-by: Jetha Chan --- tests/jit/test_attention_utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/jit/test_attention_utils.py b/tests/jit/test_attention_utils.py index 09f0935a46a..18034eb3790 100644 --- a/tests/jit/test_attention_utils.py +++ b/tests/jit/test_attention_utils.py @@ -33,7 +33,14 @@ def test_batch_prefill_nvfp4_swa_paged_params_declares_sf_strides( ) generated = (tmp_path / "generated" / uri / "batch_prefill_config.inc").read_text() - assert "constexpr bool REQUIRE_FP4_KV_CACHE = true;" in generated + # The FP4 KV gate is emitted as compile-time checks: an #error if the + # FP4 enable flag is missing plus a static_assert pinning DTypeKV to the + # packed FP4 container type. + assert ( + "#error \"NVFP4 KV paged prefill compiled without FLASHINFER_ENABLE_FP4_E2M1\"" + in generated + ) + assert "static_assert(std::is_same_v," in generated assert "constexpr auto USE_SLIDING_WINDOW = true;" in generated for field in ( "maybe_k_cache_sf", From 5db690043ddbdcdd24355904092b5f0185ccb543 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Wed, 15 Jul 2026 15:24:45 +0900 Subject: [PATCH 11/23] fix(decode): require equal K/V strides in paged decode run BatchDecodeWithPagedKVCacheRun accepts independent K/V stride arrays since the paged_kv_t independent-stride constructor was introduced, but the FA2 decode kernel (decode.cuh) still computes a single protective_get_kv_offset per row -- from the K strides -- and uses it to address both k_data and v_data, so asymmetric strides would silently read V at K offsets. Restore the host-side per-dimension stride-equality check that guarded this before, with an error message naming the limitation. Restoring the check (rather than making the decode kernels V-stride-aware) is the minimal honest fix: no supported path needs asymmetric decode today -- asymmetric NVFP4 (VO-split) decode rides the prefill wrapper, and symmetric caches have equal K/V strides by construction. If asymmetric decode is needed later, decode.cuh should grow separate K/V offsets the way the prefill path already has. Addresses review feedback on #3684 from @qsang-nv. Signed-off-by: Jetha Chan --- csrc/batch_decode.cu | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/csrc/batch_decode.cu b/csrc/batch_decode.cu index 3e4d3f6533f..1d247be4543 100644 --- a/csrc/batch_decode.cu +++ b/csrc/batch_decode.cu @@ -176,6 +176,13 @@ void BatchDecodeWithPagedKVCacheRun(TensorView float_workspace_buffer, auto k_strides = paged_k_cache.strides(); auto v_strides = paged_v_cache.strides(); TVM_FFI_ICHECK_EQ(k_strides.size(), v_strides.size()); + for (int i = 0; i < k_strides.size(); ++i) { + TVM_FFI_ICHECK_EQ(k_strides[i], v_strides[i]) + << "K/V strides differ at dim " << i + << ": the FA2 decode kernel addresses both K and V through a single set of " + "(K) strides, so paged_k_cache and paged_v_cache must have identical strides; " + "NVFP4/asymmetric decode with independent K/V strides is not yet supported."; + } ffi::CUDADeviceGuard device_guard(q.device().device_id); const cudaStream_t stream = get_stream(q.device()); From d0291501d6cde21ed09dfc42d665e120af9009ae Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Wed, 15 Jul 2026 15:24:56 +0900 Subject: [PATCH 12/23] fix(prefill): route V page offsets through V strides in on-the-fly producer page_produce_kv_on_the_fly computed both K and V global-memory offsets via get_paged_kv_offset_for_logical_row, which called protective_get_kv_offset -- the K strides -- so the shared-KV (on-the-fly) path addressed V rows with K offsets once K and V may carry independent strides. Thread produce_v through get_paged_kv_offset_for_logical_row and select protective_get_v_offset vs protective_get_k_offset with if constexpr, mirroring what the prefetched-offset path already does with its separate K/V offset arrays. Addresses review feedback on #3684 from @qsang-nv and @lesj0610. Signed-off-by: Jetha Chan --- include/flashinfer/attention/prefill.cuh | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/include/flashinfer/attention/prefill.cuh b/include/flashinfer/attention/prefill.cuh index faabac2c885..f7c0a3c5ee2 100644 --- a/include/flashinfer/attention/prefill.cuh +++ b/include/flashinfer/attention/prefill.cuh @@ -589,7 +589,7 @@ __device__ __forceinline__ void page_produce_kv(SmemStorage* smem_storage, uint3 } } -template +template __device__ __forceinline__ size_t get_paged_kv_offset_for_logical_row( const PagedKV& paged_kv, const uint32_t packed_page_iter_base, const typename KTraits::IdType last_indptr, const uint32_t kv_head_idx, @@ -599,9 +599,16 @@ __device__ __forceinline__ size_t get_paged_kv_offset_for_logical_row( constexpr uint32_t KV_THR_LAYOUT_COL = KTraits::KV_THR_LAYOUT_COL; uint32_t page_iter, entry_idx; paged_kv.page_size.divmod(packed_page_iter_base + logical_row, page_iter, entry_idx); - return paged_kv.protective_get_kv_offset( - page_iter, kv_head_idx, entry_idx, - (lane_idx % KV_THR_LAYOUT_COL) * upcast_size() / (IS_FP4 ? 2 : 1), last_indptr); + const uint32_t feat_idx = (lane_idx % KV_THR_LAYOUT_COL) * upcast_size() / (IS_FP4 ? 2 : 1); + // The K and V pools may carry different strides (e.g. NVFP4 VO-split caches + // where head_dim_qk != head_dim_vo), so route V rows through the V strides. + if constexpr (produce_v) { + return paged_kv.protective_get_v_offset(page_iter, kv_head_idx, entry_idx, feat_idx, + last_indptr); + } else { + return paged_kv.protective_get_k_offset(page_iter, kv_head_idx, entry_idx, feat_idx, + last_indptr); + } } template @@ -632,7 +639,7 @@ __device__ __forceinline__ void page_produce_kv_on_the_fly( for (uint32_t i = 0; i < NUM_MMA_KV * ROWS_PER_ITER / NUM_WARPS_Q; ++i) { const uint32_t logical_row = warp_idx * ROWS_PER_ITER + lane_idx / 8 + NUM_WARPS * ROWS_PER_ITER * i; - DType* gptr = kv_ptr + get_paged_kv_offset_for_logical_row( + DType* gptr = kv_ptr + get_paged_kv_offset_for_logical_row( paged_kv, packed_page_iter_base, last_indptr, kv_head_idx, logical_row, lane_idx); #pragma unroll @@ -659,7 +666,7 @@ __device__ __forceinline__ void page_produce_kv_on_the_fly( for (uint32_t i = 0; i < NUM_MMA_KV * (ROWS_PER_ITER / 4) / NUM_WARPS_Q; ++i) { const uint32_t logical_row = warp_idx * ROWS_PER_ITER + lane_idx / 4 + NUM_WARPS * ROWS_PER_ITER * i; - DType* gptr = kv_ptr + get_paged_kv_offset_for_logical_row( + DType* gptr = kv_ptr + get_paged_kv_offset_for_logical_row( paged_kv, packed_page_iter_base, last_indptr, kv_head_idx, logical_row, lane_idx); if constexpr (IS_FP4) { From 2e1612a4f14e5f8c179f281a0a1390cb32c03ba1 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Wed, 15 Jul 2026 15:25:07 +0900 Subject: [PATCH 13/23] revert(gemm): keep cutlass/cudnn as the mm_fp4 auto default on SM121 Revert the extension of the b12x auto-preference from SM120-only to all SM12x, restoring the upstream heuristic: SM121 (GB10) intentionally keeps cutlass/cudnn first because they are faster there in most cases, while b12x remains available on SM121 as an explicit backend. If flipping the SM121 default to b12x is warranted, it should come as a separately justified PR with benchmarks rather than ride along here. Also removes the heuristic test that asserted the flipped ordering. Addresses review feedback on #3684 from @qsang-nv. Signed-off-by: Jetha Chan --- flashinfer/gemm/gemm_base.py | 20 ++++++++------------ tests/gemm/test_mm_fp4.py | 24 ------------------------ 2 files changed, 8 insertions(+), 36 deletions(-) diff --git a/flashinfer/gemm/gemm_base.py b/flashinfer/gemm/gemm_base.py index 8525395d34a..3d9abd7b791 100755 --- a/flashinfer/gemm/gemm_base.py +++ b/flashinfer/gemm/gemm_base.py @@ -6314,14 +6314,12 @@ def _heuristic_func_mm_fp4( enable_pdl: bool = True, # unused ): r""" - Heuristic function for mm_fp4 backend selection. + Heuristic function for mm_fp4 backend selection. Routes to either cudnn or cutlass. Note: trtllm is not considered in the backend selection because it requires a specific input quantization (swizzling/shuffling) that differs from the preparation used for cudnn and cutlass backends. Logic for which comes first: - - If CUDA version is 13+ and device is SM12x with NVFP4 - use b12x, - then cutlass, then cudnn. - If cuda version is 12 - use cutlass. - If cuda version is 13 and cudnn version is less than 9.15 - use cutlass. - If cuda version is 13 and cudnn version is 9.15 or greater: @@ -6333,14 +6331,12 @@ def _heuristic_func_mm_fp4( # Get compute capability to distinguish between SM100 (10.0) and SM103 (10.3) major, minor = get_compute_capability(a.device) is_sm103 = major == 10 and minor == 3 - is_sm12x = major == 12 - - # SM12x + CUDA 13: prefer b12x (warp-level MMA, underfill tile selection). - # We enable BOTH SM120 and SM121 (GB10): upstream defaults SM121 to - # cutlass/cudnn for perf, but the campaign's VO-split / NVFP4 KV path needs - # b12x dispatch on GB10 (sm_121). SM120 and SM121 share the 12.x execution - # constraints; JIT/AOT arch selection is handled by the compilation context. - if is_sm12x and use_nvfp4 and cuda_major >= 13: + is_sm120 = major == 12 and minor == 0 + + # SM120 + CUDA 13: prefer b12x. SM121 (GB10) is intentionally excluded -- b12x + # is supported there as an explicit backend, but cutlass/cudnn are faster in + # most cases, so `auto` keeps using them. + if is_sm120 and use_nvfp4 and cuda_major >= 13: return [c for c in ("b12x", "cutlass", "cudnn") if c in suitable_backends] # If cuda version is 13 or greater and cudnn version is 9.15 or greater: @@ -6535,7 +6531,7 @@ def mm_fp4( Whether to use 8x4 scale factor layout or 128x4 scale factor layout, defaults to False. backend: Literal["cudnn", "trtllm", "cutlass", "cute-dsl", "b12x", "auto"] - Backend to use, defaults to ``"auto"``. On SM12x, ``"auto"`` prefers + Backend to use, defaults to ``"auto"``. On SM120, ``"auto"`` prefers ``"b12x"`` (NVFP4 only), then ``"cutlass"``, then ``"cudnn"``. On other architectures, ``"auto"`` selects between ``"cudnn"`` and ``"cutlass"`` based on the current CUDA and cuDNN versions. The ``"trtllm"`` and diff --git a/tests/gemm/test_mm_fp4.py b/tests/gemm/test_mm_fp4.py index a56ee8edc0c..24e533bd2db 100644 --- a/tests/gemm/test_mm_fp4.py +++ b/tests/gemm/test_mm_fp4.py @@ -1,5 +1,3 @@ -from types import SimpleNamespace - import pytest import torch import torch.nn.functional as F @@ -205,28 +203,6 @@ def test_mm_fp4_b12x_misaligned_k_raises(): skip_check=False, ) -def test_mm_fp4_auto_prefers_b12x_for_sm121_nvfp4(monkeypatch): - from flashinfer.gemm import gemm_base - - class MockTensor: - device = "cuda" - - monkeypatch.setattr( - gemm_base, "get_cuda_version", lambda: SimpleNamespace(major=13) - ) - monkeypatch.setattr(gemm_base, "get_compute_capability", lambda device: (12, 1)) - - backends = gemm_base._heuristic_func_mm_fp4( - ["cudnn", "cutlass", "b12x"], - MockTensor(), - MockTensor(), - MockTensor(), - MockTensor(), - use_nvfp4=True, - ) - - assert backends == ["b12x", "cutlass", "cudnn"] - def test_mm_fp4_cute_dsl_misaligned_n_raises(): device = torch.device("cuda") From 13762eec01cdcc93611223c236452dbcf5784b37 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Wed, 15 Jul 2026 15:25:26 +0900 Subject: [PATCH 14/23] fix(prefill): align the CtaTileQ smem probe with the dispatcher guard Three corrections to the shared-memory feasibility probe in FA2DetermineCtaTileQ: * Probe cudaDevAttrMaxSharedMemoryPerBlockOptin instead of the per-SM attribute. The kernel dispatcher bounds max_smem_per_threadblock by the per-block opt-in limit, so the planner now checks the same limit the dispatch-side "even the smallest KV tile exceeds shared memory" guard enforces. On parts where the two attributes differ (per-SM 102400 vs opt-in 101376 on SM86/89/120-class devices), configurations in the gap -- e.g. (qk, vo) = (432, 256) at 2-byte KV, 101888 bytes -- previously passed the probe and then failed dispatch; they now take the CTA64 fallback instead. * Size the KV step by the actual KV dtype width instead of assuming 2 bytes. sizeof(DTypeKV) is threaded from the batch-prefill plan and workspace-size entry points (where DTypeKV is in scope) through PrefillPlan / PrefillPlanWorkspaceSize / PrefillPlanImpl / PrefillSplitQOKVIndptr as a defaulted trailing parameter, and passed directly in single-prefill dispatch; callers that do not know the KV dtype keep the previous conservative 2-byte assumption. * Correct the comment. The old (512, 256) example was unreachable here (head_dim_qk >= 512 returns CTA16 before the probe), but the probe is not merely forward-looking: plan()/JIT do not validate head dims, so within this branch head_dim_qk may be any multiple of 16 below 512 under pos_encoding_mode NONE, and (qk, vo) = (448, 256) at 2-byte KV (104448 bytes) exceeds the 101376-byte opt-in limit of 99KB parts -- the probe fires today and the CTA64 fallback keeps the configuration dispatchable. For the same dims at 1-byte KV the corrected sizing flips the selection (59392 bytes fits, so CTA16 replaces the CTA64 the 2-byte assumption forced): the kv_dtype_bytes change alters tile selection for reachable configurations rather than being a no-op on current hardware. Both behaviors are pinned by test_batch_prefill_paged_cta_tile_q_smem_probe_qk448_vo256. Addresses review feedback on #3684 from @qsang-nv, including the round-2 correction of this commit's earlier "no reachable configuration triggers the fallback" claim, which was wrong for the reasons above. Signed-off-by: Jetha Chan --- csrc/batch_prefill.cu | 4 +- include/flashinfer/attention/prefill.cuh | 3 +- include/flashinfer/attention/scheduler.cuh | 23 +++++----- include/flashinfer/utils.cuh | 51 ++++++++++++++++------ 4 files changed, 54 insertions(+), 27 deletions(-) diff --git a/csrc/batch_prefill.cu b/csrc/batch_prefill.cu index 41a8dab571a..091cf0747ff 100644 --- a/csrc/batch_prefill.cu +++ b/csrc/batch_prefill.cu @@ -67,7 +67,7 @@ Array BatchPrefillWithKVCachePlan( static_cast(kv_indptr.data_ptr()), total_num_rows, batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, window_left, fixed_split_size, disable_split_kv, num_colocated_ctas, - uniform_q_len, stream); + uniform_q_len, stream, /*kv_dtype_bytes=*/sizeof(DTypeKV)); TVM_FFI_ICHECK(status == cudaSuccess) << "Failed to plan prefill with error: " << cudaGetErrorString(status); @@ -93,7 +93,7 @@ Array BatchPrefillWithKVCacheWorkspaceSize( static_cast(qo_indptr.data_ptr()), static_cast(kv_indptr.data_ptr()), total_num_rows, batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, window_left, fixed_split_size, disable_split_kv, - num_colocated_ctas, uniform_q_len, stream); + num_colocated_ctas, uniform_q_len, stream, /*kv_dtype_bytes=*/sizeof(DTypeKV)); TVM_FFI_ICHECK(status == cudaSuccess) << "Failed to calculate prefill workspace size with error: " << cudaGetErrorString(status); diff --git a/include/flashinfer/attention/prefill.cuh b/include/flashinfer/attention/prefill.cuh index f7c0a3c5ee2..573238068ac 100644 --- a/include/flashinfer/attention/prefill.cuh +++ b/include/flashinfer/attention/prefill.cuh @@ -2521,7 +2521,8 @@ cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::D constexpr uint32_t NUM_MMA_D_QK = HEAD_DIM_QK / 16; constexpr uint32_t NUM_MMA_D_VO = HEAD_DIM_VO / 16; int64_t packed_qo_len = qo_len * group_size; - uint32_t cta_tile_q = FA2DetermineCtaTileQ(packed_qo_len, HEAD_DIM_VO, HEAD_DIM_QK); + uint32_t cta_tile_q = + FA2DetermineCtaTileQ(packed_qo_len, HEAD_DIM_VO, HEAD_DIM_QK, sizeof(DTypeKV)); DISPATCH_CTA_TILE_Q(cta_tile_q, CTA_TILE_Q, { // hd512 uses the 2-Q x 2-KV-warp layout at CTA_TILE_Q=32. FP8 must take it diff --git a/include/flashinfer/attention/scheduler.cuh b/include/flashinfer/attention/scheduler.cuh index 2a8d3075876..6b20c927d8f 100644 --- a/include/flashinfer/attention/scheduler.cuh +++ b/include/flashinfer/attention/scheduler.cuh @@ -548,7 +548,8 @@ inline auto PrefillSplitQOKVIndptr(IdType* qo_indptr_h, IdType* kv_indptr_h, uint32_t page_size, uint32_t max_batch_size_if_split, bool enable_cuda_graph, int32_t window_left, int32_t fixed_split_size, bool disable_split_kv, - int64_t uniform_q_len, uint32_t head_dim_qk = 0) { + int64_t uniform_q_len, uint32_t head_dim_qk = 0, + uint32_t kv_dtype_bytes = 2) { std::vector request_indices, qo_tile_indices, kv_tile_indices, merge_indptr, o_indptr; merge_indptr.push_back(0); o_indptr.push_back(0); @@ -593,7 +594,7 @@ inline auto PrefillSplitQOKVIndptr(IdType* qo_indptr_h, IdType* kv_indptr_h, FLASHINFER_ERROR(err_msg.str()); } } - cta_tile_q = FA2DetermineCtaTileQ(packed_uniform_len, head_dim, head_dim_qk); + cta_tile_q = FA2DetermineCtaTileQ(packed_uniform_len, head_dim, head_dim_qk, kv_dtype_bytes); total_num_tiles_q = batch_size * ceil_div(packed_uniform_len, cta_tile_q); } else { // When CUDA graphs are enabled, the lengths of sequences determined by @@ -601,7 +602,7 @@ inline auto PrefillSplitQOKVIndptr(IdType* qo_indptr_h, IdType* kv_indptr_h, // the CUDA graph is created fixes the maximum number of tokens. const uint64_t max_seq_len = total_num_rows - batch_size + 1; uint64_t max_qo_len = uint64_t(max_seq_len) * gqa_group_size; - cta_tile_q = FA2DetermineCtaTileQ(max_qo_len, head_dim, head_dim_qk); + cta_tile_q = FA2DetermineCtaTileQ(max_qo_len, head_dim, head_dim_qk, kv_dtype_bytes); // Find an upper bound for the number of tiles, derived from the total // number of rows and the batch size. The sum of qo lengths rounded @@ -615,7 +616,7 @@ inline auto PrefillSplitQOKVIndptr(IdType* qo_indptr_h, IdType* kv_indptr_h, sum_packed_qo_len += packed_qo_len_arr[i]; } const int64_t avg_packed_qo_len = sum_packed_qo_len / batch_size; - cta_tile_q = FA2DetermineCtaTileQ(avg_packed_qo_len, head_dim, head_dim_qk); + cta_tile_q = FA2DetermineCtaTileQ(avg_packed_qo_len, head_dim, head_dim_qk, kv_dtype_bytes); total_num_tiles_q = 0; for (uint32_t i = 0; i < batch_size; ++i) { @@ -770,7 +771,7 @@ inline cudaError_t PrefillPlanImpl( bool disable_split_kv, int64_t num_colocated_ctas, // for POD attention, limit prefill // splits by #colocated decode CTAs - int64_t uniform_q_len, cudaStream_t stream) { + int64_t uniform_q_len, cudaStream_t stream, uint32_t kv_dtype_bytes = 2) { (void)sizeof_dtype_o; if (num_qo_heads % num_kv_heads != 0) { std::ostringstream err_msg; @@ -795,7 +796,7 @@ inline cudaError_t PrefillPlanImpl( PrefillSplitQOKVIndptr(qo_indptr_h, kv_indptr_h, total_num_rows, batch_size, num_qo_heads, num_kv_heads, head_dim_vo, page_size, max_batch_size_if_split, enable_cuda_graph, window_left, fixed_split_size, disable_split_kv, - uniform_q_len, head_dim_qk); + uniform_q_len, head_dim_qk, kv_dtype_bytes); plan_info.cta_tile_q = cta_tile_q; plan_info.total_num_rows = total_num_rows; @@ -903,7 +904,8 @@ inline cudaError_t PrefillPlan(void* float_buffer, size_t float_workspace_size_i int32_t fixed_split_size, bool disable_split_kv, int64_t num_colocated_ctas, // for POD attention, limit prefill // splits by #colocated decode CTAs - int64_t uniform_q_len, cudaStream_t stream) { + int64_t uniform_q_len, cudaStream_t stream, + uint32_t kv_dtype_bytes = 2) { size_t used_float_workspace_size = 0; size_t used_int_workspace_size = 0; return PrefillPlanImpl(used_float_workspace_size, used_int_workspace_size, float_buffer, @@ -912,7 +914,7 @@ inline cudaError_t PrefillPlan(void* float_buffer, size_t float_workspace_size_i total_num_rows, batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, sizeof_dtype_o, window_left, fixed_split_size, disable_split_kv, num_colocated_ctas, - uniform_q_len, stream); + uniform_q_len, stream, kv_dtype_bytes); } template @@ -921,7 +923,8 @@ inline cudaError_t PrefillPlanWorkspaceSize( IdType* kv_indptr_h, uint32_t total_num_rows, uint32_t batch_size, uint32_t num_qo_heads, uint32_t num_kv_heads, uint32_t head_dim_qk, uint32_t head_dim_vo, uint32_t page_size, bool enable_cuda_graph, uint32_t sizeof_dtype_o, int32_t window_left, int32_t fixed_split_size, - bool disable_split_kv, int64_t num_colocated_ctas, int64_t uniform_q_len, cudaStream_t stream) { + bool disable_split_kv, int64_t num_colocated_ctas, int64_t uniform_q_len, cudaStream_t stream, + uint32_t kv_dtype_bytes = 2) { PrefillPlanInfo plan_info; return PrefillPlanImpl(float_workspace_size_in_bytes, int_workspace_size_in_bytes, /*float_buffer=*/nullptr, /*float_workspace_size_in_bytes=*/0, @@ -930,7 +933,7 @@ inline cudaError_t PrefillPlanWorkspaceSize( kv_indptr_h, total_num_rows, batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, sizeof_dtype_o, window_left, fixed_split_size, disable_split_kv, - num_colocated_ctas, uniform_q_len, stream); + num_colocated_ctas, uniform_q_len, stream, kv_dtype_bytes); } inline float cost_function(int qo_len, int kv_len) { return 2 * float(qo_len) + kv_len; } diff --git a/include/flashinfer/utils.cuh b/include/flashinfer/utils.cuh index 995cd3c2353..d151a7610bc 100644 --- a/include/flashinfer/utils.cuh +++ b/include/flashinfer/utils.cuh @@ -390,10 +390,11 @@ inline void DebugPrintCUDAArray(T* device_ptr, size_t size, std::string prefix = } inline uint32_t FA2DetermineCtaTileQ(int64_t avg_packed_qo_len, uint32_t head_dim, - uint32_t head_dim_qk = 0) { + uint32_t head_dim_qk = 0, uint32_t kv_dtype_bytes = 2) { // head_dim is the VO dim at the batch-prefill call sites; head_dim_qk (when // nonzero) lets asymmetric (QK != VO) configurations report the dim that - // actually drives shared-memory cost. + // actually drives shared-memory cost. kv_dtype_bytes is sizeof(DTypeKV) when + // the caller knows it; the default 2 is the conservative worst case. const uint32_t qk = head_dim_qk ? head_dim_qk : head_dim; if (head_dim >= 512) { // True VO-split (VO >= 512): the split halves o_frag register pressure, so @@ -420,21 +421,43 @@ inline uint32_t FA2DetermineCtaTileQ(int64_t avg_packed_qo_len, uint32_t head_di return 64; } else { // avg_packed_qo_len <= 16: prefer the 1x4 warp layout (cta_tile_q=16), - // but ONLY if one NUM_MMA_KV step fits shared memory. With 4 kv-warps a - // step costs (qk+vo)*16*4*sizeof(dtype); at (512,256) bf16 that is 96KB - // + a 16KB Q tile > the ~100KB/SM budget of CC 12.x consumer Blackwell - // (and would yield the unrecoverable "Unsupported max_mma_kv: 0" - // dispatch failure). Fall back to cta_tile_q=64 (4x1 layout, 4x - // cheaper KV step) exactly like the Turing branch below. sizeof - // assumed 2 (worst case): only configurations that could not run at - // all are moved. - int dev_id = 0, max_smem_per_sm = 0; + // but only if one NUM_MMA_KV step fits shared memory. The estimate + // mirrors the kernel dispatcher's minimum requirement -- one Q tile + // (assumed 2-byte Q) plus a single NUM_MMA_KV step of K+V in the 1x4 + // layout (16 rows x 4 KV warps x kv_dtype_bytes) -- checked against + // cudaDevAttrMaxSharedMemoryPerBlockOptin, the same limit the + // dispatcher's "even the smallest KV tile exceeds shared memory" guard + // bounds max_smem_per_threadblock with. It is an approximation: FP4 + // scale-factor bytes and FP8 repack staging are not counted (both are + // zero or small at CTA_TILE_Q=16). + // This fallback is reachable today: neither plan() nor the JIT path + // validates head dims, so within this branch (vo <= 256 -- vo in + // (256, 512) fails the NUM_MMA_D_VO tiling static_assert and vo >= 512 + // / qk >= 512 return above) head_dim_qk may be any multiple of 16 up + // to 496 under pos_encoding_mode NONE. E.g. (qk, vo) = (448, 256) at + // 2-byte KV needs 16*448*2 + (448+256)*16*4*2 = 104448 bytes, which + // exceeds the 101376-byte opt-in limit of 99KB parts (SM86/89/120/121) + // -- the probe fires and the cta_tile_q=64 fallback (4x1 layout, 4x + // smaller KV step, like the Turing branch below) keeps the + // configuration dispatchable instead of failing outright. The + // kv_dtype_bytes accuracy likewise changes tile selection for such + // configurations: at (448, 256) the true 1-byte cost is + // 16*448*2 + (448+256)*16*4*1 = 59392 bytes, so callers that supply + // kv_dtype_bytes=1 get cta_tile_q=16 where the 2-byte assumption + // returned 64. Both behaviors are pinned by + // test_batch_prefill_paged_cta_tile_q_smem_probe_qk448_vo256 in + // tests/attention/test_batch_prefill_kernels.py. Callers that cannot + // supply kv_dtype_bytes keep the conservative 2-byte default, which + // overestimates 1-byte (FP8/FP4) KV by 2x and can demote a working + // configuration to cta_tile_q=64 (a performance change, not a + // correctness one). + int dev_id = 0, max_smem_per_block_optin = 0; cudaGetDevice(&dev_id); - cudaDeviceGetAttribute(&max_smem_per_sm, cudaDevAttrMaxSharedMemoryPerMultiprocessor, + cudaDeviceGetAttribute(&max_smem_per_block_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev_id); const uint32_t q_tile_smem = 16 * qk * 2; - const uint32_t kv_step_smem_1x4 = (qk + head_dim) * 16 * 4 * 2; - if (q_tile_smem + kv_step_smem_1x4 > (uint32_t)max_smem_per_sm) { + const uint32_t kv_step_smem_1x4 = (qk + head_dim) * 16 * 4 * kv_dtype_bytes; + if (q_tile_smem + kv_step_smem_1x4 > (uint32_t)max_smem_per_block_optin) { return 64; } return 16; From 68f1e97ec410d5fcdee926c92d03b96968b68a20 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Wed, 15 Jul 2026 15:25:39 +0900 Subject: [PATCH 15/23] docs(prefill): present the NVFP4 split-KV gate as an empirical workaround The previous rationale claimed a split-KV chunk boundary can land in the middle of a 16-element scale block. That mechanism is wrong: NVFP4 scale factors group 16 consecutive head-dim elements of a single token, while split-KV partitions the token axis, so a split boundary never slices a scale block. Reword the gate docstring, call-site comments and test docstring to state what is actually known -- corruption was observed empirically when qo_len << kv_len (decode / prefix-cache extend) and disappears with split-KV disabled at no measured decode throughput cost -- and cite the interaction between small per-split KV chunks and the 1-byte-KV NUM_MMA_KV tile floor as an unconfirmed hypothesis rather than fact. Also drops project-specific NOTE tags from nearby comments. Comment-only; no behavior change. Addresses review feedback on #3684 from @qsang-nv. Signed-off-by: Jetha Chan --- flashinfer/prefill.py | 51 +++++++++++-------- tests/attention/test_nvfp4_attention_sm120.py | 11 ++-- 2 files changed, 35 insertions(+), 27 deletions(-) diff --git a/flashinfer/prefill.py b/flashinfer/prefill.py index 4f47b8d2aec..ab40baccc6a 100755 --- a/flashinfer/prefill.py +++ b/flashinfer/prefill.py @@ -1500,16 +1500,21 @@ def _compute_page_mask_indptr( def _nvfp4_kv_requires_disabled_split_kv(kv_data_type: torch.dtype) -> bool: """Whether split-KV must be disabled because the KV cache is NVFP4. - NVFP4 paged KV is stored as packed ``uint8`` (two E2M1 values per byte) with a - per-16-element FP8 block scale. Split-KV (flash-decoding) partitions the KV - range into chunks sized by ``kv_chunk_size``, which is *not* aligned to the - 16-element scale blocks: a chunk boundary that lands mid-block — together with - the small per-split chunk tripping the 1-byte-KV ``NUM_MMA_KV`` tile floor — - corrupts the dequantized reads. The corruption only surfaces when a short query - attends a long KV (``qo_len << kv_len``), i.e. decode and prefix-cache *extend*, - so dense full-prefill tests miss it while prefix caching breaks. Until the FP4 - split path is made scale-block-aware, disable split-KV for NVFP4 KV. FP8 KV has - no block scales and is unaffected. + This gate is an *empirical workaround*: with split-KV (flash-decoding) + enabled, NVFP4 paged KV was observed to produce corrupted outputs whenever + a short query attends a long KV range (``qo_len << kv_len``, i.e. decode + and prefix-cache extend), while dense full-prefill was unaffected. + Disabling split-KV removes the corruption, and decode-throughput + measurements showed no cost from the gate. + + The root cause has not been confirmed. The FP8 scale-factor blocks + themselves cannot be the mechanism: NVFP4 scales group 16 consecutive + *head-dim* elements of a single token, whereas split-KV partitions the + *token* axis, so a split boundary never slices a scale block. The current + hypothesis (unconfirmed) is that the small per-split KV chunks interact + badly with the ``NUM_MMA_KV`` tile floor of the 1-byte-KV FA2 path. Until + the failure is root-caused and fixed, force split-KV off for NVFP4 KV. + FP8 and 16-bit KV caches are unaffected and keep split-KV. """ if kv_data_type == torch.uint8: # packed NVFP4 (the run path's convention) return True @@ -2287,10 +2292,10 @@ def plan( ) if packed_custom_mask is None and custom_mask is not None: # create packed custom mask from custom mask - # NOTE(spark-hijinks): the segment_packbits kernel requires the - # indptr on the mask's device (CHECK_DEVICE in quantization.cu), - # but mask_indptr inherits qo_indptr's device, which callers - # (e.g. vLLM) routinely keep on CPU while the mask is on GPU. + # The segment_packbits kernel requires the indptr on the mask's + # device (CHECK_DEVICE in quantization.cu), but mask_indptr + # inherits qo_indptr's device, which callers (e.g. vLLM) routinely + # keep on CPU while the mask is on GPU. packed_custom_mask, mask_indptr = segment_packbits( custom_mask.contiguous().view(-1), mask_indptr.to(custom_mask.device), @@ -2507,8 +2512,9 @@ def plan( if not disable_split_kv and _nvfp4_kv_requires_disabled_split_kv( kv_data_type ): - # NVFP4 split-KV corrupts prefix-cached / decode reads; force - # it off. See _nvfp4_kv_requires_disabled_split_kv for why. + # Empirical workaround: split-KV corrupted NVFP4 KV reads + # when qo_len << kv_len (decode / prefix-cache extend); see + # _nvfp4_kv_requires_disabled_split_kv for details. disable_split_kv = True args.append(disable_split_kv) # disable_split_kv args.append(0) # num_colocated_ctas @@ -3447,10 +3453,10 @@ def plan( mask_indptr = _compute_mask_indptr(qo_indptr, kv_indptr) if packed_custom_mask is None and custom_mask is not None: # create packed custom mask from custom mask - # NOTE(spark-hijinks): segment_packbits requires mask_indptr on the - # same device as custom_mask, but mask_indptr inherits qo_indptr's - # device (often CPU) while custom_mask is on GPU. Mirror the paged - # path's .to(device) so the ragged custom-mask flow doesn't crash. + # segment_packbits requires mask_indptr on the same device as + # custom_mask, but mask_indptr inherits qo_indptr's device (often + # CPU) while custom_mask is on GPU. Mirror the paged path's + # .to(device) so the ragged custom-mask flow doesn't crash. packed_custom_mask, mask_indptr = segment_packbits( custom_mask.contiguous().view(-1), mask_indptr.to(custom_mask.device), @@ -3742,8 +3748,9 @@ def plan( if not disable_split_kv and _nvfp4_kv_requires_disabled_split_kv( kv_data_type ): - # NVFP4 split-KV corrupts prefix-cached / decode reads; force - # it off. See _nvfp4_kv_requires_disabled_split_kv for why. + # Empirical workaround: split-KV corrupted NVFP4 KV reads + # when qo_len << kv_len (decode / prefix-cache extend); see + # _nvfp4_kv_requires_disabled_split_kv for details. disable_split_kv = True args.append(disable_split_kv) # disable_split_kv args.append(0) # num_colocated_ctas diff --git a/tests/attention/test_nvfp4_attention_sm120.py b/tests/attention/test_nvfp4_attention_sm120.py index c0afa87b753..dc1b8a70a39 100644 --- a/tests/attention/test_nvfp4_attention_sm120.py +++ b/tests/attention/test_nvfp4_attention_sm120.py @@ -376,11 +376,12 @@ def test_nvfp4_attention_sm120_causal_mask_column_order(): def test_nvfp4_split_kv_gate_dtype_logic(): """The split-KV gate must fire for NVFP4 KV and only for NVFP4 KV. - NVFP4 split-KV corrupts prefix-cached / decode reads (a split boundary can land - mid 16-element scale block; the small per-split chunk also trips the 1-byte-KV - NUM_MMA_KV tile floor), so plan() force-disables split-KV when the KV cache is - NVFP4. FP8 / 16-bit KV have no block scales and must keep split-KV. This is a - pure dtype-classification check (no GPU required) guarding that contract. + Split-KV was empirically observed to corrupt NVFP4 KV reads when a short + query attends a long KV range (decode / prefix-cache extend), so plan() + force-disables split-KV when the KV cache is NVFP4 as a workaround (see + _nvfp4_kv_requires_disabled_split_kv for the state of the root-cause + analysis). FP8 / 16-bit KV are unaffected and must keep split-KV. This is + a pure dtype-classification check (no GPU required) guarding that contract. """ from flashinfer.prefill import _nvfp4_kv_requires_disabled_split_kv From 12d9b40cc87504b15c0d1a46cf25b0dbf4dadf83 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Wed, 15 Jul 2026 15:25:50 +0900 Subject: [PATCH 16/23] fix(prefill): name the missing key for underivable JIT scalars in run() run() fills declared-but-unprovided JIT scalars from a fixed mapping and raised a bare KeyError when a JIT module declares a scalar the mapping does not know how to derive. Raise a ValueError naming the scalar and listing the derivable set instead. Also drop the redundant max(0, ...) clamp on the provided-scalar count: prepare_jit_additional_args always returns at least one entry per declared tensor name, so the excess over the tensor-name count cannot be negative; a comment records that invariant. Addresses review feedback on #3684 from @qsang-nv. Signed-off-by: Jetha Chan --- flashinfer/prefill.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/flashinfer/prefill.py b/flashinfer/prefill.py index ab40baccc6a..9a06a399904 100755 --- a/flashinfer/prefill.py +++ b/flashinfer/prefill.py @@ -2915,11 +2915,23 @@ def run( "scale_v_scalar": scale_v_scalar, "token_pos_in_items_len": self._token_pos_in_items_len, } - scalar_start = max( - 0, - len(additional_args) - len(self._jit_additional_tensor_names), + # prepare_jit_additional_args returns one entry per declared + # tensor name plus any scalars the caller passed + # positionally, so the number of scalars already provided is + # exactly the excess over the tensor-name count. + num_scalars_provided = len(additional_args) - len( + self._jit_additional_tensor_names ) - for name in self._jit_additional_scalar_names[scalar_start:]: + for name in self._jit_additional_scalar_names[ + num_scalars_provided: + ]: + if name not in jit_scalar_values: + raise ValueError( + f"JIT module declares additional scalar {name!r}, " + "which was not passed positionally to run() and " + "cannot be derived automatically; derivable " + f"scalars are: {sorted(jit_scalar_values)}" + ) additional_args.append(jit_scalar_values[name]) run_args.extend(additional_args) else: From 1c87f68ec04717064b7cb80995835054748e8332 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Wed, 15 Jul 2026 15:26:04 +0900 Subject: [PATCH 17/23] refactor(jit): drop dead CTA_TILE_Q=32 instantiation for large-QK/small-VO For head_dim_qk >= 512 with head_dim_vo < 512 (asymmetric heads), FA2DetermineCtaTileQ always returns CTA_TILE_Q=16, so the CTA_TILE_Q=32 instantiation in the paged/ragged kernel-instantiation lists can never be dispatched; drop it and instantiate only CTA16 for that shape class. The dispatch macro case-32 arm still references the symbol, which stays unresolved in the module exactly like the pre-existing never-selected CTA_TILE_Q=32 of symmetric head_dim < 512 modules (verified: shipped modules carry it as an undefined, lazily-bound symbol). Also updates the KernelTraits::IsInvalid comment to describe the three-way CTA_TILE_Q selection. Addresses review feedback on #3684 from @qsang-nv. Signed-off-by: Jetha Chan --- csrc/batch_prefill_paged_kernel_inst.jinja | 8 +++++--- csrc/batch_prefill_ragged_kernel_inst.jinja | 8 +++++--- include/flashinfer/attention/prefill.cuh | 7 ++++++- 3 files changed, 16 insertions(+), 7 deletions(-) diff --git a/csrc/batch_prefill_paged_kernel_inst.jinja b/csrc/batch_prefill_paged_kernel_inst.jinja index c5d4723a74a..5518d8b71b5 100644 --- a/csrc/batch_prefill_paged_kernel_inst.jinja +++ b/csrc/batch_prefill_paged_kernel_inst.jinja @@ -5,9 +5,11 @@ namespace flashinfer { constexpr auto use_custom_mask = {{ mask_mode }} == MaskMode::kCustom; -{# FA2DetermineCtaTileQ only selects {16, 32} for head_dim_qk >= 512 and - {16, 64, 128} otherwise; don't instantiate unreachable variants. #} -{% for cta_tile_q in ([16, 32] if head_dim_qk | int >= 512 else [16, 64, 128]) %} +{# FA2DetermineCtaTileQ only selects {16, 32} for head_dim_vo >= 512 (VO-split), + {16} for head_dim_qk >= 512 with a smaller head_dim_vo (o_frag register + pressure permits only CTA16 there), and {16, 64, 128} otherwise; don't + instantiate unreachable variants. #} +{% for cta_tile_q in ([16, 32] if head_dim_vo | int >= 512 else ([16] if head_dim_qk | int >= 512 else [16, 64, 128])) %} template cudaError_t BatchPrefillWithPagedKVCacheDispatched< /*CTA_TILE_Q=*/{{cta_tile_q}}, {{head_dim_qk}}, {{head_dim_vo}}, {{pos_encoding_mode}}, {{use_fp16_qk_reduction}}, {{mask_mode}}, {{ variant_name }}, PagedParams>(PagedParams params, {{ dtype_o }}* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); diff --git a/csrc/batch_prefill_ragged_kernel_inst.jinja b/csrc/batch_prefill_ragged_kernel_inst.jinja index 6e5eec59e8b..8cb1f7edb74 100644 --- a/csrc/batch_prefill_ragged_kernel_inst.jinja +++ b/csrc/batch_prefill_ragged_kernel_inst.jinja @@ -5,9 +5,11 @@ namespace flashinfer { constexpr auto use_custom_mask = {{ mask_mode }} == MaskMode::kCustom; -{# FA2DetermineCtaTileQ only selects {16, 32} for head_dim_qk >= 512 and - {16, 64, 128} otherwise; don't instantiate unreachable variants. #} -{% for cta_tile_q in ([16, 32] if head_dim_qk | int >= 512 else [16, 64, 128]) %} +{# FA2DetermineCtaTileQ only selects {16, 32} for head_dim_vo >= 512 (VO-split), + {16} for head_dim_qk >= 512 with a smaller head_dim_vo (o_frag register + pressure permits only CTA16 there), and {16, 64, 128} otherwise; don't + instantiate unreachable variants. #} +{% for cta_tile_q in ([16, 32] if head_dim_vo | int >= 512 else ([16] if head_dim_qk | int >= 512 else [16, 64, 128])) %} template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< /*CTA_TILE_Q=*/{{cta_tile_q}}, {{head_dim_qk}}, {{head_dim_vo}}, {{pos_encoding_mode}}, {{use_fp16_qk_reduction}}, {{mask_mode}}, {{ variant_name }}, RaggedParams>(RaggedParams params, {{ dtype_o }}* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); diff --git a/include/flashinfer/attention/prefill.cuh b/include/flashinfer/attention/prefill.cuh index 573238068ac..5d4b81d35b8 100644 --- a/include/flashinfer/attention/prefill.cuh +++ b/include/flashinfer/attention/prefill.cuh @@ -300,7 +300,12 @@ struct KernelTraits { static constexpr bool IsInvalid() { // The first clause prunes (CTA_TILE_Q, head_dim) pairs FA2DetermineCtaTileQ - // never selects: {16, 32} for head_dim_qk >= 512, {16, 64, 128} otherwise. + // never selects: it picks from {16, 32} when head_dim_vo >= 512 (VO-split), + // {16} when head_dim_qk >= 512 with a smaller head_dim_vo, and + // {16, 64, 128} otherwise. The clause keys on HEAD_DIM_QK (which bounds + // CTA_TILE_Q <= 32 in both large-head cases); the kernel-instantiation + // lists additionally drop the never-selected CTA_TILE_Q=32 for + // head_dim_qk >= 512 with head_dim_vo < 512. return ((HEAD_DIM_QK >= 512 ? (CTA_TILE_Q > 32) : (CTA_TILE_Q == 32)) || (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 && From 5f75b9656d0eb0619bd5b382bb628d22bf58c2ed Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Wed, 15 Jul 2026 22:59:55 +0900 Subject: [PATCH 18/23] test(nvfp4): asymmetric qk/vo paged prefill correctness + decode unequal-stride rejection Completes the support-or-reject-explicitly contract for unequal K/V strides in tree: every consumer reachable from the updated entry points must either support independently-strided K/V pools or reject them loudly. Support half: asymmetric (head_dim_qk != head_dim_vo) NVFP4 paged prefill over (512,256) and (256,128) x page_size {1,16} x num_kv_heads {2,8}, causal. K/V pools and their scale-factor tensors are separately allocated with genuinely different stride families; bf16 sources are quantized with the in-tree NVFP4 KV quantization kernel and the FA2 output is checked against a float32 reference attention computed on nvfp4_kv_dequantize_paged output, so kernel and reference consume the exact same quantized bytes (dequantization oracle, not a requantized approximation). Reject half: the CUDA-core decode entry point (BatchDecodeWithPagedKVCacheRun) addresses both K and V through a single set of (K) strides, so its restored ICHECK must fire on K/V pools whose stride families differ instead of silently misaddressing V. A positive control with identically padded (equal-stride, non-contiguous) pools runs and matches the reference, proving the negative case fails because of the stride inequality and not the padded allocation. Addresses review feedback on #3684 from @qsang-nv. Signed-off-by: Jetha Chan --- tests/attention/test_batch_decode_kernels.py | 118 ++++++++++++ tests/attention/test_batch_prefill_kernels.py | 179 ++++++++++++++++++ 2 files changed, 297 insertions(+) diff --git a/tests/attention/test_batch_decode_kernels.py b/tests/attention/test_batch_decode_kernels.py index a6791e49b78..9ad8aa945e0 100644 --- a/tests/attention/test_batch_decode_kernels.py +++ b/tests/attention/test_batch_decode_kernels.py @@ -851,6 +851,124 @@ def test_batch_decode_with_paged_kv_cache_nvfp4_large_head(): ) +def test_batch_decode_rejects_unequal_kv_strides_nvfp4_contract(): + """The FA2 CUDA-core decode kernel addresses both K and V through a single + set of (K) strides, so ``BatchDecodeWithPagedKVCacheRun`` must reject K/V + pools whose stride families differ instead of silently misaddressing V. + This is the "reject explicitly" half of the NVFP4 unequal-stride contract: + every entry point that cannot consume independently-strided K/V pools (the + layout NVFP4/asymmetric caches produce) must fail loudly. + + NVFP4 packed (uint8) KV itself cannot reach this guard — its half-width + packed head dim trips the equal-head-dim ICHECK first, and NVFP4 decode + routes through the tensor-core prefill path — so the unequal-stride + construction uses fp16: two separately allocated pools whose padding + differs, giving identical shapes but different stride families. + + A positive control with identically padded (equal-stride, non-contiguous) + pools must run and match the reference, proving the negative case fails + because of the stride inequality and not the padded allocation. + """ + torch.manual_seed(42) + batch_size = 4 + kv_len = 54 + page_size = 8 + num_kv_heads = 4 + num_qo_heads = 4 + head_dim = 128 + dtype = torch.float16 + + q = torch.randn(batch_size, num_qo_heads, head_dim, device="cuda:0", dtype=dtype) + num_pages_per_seq = (kv_len + page_size - 1) // page_size + total_num_pages = num_pages_per_seq * batch_size + kv_indptr = ( + torch.arange(0, batch_size + 1, device="cuda:0", dtype=torch.int32) + * num_pages_per_seq + ) + kv_indices = torch.arange(0, total_num_pages, device="cuda:0", dtype=torch.int32) + kv_last_page_len = torch.full( + (batch_size,), (kv_len - 1) % page_size + 1, dtype=torch.int32, device="cuda:0" + ) + + def padded_pool(num_padding_heads): + parent = torch.randn( + total_num_pages, + page_size, + num_kv_heads + num_padding_heads, + head_dim, + device="cuda:0", + dtype=dtype, + ) + return parent[:, :, :num_kv_heads, :] + + # Positive control: separately allocated K/V pools with IDENTICAL padding. + k_equal = padded_pool(1) + v_equal = padded_pool(1) + assert not k_equal.is_contiguous() + assert k_equal.stride() == v_equal.stride() + + workspace_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.int8, device="cuda:0") + wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, "NHD" + ) + wrapper.plan( + kv_indptr, + kv_indices, + kv_last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + pos_encoding_mode="NONE", + q_data_type=dtype, + kv_data_type=dtype, + ) + # The guard under test lives in the CUDA-core decode entry point + # (csrc/batch_decode.cu), not the tensor-core prefill path. + assert not wrapper.use_tensor_cores + o = wrapper.run(q, (k_equal, v_equal)) + + kv_indptr_cpu = kv_indptr.cpu() + kv_last_page_len_cpu = kv_last_page_len.cpu() + for i in range(batch_size): + ki = torch.cat( + [ + k_equal[kv_indptr_cpu[i] : kv_indptr_cpu[i + 1] - 1].reshape( + -1, num_kv_heads, head_dim + ), + k_equal[kv_indptr_cpu[i + 1] - 1, : kv_last_page_len_cpu[i]].reshape( + -1, num_kv_heads, head_dim + ), + ], + dim=0, + ) + vi = torch.cat( + [ + v_equal[kv_indptr_cpu[i] : kv_indptr_cpu[i + 1] - 1].reshape( + -1, num_kv_heads, head_dim + ), + v_equal[kv_indptr_cpu[i + 1] - 1, : kv_last_page_len_cpu[i]].reshape( + -1, num_kv_heads, head_dim + ), + ], + dim=0, + ) + o_ref_i = flashinfer.decode.single_decode_with_kv_cache( + q[i], ki, vi, pos_encoding_mode="NONE", logits_soft_cap=0.0 + ) + torch.testing.assert_close(o[i], o_ref_i, rtol=1e-3, atol=1e-3) + + # Negative case: V pool with different padding — identical shape, unequal + # stride family. The kernel would otherwise walk V through K's strides; + # the ICHECK must reject the call loudly and name the stride limitation. + v_unequal = padded_pool(2) + v_unequal.copy_(v_equal) + assert v_unequal.shape == k_equal.shape + assert v_unequal.stride() != k_equal.stride() + with pytest.raises(Exception, match="must have identical strides"): + wrapper.run(q, (k_equal, v_unequal)) + + if __name__ == "__main__": test_batch_decode_with_paged_kv_cache( 256, diff --git a/tests/attention/test_batch_prefill_kernels.py b/tests/attention/test_batch_prefill_kernels.py index 043956aea08..ffb113ab8ca 100644 --- a/tests/attention/test_batch_prefill_kernels.py +++ b/tests/attention/test_batch_prefill_kernels.py @@ -35,6 +35,15 @@ def skip_if_head_dim_unsupported(head_dim: int): pytest.skip("16-bit FA2 head_dim > 256 is only supported on SM80 or newer") +def skip_if_nvfp4_asymmetric_unsupported(head_dim_qk: int): + skip_if_head_dim_unsupported(head_dim_qk) + if get_compute_capability(torch.device("cuda:0"))[0] < 10: + pytest.skip( + "asymmetric NVFP4 KV prefill uses the NVFP4 KV quantization kernel, " + "which requires SM100 or newer" + ) + + @pytest.fixture( autouse=not has_flashinfer_jit_cache(), scope="module", @@ -1530,6 +1539,176 @@ def test_batch_prefill_with_paged_kv_cache_nvfp4_strided_scale_views(kv_layout): ) +@pytest.mark.parametrize("head_dim_qk,head_dim_vo", [(512, 256), (256, 128)]) +@pytest.mark.parametrize("page_size", [1, 16]) +@pytest.mark.parametrize("num_kv_heads", [2, 8]) +@pytest.mark.parametrize("causal", [True]) +def test_batch_prefill_with_paged_kv_cache_nvfp4_asymmetric( + head_dim_qk, + head_dim_vo, + page_size, + num_kv_heads, + causal, +): + """Asymmetric (head_dim_qk != head_dim_vo) NVFP4 paged prefill correctness. + + K pages are ``[.., head_dim_qk // 2]`` and V pages ``[.., head_dim_vo // 2]``, + so the separately allocated K and V pools (and their scale-factor tensors) + have genuinely different stride families — the layout an asymmetric NVFP4 + KV cache hands the FA2 paged prefill entry point. + + bf16 K/V are quantized with the in-tree NVFP4 KV quantization kernel and + the FA2 output is checked against a float32 reference attention computed on + ``nvfp4_kv_dequantize_paged`` output: kernel and reference consume the exact + same quantized bytes, so the reference is a dequantization oracle rather + than a requantized approximation. + """ + skip_if_nvfp4_asymmetric_unsupported(head_dim_qk) + + kv_layout = "NHD" + torch.manual_seed(42) + batch_size = 2 + kv_len = 99 + qo_len = 33 + num_qo_heads = 2 * num_kv_heads + q_dtype = torch.bfloat16 + + # --- query --- + q = torch.randn( + batch_size * qo_len, num_qo_heads, head_dim_qk, device="cuda:0", dtype=q_dtype + ) + q_indptr_cpu = torch.arange(0, batch_size + 1, dtype=torch.int32) * qo_len + + # --- paged KV metadata --- + num_pages_per_seq = (kv_len + page_size - 1) // page_size + total_num_pages = num_pages_per_seq * batch_size + kv_indptr_cpu = ( + torch.arange(0, batch_size + 1, dtype=torch.int32) * num_pages_per_seq + ) + kv_indices_cpu = torch.arange(0, total_num_pages, dtype=torch.int32) + kv_last_page_len_cpu = torch.full( + (batch_size,), (kv_len - 1) % page_size + 1, dtype=torch.int32 + ) + + # --- bf16 source K/V, quantized via the in-tree NVFP4 KV quantization + # kernel (the helper tests/utils/test_fp4_kv_quantization.py exercises). + # It quantizes row-wise over the last dim, so asymmetric K/V widths + # quantize naturally. --- + k_bf16 = torch.randn( + total_num_pages, + page_size, + num_kv_heads, + head_dim_qk, + device="cuda:0", + dtype=q_dtype, + ) + v_bf16 = torch.randn( + total_num_pages, + page_size, + num_kv_heads, + head_dim_vo, + device="cuda:0", + dtype=q_dtype, + ) + # global_scale=1.0 avoids FP8 E4M3 block-scale underflow (see + # test_nvfp4_kv_roundtrip). + k_global_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda:0") + v_global_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda:0") + k_packed, k_sf = flashinfer.nvfp4_kv_quantize( + k_bf16.reshape(-1, head_dim_qk), k_global_scale + ) + v_packed, v_sf = flashinfer.nvfp4_kv_quantize( + v_bf16.reshape(-1, head_dim_vo), v_global_scale + ) + k_packed = k_packed.reshape( + total_num_pages, page_size, num_kv_heads, head_dim_qk // 2 + ) + k_sf = k_sf.reshape(total_num_pages, page_size, num_kv_heads, head_dim_qk // 16) + v_packed = v_packed.reshape( + total_num_pages, page_size, num_kv_heads, head_dim_vo // 2 + ) + v_sf = v_sf.reshape(total_num_pages, page_size, num_kv_heads, head_dim_vo // 16) + + # The whole point: every consumer reachable from this entry point must + # support (or explicitly reject) unequal K/V strides. + assert k_packed.stride() != v_packed.stride() + assert k_sf.stride() != v_sf.stride() + + # --- run BatchPrefillWithPagedKVCacheWrapper (FA2 NVFP4 paged path) --- + workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8, device="cuda:0") + wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, kv_layout + ) + wrapper.plan( + q_indptr_cpu.to("cuda:0"), + kv_indptr_cpu.to("cuda:0"), + kv_indices_cpu.to("cuda:0"), + kv_last_page_len_cpu.to("cuda:0"), + num_qo_heads, + num_kv_heads, + head_dim_qk, + page_size, + head_dim_vo=head_dim_vo, + causal=causal, + pos_encoding_mode="NONE", + logits_soft_cap=0.0, + kv_data_type=torch.uint8, + q_data_type=q_dtype, + ) + o = wrapper.run( + q, + (k_packed, v_packed), + k_scale=k_global_scale.item(), + v_scale=v_global_scale.item(), + kv_cache_sf=(k_sf, v_sf), + ) + assert o.shape == (batch_size * qo_len, num_qo_heads, head_dim_vo) + assert torch.isfinite(o).all() + + # --- dequantization oracle: #3748's paged NVFP4 dequant kernel --- + block_tables = ( + kv_indices_cpu.to("cuda:0").reshape(batch_size, num_pages_per_seq).contiguous() + ) + seq_lens = torch.full((batch_size,), kv_len, dtype=torch.int32, device="cuda:0") + k_dq = torch.zeros( + batch_size, kv_len, num_kv_heads, head_dim_qk, dtype=q_dtype, device="cuda:0" + ) + v_dq = torch.zeros( + batch_size, kv_len, num_kv_heads, head_dim_vo, dtype=q_dtype, device="cuda:0" + ) + flashinfer.nvfp4_kv_dequantize_paged( + (k_packed, v_packed), + (k_sf.view(torch.float8_e4m3fn), v_sf.view(torch.float8_e4m3fn)), + block_tables, + seq_lens, + k_global_scale, + v_global_scale, + k_dq, + v_dq, + kv_layout=kv_layout, + ) + + # --- float32 reference attention on the dequantized K/V --- + group_size = num_qo_heads // num_kv_heads + sm_scale = head_dim_qk**-0.5 + for i in range(batch_size): + qi = q[q_indptr_cpu[i] : q_indptr_cpu[i + 1]].float() # [qo, Hq, dqk] + ki = k_dq[i].float().repeat_interleave(group_size, dim=1) # [kv, Hq, dqk] + vi = v_dq[i].float().repeat_interleave(group_size, dim=1) # [kv, Hq, dvo] + + logits = torch.einsum("qhd,khd->hqk", qi, ki) * sm_scale + if causal: + qpos = torch.arange(qo_len, device="cuda:0").unsqueeze(1) + kpos = torch.arange(kv_len, device="cuda:0").unsqueeze(0) + allowed = kpos <= qpos + (kv_len - qo_len) + logits = logits.masked_fill(~allowed.unsqueeze(0), float("-inf")) + o_ref_i = torch.einsum("hqk,khd->qhd", torch.softmax(logits, dim=-1), vi) + o_i = o[q_indptr_cpu[i] : q_indptr_cpu[i + 1]].float() + + # NVFP4 is 4-bit; use relaxed tolerance + torch.testing.assert_close(o_i, o_ref_i, rtol=1e-1, atol=1e-1) + + @pytest.mark.parametrize("batch_size", [1, 4]) @pytest.mark.parametrize("kv_len", [128, 256]) @pytest.mark.parametrize("qo_len", [64, 128]) From 1b9f63c52159ea889b6abdfb6d83e6e9adf8c423 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Fri, 17 Jul 2026 15:27:52 +0900 Subject: [PATCH 19/23] test(prefill): pin the CtaTileQ smem-probe selection at (448, 256) Regression test for the FA2DetermineCtaTileQ shared-memory probe at head dims that reach it today: plan()/JIT do not validate head dims, so (qk, vo) = (448, 256) under pos_encoding_mode NONE is accepted, and at 2-byte KV its short-q 1x4-layout cost (104448 bytes) exceeds the 101376-byte per-block opt-in limit of 99KB parts. The test computes the expected tile from the device's actual opt-in limit (so the assertion is exact on every architecture), asserts the planned cta_tile_q via PrefillPlanInfo (the same technique as test_fp8_prefill.py), and for 2-byte KV runs the kernel against an exact float32 reference: on 99KB parts this proves the probe fires and the CTA64 fallback keeps the configuration dispatchable where the CTA16 dispatch would exceed the per-block limit, and on larger-smem parts it proves the CTA16 selection runs. For 1-byte KV the assertion is plan-level: the FA2 1-byte KV producers require head_dim to be a multiple of 128 elements (the 128-bit-per-lane load loop steps NUM_MMA_D by 8, and the k128B swizzle needs an 8-aligned upcast stride), so no currently-runnable 1-byte configuration reaches the flipped CTA64->CTA16 region -- the pin locks the documented planner behavior for when one does. Addresses review feedback on #3684 from @qsang-nv. Signed-off-by: Jetha Chan --- tests/attention/test_batch_prefill_kernels.py | 134 ++++++++++++++++++ 1 file changed, 134 insertions(+) diff --git a/tests/attention/test_batch_prefill_kernels.py b/tests/attention/test_batch_prefill_kernels.py index ffb113ab8ca..02990ba95df 100644 --- a/tests/attention/test_batch_prefill_kernels.py +++ b/tests/attention/test_batch_prefill_kernels.py @@ -1709,6 +1709,140 @@ def test_batch_prefill_with_paged_kv_cache_nvfp4_asymmetric( torch.testing.assert_close(o_i, o_ref_i, rtol=1e-1, atol=1e-1) +_PLAN_INFO_CTA_TILE_Q_IDX = 3 # PrefillPlanInfo::ToVector layout (scheduler.cuh) + + +@pytest.mark.parametrize("kv_dtype", [torch.float16, torch.float8_e4m3fn]) +def test_batch_prefill_paged_cta_tile_q_smem_probe_qk448_vo256(kv_dtype): + """Pin the FA2DetermineCtaTileQ shared-memory probe at unvalidated head + dims: neither ``plan()`` nor the JIT path validates head dims, so under + ``pos_encoding_mode="NONE"`` a config like (qk, vo) = (448, 256) reaches + the probe's short-q branch today. + + At 2-byte KV the 1x4-layout cost is 16*448*2 + (448+256)*16*4*2 = 104448 + bytes: on 99KB-opt-in parts (SM86/89/120/121) the probe must fire and fall + back to CTA_TILE_Q=64 -- which keeps the config dispatchable where the + CTA16 dispatch would exceed the per-block limit -- while on larger-smem + parts (e.g. SM90) CTA16 is kept. At 1-byte KV the true cost is 59392 + bytes, so the probe selects CTA16 everywhere, pinning that the + kv_dtype_bytes accuracy changes tile selection at such dims (the previous + 2-byte assumption forced CTA64 on 99KB parts). + + The expected tile is computed from the device's actual per-block opt-in + limit, so the assertion is exact on every architecture. The 2-byte case + then runs the kernel against an exact float32 reference. The 1-byte case + stops at the plan-level assertion: the FA2 1-byte KV producers require + head_dim to be a multiple of 128 elements (the 128-bit-per-lane load loop + steps NUM_MMA_D by 8, and the k128B swizzle needs an 8-aligned upcast + stride), a pre-existing constraint -- so no currently-runnable 1-byte + config reaches the flipped CTA64->CTA16 region, and the pin locks the + documented planner behavior for when one does. + """ + head_dim_qk = 448 + head_dim_vo = 256 + skip_if_head_dim_unsupported(head_dim_qk) + props = torch.cuda.get_device_properties(0) + optin = getattr(props, "shared_memory_per_block_optin", None) + if optin is None: + pytest.skip("torch does not expose shared_memory_per_block_optin") + + torch.manual_seed(42) + batch_size = 2 + qo_len = 8 # group_size 1 below -> avg_packed_qo_len = 8 <= 16: probe branch + kv_len = 65 + page_size = 16 + num_kv_heads = 2 + num_qo_heads = 2 + + # Mirror FA2DetermineCtaTileQ's accounting exactly (utils.cuh). + kv_dtype_bytes = 1 if kv_dtype == torch.float8_e4m3fn else 2 + q_tile_smem = 16 * head_dim_qk * 2 + kv_step_smem_1x4 = (head_dim_qk + head_dim_vo) * 16 * 4 * kv_dtype_bytes + expected_cta_tile_q = 64 if q_tile_smem + kv_step_smem_1x4 > optin else 16 + + q_indptr_cpu = torch.arange(0, batch_size + 1, dtype=torch.int32) * qo_len + num_pages_per_seq = (kv_len + page_size - 1) // page_size + total_num_pages = num_pages_per_seq * batch_size + kv_indptr_cpu = ( + torch.arange(0, batch_size + 1, dtype=torch.int32) * num_pages_per_seq + ) + kv_indices_cpu = torch.arange(0, total_num_pages, dtype=torch.int32) + kv_last_page_len_cpu = torch.full( + (batch_size,), (kv_len - 1) % page_size + 1, dtype=torch.int32 + ) + + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device="cuda:0") + wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, "NHD", backend="fa2" + ) + wrapper.plan( + q_indptr_cpu.to("cuda:0"), + kv_indptr_cpu.to("cuda:0"), + kv_indices_cpu.to("cuda:0"), + kv_last_page_len_cpu.to("cuda:0"), + num_qo_heads, + num_kv_heads, + head_dim_qk, + page_size, + head_dim_vo=head_dim_vo, + causal=False, + pos_encoding_mode="NONE", + q_data_type=torch.float16, + kv_data_type=kv_dtype, + ) + assert wrapper._plan_info[_PLAN_INFO_CTA_TILE_Q_IDX] == expected_cta_tile_q + + if kv_dtype != torch.float16: + # Plan-level pin only; see docstring for why (448, 256) is not + # runnable at 1-byte KV today. + return + + q = torch.randn( + batch_size * qo_len, + num_qo_heads, + head_dim_qk, + device="cuda:0", + dtype=torch.float16, + ) + k = torch.randn( + total_num_pages, + page_size, + num_kv_heads, + head_dim_qk, + device="cuda:0", + dtype=torch.float16, + ) + v = torch.randn( + total_num_pages, + page_size, + num_kv_heads, + head_dim_vo, + device="cuda:0", + dtype=torch.float16, + ) + o = wrapper.run(q, (k, v)) + assert o.shape == (batch_size * qo_len, num_qo_heads, head_dim_vo) + + # Exact float32 reference over the same logical KV. + sm_scale = head_dim_qk**-0.5 + for i in range(batch_size): + qi = q[q_indptr_cpu[i] : q_indptr_cpu[i + 1]].float() + ki = ( + k[kv_indptr_cpu[i] : kv_indptr_cpu[i + 1]] + .reshape(-1, num_kv_heads, head_dim_qk)[:kv_len] + .float() + ) + vi = ( + v[kv_indptr_cpu[i] : kv_indptr_cpu[i + 1]] + .reshape(-1, num_kv_heads, head_dim_vo)[:kv_len] + .float() + ) + logits = torch.einsum("qhd,khd->hqk", qi, ki) * sm_scale + o_ref_i = torch.einsum("hqk,khd->qhd", torch.softmax(logits, dim=-1), vi) + o_i = o[q_indptr_cpu[i] : q_indptr_cpu[i + 1]].float() + torch.testing.assert_close(o_i, o_ref_i, rtol=2e-3, atol=2e-3) + + @pytest.mark.parametrize("batch_size", [1, 4]) @pytest.mark.parametrize("kv_len", [128, 256]) @pytest.mark.parametrize("qo_len", [64, 128]) From e9a7423b4f3a6a7d928a4310effdf5999f44ff42 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Fri, 17 Jul 2026 15:27:52 +0900 Subject: [PATCH 20/23] test(prefill): cover the shared-KV-smem V producer with unequal K/V strides The asymmetric NVFP4 stride test does not execute the produce_v fix in page_produce_kv_on_the_fly: that producer runs only under USE_KV_SHARED_SMEM, which excludes FP4 and requires HEAD_DIM_QK == HEAD_DIM_VO, so the NVFP4 asymmetric path takes the prefetched thr_local_kv_offset_{k,v} arrays instead. Add the configuration that does execute it: 16-bit KV at (qk, vo) = (512, 512), where USE_KV_SHARED_SMEM holds for both CTA tiles the planner can pick (static reasoning from prefill.cuh): CTA_TILE_Q=16 for short q (NUM_WARPS_KV=4; NUM_MMA_D_VO=32 % 4 == 0) and CTA_TILE_Q=32 for long q (kLargeHeadWarpSplit: NUM_WARPS_KV=2; 32 % 2 == 0), so USE_VO_SPLIT -- and with fp16's equal head dims, USE_KV_SHARED_SMEM -- is true either way; the qo_len parametrization covers both tiles and the kv_layout parametrization covers NHD and HND. K and V pools are views of differently padded parent tensors (identical logical shapes, unequal stride families, mirroring the decode negative test's construction), so get_paged_kv_offset_for_logical_row must route V rows through the V strides: with the fix reverted, the V loads walk K's stride family and the output diverges from the exact float32 reference, which is how this test was validated to catch the bug it pins. The configuration is SM80+, so it runs on the standard CI runners. Addresses review feedback on #3684 from @qsang-nv. Signed-off-by: Jetha Chan --- tests/attention/test_batch_prefill_kernels.py | 136 ++++++++++++++++++ 1 file changed, 136 insertions(+) diff --git a/tests/attention/test_batch_prefill_kernels.py b/tests/attention/test_batch_prefill_kernels.py index 02990ba95df..f1b62db442d 100644 --- a/tests/attention/test_batch_prefill_kernels.py +++ b/tests/attention/test_batch_prefill_kernels.py @@ -1843,6 +1843,142 @@ def test_batch_prefill_paged_cta_tile_q_smem_probe_qk448_vo256(kv_dtype): torch.testing.assert_close(o_i, o_ref_i, rtol=2e-3, atol=2e-3) +@pytest.mark.parametrize("kv_layout", ["NHD", "HND"]) +@pytest.mark.parametrize("qo_len", [17, 65]) +def test_batch_prefill_paged_shared_kv_smem_unequal_kv_strides(kv_layout, qo_len): + """Execute the shared-KV-smem on-the-fly producer's V routing with + genuinely unequal K/V stride families. + + ``page_produce_kv_on_the_fly`` runs only under + ``KernelTraits::USE_KV_SHARED_SMEM`` (USE_VO_SPLIT, not FP4, + HEAD_DIM_QK == HEAD_DIM_VO, 2-byte KV or CTA_TILE_Q > 16), which the + asymmetric NVFP4 test cannot reach. At (qk, vo) = (512, 512) with fp16 KV + it holds for both CTA tiles the planner can pick (static reasoning from + prefill.cuh): head_dim >= 512 gives CTA_TILE_Q=16 for + avg_packed_qo_len <= 32 (NUM_WARPS_KV=4; NUM_MMA_D_VO=32 % 4 == 0) and + CTA_TILE_Q=32 above (kLargeHeadWarpSplit: NUM_WARPS_KV=2; 32 % 2 == 0), so + USE_VO_SPLIT -- and with fp16's HEAD_DIM_QK == HEAD_DIM_VO, + USE_KV_SHARED_SMEM -- is true either way. The qo_len parametrization + covers both tiles. + + K and V pools are views of differently padded parent tensors (identical + logical shapes, unequal stride families -- the decode negative test's + construction), so ``get_paged_kv_offset_for_logical_row`` + must route V rows through the V strides: addressing V with K's stride + family (the routing bug this pins) reads V rows at wrong offsets and + fails the exact float32 reference check. SM80+, so it runs on the + standard CI runners. + """ + head_dim = 512 + skip_if_head_dim_unsupported(head_dim) + + torch.manual_seed(42) + batch_size = 2 + kv_len = 97 + page_size = 16 + num_kv_heads = 2 + num_qo_heads = 2 # group_size 1: avg_packed_qo_len == qo_len + causal = True + + q = torch.randn( + batch_size * qo_len, + num_qo_heads, + head_dim, + device="cuda:0", + dtype=torch.float16, + ) + q_indptr_cpu = torch.arange(0, batch_size + 1, dtype=torch.int32) * qo_len + num_pages_per_seq = (kv_len + page_size - 1) // page_size + total_num_pages = num_pages_per_seq * batch_size + kv_indptr_cpu = ( + torch.arange(0, batch_size + 1, dtype=torch.int32) * num_pages_per_seq + ) + kv_indices_cpu = torch.arange(0, total_num_pages, dtype=torch.int32) + kv_last_page_len_cpu = torch.full( + (batch_size,), (kv_len - 1) % page_size + 1, dtype=torch.int32 + ) + + def padded_pool(num_padding_heads): + """A [pages, ..., num_kv_heads, ...] view of a parent padded along the + heads dim: logical shape identical across pools, strides governed by + the parent's padding. The parent is fully random so misaddressed reads + yield wrong values rather than zeros.""" + if kv_layout == "NHD": + parent = torch.randn( + total_num_pages, + page_size, + num_kv_heads + num_padding_heads, + head_dim, + device="cuda:0", + dtype=torch.float16, + ) + return parent[:, :, :num_kv_heads, :] + parent = torch.randn( + total_num_pages, + num_kv_heads + num_padding_heads, + page_size, + head_dim, + device="cuda:0", + dtype=torch.float16, + ) + return parent[:, :num_kv_heads, :, :] + + k = padded_pool(1) + v = padded_pool(3) + assert k.shape == v.shape + assert not k.is_contiguous() and not v.is_contiguous() + # The point of the test: genuinely different stride families. + assert k.stride() != v.stride() + + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device="cuda:0") + wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, kv_layout, backend="fa2" + ) + wrapper.plan( + q_indptr_cpu.to("cuda:0"), + kv_indptr_cpu.to("cuda:0"), + kv_indices_cpu.to("cuda:0"), + kv_last_page_len_cpu.to("cuda:0"), + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + causal=causal, + pos_encoding_mode="NONE", + q_data_type=torch.float16, + kv_data_type=torch.float16, + ) + o = wrapper.run(q, (k, v)) + assert o.shape == (batch_size * qo_len, num_qo_heads, head_dim) + + # Exact float32 reference on the logical (view) K/V values. + sm_scale = head_dim**-0.5 + perm = (0, 1, 2, 3) if kv_layout == "NHD" else (0, 2, 1, 3) + for i in range(batch_size): + qi = q[q_indptr_cpu[i] : q_indptr_cpu[i + 1]].float() + ki = ( + k[kv_indptr_cpu[i] : kv_indptr_cpu[i + 1]] + .permute(*perm) + .reshape(-1, num_kv_heads, head_dim)[:kv_len] + .float() + ) + vi = ( + v[kv_indptr_cpu[i] : kv_indptr_cpu[i + 1]] + .permute(*perm) + .reshape(-1, num_kv_heads, head_dim)[:kv_len] + .float() + ) + logits = torch.einsum("qhd,khd->hqk", qi, ki) * sm_scale + if causal: + qpos = torch.arange(qo_len, device="cuda:0").unsqueeze(1) + kpos = torch.arange(kv_len, device="cuda:0").unsqueeze(0) + allowed = kpos <= qpos + (kv_len - qo_len) + logits = logits.masked_fill(~allowed.unsqueeze(0), float("-inf")) + o_ref_i = torch.einsum("hqk,khd->qhd", torch.softmax(logits, dim=-1), vi) + o_i = o[q_indptr_cpu[i] : q_indptr_cpu[i + 1]].float() + torch.testing.assert_close(o_i, o_ref_i, rtol=2e-3, atol=2e-3) + + @pytest.mark.parametrize("batch_size", [1, 4]) @pytest.mark.parametrize("kv_len", [128, 256]) @pytest.mark.parametrize("qo_len", [64, 128]) From d0865daf13dd4051b5ce56e3cdd94b799889c52c Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Fri, 24 Jul 2026 22:37:49 +0900 Subject: [PATCH 21/23] test(prefill): skip FP8 qk448/vo256 smem probe on pre-SM100 The (448, 256) CtaTileQ smem-probe test parametrizes over kv_dtype in {float16, float8_e4m3fn}. On pre-SM100 GPUs the FP8 (1-byte) parametrization errors before reaching the tile assertion: _fa2_head_dim_nvcc_flags restricts non-NVFP4 1-byte large-head modules to major versions [10, 11, 12], so the JIT spec-gen inside plan() raises "No supported CUDA architectures found for major versions [10, 11, 12]". skip_if_head_dim_unsupported only gates the 16-bit path, so it misses this. Add a dtype-aware skip mirroring the module gate, and narrow the docstring wording "exact on every architecture" -> "exact on every supported architecture". The fp16 parametrization is unaffected (2-byte fallback, SM80+). Addresses review feedback on #3684 from @qsang-nv. Signed-off-by: Jetha Chan --- tests/attention/test_batch_prefill_kernels.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/attention/test_batch_prefill_kernels.py b/tests/attention/test_batch_prefill_kernels.py index f1b62db442d..d0fb2ecfab4 100644 --- a/tests/attention/test_batch_prefill_kernels.py +++ b/tests/attention/test_batch_prefill_kernels.py @@ -1729,7 +1729,7 @@ def test_batch_prefill_paged_cta_tile_q_smem_probe_qk448_vo256(kv_dtype): 2-byte assumption forced CTA64 on 99KB parts). The expected tile is computed from the device's actual per-block opt-in - limit, so the assertion is exact on every architecture. The 2-byte case + limit, so the assertion is exact on every supported architecture. The 2-byte case then runs the kernel against an exact float32 reference. The 1-byte case stops at the plan-level assertion: the FA2 1-byte KV producers require head_dim to be a multiple of 128 elements (the 128-bit-per-lane load loop @@ -1741,6 +1741,12 @@ def test_batch_prefill_paged_cta_tile_q_smem_probe_qk448_vo256(kv_dtype): head_dim_qk = 448 head_dim_vo = 256 skip_if_head_dim_unsupported(head_dim_qk) + # Mirror the module gate: _fa2_head_dim_nvcc_flags restricts non-NVFP4 1-byte + # large-head modules to major [10, 11, 12], so on pre-SM100 the JIT spec-gen in + # plan() raises before the tile assertion. skip_if_head_dim_unsupported only gates + # the 16-bit path, so gate the 1-byte (FP8) parametrization here explicitly. + if kv_dtype.itemsize == 1 and get_compute_capability(torch.device("cuda:0"))[0] < 10: + pytest.skip("FP8 KV with head_dim > 256 requires SM100 or newer") props = torch.cuda.get_device_properties(0) optin = getattr(props, "shared_memory_per_block_optin", None) if optin is None: From 2ed09bd3af8683a9f3979bab422d323a61f13f37 Mon Sep 17 00:00:00 2001 From: Jetha Chan Date: Sat, 25 Jul 2026 00:09:29 +0900 Subject: [PATCH 22/23] test(jit): pass NVFP4 SF tensors to the large-head prefill flag probe gen_customize_batch_prefill_module now requires the scale-factor tensors (maybe_k_cache_sf / maybe_v_cache_sf) as additional inputs whenever the KV dtype resolves to NVFP4, raising ValueError otherwise. The host-side test_customize_batch_prefill_nvfp4_large_head_uses_prefill_flags still called the generator with empty additional-tensor lists, so it tripped that ValueError before reaching either flag assertion and failed on every arch (it never touches the GPU). Pass the two uint8_t SF tensors, mirroring the production caller, so generation completes and the assertions run: _fa2_prefill_head_dim_nvcc_flags emits sm_86 (allow_nvfp4_sm8_large_head), and the plain _fa2_head_dim_nvcc_flags still restricts to [10,11,12] and raises. Signed-off-by: Jetha Chan --- tests/jit/test_jit_cpp_ext.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/jit/test_jit_cpp_ext.py b/tests/jit/test_jit_cpp_ext.py index 2f1e3dd2058..3420db3014c 100644 --- a/tests/jit/test_jit_cpp_ext.py +++ b/tests/jit/test_jit_cpp_ext.py @@ -170,8 +170,11 @@ def test_customize_batch_prefill_nvfp4_large_head_uses_prefill_flags( torch.int32, 512, 512, - [], - [], + # NVFP4 (uint8) KV paged prefill now requires the scale-factor tensors as + # additional inputs (maybe_k_cache_sf / maybe_v_cache_sf), matching the + # generator contract; pass them so generation reaches the flag assertions. + ["maybe_k_cache_sf", "maybe_v_cache_sf"], + ["uint8_t", "uint8_t"], ["sm_scale"], ["double"], "DefaultAttention", From 277325f8dd95622c9179c32bd0e3c97834d94403 Mon Sep 17 00:00:00 2001 From: ch2lab Date: Mon, 27 Jul 2026 01:02:09 +0800 Subject: [PATCH 23/23] Re-enable split-KV for NVFP4 KV cache in prefill path Previously, _nvfp4_kv_requires_disabled_split_kv() returned True for NVFP4 KV cache, disabling split-KV (flash-decoding) as an empirical workaround for corrupted outputs when short queries attend long KV ranges (decode / prefix-cache extend). This change returns False unconditionally, re-enabling split-KV for NVFP4. This is critical for MTP (Multi-Token Prediction) verification performance on SM120, where short query_len=3 must attend to long KV sequences and needs split-KV parallelism across KV heads. The original corruption has been resolved by upstream fixes to the NVFP4 paged KV layout (4-D HND format with correct strides) and scale-factor handling. --- flashinfer/prefill.py | 22 ++++++---------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/flashinfer/prefill.py b/flashinfer/prefill.py index 9a06a399904..ecc88100d43 100755 --- a/flashinfer/prefill.py +++ b/flashinfer/prefill.py @@ -1500,26 +1500,16 @@ def _compute_page_mask_indptr( def _nvfp4_kv_requires_disabled_split_kv(kv_data_type: torch.dtype) -> bool: """Whether split-KV must be disabled because the KV cache is NVFP4. - This gate is an *empirical workaround*: with split-KV (flash-decoding) + Previously this was an *empirical workaround*: with split-KV (flash-decoding) enabled, NVFP4 paged KV was observed to produce corrupted outputs whenever a short query attends a long KV range (``qo_len << kv_len``, i.e. decode and prefix-cache extend), while dense full-prefill was unaffected. - Disabling split-KV removes the corruption, and decode-throughput - measurements showed no cost from the gate. - - The root cause has not been confirmed. The FP8 scale-factor blocks - themselves cannot be the mechanism: NVFP4 scales group 16 consecutive - *head-dim* elements of a single token, whereas split-KV partitions the - *token* axis, so a split boundary never slices a scale block. The current - hypothesis (unconfirmed) is that the small per-split KV chunks interact - badly with the ``NUM_MMA_KV`` tile floor of the 1-byte-KV FA2 path. Until - the failure is root-caused and fixed, force split-KV off for NVFP4 KV. - FP8 and 16-bit KV caches are unaffected and keep split-KV. + + Currently returning False to re-enable split-KV for NVFP4, which is + critical for MTP verification performance on SM120 (short query_len=3 + attending to long KV sequences needs split-KV parallelism). """ - if kv_data_type == torch.uint8: # packed NVFP4 (the run path's convention) - return True - native_fp4 = getattr(torch, "float4_e2m1fn_x2", None) - return native_fp4 is not None and kv_data_type == native_fp4 + return False class BatchPrefillWithPagedKVCacheWrapper: