diff --git a/tests/kernels/attention/test_flashinfer.py b/tests/kernels/attention/test_flashinfer.py index 109f59bc8508..36c0b9bd2663 100644 --- a/tests/kernels/attention/test_flashinfer.py +++ b/tests/kernels/attention/test_flashinfer.py @@ -2,10 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from itertools import product +from types import SimpleNamespace + import pytest from vllm.platforms import current_platform -from vllm.utils.torch_utils import set_random_seed +from vllm.utils.torch_utils import ( + nvfp4_kv_cache_full_dim, + nvfp4_kv_cache_split_views, + set_random_seed, +) try: import flashinfer @@ -16,6 +23,15 @@ ) import torch +from flashinfer import get_seq_lens + +from vllm.platforms.interface import DeviceCapability +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheLayout, + KVQuantMode, + compute_layer_kv_cache_shape_bytes, +) NUM_HEADS = [(32, 8), (6, 1)] HEAD_SIZES = [128, 256] @@ -26,6 +42,19 @@ SLIDING_WINDOWS = [None, 64] +_TEST_KV_LAYOUTS = {"NHD": KVCacheLayout.LBNHC, "HND": KVCacheLayout.LBHNC} + + +def _patch_impl_kv_cache_layout(monkeypatch, flashinfer_backend, name: str): + """The impl reads its layout from cache_config via a property now; the + module-level get_kv_cache_layout() these tests used to patch is gone.""" + monkeypatch.setattr( + flashinfer_backend.FlashInferImpl, + "kv_cache_layout", + property(lambda self: _TEST_KV_LAYOUTS[name]), + ) + + def ref_paged_attn( query: torch.Tensor, key_cache: torch.Tensor, @@ -146,6 +175,1315 @@ def _make_cg_decode_wrapper( ) +def test_flashinfer_backend_accepts_nvfp4_kv_cache() -> None: + from vllm.v1.attention.backends.flashinfer import FlashInferBackend + + invalid_reasons = FlashInferBackend.validate_configuration( + head_size=128, + dtype=torch.bfloat16, + kv_cache_dtype="nvfp4", + block_size=16, + use_mla=False, + has_sink=False, + use_sparse=False, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=DeviceCapability(8, 6), + attn_type="decoder", + ) + + assert invalid_reasons == [] + + +def _make_flashinfer_q_dtype_builder( + *, + cache_dtype: str, + model_dtype: torch.dtype = torch.bfloat16, + disable_q_quantization: bool = False, +): + from vllm.v1.attention.backends.flashinfer import FlashInferMetadataBuilder + + builder = FlashInferMetadataBuilder.__new__(FlashInferMetadataBuilder) + builder.cache_dtype = cache_dtype + builder.is_kvcache_nvfp4 = cache_dtype == "nvfp4" + builder.kv_cache_dtype = model_dtype + builder.kv_cache_spec = SimpleNamespace(dtype=model_dtype) + builder.model_config = SimpleNamespace(dtype=model_dtype) + builder.vllm_config = SimpleNamespace( + attention_config=SimpleNamespace( + disable_flashinfer_q_quantization=disable_q_quantization + ) + ) + return builder + + +@pytest.mark.parametrize("is_prefill", [True, False]) +def test_flashinfer_nvfp4_native_q_dtype_uses_model_dtype( + is_prefill: bool, +) -> None: + builder = _make_flashinfer_q_dtype_builder(cache_dtype="nvfp4") + + q_dtype = builder.get_q_data_type( + is_prefill=is_prefill, + use_trtllm_gen=False, + ) + + assert q_dtype == torch.bfloat16 + + +@pytest.mark.parametrize("is_prefill", [True, False]) +def test_flashinfer_nvfp4_trtllm_gen_q_dtype_uses_fp8(is_prefill: bool) -> None: + from vllm.v1.attention.backends.flashinfer import FP8_DTYPE + + builder = _make_flashinfer_q_dtype_builder(cache_dtype="nvfp4") + + q_dtype = builder.get_q_data_type( + is_prefill=is_prefill, + use_trtllm_gen=True, + ) + + assert q_dtype == FP8_DTYPE + + +def test_flashinfer_q_quantization_disable_overrides_nvfp4_trtllm_gen() -> None: + builder = _make_flashinfer_q_dtype_builder( + cache_dtype="nvfp4", + disable_q_quantization=True, + ) + + q_dtype = builder.get_q_data_type( + is_prefill=False, + use_trtllm_gen=True, + ) + + assert q_dtype == torch.bfloat16 + + +def _nvfp4_spec(head_size: int, head_size_v: int) -> FullAttentionSpec: + return FullAttentionSpec( + block_size=16, + num_kv_heads=2, + head_size=head_size, + head_size_v=head_size_v, + dtype=torch.bfloat16, + kv_quant_mode=KVQuantMode.NVFP4, + ) + + +def test_flashinfer_nvfp4_mixed_head_spec_packs_k_and_v_in_one_slot() -> None: + from vllm.utils.torch_utils import get_dtype_size + from vllm.v1.attention.backends.flashinfer import FlashInferBackend + + spec = FlashInferBackend.customize_spec(_nvfp4_spec(256, 512)) + + assert spec.num_head_slots == 2 + assert spec.state_content_bytes == ( + nvfp4_kv_cache_full_dim(256) + nvfp4_kv_cache_full_dim(512) + ) * get_dtype_size(torch.bfloat16) + + +def test_flashinfer_nvfp4_same_head_spec_uses_separate_head_groups() -> None: + from vllm.utils.torch_utils import get_dtype_size + from vllm.v1.attention.backends.flashinfer import FlashInferBackend + + spec = FlashInferBackend.customize_spec(_nvfp4_spec(256, 256)) + + assert spec.num_head_slots == 4 + assert spec.state_content_bytes == nvfp4_kv_cache_full_dim(256) * get_dtype_size( + torch.bfloat16 + ) + + +@pytest.mark.parametrize( + ("quant_mode_name", "head_size_v", "expected_shape"), + [ + pytest.param("NONE", None, (1392, 4, 32, 2 * 256 * 2), id="auto"), + pytest.param("FP8_PER_TENSOR", None, (1392, 4, 32, 2 * 256 * 2), id="fp8"), + pytest.param( + "NVFP4", + 256, + (1392, 8, 32, nvfp4_kv_cache_full_dim(256) * 2), + id="nvfp4-same-head", + ), + pytest.param( + "NVFP4", + 128, + ( + 1392, + 4, + 32, + (nvfp4_kv_cache_full_dim(256) + nvfp4_kv_cache_full_dim(128)) * 2, + ), + id="nvfp4-mixed-head", + ), + ], +) +def test_flashinfer_kv_cache_byte_shape( + quant_mode_name: str, + head_size_v: int | None, + expected_shape: tuple[int, ...], +) -> None: + """NVFP4 splits K and V across head slots, or packs both into one slot when + their head sizes differ. The trailing dimension counts bytes.""" + from vllm.v1.attention.backends.flashinfer import FlashInferBackend + + spec = FlashInferBackend.customize_spec( + FullAttentionSpec( + block_size=32, + num_kv_heads=4, + head_size=256, + head_size_v=head_size_v, + dtype=torch.bfloat16, + kv_quant_mode=KVQuantMode[quant_mode_name], + ) + ) + + assert compute_layer_kv_cache_shape_bytes(spec, 1392) == expected_shape + + +def test_flashinfer_cascade_passes_kv_tuple(monkeypatch) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + impl = flashinfer_backend.FlashInferImpl.__new__(flashinfer_backend.FlashInferImpl) + impl.bmm1_scale = None + impl.bmm2_scale = None + impl.scale = 1.0 + impl.kv_cache_dtype = "auto" + impl.is_kvcache_nvfp4 = False + impl.head_size = 8 + impl.num_kv_heads = 1 + _patch_impl_kv_cache_layout(monkeypatch, flashinfer_backend, "NHD") + + seen_kv_caches: list[tuple[torch.Tensor, torch.Tensor]] = [] + + def run_cascade( + query: torch.Tensor, paged_kv_cache: tuple[torch.Tensor, torch.Tensor] + ) -> torch.Tensor: + seen_kv_caches.append(paged_kv_cache) + return torch.zeros_like(query) + + attn_metadata = SimpleNamespace( + num_actual_tokens=2, + prefill=None, + decode=None, + use_cascade=True, + cascade_wrapper=SimpleNamespace(run=run_cascade), + ) + query = torch.ones((2, 1, 8)) + key = torch.ones_like(query) + value = torch.ones_like(query) + kv_cache = torch.empty((3, 1, 4, 16)) + output = torch.ones_like(query) + + result = impl.forward( + layer=SimpleNamespace(), + query=query, + key=key, + value=value, + kv_cache=kv_cache, + attn_metadata=attn_metadata, + output=output, + ) + + assert len(seen_kv_caches) == 1 + k_cache, v_cache = seen_kv_caches[0] + assert k_cache.shape == (3, 4, 1, 8) + assert v_cache.shape == (3, 4, 1, 8) + assert k_cache.untyped_storage().data_ptr() == kv_cache.untyped_storage().data_ptr() + assert v_cache.untyped_storage().data_ptr() == kv_cache.untyped_storage().data_ptr() + assert torch.count_nonzero(result) == 0 + + +def _storage_offsets(tensor: torch.Tensor) -> set[int]: + return { + tensor.storage_offset() + + sum(idx * stride for idx, stride in zip(indices, tensor.stride())) + for indices in product(*(range(size) for size in tensor.shape)) + } + + +@pytest.mark.parametrize( + "shape", + [ + pytest.param((2, 4, 3), id="NHD"), + pytest.param((2, 3, 4), id="HND"), + ], +) +def test_nvfp4_kv_cache_split_views_mixed_packed_layout( + shape: tuple[int, int, int], +) -> None: + head_size = 64 + head_size_v = 128 + k_full_dim = nvfp4_kv_cache_full_dim(head_size) + v_full_dim = nvfp4_kv_cache_full_dim(head_size_v) + num_pages, dim_1, dim_2 = shape + kv_cache = torch.empty( + num_pages, dim_1, dim_2, k_full_dim + v_full_dim, dtype=torch.uint8 + ) + + (k_data, v_data), (k_scales, v_scales) = nvfp4_kv_cache_split_views( + kv_cache, head_size, head_size_v + ) + + assert k_data.shape == (num_pages, dim_1, dim_2, head_size // 2) + assert k_scales.shape == (num_pages, dim_1, dim_2, head_size // 16) + assert v_data.shape == (num_pages, dim_1, dim_2, head_size_v // 2) + assert v_scales.shape == (num_pages, dim_1, dim_2, head_size_v // 16) + + page_items = dim_1 * dim_2 + assert k_data.storage_offset() == kv_cache.storage_offset() + assert k_scales.storage_offset() == kv_cache.storage_offset() + page_items * ( + head_size // 2 + ) + assert v_data.storage_offset() == kv_cache.storage_offset() + page_items * ( + k_full_dim + ) + assert v_scales.storage_offset() == kv_cache.storage_offset() + page_items * ( + k_full_dim + head_size_v // 2 + ) + + offset_sets = [ + _storage_offsets(k_data), + _storage_offsets(k_scales), + _storage_offsets(v_data), + _storage_offsets(v_scales), + ] + assert len(set().union(*offset_sets)) == sum( + len(offsets) for offsets in offset_sets + ) + + +def test_nvfp4_kv_cache_split_views_rejects_incompatible_strides() -> None: + head_size = 64 + full_dim = nvfp4_kv_cache_full_dim(head_size) + storage = torch.empty(1000, dtype=torch.uint8) + kv_side = torch.as_strided( + storage, + (2, 4, 3, full_dim), + (500, 100, 37, 1), + ) + + with pytest.raises(ValueError, match="strides are not compatible"): + nvfp4_kv_cache_split_views(kv_side, head_size) + + +def test_nvfp4_kv_cache_split_views_accepts_size_one_strides() -> None: + head_size = 64 + full_dim = nvfp4_kv_cache_full_dim(head_size) + storage = torch.empty(256, dtype=torch.uint8) + kv_side = torch.as_strided( + storage, + (2, 1, 3, full_dim), + (128, 7, full_dim, 1), + ) + + (data,), (scale,) = nvfp4_kv_cache_split_views(kv_side, head_size) + + assert data.shape == (2, 1, 3, head_size // 2) + assert scale.shape == (2, 1, 3, head_size // 16) + + +def test_nvfp4_kv_cache_split_views_rejects_mixed_incompatible_strides() -> None: + head_size = 64 + head_size_v = 128 + k_full_dim = nvfp4_kv_cache_full_dim(head_size) + v_full_dim = nvfp4_kv_cache_full_dim(head_size_v) + full_dim = k_full_dim + v_full_dim + storage = torch.empty(1000, dtype=torch.uint8) + kv_cache = torch.as_strided( + storage, + (2, 4, 3, full_dim), + (500, 100, 37, 1), + ) + + with pytest.raises(ValueError, match="strides are not compatible"): + nvfp4_kv_cache_split_views(kv_cache, head_size, head_size_v) + + +def test_nvfp4_kv_cache_split_views_rejects_mixed_side_slice() -> None: + head_size = 64 + head_size_v = 128 + k_full_dim = nvfp4_kv_cache_full_dim(head_size) + v_full_dim = nvfp4_kv_cache_full_dim(head_size_v) + kv_cache = torch.empty((2, 4, 3, k_full_dim + v_full_dim), dtype=torch.uint8) + + with pytest.raises(ValueError, match="strides are not compatible"): + nvfp4_kv_cache_split_views(kv_cache[..., :k_full_dim], head_size) + + +def test_nvfp4_kv_cache_split_views_rejects_non_inferable_dim() -> None: + kv_side = torch.empty((2, 4, 3, 10), dtype=torch.uint8) + + with pytest.raises(ValueError, match="last dimension cannot be inferred"): + nvfp4_kv_cache_split_views(kv_side) + + +@pytest.mark.parametrize( + ("head_dim", "head_dim_v", "uses_alibi", "expected_calls"), + [ + (128, 128, True, 1), + (128, 128, False, 0), + (256, 128, True, 0), + (128, 256, True, 0), + (512, 512, True, 0), + ], +) +def test_flashinfer_nvfp4_fa2_prefill_reservation_requires_matching_head_dims( + monkeypatch, + head_dim: int, + head_dim_v: int, + uses_alibi: bool, + expected_calls: int, +) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + FlashInferMetadataBuilder = flashinfer_backend.FlashInferMetadataBuilder + builder = FlashInferMetadataBuilder.__new__(FlashInferMetadataBuilder) + builder.is_kvcache_nvfp4 = True + builder.head_dim = head_dim + builder.head_dim_v = head_dim_v + builder.use_dcp = False + builder.model_config = SimpleNamespace( + max_model_len=1024, + dtype=torch.float16, + uses_alibi=uses_alibi, + ) + builder.vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace( + enable_chunked_prefill=True, + max_num_batched_tokens=8, + ) + ) + builder.num_kv_heads = 2 + + calls = [] + + class FakeWorkspaceManager: + def get_simultaneous(self, *specs): + calls.append(specs) + + monkeypatch.setattr( + flashinfer_backend, "_is_flash_attn_varlen_func_available", lambda: True + ) + monkeypatch.setattr( + flashinfer_backend.flashinfer, + "nvfp4_kv_dequantize_paged", + object(), + raising=False, + ) + monkeypatch.setattr( + flashinfer_backend, "is_workspace_manager_initialized", lambda: True + ) + monkeypatch.setattr( + flashinfer_backend, "current_workspace_manager", FakeWorkspaceManager + ) + + builder._reserve_nvfp4_fa2_prefill_workspace(can_use_trtllm=False) + + assert len(calls) == expected_calls + if expected_calls: + assert calls[0] == ( + ((1024, 2, head_dim), torch.float16), + ((1024, 2, head_dim_v), torch.float16), + ) + + +def _make_nvfp4_fa2_prefill_impl( + flashinfer_backend, + *, + head_size: int = 4, + head_size_v: int | None = None, + alibi: bool = False, +): + impl = flashinfer_backend.FlashInferImpl.__new__(flashinfer_backend.FlashInferImpl) + impl.is_kvcache_nvfp4 = True + impl._nvfp4_paged_dequant = None + impl.fa_version = 2 + impl.head_size = head_size + impl.head_size_v = head_size if head_size_v is None else head_size_v + impl.dcp_world_size = 1 + impl._nvfp4_fa2_cu_q = torch.zeros(2, dtype=torch.int32) + impl._nvfp4_fa2_cu_k = torch.zeros(2, dtype=torch.int32) + impl.alibi_slopes = torch.ones(1) if alibi else None + impl.num_kv_heads = 1 + impl.sinks = None + return impl + + +@pytest.mark.parametrize( + ("query_lens", "seq_lens"), + [ + ([15, 3], [15, 16]), + ([3, 17], [16, 17]), + ([31, 3], [31, 32]), + ([3, 33], [32, 33]), + ], +) +def test_nvfp4_fa2_prefill_mixed_runs_native_then_overwrites_first_chunks( + monkeypatch, + query_lens: list[int], + seq_lens: list[int], +) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + monkeypatch.setattr( + flashinfer_backend, "_is_flash_attn_varlen_func_available", lambda: True + ) + impl = _make_nvfp4_fa2_prefill_impl(flashinfer_backend) + events: list[tuple[str, int]] = [] + + query_start_loc = [0] + for query_len in query_lens: + query_start_loc.append(query_start_loc[-1] + query_len) + num_tokens = query_start_loc[-1] + query = torch.zeros((num_tokens, 1, 4), dtype=torch.float16) + key = torch.zeros_like(query) + value = torch.zeros_like(query) + for req_idx, start in enumerate(query_start_loc[:-1]): + end = query_start_loc[req_idx + 1] + query[start:end].fill_(req_idx + 1) + + def fake_flash_attn(**kwargs): + q = kwargs["q"] + marker = int(q[0, 0, 0].item()) + events.append(("fa", marker)) + assert kwargs["k"].shape[0] == q.shape[0] + assert kwargs["max_seqlen_q"] == q.shape[0] + assert kwargs["max_seqlen_k"] == q.shape[0] + return torch.full_like(q, marker + 10) + + impl._flash_attn_varlen = fake_flash_attn + + class FakeNativeWrapper: + def run(self, *args, **kwargs): + events.append(("native", 0)) + kwargs["out"].fill_(-1) + + wrapper = FakeNativeWrapper() + prefill = flashinfer_backend.FIPrefill( + wrapper=wrapper, + block_tables=torch.zeros((2, 3), dtype=torch.int32), + seq_lens=torch.tensor(seq_lens, dtype=torch.int32), + seq_lens_cpu=torch.tensor(seq_lens, dtype=torch.int32), + query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32), + query_start_loc_cpu=torch.tensor(query_start_loc, dtype=torch.int32), + ) + metadata = SimpleNamespace(causal=True, use_cascade=False, prefill=prefill) + layer = SimpleNamespace( + _q_scale_float=1.0, + _k_scale_float=1.0, + _v_scale_float=1.0, + ) + full_output = torch.full((num_tokens + 2, 1, 4), 99, dtype=torch.float16) + output = full_output[2:] + + handled = impl._run_nvfp4_fa2_prefill( + layer, + wrapper, + query, + key, + value, + output, + metadata, + (torch.empty(0),), + (torch.empty(0),), + ) + + first_req_idx = query_lens.index( + next(q_len for q_len, seq_len in zip(query_lens, seq_lens) if q_len == seq_len) + ) + assert handled + assert events == [("native", 0), ("fa", first_req_idx + 1)] + assert torch.all(full_output[:2] == 99) + for req_idx, start in enumerate(query_start_loc[:-1]): + end = query_start_loc[req_idx + 1] + expected = req_idx + 11 if query_lens[req_idx] == seq_lens[req_idx] else -1 + assert torch.all(output[start:end] == expected) + + +@pytest.mark.parametrize( + ("query_lens", "seq_lens", "expected_handled", "expected_fa_calls"), + [ + ([15, 17], [15, 17], True, 2), + ([3, 5], [16, 33], False, 0), + ], +) +def test_nvfp4_fa2_prefill_routes_uniform_chunk_kinds( + monkeypatch, + query_lens: list[int], + seq_lens: list[int], + expected_handled: bool, + expected_fa_calls: int, +) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + monkeypatch.setattr( + flashinfer_backend, "_is_flash_attn_varlen_func_available", lambda: True + ) + impl = _make_nvfp4_fa2_prefill_impl(flashinfer_backend) + query_start_loc = [0] + for query_len in query_lens: + query_start_loc.append(query_start_loc[-1] + query_len) + num_tokens = query_start_loc[-1] + query = torch.zeros((num_tokens, 1, 4), dtype=torch.float16) + output = torch.full_like(query, -1) + fa_calls = 0 + + def fake_flash_attn(**kwargs): + nonlocal fa_calls + fa_calls += 1 + return torch.ones_like(kwargs["q"]) + + impl._flash_attn_varlen = fake_flash_attn + + class FakeNativeWrapper: + def run(self, *args, **kwargs): + raise AssertionError("native prefill must be dispatched by the caller") + + wrapper = FakeNativeWrapper() + prefill = flashinfer_backend.FIPrefill( + wrapper=wrapper, + block_tables=torch.zeros((2, 3), dtype=torch.int32), + seq_lens=torch.tensor(seq_lens, dtype=torch.int32), + seq_lens_cpu=torch.tensor(seq_lens, dtype=torch.int32), + query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32), + query_start_loc_cpu=torch.tensor(query_start_loc, dtype=torch.int32), + ) + metadata = SimpleNamespace(causal=True, use_cascade=False, prefill=prefill) + layer = SimpleNamespace( + _q_scale_float=1.0, + _k_scale_float=1.0, + _v_scale_float=1.0, + ) + + handled = impl._run_nvfp4_fa2_prefill( + layer, + wrapper, + query, + query, + query, + output, + metadata, + (torch.empty(0),), + (torch.empty(0),), + ) + + assert handled is expected_handled + assert fa_calls == expected_fa_calls + + +@pytest.mark.parametrize( + ("query_lens", "seq_lens"), + [ + ([15, 17], [15, 17]), + ([3, 5], [16, 17]), + ([3, 15, 3], [16, 15, 17]), + ([31, 3, 33], [31, 32, 33]), + ], +) +@torch.inference_mode() +def test_nvfp4_fa2_prefill_routing_matches_legacy_dequant_fa2( + query_lens: list[int], + seq_lens: list[int], +) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + torch.manual_seed(0) + dtype = torch.bfloat16 + page_size = 16 + num_qo_heads = 4 + num_kv_heads = 1 + head_size = 128 + num_reqs = len(query_lens) + pages_per_req = [(seq_len + page_size - 1) // page_size for seq_len in seq_lens] + num_pages = sum(pages_per_req) + + kv_data = tuple( + torch.empty( + (num_pages, page_size, num_kv_heads, head_size // 2), + dtype=torch.uint8, + device="cuda", + ) + for _ in range(2) + ) + kv_scales = tuple( + torch.empty( + (num_pages, page_size, num_kv_heads, head_size // 16), + dtype=torch.float8_e4m3fn, + device="cuda", + ) + for _ in range(2) + ) + block_tables = torch.full( + (num_reqs, max(pages_per_req)), + -1, + dtype=torch.int32, + device="cuda", + ) + slot_mapping: list[int] = [] + raw_keys: list[torch.Tensor] = [] + raw_values: list[torch.Tensor] = [] + page_start = 0 + for req_idx, (seq_len, num_req_pages) in enumerate(zip(seq_lens, pages_per_req)): + req_pages = torch.arange( + page_start, + page_start + num_req_pages, + dtype=torch.int32, + device="cuda", + ) + block_tables[req_idx, :num_req_pages] = req_pages + raw_keys.append( + torch.randn((seq_len, num_kv_heads, head_size), dtype=dtype, device="cuda") + * 0.2 + ) + raw_values.append(torch.randn_like(raw_keys[-1]) * 0.2) + slot_mapping.extend( + int(req_pages[token_idx // page_size].item()) * page_size + + token_idx % page_size + for token_idx in range(seq_len) + ) + page_start += num_req_pages + + global_scale = torch.ones((), dtype=torch.float32, device="cuda") + flashinfer.nvfp4_quantize_append_paged_kv_cache_with_slot_mapping( + torch.cat(raw_keys), + torch.cat(raw_values), + torch.tensor(slot_mapping, dtype=torch.int64, device="cuda"), + kv_data, + kv_scales, + global_scale, + global_scale, + "NHD", + ) + + query_start_loc = [0] + for query_len in query_lens: + query_start_loc.append(query_start_loc[-1] + query_len) + query = ( + torch.randn( + (query_start_loc[-1], num_qo_heads, head_size), + dtype=dtype, + device="cuda", + ) + * 0.2 + ) + current_key = torch.cat( + [key[-query_len:] for key, query_len in zip(raw_keys, query_lens)] + ) + current_value = torch.cat( + [value[-query_len:] for value, query_len in zip(raw_values, query_lens)] + ) + + paged_kv_indptr = [0] + paged_kv_indices: list[int] = [] + for req_idx, num_req_pages in enumerate(pages_per_req): + paged_kv_indices.extend(block_tables[req_idx, :num_req_pages].tolist()) + paged_kv_indptr.append(len(paged_kv_indices)) + workspace = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda") + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(workspace, "NHD") + wrapper.plan( + torch.tensor(query_start_loc, dtype=torch.int32), + torch.tensor(paged_kv_indptr, dtype=torch.int32), + torch.tensor(paged_kv_indices, dtype=torch.int32, device="cuda"), + torch.tensor( + [seq_len % page_size or page_size for seq_len in seq_lens], + dtype=torch.int32, + ), + num_qo_heads, + num_kv_heads, + head_size, + page_size, + causal=True, + q_data_type=dtype, + kv_data_type=torch.uint8, + o_data_type=dtype, + ) + + impl = _make_nvfp4_fa2_prefill_impl( + flashinfer_backend, + head_size=head_size, + ) + impl.fa_version = 2 + impl.scale = head_size**-0.5 + impl.sliding_window = (-1, -1) + impl.logits_soft_cap = None + impl._nvfp4_fa2_cu_q = torch.zeros(2, dtype=torch.int32, device="cuda") + impl._nvfp4_fa2_cu_k = torch.zeros(2, dtype=torch.int32, device="cuda") + prefill = flashinfer_backend.FIPrefill( + wrapper=wrapper, + block_tables=block_tables, + seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), + seq_lens_cpu=torch.tensor(seq_lens, dtype=torch.int32), + query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32, device="cuda"), + query_start_loc_cpu=torch.tensor(query_start_loc, dtype=torch.int32), + ) + metadata = SimpleNamespace(causal=True, use_cascade=False, prefill=prefill) + layer = SimpleNamespace( + _q_scale_float=1.0, + _k_scale_float=1.0, + _v_scale_float=1.0, + ) + full_output = torch.full( + (query.shape[0] + 2, num_qo_heads, head_size), + 99, + dtype=dtype, + device="cuda", + ) + output = full_output[2:] + handled = impl._run_nvfp4_fa2_prefill( + layer, + wrapper, + query, + current_key, + current_value, + output, + metadata, + kv_data, + kv_scales, + ) + if not handled: + impl._run_native_nvfp4_prefill( + layer, + wrapper, + query, + output, + kv_data, + kv_scales, + ) + + dequant_key = torch.empty( + (num_reqs, max(seq_lens), num_kv_heads, head_size), + dtype=dtype, + device="cuda", + ) + dequant_value = torch.empty_like(dequant_key) + flashinfer.nvfp4_kv_dequantize_paged( + kv_data, + kv_scales, + block_tables, + prefill.seq_lens, + global_scale, + global_scale, + dequant_key, + dequant_value, + "NHD", + ) + reference_parts = [] + for req_idx, (query_len, seq_len) in enumerate(zip(query_lens, seq_lens)): + q_start = query_start_loc[req_idx] + q_end = query_start_loc[req_idx + 1] + if query_len == seq_len: + key = current_key[q_start:q_end] + value = current_value[q_start:q_end] + else: + key = dequant_key[req_idx, :seq_len] + value = dequant_value[req_idx, :seq_len] + reference_parts.append( + impl._flash_attn_varlen( + query[q_start:q_end], + key, + value, + torch.tensor([0, query_len], dtype=torch.int32, device="cuda"), + torch.tensor([0, seq_len], dtype=torch.int32, device="cuda"), + query_len, + seq_len, + True, + ) + ) + + reference = torch.cat(reference_parts) + torch.testing.assert_close(output, reference, atol=2e-3, rtol=2e-2) + assert torch.all(full_output[:2] == 99) + + +@pytest.mark.parametrize("guard", ["mm_prefix", "mm_wrapper", "custom_causal"]) +def test_nvfp4_fa2_prefill_keeps_custom_mask_on_native_wrapper( + monkeypatch, guard: str +) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + monkeypatch.setattr( + flashinfer_backend, "_is_flash_attn_varlen_func_available", lambda: True + ) + impl = _make_nvfp4_fa2_prefill_impl(flashinfer_backend) + + class FakeNativeWrapper: + pass + + wrapper = FakeNativeWrapper() + prefill = flashinfer_backend.FIPrefill( + wrapper=wrapper, + block_tables=torch.zeros((1, 1), dtype=torch.int32), + seq_lens=torch.tensor([1], dtype=torch.int32), + seq_lens_cpu=torch.tensor([1], dtype=torch.int32), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32), + ) + metadata = SimpleNamespace(causal=True, use_cascade=False, prefill=prefill) + if guard == "mm_prefix": + metadata.mm_prefix_range = {0: [(0, 1)]} + elif guard == "mm_wrapper": + prefill.mm_wrapper = object() + else: + metadata.causal = torch.ones(1, dtype=torch.bool) + + query = torch.zeros((1, 1, 4), dtype=torch.float16) + handled = impl._run_nvfp4_fa2_prefill( + SimpleNamespace(), + wrapper, + query, + query, + query, + torch.empty_like(query), + metadata, + (torch.empty(0),), + (torch.empty(0),), + ) + + assert not handled + + +@pytest.mark.parametrize("head_size,head_size_v", [(256, 128), (128, 256)]) +def test_nvfp4_fa2_prefill_keeps_asymmetric_heads_on_native_fallback( + monkeypatch, head_size: int, head_size_v: int +) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + monkeypatch.setattr( + flashinfer_backend, "_is_flash_attn_varlen_func_available", lambda: True + ) + impl = _make_nvfp4_fa2_prefill_impl( + flashinfer_backend, + head_size=head_size, + head_size_v=head_size_v, + ) + query = torch.zeros((1, 1, head_size), dtype=torch.float16) + key = torch.zeros_like(query) + value = torch.zeros((1, 1, head_size_v), dtype=torch.float16) + prefill = flashinfer_backend.FIPrefill( + wrapper=SimpleNamespace(), + block_tables=torch.zeros((1, 1), dtype=torch.int32), + seq_lens=torch.tensor([1], dtype=torch.int32), + seq_lens_cpu=torch.tensor([1], dtype=torch.int32), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32), + ) + + assert not impl._run_nvfp4_fa2_prefill( + SimpleNamespace(), + prefill.wrapper, + query, + key, + value, + torch.empty_like(query), + SimpleNamespace(causal=True, use_cascade=False, prefill=prefill), + (torch.empty(0),), + (torch.empty(0),), + ) + + +def test_nvfp4_fa2_prefill_keeps_alibi_on_legacy_scratch(monkeypatch) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + _patch_impl_kv_cache_layout(monkeypatch, flashinfer_backend, "NHD") + monkeypatch.setattr( + flashinfer_backend, "_is_flash_attn_varlen_func_available", lambda: True + ) + monkeypatch.setattr( + flashinfer_backend, "is_workspace_manager_initialized", lambda: True + ) + impl = _make_nvfp4_fa2_prefill_impl(flashinfer_backend, alibi=True) + events: list[str] = [] + + class FakeWorkspaceManager: + def get_simultaneous(self, *specs): + events.append("workspace") + return tuple(torch.empty(shape, dtype=dtype) for shape, dtype in specs) + + monkeypatch.setattr( + flashinfer_backend, "current_workspace_manager", FakeWorkspaceManager + ) + + def fake_paged_dequant(*args, **kwargs): + events.append("dequant") + args[6].zero_() + args[7].zero_() + + impl._nvfp4_paged_dequant = fake_paged_dequant + + def fake_flash_attn(**kwargs): + events.append("fa") + return torch.ones_like(kwargs["q"]) + + impl._flash_attn_varlen = fake_flash_attn + + class FakeNativeWrapper: + def run(self, *args, **kwargs): + raise AssertionError("ALiBi continuation must keep the legacy path") + + wrapper = FakeNativeWrapper() + prefill = flashinfer_backend.FIPrefill( + wrapper=wrapper, + block_tables=torch.zeros((1, 1), dtype=torch.int32), + seq_lens=torch.tensor([16], dtype=torch.int32), + seq_lens_cpu=torch.tensor([16], dtype=torch.int32), + query_start_loc=torch.tensor([0, 3], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 3], dtype=torch.int32), + ) + query = torch.zeros((3, 1, 4), dtype=torch.float16) + output = torch.empty_like(query) + + assert impl._run_nvfp4_fa2_prefill( + SimpleNamespace(_k_scale=1.0, _v_scale=1.0), + wrapper, + query, + query, + query, + output, + SimpleNamespace(causal=True, use_cascade=False, prefill=prefill), + (torch.empty(0),), + (torch.empty(0),), + ) + assert events == ["workspace", "dequant", "fa"] + assert torch.all(output == 1) + + +def test_flashinfer_impl_caches_nvfp4_slot_mapping_writer(monkeypatch) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + def fake_slot_writer(*args, **kwargs): + pass + + monkeypatch.setattr( + flashinfer_backend.flashinfer, + "nvfp4_quantize_append_paged_kv_cache_with_slot_mapping", + fake_slot_writer, + raising=False, + ) + monkeypatch.setattr( + flashinfer_backend.current_platform, + "is_device_capability_family", + lambda family: False, + ) + monkeypatch.setattr( + flashinfer_backend, + "can_use_trtllm_attention", + lambda num_heads, num_kv_heads, is_prefill=False: False, + ) + + impl = flashinfer_backend.FlashInferImpl( + num_heads=1, + head_size=128, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="nvfp4", + ) + + assert impl._nvfp4_slot_writer is fake_slot_writer + + +@pytest.mark.parametrize( + ("prefill_ok", "decode_ok", "expected_native"), + [ + (True, True, True), + (False, True, False), + (True, False, False), + (False, False, False), + ], +) +def test_flashinfer_impl_gates_native_nvfp4_update_on_trtllm_availability( + monkeypatch, prefill_ok: bool, decode_ok: bool, expected_native: bool +) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + def fake_slot_writer(*args, **kwargs): + pass + + def fake_can_use_trtllm_attention( + num_heads: int, num_kv_heads: int, is_prefill: bool = False + ) -> bool: + return prefill_ok if is_prefill else decode_ok + + monkeypatch.setattr( + flashinfer_backend.flashinfer, + "nvfp4_quantize_append_paged_kv_cache_with_slot_mapping", + fake_slot_writer, + raising=False, + ) + monkeypatch.setattr( + flashinfer_backend.current_platform, + "is_device_capability_family", + lambda family: family == 120, + ) + monkeypatch.setattr( + flashinfer_backend, + "can_use_trtllm_attention", + fake_can_use_trtllm_attention, + ) + + impl = flashinfer_backend.FlashInferImpl( + num_heads=1, + head_size=128, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="nvfp4", + ) + + assert impl.use_native_nvfp4_kv_cache_update is expected_native + if expected_native: + assert impl._nvfp4_slot_writer is None + else: + assert impl._nvfp4_slot_writer is fake_slot_writer + + +def test_flashinfer_dcp_prefill_forwards_nvfp4_scales(monkeypatch) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + seen: dict[str, object] = {} + + class FakeContext: + def run(self, *args, **kwargs): + seen["kv_cache"] = args[1] + seen["kv_cache_sf"] = kwargs["kv_cache_sf"] + return torch.zeros_like(args[0]), torch.zeros(args[0].shape[:2]) + + class FakeNewTokens: + def run(self, query, *args, **kwargs): + return torch.zeros_like(query), torch.zeros(query.shape[:2]) + + group = SimpleNamespace(all_gather=lambda tensor, dim: tensor) + monkeypatch.setattr(flashinfer_backend, "get_dcp_group", lambda: group) + monkeypatch.setattr(flashinfer_backend, "merge_attn_states", lambda *args: None) + + wrapper = flashinfer_backend.BatchDCPPrefillWrapper.__new__( + flashinfer_backend.BatchDCPPrefillWrapper + ) + wrapper._context = FakeContext() + wrapper._new_tokens = FakeNewTokens() + wrapper._dcp_combine = lambda output, lse, group, return_lse: (output, lse) + + query = torch.empty((2, 1, 8)) + kv_cache = (torch.empty(0), torch.empty(0)) + kv_cache_sf = (torch.empty(0), torch.empty(0)) + wrapper.run( + SimpleNamespace(_k_scale_float=1.0, _v_scale_float=1.0), + query, + kv_cache, + query, + query, + torch.empty_like(query), + kv_cache_sf=kv_cache_sf, + ) + + assert seen["kv_cache"] is kv_cache + assert seen["kv_cache_sf"] is kv_cache_sf + + +def test_flashinfer_dcp_rejects_mixed_head_nvfp4() -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + builder = flashinfer_backend.FlashInferMetadataBuilder.__new__( + flashinfer_backend.FlashInferMetadataBuilder + ) + builder.use_dcp = True + builder.is_kvcache_nvfp4 = True + builder.head_dim = 64 + builder.head_dim_v = 128 + + with pytest.raises(NotImplementedError, match="different K/V head dimensions"): + builder.build(0, SimpleNamespace()) + + +def test_flashinfer_impl_caches_nvfp4_kv_cache_views(monkeypatch) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + def fake_slot_writer(*args, **kwargs): + pass + + monkeypatch.setattr( + flashinfer_backend.flashinfer, + "nvfp4_quantize_append_paged_kv_cache_with_slot_mapping", + fake_slot_writer, + raising=False, + ) + monkeypatch.setattr( + flashinfer_backend.current_platform, + "is_device_capability_family", + lambda family: False, + ) + monkeypatch.setattr( + flashinfer_backend, + "can_use_trtllm_attention", + lambda num_heads, num_kv_heads, is_prefill=False: False, + ) + _patch_impl_kv_cache_layout(monkeypatch, flashinfer_backend, "NHD") + + head_size = 64 + head_size_v = 128 + impl = flashinfer_backend.FlashInferImpl( + num_heads=1, + head_size=head_size, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="nvfp4", + head_size_v=head_size_v, + ) + + full_dim = nvfp4_kv_cache_full_dim(head_size) + nvfp4_kv_cache_full_dim(head_size_v) + kv_cache = torch.empty((2, 1, 4, full_dim), dtype=torch.uint8) + + views = impl._get_nvfp4_kv_cache_views(kv_cache) + cached_views = impl._get_nvfp4_kv_cache_views(kv_cache) + + assert cached_views is views + assert cached_views.kv_cache is views.kv_cache + assert cached_views.data[0] is views.data[0] + assert cached_views.block_scales[0] is views.block_scales[0] + assert views.data[0].shape == (2, 4, 1, head_size // 2) + assert views.data[1].shape == (2, 4, 1, head_size_v // 2) + + rebound_kv_cache = torch.empty_like(kv_cache) + rebound_views = impl._get_nvfp4_kv_cache_views(rebound_kv_cache) + + assert rebound_views is not views + assert rebound_views.data[0].data_ptr() == rebound_kv_cache.data_ptr() + + +@pytest.mark.parametrize("cache_layout", ["NHD", "HND"]) +def test_flashinfer_impl_same_head_nvfp4_views_cover_compact_pages( + monkeypatch, cache_layout: str +) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + _patch_impl_kv_cache_layout(monkeypatch, flashinfer_backend, cache_layout) + + num_blocks = 2 + num_kv_heads = 3 + block_size = 16 + head_size = 128 + full_dim = nvfp4_kv_cache_full_dim(head_size) + logical_shape = (num_blocks, 2 * num_kv_heads, block_size, full_dim) + stride_order = (0, 2, 1, 3) if cache_layout == "NHD" else (0, 1, 2, 3) + physical_shape = tuple(logical_shape[i] for i in stride_order) + physical_cache = torch.empty(physical_shape, dtype=torch.uint8) + inverse_order = tuple(stride_order.index(i) for i in range(4)) + kv_cache = physical_cache.permute(*inverse_order) + + impl = flashinfer_backend.FlashInferImpl.__new__(flashinfer_backend.FlashInferImpl) + impl.head_size = head_size + impl.head_size_v = head_size + impl.num_kv_heads = num_kv_heads + impl.kv_cache_dtype = "nvfp4" + impl.use_native_nvfp4_kv_cache_update = False + impl._nvfp4_kv_cache_view_key = None + impl._nvfp4_kv_cache_views = None + + views = impl._get_nvfp4_kv_cache_views(kv_cache) + cached_views = impl._get_nvfp4_kv_cache_views(kv_cache) + + assert cached_views is views + if cache_layout == "NHD": + expected_prefix = (num_blocks, block_size, num_kv_heads) + else: + expected_prefix = (num_blocks, num_kv_heads, block_size) + assert views.data[0].shape == (*expected_prefix, head_size // 2) + assert views.data[1].shape == (*expected_prefix, head_size // 2) + assert views.block_scales[0].shape == (*expected_prefix, head_size // 16) + assert views.block_scales[1].shape == (*expected_prefix, head_size // 16) + + offset_sets = [ + _storage_offsets(views.data[0]), + _storage_offsets(views.block_scales[0]), + _storage_offsets(views.data[1]), + _storage_offsets(views.block_scales[1]), + ] + assert len(set().union(*offset_sets)) == sum( + len(offsets) for offsets in offset_sets + ) + assert len(set().union(*offset_sets)) == physical_cache.numel() + + +def test_flashinfer_impl_requires_nvfp4_slot_mapping_writer(monkeypatch) -> None: + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + monkeypatch.delattr( + flashinfer_backend.flashinfer, + "nvfp4_quantize_append_paged_kv_cache_with_slot_mapping", + raising=False, + ) + monkeypatch.setattr( + flashinfer_backend.current_platform, + "is_device_capability_family", + lambda family: False, + ) + monkeypatch.setattr( + flashinfer_backend, + "can_use_trtllm_attention", + lambda num_heads, num_kv_heads, is_prefill=False: False, + ) + + with pytest.raises(RuntimeError, match="NVFP4 slot-mapping KV cache update"): + flashinfer_backend.FlashInferImpl( + num_heads=1, + head_size=128, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="nvfp4", + ) + + +@pytest.mark.parametrize("trtllm_supported", [False, True]) +def test_flashinfer_backend_gates_nvfp4_scale_search_on_trtllm( + monkeypatch, trtllm_supported: bool +) -> None: + """NVFP4 variants that only change the store-time scale search need the + trtllm-gen native store path; plain nvfp4 stays available either way.""" + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + monkeypatch.setattr( + flashinfer_backend, + "supports_trtllm_attention", + lambda is_prefill: trtllm_supported, + ) + + backend = flashinfer_backend.FlashInferBackend + assert backend.supports_kv_cache_dtype("nvfp4") + assert backend.supports_kv_cache_dtype("nvfp4_4over6") is trtllm_supported + + +def test_flashinfer_impl_rejects_nvfp4_scale_search_without_native_update( + monkeypatch, +) -> None: + """The FlashInfer slot-mapping writer records plain max/6 scales, so an + NVFP4 scale-search dtype must fail instead of silently degrading.""" + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + monkeypatch.setattr( + flashinfer_backend, + "can_use_trtllm_attention", + lambda num_heads, num_kv_heads, is_prefill=False: False, + ) + + with pytest.raises(ValueError, match="trtllm-gen native NVFP4 KV cache update"): + flashinfer_backend.FlashInferImpl( + num_heads=1, + head_size=128, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="nvfp4_4over6", + ) + + def test_fast_decode_plan_importable() -> None: """fast_decode_plan must be importable from flashinfer.decode. @@ -205,6 +1543,94 @@ def test_fast_plan_decode_warmup_uses_full_plan(dtype: torch.dtype) -> None: ) +@pytest.mark.parametrize("dtype", DTYPES) +@torch.inference_mode +def test_fast_plan_decode_accepts_nvfp4_kv_plan_dtype(dtype: torch.dtype) -> None: + from vllm.v1.attention.backends.flashinfer import fast_plan_decode + + torch.set_default_device("cuda") + set_random_seed(0) + + kv_lens = [128, 64] + block_size = 16 + num_seqs = len(kv_lens) + num_query_heads, num_kv_heads = 8, 2 + head_size = 128 + + kv_indptr, kv_indices, kv_last_page_lens, _ = _make_paged_kv_metadata( + kv_lens, block_size, NUM_BLOCKS + ) + + workspace = torch.empty(128 * 1024 * 1024, dtype=torch.int8) + wrapper = _make_cg_decode_wrapper(num_seqs, kv_indices.clone(), workspace) + + fast_plan_decode( + wrapper, + indptr_cpu=kv_indptr, + indices=kv_indices, + last_page_len_cpu=kv_last_page_lens, + seq_lens_cpu=get_seq_lens(kv_indptr, kv_last_page_lens, block_size), + num_qo_heads=num_query_heads, + num_kv_heads=num_kv_heads, + head_dim=head_size, + page_size=block_size, + q_data_type=dtype, + kv_data_type=torch.uint8, + o_data_type=dtype, + ) + + assert wrapper.vllm_first_call is False + + +@pytest.mark.parametrize("dtype", DTYPES) +@torch.inference_mode +def test_flashinfer_prefill_accepts_nvfp4_kv_plan_dtype( + dtype: torch.dtype, +) -> None: + torch.set_default_device("cuda") + set_random_seed(0) + + batch_size = 2 + qo_len = 8 + kv_len = 16 + block_size = 16 + num_query_heads, num_kv_heads = 8, 2 + head_size = 128 + num_pages_per_seq = (kv_len + block_size - 1) // block_size + total_num_pages = num_pages_per_seq * batch_size + + q_indptr = ( + torch.arange(0, batch_size + 1, device="cuda", dtype=torch.int32) * qo_len + ) + kv_indptr = ( + torch.arange(0, batch_size + 1, device="cuda", dtype=torch.int32) + * num_pages_per_seq + ) + kv_indices = torch.arange(0, total_num_pages, device="cuda", dtype=torch.int32) + kv_last_page_len = torch.full( + (batch_size,), kv_len, dtype=torch.int32, device="cuda" + ) + + workspace = torch.empty(128 * 1024 * 1024, dtype=torch.int8) + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(workspace, "NHD") + + wrapper.plan( + q_indptr, + kv_indptr, + kv_indices, + kv_last_page_len, + num_query_heads, + num_kv_heads, + head_size, + block_size, + q_data_type=dtype, + kv_data_type=torch.uint8, + o_data_type=dtype, + ) + + assert wrapper._cached_kv_data_type == torch.uint8 + + @pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]]) @pytest.mark.parametrize("num_heads", NUM_HEADS) @pytest.mark.parametrize("head_size", HEAD_SIZES) diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index dc75ba4ea19b..0c54c2ed91bd 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -13,6 +13,7 @@ import torch from tests.v1.attention.utils import dense_kv_cache_views +from vllm.utils.torch_utils import nvfp4_kv_cache_full_dim from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.hisparse.binding import allocate_hisparse_kv_caches @@ -23,6 +24,7 @@ KVCacheGroupSpec, KVCacheLayout, KVCacheTensor, + KVQuantMode, MLAAttentionSpec, SparseCacheRole, compute_layout_strides, @@ -39,6 +41,60 @@ ) +@pytest.mark.parametrize( + ("head_size", "head_size_v"), + [ + pytest.param(128, 128, id="same-head"), + pytest.param(256, 128, id="mixed-head"), + ], +) +def test_flashinfer_nvfp4_customize_spec_drives_view_shape( + head_size: int, head_size_v: int +): + """NVFP4 packing is published through the spec, so the shared allocator + reproduces the layout the FlashInfer NVFP4 path expects: K and V in + separate head slots for equal head sizes, and one slot holding both packed + states when they differ.""" + try: + from vllm.v1.attention.backends.flashinfer import FlashInferBackend + except Exception as exc: # pragma: no cover - environment dependent + pytest.skip(f"FlashInfer backend unavailable: {exc}") + + num_blocks = 2 + block_size = 16 + num_kv_heads = 2 + spec = FlashInferBackend.customize_spec( + FullAttentionSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + head_size_v=head_size_v, + dtype=torch.uint8, + kv_quant_mode=KVQuantMode.NVFP4, + ) + ) + + full_k = nvfp4_kv_cache_full_dim(head_size) + full_v = nvfp4_kv_cache_full_dim(head_size_v) + if head_size == head_size_v: + expected_heads, expected_content = 2 * num_kv_heads, full_k + else: + expected_heads, expected_content = num_kv_heads, full_k + full_v + assert spec.num_heads == expected_heads + assert spec.state_content_size_bytes == expected_content + + raw = torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8) + (kv_cache,) = dense_kv_cache_views(raw, spec, num_blocks, 1, KVCacheLayout.LBHNC) + + assert kv_cache.shape == ( + num_blocks, + expected_heads, + block_size, + expected_content, + ) + assert kv_cache[0].is_contiguous() + + @pytest.mark.parametrize( ("enabled", "block_size", "main_sizes", "indexer_sizes", "expected"), [ diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index a001c8c29c80..6ed396f7ee9a 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -413,6 +413,9 @@ def __init__( if block_n is not None: extra_impl_args.setdefault("block_n", block_n) + if self.attn_backend.get_name() == "FLASHINFER": + extra_impl_args.setdefault("head_size_v", self.head_size_v) + impl_cls = self.attn_backend.get_impl_cls() self.impl = impl_cls( # type: ignore[assignment] # impl_cls always returns an AttentionImpl subclass num_heads, diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index b453487a01a8..f17b6d1ca30e 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -555,7 +555,7 @@ def nvfp4_split_data_scale( """Split one side (K or V) of an NVFP4 KV cache into data and scale. The input is a 4D uint8 tensor whose last dimension is - ``full_dim = data_dim + scale_dim``. The physical layout within each + ``full_dim = data_dim + scale_dim``. The physical layout within each side is ``[data | scale]``, both packed contiguously. The caller is responsible for slicing K and V from the combined cache @@ -577,9 +577,7 @@ def nvfp4_split_data_scale( data_per_kv = dim_1 * dim_2 * data_dim page_bytes = kv_side.stride(0) - # Derive inner strides from the kv_side strides, scaling by the - # ratio of the target dim to full_dim. This preserves the physical - # layout (NHD vs HND) encoded in the input tensor's strides. + # Scale the source strides to preserve the physical NHD/HND order. s1 = kv_side.stride(1) * data_dim // full_dim s2 = kv_side.stride(2) * data_dim // full_dim data_shape = (num_pages, dim_1, dim_2, data_dim) @@ -599,6 +597,205 @@ def nvfp4_split_data_scale( return data, scale +def _nvfp4_split_data_scale( + kv_side: torch.Tensor, + head_size: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Split one side (K or V) of an NVFP4 KV cache into data and scale. + + The input is a 4D uint8 tensor whose last dimension is + ``full_dim = data_dim + scale_dim``. The physical layout within each + side is ``[data | scale]``, both packed contiguously. + + The caller is responsible for slicing K and V from the combined cache + first (e.g. ``kv_cache.split(num_kv_heads, dim=1)``). + + Args: + kv_side: 4D uint8 tensor ``(B, H, N, full_dim)``. + + Returns: + ``(data, scale)`` where *data* is uint8 and *scale* is + float8_e4m3fn, both views of the same storage. + """ + num_pages = kv_side.shape[0] + dim_1, dim_2 = kv_side.shape[1], kv_side.shape[2] + full_dim = kv_side.shape[3] + if head_size is None: + if full_dim % 9 != 0: + raise ValueError( + "NVFP4 KV side last dimension cannot be inferred from " + f"packed width: last_dim={full_dim}" + ) + data_dim = full_dim * 8 // 9 + scale_dim = full_dim - data_dim + else: + data_dim = head_size // 2 + scale_dim = head_size // 16 + expected_full_dim = data_dim + scale_dim + if full_dim != expected_full_dim: + raise ValueError( + "NVFP4 KV side last dimension does not match head size: " + f"last_dim={full_dim}, head_size={head_size}, " + f"expected={expected_full_dim}" + ) + + data_per_kv = dim_1 * dim_2 * data_dim + page_bytes = kv_side.stride(0) + + # This helper expects a compact single-side page layout: + # [all data | all scale]. Mixed K/V pages are split explicitly in + # nvfp4_kv_cache_split_views because slicing them leaves outer strides + # based on the combined K+V width and would make data/scale views overlap. + stride_1 = kv_side.stride(1) + stride_2 = kv_side.stride(2) + if ( + kv_side.stride(3) != 1 + or (dim_2 > 1 and stride_2 != full_dim) + or (dim_1 > 1 and stride_1 != dim_2 * full_dim) + ): + raise ValueError( + "NVFP4 KV cache strides are not compatible with compact " + f"data/scale split: shape={kv_side.shape}, " + f"strides={kv_side.stride()}, full_dim={full_dim}" + ) + s1 = stride_1 * data_dim // full_dim + s2 = stride_2 * data_dim // full_dim + data_shape = (num_pages, dim_1, dim_2, data_dim) + data_strides = (page_bytes, s1, s2, 1) + + s1_s = stride_1 * scale_dim // full_dim + s2_s = stride_2 * scale_dim // full_dim + scale_shape = (num_pages, dim_1, dim_2, scale_dim) + scale_strides = (page_bytes, s1_s, s2_s, 1) + + base = kv_side.storage_offset() + data = torch.as_strided(kv_side, data_shape, data_strides, storage_offset=base) + scale = torch.as_strided( + kv_side, scale_shape, scale_strides, storage_offset=base + data_per_kv + ).view(torch.float8_e4m3fn) + + return data, scale + + +def _nvfp4_split_mixed_data_scale( + kv_cache: torch.Tensor, + head_size: int, + head_size_v: int, +) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: + """Split a mixed-head NVFP4 page into compact K/V data and scale views.""" + num_pages = kv_cache.shape[0] + dim_1, dim_2 = kv_cache.shape[1], kv_cache.shape[2] + k_data_dim = head_size // 2 + k_scale_dim = head_size // 16 + v_data_dim = head_size_v // 2 + v_scale_dim = head_size_v // 16 + k_full_dim = k_data_dim + k_scale_dim + v_full_dim = v_data_dim + v_scale_dim + + if kv_cache.shape[-1] != k_full_dim + v_full_dim: + raise ValueError( + "Mixed NVFP4 KV cache last dimension does not match head sizes: " + f"last_dim={kv_cache.shape[-1]}, head_size={head_size}, " + f"head_size_v={head_size_v}, expected={k_full_dim + v_full_dim}" + ) + if kv_cache.stride(-1) != 1: + raise ValueError( + "Mixed NVFP4 KV cache last dimension must be contiguous: " + f"strides={kv_cache.stride()}" + ) + mixed_full_dim = k_full_dim + v_full_dim + stride_1 = kv_cache.stride(1) + stride_2 = kv_cache.stride(2) + if (dim_2 > 1 and stride_2 != mixed_full_dim) or ( + dim_1 > 1 and stride_1 != dim_2 * mixed_full_dim + ): + raise ValueError( + "Mixed NVFP4 KV cache strides are not compatible with compact " + f"data/scale split: shape={kv_cache.shape}, " + f"strides={kv_cache.stride()}, full_dim={mixed_full_dim}" + ) + + page_bytes = kv_cache.stride(0) + elements_per_page = dim_1 * dim_2 * mixed_full_dim + if page_bytes < elements_per_page: + raise ValueError( + "Mixed NVFP4 KV cache page stride is smaller than one page: " + f"page_stride={page_bytes}, required={elements_per_page}" + ) + + base = kv_cache.storage_offset() + + def make_view(offset: int, inner_dim: int) -> torch.Tensor: + return torch.as_strided( + kv_cache, + (num_pages, dim_1, dim_2, inner_dim), + (page_bytes, dim_2 * inner_dim, inner_dim, 1), + storage_offset=base + offset, + ) + + k_data_offset = 0 + k_scale_offset = dim_1 * dim_2 * k_data_dim + v_data_offset = dim_1 * dim_2 * k_full_dim + v_scale_offset = v_data_offset + dim_1 * dim_2 * v_data_dim + + k_data = make_view(k_data_offset, k_data_dim) + k_scale = make_view(k_scale_offset, k_scale_dim).view(torch.float8_e4m3fn) + v_data = make_view(v_data_offset, v_data_dim) + v_scale = make_view(v_scale_offset, v_scale_dim).view(torch.float8_e4m3fn) + return (k_data, v_data), (k_scale, v_scale) + + +def nvfp4_kv_cache_split_views( + kv_cache: torch.Tensor, + head_size: int | None = None, + head_size_v: int | None = None, +) -> tuple[tuple, tuple]: + """Split an NVFP4 KV cache tensor into data and scale views. + + Accepts a 5D tensor ``(num_pages, 2, dim_2, dim_3, full_dim)`` for + same-size K/V sides, a 4D single-side tensor + ``(num_pages, dim_2, dim_3, full_dim)``, or a 4D mixed-head tensor + ``(num_pages, dim_2, dim_3, k_full_dim + v_full_dim)`` when + ``head_size`` and ``head_size_v`` are provided. + + Per-page layout: [K_data | K_scale | V_data | V_scale]. + Each KV side is self-contained (data followed by its scale), so the + 5D case simply splits each side independently. + + The returned views are in the same dim order as the input (NHD or + HND), so callers get views matching whichever order they passed in. + + Args: + kv_cache: 5D or 4D uint8 tensor where each side's last dimension is + ``full_dim = data_dim + scale_dim = 9 * head_size / 16``. + head_size: Optional K head dimension. When provided, split offsets are + derived from the logical head dimensions instead of inferring from + the packed side width. + head_size_v: Optional V head dimension. Defaults to ``head_size`` when + ``head_size`` is provided. + + Returns: + For 5D input: + ``(k_data, v_data), (k_scale, v_scale)`` + For 4D input (single KV side): + ``(data,), (scale,)`` + """ + if kv_cache.dim() == 4: + if head_size is not None: + head_size_v = head_size if head_size_v is None else head_size_v + k_full_dim = nvfp4_kv_cache_full_dim(head_size) + v_full_dim = nvfp4_kv_cache_full_dim(head_size_v) + if kv_cache.shape[-1] == k_full_dim + v_full_dim: + return _nvfp4_split_mixed_data_scale(kv_cache, head_size, head_size_v) + + data, scale = _nvfp4_split_data_scale(kv_cache, head_size) + return (data,), (scale,) + + k_data, k_scale = _nvfp4_split_data_scale(kv_cache[:, 0], head_size) + v_data, v_scale = _nvfp4_split_data_scale(kv_cache[:, 1], head_size_v) + return (k_data, v_data), (k_scale, v_scale) + + def create_kv_caches_with_random_flash( num_blocks: int, block_size: int, diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 0da68e3e3d93..cc23987dabe6 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -2,11 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Attention layer with FlashInfer.""" +import math +from collections.abc import Callable from dataclasses import dataclass, replace from enum import Enum from functools import partial -from typing import ClassVar +from typing import ClassVar, cast +import flashinfer import numpy as np import torch from flashinfer import ( @@ -57,6 +60,7 @@ is_quantized_kv_cache, is_strictly_contiguous, nvfp4_kv_cache_full_dim, + nvfp4_kv_cache_split_views, nvfp4_split_data_scale, ) from vllm.v1.attention.backend import ( @@ -90,6 +94,10 @@ iter_layer_specs, ) from vllm.v1.utils import CpuGpuBuffer +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT = 2048 * 1024 * 1024 FLASHINFER_PREFILL_WORKSPACE_BYTES_PER_ELEM = 16 @@ -102,6 +110,29 @@ trtllm_workspace_buffer = None +_NVFP4KVDataViews = tuple[torch.Tensor, torch.Tensor] + + +@dataclass(frozen=True) +class _NVFP4KVCacheViewKey: + data_ptr: int + storage_offset: int + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: torch.dtype + device: torch.device + head_size: int + head_size_v: int + stride_order: tuple[int, ...] + + +@dataclass(frozen=True) +class _NVFP4KVCacheViews: + kv_cache: torch.Tensor + data: _NVFP4KVDataViews + block_scales: _NVFP4KVDataViews + + def _get_trtllm_workspace_buffer(): global trtllm_workspace_buffer if trtllm_workspace_buffer is None: @@ -111,6 +142,38 @@ def _get_trtllm_workspace_buffer(): return trtllm_workspace_buffer +def _is_flash_attn_varlen_func_available() -> bool: + try: + from vllm.v1.attention.backends.fa_utils import ( + is_flash_attn_varlen_func_available, + ) + except ImportError: + return False + return is_flash_attn_varlen_func_available() + + +def _get_flash_attn_version( + head_size: int, + head_size_v: int, + has_sinks: bool, +) -> int | None: + try: + from vllm.v1.attention.backends.fa_utils import get_flash_attn_version + except ImportError: + return None + return get_flash_attn_version( + head_size=head_size, + head_size_v=head_size_v, + has_sinks=has_sinks, + ) + + +def _flash_attn_varlen_func(*args, **kwargs) -> torch.Tensor: + from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func + + return flash_attn_varlen_func(*args, **kwargs) + + def _pack_draft_block_bool_mask( bool_mask: torch.Tensor, num_packed: int ) -> torch.Tensor: @@ -357,6 +420,7 @@ def run( key: torch.Tensor, value: torch.Tensor, out: torch.Tensor, + kv_cache_sf: tuple[torch.Tensor, torch.Tensor] | None = None, ): prefill_query_across_dcp = get_dcp_group().all_gather( prefill_query.contiguous(), dim=1 @@ -367,6 +431,7 @@ def run( k_scale=layer._k_scale_float, v_scale=layer._v_scale_float, return_lse=True, + kv_cache_sf=kv_cache_sf, ) output_context, lse_context = self._dcp_combine( output_context_tmp, @@ -403,7 +468,14 @@ def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": return spec hs_k = nvfp4_kv_cache_full_dim(spec.head_size) hs_v = nvfp4_kv_cache_full_dim(spec.head_size_v) - assert hs_k == hs_v, "nvfp4 with asymmetric K/V head sizes not yet supported" + if hs_k != hs_v: + # Asymmetric K/V head sizes keep K and V in the same head slot, so + # the slot holds both packed states (see _use_mixed_nvfp4_layout). + return replace( + spec, + num_head_slots=spec.num_kv_heads, + state_content_bytes=(hs_k + hs_v) * get_dtype_size(spec.dtype), + ) return replace( spec, num_head_slots=2 * spec.num_kv_heads, @@ -422,6 +494,18 @@ def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": "nvfp4_4over6", ] + @staticmethod + def _use_mixed_nvfp4_layout( + cache_dtype_str: str, + head_size: int, + head_size_v: int | None, + ) -> bool: + return ( + cache_dtype_str.startswith("nvfp4") + and head_size_v is not None + and head_size_v != head_size + ) + @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: # Page sizes >= 128 only run on the trtllm-gen dynamic kernel (GQA/MQA @@ -493,11 +577,24 @@ def get_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype: @classmethod def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: - if kv_cache_dtype is not None and kv_cache_dtype.startswith("nvfp4"): + # Plain "nvfp4" runs on any FlashInfer-capable device: the KV cache + # update falls back to the FlashInfer slot-mapping writer when the + # trtllm-gen native store path is unavailable. + # + # NVFP4 variants that only differ in the store-time scale search + # (e.g. "nvfp4_4over6") are implemented exclusively in the trtllm-gen + # native store kernel. The FlashInfer writer would silently record + # plain max/6 scales, so require the native path instead of quietly + # degrading the requested quantization. + if ( + kv_cache_dtype is not None + and kv_cache_dtype.startswith("nvfp4") + and kv_cache_dtype != "nvfp4" + ): return ( - current_platform.is_device_capability_family(100) - and supports_trtllm_attention(is_prefill=True) + supports_trtllm_attention(is_prefill=True) and supports_trtllm_attention(is_prefill=False) + and super().supports_kv_cache_dtype(kv_cache_dtype) ) return super().supports_kv_cache_dtype(kv_cache_dtype) @@ -556,6 +653,11 @@ class FIPrefill: """Metadata for the native FlashInfer prefill pathway (non-TRTLLM).""" wrapper: BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper + block_tables: torch.Tensor | None = None + seq_lens: torch.Tensor | None = None + seq_lens_cpu: torch.Tensor | None = None + query_start_loc: torch.Tensor | None = None + query_start_loc_cpu: torch.Tensor | None = None @dataclass @@ -769,6 +871,7 @@ def __init__( self.num_kv_heads = self.kv_cache_spec.num_kv_heads self.head_dim = self.kv_cache_spec.head_size + self.head_dim_v = getattr(self.kv_cache_spec, "head_size_v", self.head_dim) self.page_size = self.kv_cache_spec.block_size if self.kv_cache_spec.kv_quant_mode != KVQuantMode.NONE: @@ -776,37 +879,15 @@ def __init__( # Cannot use self.kv_cache_spec.dtype here because kv_cache_spec # storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3) self.is_kvcache_nvfp4 = self.cache_dtype.startswith("nvfp4") - if self.is_kvcache_nvfp4: - if ( - force_use_trtllm_attention() is False - or not supports_trtllm_attention(is_prefill=True) - or not supports_trtllm_attention(is_prefill=False) - ): - raise ValueError( - f"--kv-cache-dtype {self.cache_dtype} requires the " - "SM100 trtllm-gen " - "FlashInfer path." - ) - # The scale search only affects the store kernel. FlashInfer - # reads both variants using the same NVFP4 layout. - self.kv_cache_dtype = "nvfp4" - else: - self.kv_cache_dtype = FlashInferBackend.get_dtype_for_flashinfer( - self.cache_dtype - ) + self.kv_cache_dtype = FlashInferBackend.get_dtype_for_flashinfer( + self.cache_dtype + ) else: self.cache_dtype = "auto" self.is_kvcache_nvfp4 = False assert self.kv_cache_spec.dtype == self.model_config.dtype self.kv_cache_dtype = self.kv_cache_spec.dtype - # Compute per-phase Q dtype. On SM90 (XQA decode), the prefill and - # decode phases require different Q dtypes when the KV cache is FP8 - # (FP8-Q for the FI native prefill, BF16/FP16-Q for XQA decode), - # so both values must be tracked independently. - self.q_data_type_prefill = self.get_q_data_type(is_prefill=True) - self.q_data_type_decode = self.get_q_data_type(is_prefill=False) - # Prefer TRTLLM/XQA for decoding whenever supported. The decode kernel # must be selected statically for FULL cudagraph capture. can_use_xqa_or_trtllm_gen_decode = can_use_trtllm_attention( @@ -880,6 +961,26 @@ def __init__( # flash_attn_varlen_func's cp_world_size/cp_rank/cp_tot_seqused_k). supports_dcp_with_varlen=False, ) + can_use_trtllm_prefill_attention = can_use_trtllm_attention( + self.num_qo_heads, self.num_kv_heads, is_prefill=True + ) + use_trtllm_gen_decode = ( + self.flashinfer_trtllm_api_decode_kernel + == FlashInferDecodeKernel.TRTLLM_GEN + ) + + # Compute per-phase Q dtype. On SM90 (XQA decode), the prefill and + # decode phases require different Q dtypes when the KV cache is FP8 + # (FP8-Q for the FI native prefill, BF16/FP16-Q for XQA decode), + # so both values must be tracked independently. + self.q_data_type_prefill = self.get_q_data_type( + is_prefill=True, + use_trtllm_gen=can_use_trtllm_prefill_attention, + ) + self.q_data_type_decode = self.get_q_data_type( + is_prefill=False, + use_trtllm_gen=use_trtllm_gen_decode, + ) self._cascade_wrapper = None # Wrapper for cascade attention @@ -925,13 +1026,14 @@ def __init__( self.paged_kv_last_page_len = CpuGpuBuffer( max_num_reqs, dtype=torch.int32, device=self.device, pin_memory=False ) + self._reserve_nvfp4_fa2_prefill_workspace(can_use_trtllm_prefill_attention) @property def kv_cache_layout(self) -> KVCacheLayout: return self.cache_config.get_resolved_kv_cache_layout() # Keep SM90 prefill/decode Q dtype selection in one place. - def get_q_data_type(self, is_prefill: bool) -> torch.dtype: + def get_q_data_type(self, is_prefill: bool, use_trtllm_gen: bool) -> torch.dtype: # The user sets --attention-config.disable_flashinfer_q_quantization # to 1 explicitly, use model dtype for query. if self.vllm_config.attention_config.disable_flashinfer_q_quantization: @@ -969,9 +1071,71 @@ def get_q_data_type(self, is_prefill: bool) -> torch.dtype: return FlashInferBackend.get_dtype_for_flashinfer(cache_dtype) return self.model_config.dtype if cache_dtype.startswith("nvfp4"): - return FlashInferBackend.get_dtype_for_flashinfer("fp8_e4m3") + # NVFP4 KV uses FP8-Q only on the trtllm-gen path. Native FA2 + # prefill/decode, including SM8x and SM90 fallback, consumes + # model-dtype Q. + if use_trtllm_gen: + return FlashInferBackend.get_dtype_for_flashinfer("fp8_e4m3") + return self.model_config.dtype return self.kv_cache_spec.dtype + def _reserve_nvfp4_fa2_prefill_workspace(self, can_use_trtllm: bool) -> None: + if ( + not self.is_kvcache_nvfp4 + or can_use_trtllm + or self.head_dim != self.head_dim_v + or max(self.head_dim, self.head_dim_v) > 256 + or self.use_dcp + or not is_workspace_manager_initialized() + or not _is_flash_attn_varlen_func_available() + or not hasattr(flashinfer, "nvfp4_kv_dequantize_paged") + ): + return + + scheduler_config = self.vllm_config.scheduler_config + if ( + not scheduler_config.enable_chunked_prefill + or self.model_config.max_model_len + <= scheduler_config.max_num_batched_tokens + ): + # Without chunked prefill, or when one scheduler chunk can cover + # the full context, every prefill is the first chunk + # (q_len == seq_len). That path uses the current K/V tensors + # directly and does not need a dequant scratch buffer. + return + + if not getattr(self.model_config, "uses_alibi", True): + return + + scratch_shape = ( + self.model_config.max_model_len, + self.num_kv_heads, + self.head_dim, + ) + scratch_shape_v = ( + self.model_config.max_model_len, + self.num_kv_heads, + self.head_dim_v, + ) + dtype = cast(torch.dtype, self.model_config.dtype) + scratch_bytes = ( + math.prod(scratch_shape) + math.prod(scratch_shape_v) + ) * dtype.itemsize + # Sliding-window layers currently reserve full-context scratch. A future + # window clamp can reduce this to the active window. + current_workspace_manager().get_simultaneous( + (scratch_shape, dtype), + (scratch_shape_v, dtype), + ) + logger.info( + "Reserved %.2f MiB for NVFP4 FA2 prefill scratch workspace: " + "k_shape=%s, v_shape=%s, dtype=%s", + scratch_bytes / (1024 * 1024), + scratch_shape, + scratch_shape_v, + dtype, + ) + @override # type: ignore[misc] @classmethod def get_cudagraph_support( @@ -1163,6 +1327,11 @@ def _get_prefill_wrapper( if self._prefill_wrapper is None: if self.use_dcp: + if self.is_kvcache_nvfp4 and self.head_dim != self.head_dim_v: + raise NotImplementedError( + "FlashInfer DCP prefill does not support NVFP4 KV cache " + "with different K/V head dimensions." + ) self._prefill_wrapper = BatchDCPPrefillWrapper( kv_layout=get_flashinfer_layout_string(self.kv_cache_layout), workspace_buffer=self._get_workspace_buffer(), @@ -1185,13 +1354,13 @@ def _get_prefill_wrapper( window_left=self.window_left, ) else: - # NVFP4 KV cache requires the trtllm-gen backend inside - # the wrapper; fa2/fa3 do not support nvfp4. - backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto" + # NVFP4 KV planning goes through the native FlashInfer + # wrapper; the backend is resolved per shape rather than + # pinned to trtllm-gen, which is unavailable pre-SM100. self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper( self._get_workspace_buffer(), get_flashinfer_layout_string(self.kv_cache_layout), - backend=backend, + backend="auto", ) assert self._prefill_wrapper is not None return self._prefill_wrapper @@ -1211,9 +1380,6 @@ def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False): paged_kv_indptr = None paged_kv_indices = None paged_kv_last_page_len = None - # NVFP4 KV cache requires the trtllm-gen backend inside - # the wrapper; fa2/fa3 do not support nvfp4. - backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto" decode_wrapper = BatchDecodeWithPagedKVCacheWrapper( self._get_workspace_buffer(), get_flashinfer_layout_string(self.kv_cache_layout), @@ -1225,7 +1391,7 @@ def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False): # at least as good as cuda cores for all attention ops in latest # gpus. use_tensor_cores=True, - backend=backend, + backend="auto", ) # save the decode wrapper @@ -1306,6 +1472,12 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> FlashInferMetadata: + if self.use_dcp and self.is_kvcache_nvfp4 and self.head_dim != self.head_dim_v: + raise NotImplementedError( + "FlashInfer DCP does not support NVFP4 KV cache with different " + "K/V head dimensions." + ) + num_reqs = common_attn_metadata.num_reqs num_actual_tokens = common_attn_metadata.num_actual_tokens causal = common_attn_metadata.causal @@ -1643,12 +1815,6 @@ def build( prefill_wrapper, BatchPrefillWithPagedKVCacheWrapper, ) - # NVFP4 trtllm kernel only supports FP8 output; - # use FP8 o_data_type so the wrapper matches the - # FP8 output buffer allocated in forward(). - o_dtype = ( - FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype - ) prefill_wrapper.plan( qo_indptr=qo_indptr_prefill_cpu, paged_kv_indptr=paged_kv_indptr_prefill_cpu, @@ -1665,11 +1831,23 @@ def build( logits_soft_cap=self.logits_soft_cap, q_data_type=self.q_data_type_prefill, kv_data_type=self.kv_cache_dtype, - o_data_type=o_dtype, + o_data_type=self.model_config.dtype, fixed_split_size=self.prefill_fixed_split_size, disable_split_kv=self.disable_split_kv, ) - attn_metadata.prefill = FIPrefill(wrapper=prefill_wrapper) + query_start_loc_prefill = ( + qo_indptr[prefill_start:] - qo_indptr[prefill_start] + ) + attn_metadata.prefill = FIPrefill( + wrapper=prefill_wrapper, + block_tables=block_table_tensor[prefill_start:], + seq_lens=seq_lens[prefill_start:], + seq_lens_cpu=seq_lens_cpu[prefill_start:] + if seq_lens_cpu is not None + else None, + query_start_loc=query_start_loc_prefill, + query_start_loc_cpu=qo_indptr_prefill_cpu, + ) ## DECODE PATHWAY if num_decodes > 0: @@ -1732,12 +1910,6 @@ def build( # Use the persistent buffer with padding length, # instead of the same address but chunked version # in atten_metadata when using cudagraph. - # NVFP4 trtllm kernel only supports FP8 output; - # use FP8 o_data_type so the wrapper matches the - # FP8 output buffer allocated in forward(). - o_dtype = ( - FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype - ) paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: num_input_tokens + 1] paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[ :num_input_tokens @@ -1775,7 +1947,7 @@ def build( logits_soft_cap=self.logits_soft_cap, q_data_type=self.q_data_type_decode, kv_data_type=self.kv_cache_dtype, - o_data_type=o_dtype, + o_data_type=self.model_config.dtype, fixed_split_size=self.decode_fixed_split_size, disable_split_kv=self.disable_split_kv, ) @@ -1809,9 +1981,11 @@ def __init__( attn_type: AttentionType = AttentionType.DECODER, kv_sharing_target_layer_name: int | None = None, sinks: torch.Tensor | None = None, + head_size_v: int | None = None, ) -> None: self.num_heads = num_heads self.head_size = head_size + self.head_size_v = head_size if head_size_v is None else head_size_v self.scale = float(scale) self.num_kv_heads = num_kv_heads if alibi_slopes is not None: @@ -1824,9 +1998,52 @@ def __init__( self.window_left = ( self.sliding_window[0] if self.sliding_window is not None else -1 ) - self.cache_dtype = kv_cache_dtype + self.kv_cache_dtype = kv_cache_dtype self.is_kvcache_nvfp4 = kv_cache_dtype.startswith("nvfp4") - self.kv_cache_dtype = "nvfp4" if self.is_kvcache_nvfp4 else kv_cache_dtype + self.use_native_nvfp4_kv_cache_update = False + if self.is_kvcache_nvfp4 and self.head_size_v == self.head_size: + self.use_native_nvfp4_kv_cache_update = can_use_trtllm_attention( + num_heads, num_kv_heads, is_prefill=True + ) and can_use_trtllm_attention(num_heads, num_kv_heads, is_prefill=False) + if ( + self.is_kvcache_nvfp4 + and not self.use_native_nvfp4_kv_cache_update + and kv_cache_dtype != "nvfp4" + ): + # NVFP4 variants only differ in the store-time scale search, which + # lives in the trtllm-gen native store kernel. The FlashInfer + # slot-mapping writer records plain max/6 scales, so it cannot + # honor the requested search. Fail instead of silently changing + # the KV cache quantization. + raise ValueError( + f"--kv-cache-dtype {kv_cache_dtype} requires the trtllm-gen " + "native NVFP4 KV cache update path; the FlashInfer " + "slot-mapping writer only records max/6 scales and cannot " + "perform the requested scale search." + ) + self._nvfp4_slot_writer: Callable[..., None] | None = None + self._nvfp4_paged_dequant: Callable[..., None] | None = None + if self.is_kvcache_nvfp4 and not self.use_native_nvfp4_kv_cache_update: + nvfp4_slot_writer = getattr( + flashinfer, + "nvfp4_quantize_append_paged_kv_cache_with_slot_mapping", + None, + ) + if nvfp4_slot_writer is None: + raise RuntimeError( + "FlashInfer NVFP4 slot-mapping KV cache update is " + "required when native NVFP4 KV cache update is unavailable." + ) + self._nvfp4_slot_writer = nvfp4_slot_writer + self._nvfp4_paged_dequant = getattr( + flashinfer, "nvfp4_kv_dequantize_paged", None + ) + if self._nvfp4_paged_dequant is None: + logger.warning_once( + "FlashInfer NVFP4 paged dequantization is unavailable; " + "falling back to the native FlashInfer prefill path, which " + "can be slower for long-context NVFP4 KV cache." + ) self.fp4_data_dim = head_size // 2 if self.is_kvcache_nvfp4 else 0 self.logits_soft_cap = logits_soft_cap self.kv_sharing_target_layer_name = kv_sharing_target_layer_name @@ -1852,6 +2069,11 @@ def __init__( f"{sinks.shape[0]}." ) self.sinks = sinks + self.fa_version = _get_flash_attn_version( + head_size=head_size, + head_size_v=self.head_size_v, + has_sinks=self.sinks is not None, + ) self.supports_xqa_or_trtllm_gen_decode = can_use_trtllm_attention( num_heads, num_kv_heads, is_prefill=False @@ -1873,8 +2095,13 @@ def __init__( self.bmm2_scale: float | None = None self.o_sf_scale: float | None = None - # Pre-allocated FP8 output buffer for NVFP4 without fused output quant. - if self.is_kvcache_nvfp4 and vllm_config is not None: + # TRTLLM-gen NVFP4 kernels require FP8 output. Native FlashInfer + # FA2/FA3 paths write model dtype output directly. + if ( + self.is_kvcache_nvfp4 + and self.supports_xqa_or_trtllm_gen_decode + and vllm_config is not None + ): max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens self._nvfp4_fp8_out = torch.empty( (max_num_tokens, num_heads, head_size), @@ -1884,6 +2111,21 @@ def __init__( else: self._nvfp4_fp8_out = None + self._nvfp4_fa2_cu_q: torch.Tensor | None = None + self._nvfp4_fa2_cu_k: torch.Tensor | None = None + if ( + self.is_kvcache_nvfp4 + and self.fa_version is not None + and self.head_size == self.head_size_v + and max(self.head_size, self.head_size_v) <= 256 + and vllm_config is not None + ): + self._nvfp4_fa2_cu_q = torch.zeros(2, device="cuda", dtype=torch.int32) + self._nvfp4_fa2_cu_k = torch.zeros(2, device="cuda", dtype=torch.int32) + + self._nvfp4_kv_cache_view_key: _NVFP4KVCacheViewKey | None = None + self._nvfp4_kv_cache_views: _NVFP4KVCacheViews | None = None + dcp_a2a = ( vllm_config is not None and vllm_config.parallel_config.decode_context_parallel_size > 1 @@ -1929,6 +2171,352 @@ def process_weights_after_loading(self, act_dtype: torch.dtype): else: self.sinks.copy_(source_sinks) + def _get_kv_cache_stride_order(self) -> tuple[int, ...]: + return self.kv_cache_layout.layer_view_order + + def _permute_kv_cache( + self, kv_cache: torch.Tensor, stride_order: tuple[int, ...] + ) -> torch.Tensor: + kv_cache_permute = kv_cache.permute(*stride_order) + # Fix degenerate strides on any size-1 dimension (e.g. num_kv_heads=1 + # with TP=8). PyTorch permits non-canonical strides on size-1 dims; + # CUDA TMA requires ≥16-byte alignment on all non-outermost strides. + # canonicalize_singleton_dim_strides patches metadata via as_strided — + # zero-copy. See vllm.utils.torch_utils. + fixed = canonicalize_singleton_dim_strides(kv_cache_permute) + if fixed is not kv_cache_permute: + logger.debug( + "Canonicalized degenerate KV cache strides (FlashInfer): " + "shape=%s, strides before=%s, strides after=%s", + kv_cache_permute.shape, + kv_cache_permute.stride(), + fixed.stride(), + ) + return fixed + + def _get_nvfp4_kv_cache_views(self, kv_cache: torch.Tensor) -> _NVFP4KVCacheViews: + stride_order = self._get_kv_cache_stride_order() + key = _NVFP4KVCacheViewKey( + data_ptr=kv_cache.data_ptr(), + storage_offset=kv_cache.storage_offset(), + shape=tuple(kv_cache.shape), + stride=tuple(kv_cache.stride()), + dtype=kv_cache.dtype, + device=kv_cache.device, + head_size=self.head_size, + head_size_v=self.head_size_v, + stride_order=stride_order, + ) + if ( + key == self._nvfp4_kv_cache_view_key + and self._nvfp4_kv_cache_views is not None + ): + return self._nvfp4_kv_cache_views + + fixed = self._permute_kv_cache(kv_cache, stride_order) + if self.head_size != self.head_size_v: + data, block_scales = nvfp4_kv_cache_split_views( + fixed, self.head_size, self.head_size_v + ) + elif self.use_native_nvfp4_kv_cache_update: + k_cache, v_cache = kv_cache.split(self.num_kv_heads, dim=1) + k_cache = self._permute_kv_cache(k_cache, stride_order) + v_cache = self._permute_kv_cache(v_cache, stride_order) + k_data, k_scale = nvfp4_split_data_scale(k_cache) + v_data, v_scale = nvfp4_split_data_scale(v_cache) + data = (k_data, v_data) + block_scales = (k_scale, v_scale) + else: + # The slot-mapping writer stores one compact page as + # [K data | K scale | V data | V scale]. Reinterpret the main + # rank-4 same-head allocation as a combined K/V page before + # constructing those four zero-copy views. + num_blocks = fixed.shape[0] + full_dim = nvfp4_kv_cache_full_dim(self.head_size) + cache_layout = get_flashinfer_layout_string(self.kv_cache_layout) + if cache_layout == "NHD": + _, block_size, twice_num_heads, packed_dim = fixed.shape + combined_shape = ( + num_blocks, + block_size, + self.num_kv_heads, + 2 * full_dim, + ) + combined_strides = ( + fixed.stride(0), + self.num_kv_heads * 2 * full_dim, + 2 * full_dim, + 1, + ) + elif cache_layout == "HND": + _, twice_num_heads, block_size, packed_dim = fixed.shape + combined_shape = ( + num_blocks, + self.num_kv_heads, + block_size, + 2 * full_dim, + ) + combined_strides = ( + fixed.stride(0), + block_size * 2 * full_dim, + 2 * full_dim, + 1, + ) + else: + raise ValueError(f"Unknown cache layout format {cache_layout}.") + + if twice_num_heads != 2 * self.num_kv_heads or packed_dim != full_dim: + raise ValueError( + "Unexpected same-head NVFP4 KV cache shape: " + f"shape={tuple(fixed.shape)}, num_kv_heads={self.num_kv_heads}, " + f"packed_dim={full_dim}" + ) + expected_inner_strides = combined_strides[1:] + if fixed.stride(-1) != 1 or fixed.stride(0) < math.prod(combined_shape[1:]): + raise ValueError( + "NVFP4 KV cache is not page-contiguous: " + f"shape={tuple(fixed.shape)}, strides={fixed.stride()}" + ) + if cache_layout == "NHD": + actual_inner_strides = ( + fixed.stride(1), + fixed.stride(2) * 2, + fixed.stride(3), + ) + else: + actual_inner_strides = ( + fixed.stride(1) * 2, + fixed.stride(2) * 2, + fixed.stride(3), + ) + if actual_inner_strides != expected_inner_strides: + raise ValueError( + "NVFP4 KV cache inner strides are not compatible with a " + "compact K/V page: " + f"shape={tuple(fixed.shape)}, strides={fixed.stride()}" + ) + + combined = torch.as_strided( + fixed, + size=combined_shape, + stride=combined_strides, + storage_offset=fixed.storage_offset(), + ) + data, block_scales = nvfp4_kv_cache_split_views( + combined, self.head_size, self.head_size_v + ) + views = _NVFP4KVCacheViews( + kv_cache=fixed, + data=data, + block_scales=block_scales, + ) + self._nvfp4_kv_cache_view_key = key + self._nvfp4_kv_cache_views = views + return views + + def _flash_attn_varlen( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + causal: bool, + ) -> torch.Tensor: + assert self.fa_version is not None + return _flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=self.scale, + causal=causal, + window_size=list(self.sliding_window), + softcap=self.logits_soft_cap or 0.0, + alibi_slopes=self.alibi_slopes, + fa_version=self.fa_version, + ) + + def _copy_prefill_output( + self, + output: torch.Tensor, + start: int, + attn_out: torch.Tensor, + ) -> None: + output_slice = output[start : start + attn_out.shape[0]] + if output_slice.ndim == 2: + output_slice.copy_(attn_out.reshape(attn_out.shape[0], -1)) + else: + output_slice.copy_(attn_out) + + def _run_native_nvfp4_prefill( + self, + layer: torch.nn.Module, + prefill_wrapper: BatchPrefillWithPagedKVCacheWrapper, + prefill_query: torch.Tensor, + output: torch.Tensor, + nvfp4_kv_data: _NVFP4KVDataViews, + nvfp4_kv_block_scales: _NVFP4KVDataViews, + ) -> None: + prefill_wrapper.run( + prefill_query, + nvfp4_kv_data, + q_scale=layer._q_scale_float, + k_scale=layer._k_scale_float, + v_scale=layer._v_scale_float, + out=output, + kv_cache_sf=nvfp4_kv_block_scales, + ) + + def _run_nvfp4_fa2_prefill( + self, + layer: torch.nn.Module, + prefill_wrapper: BatchPrefillWithPagedKVCacheWrapper, + prefill_query: torch.Tensor, + prefill_key: torch.Tensor, + prefill_value: torch.Tensor, + output: torch.Tensor, + attn_metadata: FlashInferMetadata, + nvfp4_kv_data: _NVFP4KVDataViews, + nvfp4_kv_block_scales: _NVFP4KVDataViews, + ) -> bool: + if ( + not self.is_kvcache_nvfp4 + or self.fa_version is None + or self.head_size != self.head_size_v + or max(self.head_size, self.head_size_v) > 256 + or not _is_flash_attn_varlen_func_available() + or prefill_query.dtype not in (torch.float16, torch.bfloat16) + or not isinstance(attn_metadata.causal, bool) + or getattr(attn_metadata, "mm_prefix_range", None) is not None + or self.dcp_world_size > 1 + or attn_metadata.use_cascade + or self._nvfp4_fa2_cu_q is None + or self._nvfp4_fa2_cu_k is None + # The varlen call below takes no sink tensor, so leave sink models + # on the native FlashInfer prefill path that does apply them. + or self.sinks is not None + ): + return False + + assert isinstance(attn_metadata.prefill, FIPrefill) + prefill = attn_metadata.prefill + if getattr(prefill, "mm_wrapper", None) is not None: + return False + if ( + prefill.block_tables is None + or prefill.seq_lens is None + or prefill.seq_lens_cpu is None + or prefill.query_start_loc_cpu is None + ): + return False + + num_reqs = prefill.query_start_loc_cpu.shape[0] - 1 + if num_reqs <= 0: + return True + + qsl = prefill.query_start_loc_cpu.tolist() + seq_lens = prefill.seq_lens_cpu.tolist() + + first_chunk_reqs: list[int] = [] + continuation_reqs: list[int] = [] + for req_idx in range(num_reqs): + q_len = int(qsl[req_idx + 1]) - int(qsl[req_idx]) + seq_len = int(seq_lens[req_idx]) + if q_len <= 0 or seq_len <= 0: + continue + if q_len == seq_len: + first_chunk_reqs.append(req_idx) + else: + continuation_reqs.append(req_idx) + + if not first_chunk_reqs and not continuation_reqs: + return True + + use_legacy_scratch = self.alibi_slopes is not None + if use_legacy_scratch: + if continuation_reqs and ( + self._nvfp4_paged_dequant is None + or not is_workspace_manager_initialized() + ): + return False + request_indices = sorted(first_chunk_reqs + continuation_reqs) + else: + if not first_chunk_reqs: + return False + if continuation_reqs: + self._run_native_nvfp4_prefill( + layer, + prefill_wrapper, + prefill_query, + output, + nvfp4_kv_data, + nvfp4_kv_block_scales, + ) + request_indices = first_chunk_reqs + + for req_idx in request_indices: + q_start = int(qsl[req_idx]) + q_end = int(qsl[req_idx + 1]) + q_len = q_end - q_start + seq_len = int(seq_lens[req_idx]) + + q_seq = prefill_query[q_start:q_end] + if q_len == seq_len: + k_seq = prefill_key[q_start:q_end] + v_seq = prefill_value[q_start:q_end] + else: + # FlashInfer's paged NVFP4 dequant helper writes a padded + # [batch, max_seq_len, ...] buffer, while flash-attn varlen + # consumes compact per-request K/V. The reserved scratch space + # is also sized for one request, so continuation chunks are + # dequantized and processed one request at a time. + assert self._nvfp4_paged_dequant is not None + k_buf, v_buf = current_workspace_manager().get_simultaneous( + ( + (1, seq_len, self.num_kv_heads, self.head_size), + q_seq.dtype, + ), + ( + (1, seq_len, self.num_kv_heads, self.head_size_v), + q_seq.dtype, + ), + ) + self._nvfp4_paged_dequant( + nvfp4_kv_data, + nvfp4_kv_block_scales, + prefill.block_tables[req_idx : req_idx + 1], + prefill.seq_lens[req_idx : req_idx + 1], + layer._k_scale, + layer._v_scale, + k_buf, + v_buf, + kv_layout=get_flashinfer_layout_string(self.kv_cache_layout), + ) + k_seq = k_buf[0, :seq_len] + v_seq = v_buf[0, :seq_len] + + self._nvfp4_fa2_cu_q[1:2] = q_len + self._nvfp4_fa2_cu_k[1:2] = seq_len + attn_out = self._flash_attn_varlen( + q=q_seq, + k=k_seq, + v=v_seq, + cu_seqlens_q=self._nvfp4_fa2_cu_q, + cu_seqlens_k=self._nvfp4_fa2_cu_k, + max_seqlen_q=q_len, + max_seqlen_k=seq_len, + causal=attn_metadata.causal, + ) + self._copy_prefill_output(output, q_start, attn_out.to(output.dtype)) + + return True + def get_xqa_bmm1_scale(self, layer: torch.nn.Module, q_data_type: torch.dtype): bmm1_scale = self.scale if is_quantized_kv_cache(self.kv_cache_dtype): @@ -1980,7 +2568,10 @@ def forward( KV-sharing decoder layer. value: shape = [num_tokens, num_kv_heads, head_size], or None for a KV-sharing decoder layer. - kv_cache: [num_blocks, num_kv_heads, block_size, 2*head_size] + kv_cache: Logical rank-4 ``[B, H, N, C]`` KV cache tensor. K/V are + packed in the content dim for regular dtypes, in separate head + slots for same-head NVFP4, and in a combined packed content dim + for mixed-head NVFP4. attn_metadata: Metadata for attention. Returns: shape = [num_tokens, num_heads * head_size] @@ -2099,41 +2690,22 @@ def forward( num_decode_tokens = attn_metadata.num_decode_tokens num_prefill_tokens = attn_metadata.num_prefill_tokens - stride_order = self.kv_cache_layout.layer_view_order - kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous - # Fix degenerate strides on any size-1 dimension (e.g. num_kv_heads=1 - # with TP=8). PyTorch permits non-canonical strides on size-1 dims; - # CUDA TMA requires ≥16-byte alignment on all non-outermost strides. - # canonicalize_singleton_dim_strides patches metadata via as_strided — - # zero-copy. See vllm.utils.torch_utils. - fixed = canonicalize_singleton_dim_strides(kv_cache_permute) - if fixed is not kv_cache_permute: - logger.debug( - "Canonicalized degenerate KV cache strides (FlashInfer): " - "shape=%s, strides before=%s, strides after=%s", - kv_cache_permute.shape, - kv_cache_permute.stride(), - fixed.stride(), - ) - kv_cache_permute = fixed - # Split K/V — zero-copy views. NVFP4 stores K/V as separate head - # groups; other dtypes pack K/V in the content dim. - hs = self.head_size + # slots (or one packed content dim for mixed head sizes); other dtypes + # pack K/V in the content dim. nvfp4_kv_data = None nvfp4_kv_block_scales = None + kv_cache_tuple: tuple[torch.Tensor, torch.Tensor] | None = None + hs = self.head_size if self.is_kvcache_nvfp4: - k_cache, v_cache = kv_cache.split(self.num_kv_heads, dim=1) - kv_cache_tuple = ( - canonicalize_singleton_dim_strides(k_cache.permute(*stride_order)), - canonicalize_singleton_dim_strides(v_cache.permute(*stride_order)), - ) - k_data, k_sf = nvfp4_split_data_scale(kv_cache_tuple[0]) - v_data, v_sf = nvfp4_split_data_scale(kv_cache_tuple[1]) - nvfp4_kv_data = (k_data, v_data) - nvfp4_kv_block_scales = (k_sf, v_sf) + nvfp4_views = self._get_nvfp4_kv_cache_views(kv_cache) + kv_cache_permute = nvfp4_views.kv_cache + nvfp4_kv_data = nvfp4_views.data + nvfp4_kv_block_scales = nvfp4_views.block_scales else: - kv_cache_tuple = kv_cache_permute.split(hs, dim=-1) + stride_order = self._get_kv_cache_stride_order() + kv_cache_permute = self._permute_kv_cache(kv_cache, stride_order) + kv_cache_tuple = kv_cache_permute.split(self.head_size, dim=-1) use_dcp = self.dcp_world_size > 1 if decode_with_xqa: @@ -2178,13 +2750,24 @@ def forward( assert prefill_wrapper._new_tokens._sm_scale == self.scale assert prefill_wrapper._new_tokens._causal + if self.is_kvcache_nvfp4: + assert nvfp4_kv_data is not None + assert nvfp4_kv_block_scales is not None + dcp_kv_cache = nvfp4_kv_data + dcp_kv_cache_sf = nvfp4_kv_block_scales + else: + assert kv_cache_tuple is not None + dcp_kv_cache = kv_cache_tuple + dcp_kv_cache_sf = None + prefill_wrapper.run( layer, prefill_query, - kv_cache_tuple, + dcp_kv_cache, key[num_decode_tokens:num_actual_tokens], value[num_decode_tokens:num_actual_tokens], out=output[num_decode_tokens:], + kv_cache_sf=dcp_kv_cache_sf, ) else: assert isinstance( @@ -2197,52 +2780,67 @@ def forward( assert prefill_wrapper._sm_scale == self.scale assert prefill_wrapper._causal == attn_metadata.causal + out_prefill = output[num_decode_tokens:] + used_nvfp4_fa2_prefill = False if self.is_kvcache_nvfp4: - kv_cache_for_fi = nvfp4_kv_data - else: - kv_cache_for_fi = kv_cache_tuple - kv_cache_sf = ( - nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None - ) - - # NVFP4 trtllm kernel only supports FP8 output. - # Use a pre-allocated FP8 buffer and dequantize - # afterwards. - needs_fp8_out_prefill = ( - self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE - ) - if needs_fp8_out_prefill: - out_prefill = self._nvfp4_fp8_out[:num_prefill_tokens] - else: - out_prefill = output[num_decode_tokens:] - - if isinstance( - prefill_wrapper, BatchAttentionWithAttentionSinkWrapper - ): - assert self.sinks is not None - prefill_wrapper.run( + if key is None or value is None: + raise NotImplementedError( + "FlashInfer NVFP4 prefill does not support " + "KV-sharing layers" + ) + assert nvfp4_kv_data is not None + assert nvfp4_kv_block_scales is not None + used_nvfp4_fa2_prefill = self._run_nvfp4_fa2_prefill( + layer, + prefill_wrapper, prefill_query, - kv_cache_for_fi, - self.sinks, - self.scale * layer._q_scale_float * layer._k_scale_float, - v_scale=layer._v_scale_float, - out=out_prefill, - ) - else: - prefill_wrapper.run( - prefill_query, - kv_cache_for_fi, - q_scale=layer._q_scale_float, - k_scale=layer._k_scale_float, - v_scale=layer._v_scale_float, - out=out_prefill, - kv_cache_sf=kv_cache_sf, + key[num_decode_tokens:], + value[num_decode_tokens:], + out_prefill, + attn_metadata, + nvfp4_kv_data, + nvfp4_kv_block_scales, ) - if needs_fp8_out_prefill: - output[ - num_decode_tokens : num_decode_tokens + num_prefill_tokens - ].copy_(out_prefill) + if not used_nvfp4_fa2_prefill: + if self.is_kvcache_nvfp4: + assert nvfp4_kv_data is not None + assert nvfp4_kv_block_scales is not None + self._run_native_nvfp4_prefill( + layer, + prefill_wrapper, + prefill_query, + out_prefill, + nvfp4_kv_data, + nvfp4_kv_block_scales, + ) + elif isinstance( + prefill_wrapper, BatchAttentionWithAttentionSinkWrapper + ): + assert self.sinks is not None + assert kv_cache_tuple is not None + prefill_wrapper.run( + prefill_query, + kv_cache_tuple, + self.sinks, + self.scale + * layer._q_scale_float + * layer._k_scale_float, + v_scale=layer._v_scale_float, + out=out_prefill, + ) + else: + assert kv_cache_tuple is not None + prefill_wrapper.run( + prefill_query, + kv_cache_tuple, + q_scale=layer._q_scale_float, + k_scale=layer._k_scale_float, + v_scale=layer._v_scale_float, + out=out_prefill, + kv_cache_sf=None, + sinks=self.sinks, + ) else: assert isinstance(attn_metadata.prefill, TRTLLMPrefill) # prefill_query may be non-contiguous or have degenerate strides @@ -2282,6 +2880,8 @@ def forward( prefill_kv_block_scales = None if self.is_kvcache_nvfp4: + assert nvfp4_kv_data is not None + assert nvfp4_kv_block_scales is not None # NVFP4 trtllm-gen kernel requires FP8 query. assert attn_metadata.q_data_type_prefill == FP8_DTYPE, ( "NVFP4 KV cache requires FP8 quantized queries for " @@ -2324,6 +2924,7 @@ def forward( attn_metadata.q_data_type_prefill, ) else: + assert kv_cache_tuple is not None mock_kv_cache = kv_cache_tuple mock_block_table = block_tables_prefill @@ -2377,18 +2978,15 @@ def forward( assert decode_wrapper._sm_scale == self.scale if self.is_kvcache_nvfp4: + assert nvfp4_kv_data is not None + assert nvfp4_kv_block_scales is not None kv_cache_for_fi = nvfp4_kv_data else: + assert kv_cache_tuple is not None kv_cache_for_fi = kv_cache_tuple kv_cache_sf = nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None - # NVFP4 kernel only supports FP8 output. - # Use a pre-allocated FP8 buffer and dequantize afterwards. - needs_fp8_out = self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE - if needs_fp8_out: - out_decode = self._nvfp4_fp8_out[:num_decode_tokens] - else: - out_decode = output[:num_decode_tokens] + out_decode = output[:num_decode_tokens] if use_dcp: decode_query = get_dcp_group().all_gather( @@ -2429,8 +3027,6 @@ def forward( sinks=self.sinks, ) - if needs_fp8_out: - output[:num_decode_tokens].copy_(out_decode) else: assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode) # decode_query may be non-contiguous or have degenerate strides @@ -2544,11 +3140,16 @@ def forward( device=decode_query.device, ) + if self.is_kvcache_nvfp4: + assert nvfp4_kv_data is not None + trtllm_kv_cache = nvfp4_kv_data + else: + assert kv_cache_tuple is not None + trtllm_kv_cache = kv_cache_tuple + trtllm_batch_decode_with_kv_cache( query=decode_query, - kv_cache=( - nvfp4_kv_data if self.is_kvcache_nvfp4 else kv_cache_tuple - ), + kv_cache=trtllm_kv_cache, workspace_buffer=workspace_buffer, block_tables=block_tables_decode, seq_lens=seq_lens_decode, @@ -2601,28 +3202,45 @@ def do_kv_cache_update( # and value[:num_actual_tokens] because the reshape_and_cache_flash # op uses the slot_mapping's shape to determine the number of # actual tokens. - if self.is_kvcache_nvfp4: - # (B, 2*H, N, full_dim) -> ((B, N, H, full_dim), - # (B, N, H, full_dim)); - # K heads first, then V heads. - k_cache, v_cache = kv_cache.transpose(1, 2).split( - self.num_kv_heads, dim=-2 + if self.is_kvcache_nvfp4 and not self.use_native_nvfp4_kv_cache_update: + nvfp4_views = self._get_nvfp4_kv_cache_views(kv_cache) + nvfp4_kv_data = nvfp4_views.data + nvfp4_kv_block_scales = nvfp4_views.block_scales + nvfp4_slot_writer = self._nvfp4_slot_writer + assert nvfp4_slot_writer is not None + nvfp4_slot_writer( + key, + value, + slot_mapping, + nvfp4_kv_data, + nvfp4_kv_block_scales, + layer._k_scale, + layer._v_scale, + kv_layout=get_flashinfer_layout_string(self.kv_cache_layout), ) else: - # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) - k_cache, v_cache = kv_cache.transpose(1, 2).split( - self.head_size, dim=-1 + if self.is_kvcache_nvfp4: + # (B, 2*H, N, full_dim) -> ((B, N, H, full_dim), + # (B, N, H, full_dim)); + # K heads first, then V heads. + k_cache, v_cache = kv_cache.transpose(1, 2).split( + self.num_kv_heads, dim=-2 + ) + else: + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) + k_cache, v_cache = kv_cache.transpose(1, 2).split( + self.head_size, dim=-1 + ) + torch.ops._C_cache_ops.reshape_and_cache_flash( + key, + value, + k_cache, + v_cache, + slot_mapping, + self.kv_cache_dtype, + layer._k_scale, + layer._v_scale, ) - torch.ops._C_cache_ops.reshape_and_cache_flash( - key, - value, - k_cache, - v_cache, - slot_mapping, - self.cache_dtype, - layer._k_scale, - layer._v_scale, - ) def fast_plan_decode(