diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 2a26023f740b..bac2afc828b3 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG 9cd61de38763d712bb6ce56e2a02cc2bf718c89f + GIT_TAG b273a02621d3c94bc747df67204f591b67f38fae GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index d175add40b84..21cde2fbed16 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -147,7 +147,7 @@ Priority is **1 = highest** (tried first). --8<-- "gen:priority-mla" -> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise. +> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASH_ATTN_MLA_SPARSE_FA4` (then `FLASHINFER_MLA_SPARSE`) is preferred for low query-head counts (<= 16) on SM 10.x, while `FLASHMLA_SPARSE` is preferred otherwise. > > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details. diff --git a/tests/kernels/attention/test_flashinfer_mla_decode.py b/tests/kernels/attention/test_flashinfer_mla_decode.py index c3f7b4c96f78..fc34b8725746 100644 --- a/tests/kernels/attention/test_flashinfer_mla_decode.py +++ b/tests/kernels/attention/test_flashinfer_mla_decode.py @@ -175,7 +175,15 @@ def unexpected_helper(*args, **kwargs): backend, "prepare_sparse_mla_safe_lengths", unexpected_helper ) - out, lse = impl._run_mqa_kernel(query.squeeze(1), kv_cache, indices, counts) + out, lse = impl._forward_mqa_kernel( + query.squeeze(1), + kv_cache, + indices, + counts, + layer=None, + block_size=block_size, + is_decode=True, + ) reference = trtllm_batch_decode_with_kv_cache_mla( query=query, kv_cache=kv_cache.unsqueeze(1), diff --git a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py index 8532d2f43aee..8f31bd488aa4 100644 --- a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py +++ b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py @@ -238,6 +238,7 @@ def _sparse_order(kv_cache_dtype, head_size, num_heads=32, capability=SM100): ) sparse = { "FLASH_ATTN_MLA_SPARSE", + "FLASH_ATTN_MLA_SPARSE_FA4", "FLASHMLA_SPARSE", "FLASHINFER_MLA_SPARSE", "FLASHINFER_MLA_SPARSE_SM90", @@ -276,3 +277,36 @@ def test_sm90_nope_fp8_ds_mla_resolves_to_flashmla(): for name in order[: order.index("FLASHMLA_SPARSE")]: backend_cls = AttentionBackendEnum[name].get_class() assert not backend_cls.supports_kv_cache_dtype("fp8_ds_mla"), name + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA backends") +def test_sm100_576_bf16_low_heads_falls_back_from_fa4(monkeypatch): + """With FA4 unusable, low-head bf16 selection lands on FlashInfer.""" + import vllm.v1.attention.backends.mla.flashattn_mla_sparse as fa4_sparse + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + monkeypatch.setattr(fa4_sparse, "_fa4_cute_mla_available", lambda: "no FA4") + + def reasons(name): + backend_cls = AttentionBackendEnum[name].get_class() + return backend_cls.validate_configuration( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="auto", + block_size=64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=SM100, + attn_type="decoder", + ) + + order = _sparse_order("auto", 576, num_heads=16) + with set_current_vllm_config(VllmConfig()): + # The selector takes the first candidate with no rejection reasons. + selected = next(name for name in order if not reasons(name)) + assert order[0] == "FLASH_ATTN_MLA_SPARSE_FA4", order + assert selected == "FLASHINFER_MLA_SPARSE" diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index f7c47fe6a3ec..a2e216994328 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -34,6 +34,9 @@ ) from vllm.forward_context import set_forward_context from vllm.model_executor.layers.attention import mla_attention +from vllm.model_executor.layers.attention import ( + sparse_mla_attention as sparse_mla_module, +) from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha from vllm.model_executor.layers.attention.sparse_mla_attention import ( GLOBAL_TOPK_MASK_MAX_BYTES, @@ -46,6 +49,7 @@ ) from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability # This module's end-to-end sparse-backend correctness test supports CUDA and # ROCm. The remaining tests exercise CUDA-specific FlashMLA/FlashInfer APIs @@ -71,8 +75,11 @@ flashattn_mla_sparse as flashattn_sparse_module, ) from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( + FlashAttnMLASparseFA4Backend, + FlashAttnMLASparseFA4Impl, FlashAttnMLASparseImpl, FlashAttnMLASparseMetadataBuilder, + _fa4_cute_mla_available, ) from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseImpl, @@ -227,6 +234,7 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch): req_id_per_token=torch.zeros(1, dtype=torch.int32), block_table=torch.zeros((1, 1), dtype=torch.int32), block_size=32, + num_decode_tokens=1, ) recorded_scale = None @@ -1401,6 +1409,7 @@ def test_masked_mha_workspace_fits_single_request_boundary(max_query_len, expect ("FLASHMLA_SPARSE", 8, 112 * 1024), ("FLASHINFER_MLA_SPARSE", 4, 36 * 1024), ("FLASHINFER_MLA_SPARSE", 8, 64 * 1024), + ("FLASH_ATTN_MLA_SPARSE_FA4", 4, 36 * 1024), ], ) def test_masked_mha_workspace_guards_long_routing_policy( @@ -3569,56 +3578,6 @@ def test_hisparse_keeps_speculative_rows_as_decodes(): assert builder.reorder_batch_threshold == 4 -def test_flashinfer_hisparse_decode_runs_batched_attention(): - device = torch.device("cpu") - num_tokens = 6 - q = torch.randn(num_tokens, 2, 4, device=device) - topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=device) - physical_topk = topk + 10 - valid_counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=device) - kernel_shapes: list[torch.Size] = [] - - def convert_decode(self, *args, **kwargs): # noqa: ARG001 - return physical_topk, valid_counts - - def prepare_kernel(self, *args, **kwargs): # noqa: ARG001 - pass - - def run_kernel(self, q, cache, indices, counts): # noqa: ARG001 - kernel_shapes.append(q.shape) - return q[..., :1], None - - cache_handle = SimpleNamespace( - runtime=SimpleNamespace( - hot=SimpleNamespace(attention_cache=torch.empty(1, device=device)) - ), - ) - index_group = object.__new__(HiSparseMLAIndexGroup) - index_group.caches = [cache_handle] - index_group.convert_logical_to_physical_topk = MethodType( - convert_decode, index_group - ) - impl = object.__new__(FlashInferMLASparseImpl) - impl.topk_indices_buffer = topk - impl.index_group = index_group - impl.index_group_index = 0 - impl._prepare_mqa_kernel = MethodType(prepare_kernel, impl) - impl._run_mqa_kernel = MethodType(run_kernel, impl) - metadata = SimpleNamespace(num_decode_tokens=num_tokens) - - output, lse = FlashInferMLASparseImpl.forward_mqa( - impl, - q, - torch.empty(1, device=device), - metadata, - SimpleNamespace(), - ) - - assert kernel_shapes == [q.shape] - assert output.shape == (num_tokens, 2, 1) - assert lse is None - - @pytest.mark.skipif( not current_platform.is_device_capability_family(90), reason="FlashAttention MLA requires Hopper", @@ -3690,7 +3649,7 @@ def test_flashattn_shared_indices_match_independent_conversion( convert = triton_convert_req_index_to_global_index convert_spy = MagicMock(wraps=convert) - for module in (index_group_module, flashattn_sparse_module): + for module in (index_group_module, sparse_mla_module): monkeypatch.setattr( module, "triton_convert_req_index_to_global_index", convert_spy ) @@ -3735,48 +3694,19 @@ def run(): return_valid_counts=True, ) for impl, cache, actual in zip(impls, caches, outputs): - expected = impl._run_mqa_kernel( - q_nope, q_rope, cache, physical, counts, block_size + expected, _ = impl._forward_mqa_kernel( + (q_nope, q_rope), + cache, + physical, + counts, + layer=None, + block_size=block_size, + is_decode=num_decode_tokens == num_tokens, ) torch.testing.assert_close(actual, expected, rtol=0, atol=0) stream.synchronize() -def test_flashattn_hisparse_decode_uses_index_group(): - num_tokens = 4 - q_nope = torch.empty(num_tokens, 2, 3, device=DEVICE_TYPE) - q_rope = torch.empty(num_tokens, 2, 1, device=DEVICE_TYPE) - topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE) - physical = topk.clone() - counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE) - index_group = object.__new__(HiSparseMLAIndexGroup) - index_group.convert_logical_to_physical_topk = MagicMock( - return_value=(physical, counts) - ) - index_group.physical_kv_cache = MagicMock( - return_value=torch.empty(1, device=DEVICE_TYPE) - ) - impl = object.__new__(FlashAttnMLASparseImpl) - impl.topk_indices_buffer = topk - impl.index_group = index_group - impl.index_group_index = 0 - impl._run_mqa_kernel = MagicMock(return_value=q_nope[..., :1]) - metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64) - - output, lse = FlashAttnMLASparseImpl.forward_mqa( - impl, - (q_nope, q_rope), - torch.empty(1, device=DEVICE_TYPE), - metadata, - SimpleNamespace(), - ) - - assert output.shape == (num_tokens, 2, 1) - assert lse is None - index_group.convert_logical_to_physical_topk.assert_called_once() - impl._run_mqa_kernel.assert_called_once() - - def test_flashinfer_sm120_hisparse_decode_uses_index_group(): num_tokens = 4 q = torch.empty(num_tokens, 2, 4, device=DEVICE_TYPE) @@ -3811,6 +3741,131 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group(): impl._run_mqa_kernel.assert_called_once() +def test_sparse_mqa_plain_lane_wiring(monkeypatch): + """Under DCP the plain lane converts with the DCP filter at the flat stride.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla + + num_tokens, topk_tokens, block_size, head_dim = 3, 8, 4, 6 + # Flat row stride (8) differs from block size (4): the wrong one misaddresses. + kv_cache = torch.zeros(2, 8, head_dim)[:, :block_size] + q = torch.zeros(num_tokens, 2, head_dim) + metadata = SimpleNamespace( + num_decode_tokens=num_tokens, + block_size=block_size, + block_table=torch.zeros((num_tokens, 1), dtype=torch.int32), + req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32), + cp_kv_cache_interleave_size=2, + ) + index_filter = MagicMock(return_value=(None, None)) + monkeypatch.setattr(sparse_mla, "triton_filter_and_convert_dcp_index", index_filter) + + impl = object.__new__(FlashInferMLASparseImpl) + impl.topk_indices_buffer = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32) + impl.index_group = MagicMock() + impl.index_group_index = 0 + impl.dcp_world_size = 2 + impl.dcp_rank = 1 + impl._convert_logical_to_physical_topk = MagicMock() + impl._forward_mqa_kernel = MagicMock(return_value=(q, None)) + + impl.forward_mqa(q, kv_cache, metadata, None) + + impl._convert_logical_to_physical_topk.assert_not_called() + assert index_filter.call_args.kwargs["BLOCK_STRIDE_ROWS"] == 8 + impl._forward_mqa_kernel.assert_called_once() + + +@pytest.mark.parametrize( + "num_decode_tokens,all_resident,split_q", + [ + pytest.param(0, True, False, id="resident"), + pytest.param(0, False, False, id="staged"), + pytest.param(2, True, True, id="mixed"), + pytest.param(5, True, False, id="decode"), + ], +) +def test_sparse_mqa_hisparse_lanes( + monkeypatch, num_decode_tokens, all_resident, split_q +): + """HiSparse splits a batch into at most one decode and one prefill call.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla + + num_tokens, topk_tokens, block_size, head_dim = 5, 8, 4, 6 + num_prefill = num_tokens - num_decode_tokens + # Distinct block counts, so the KV tensor a call is handed names its lane. + host = torch.zeros(2, block_size, head_dim, dtype=torch.bfloat16) + hot = torch.zeros(6, block_size, head_dim, dtype=torch.bfloat16) + staged = torch.zeros(3, block_size, head_dim, dtype=torch.bfloat16) + q = torch.arange(num_tokens, dtype=torch.float32).view(num_tokens, 1, 1) + metadata = SimpleNamespace( + num_decode_tokens=num_decode_tokens, block_size=block_size + ) + + def converted(rows, fill): + return ( + torch.full((rows, topk_tokens), fill, dtype=torch.int32), + torch.full((rows,), rows, dtype=torch.int32), + ) + + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.physical_kv_cache = MagicMock(return_value=hot) + index_group.cache = MagicMock( + return_value=SimpleNamespace(all_context_pages_resident=all_resident) + ) + # The fill value names the converter that produced a call's indices. + index_group.convert_logical_to_physical_topk = MagicMock( + return_value=converted(num_decode_tokens, 1) + if num_decode_tokens + else converted(num_tokens, 2) + ) + index_group.stage_prefill_rows = MagicMock(return_value=(staged, None, None)) + calls: list[tuple] = [] + + def convert_staged(*args, **kwargs): + # Decode rows are consumed before a later conversion can reuse their buffer. + assert len(calls) == min(num_decode_tokens, 1) + return converted(num_prefill, 3) + + monkeypatch.setattr( + sparse_mla, + "triton_convert_req_index_to_global_index", + MagicMock(side_effect=convert_staged), + ) + + def hook(rows, cache, indices, counts, *, is_decode, **kwargs): + if split_q: + assert rows[0].tolist() == rows[1].tolist() + rows = rows[0] + calls.append((rows.flatten().tolist(), cache.shape[0], indices[0, 0].item())) + return torch.full((rows.shape[0], 1, 1), 1.0 if is_decode else 2.0), None + + impl = object.__new__(FlashInferMLASparseImpl) + impl.topk_indices_buffer = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32) + impl.index_group = index_group + impl.index_group_index = 0 + impl.dcp_world_size = 1 + impl._forward_mqa_kernel = hook + + out, lse = impl.forward_mqa((q, q) if split_q else q, host, metadata, None) + + q_rows = q.flatten().tolist() + expected = [] + if num_decode_tokens: + # Decode reads the hot cache through the group's conversion. + expected.append((q_rows[:num_decode_tokens], hot.shape[0], 1)) + if num_decode_tokens == 0 and all_resident: + expected.append((q_rows, hot.shape[0], 2)) + elif num_prefill: + expected.append((q_rows[num_decode_tokens:], staged.shape[0], 3)) + assert calls == expected + if not num_prefill: + # A decode-only batch returns before it looks up the prefill cache. + index_group.cache.assert_not_called() + assert lse is None + # A lane handed the whole batch would concatenate to more rows than this. + assert out.flatten().tolist() == [1.0] * num_decode_tokens + [2.0] * num_prefill + + def test_hisparse_resident_prefill_uses_attention_block_stride(): """A batch beyond the swap-row capacity must bypass swap_in. @@ -3887,6 +3942,250 @@ def test_hisparse_mixed_single_token_batch_is_not_decode_only(): assert not cache.decode_batch +def _fa4_impl(topk_indices_buffer, **fields): + impl = object.__new__(FlashAttnMLASparseFA4Impl) + impl.__dict__.update( + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + scale=1.0, + kv_cache_dtype="auto", + dcp_world_size=1, + dcp_rank=0, + need_to_return_lse_for_decode=False, + is_nope_mla=False, + _workspace_buffer=None, + bmm1_scale=None, + bmm2_scale=None, + # A non-index-producing layer does not wait on the indexer's ready event. + index_group=SparseMLAIndexGroupBuilder(topk_indices_buffer).register_layer( + False + )[0], + index_group_index=0, + ) + impl.__dict__.update(fields) + # topk_indices_buffer is a data descriptor, so it needs its own setter. + impl.topk_indices_buffer = topk_indices_buffer + return impl + + +def _fa4_inputs(counts, *, num_heads=16, topk=128, num_blocks=4, device=DEVICE_TYPE): + """Decode inputs whose token ``i`` has ``counts[i]`` valid rows, -1 padded.""" + device = torch.device(device) + block_size, kv_lora_rank, rope_dim = 64, 512, 64 + num_tokens = len(counts) + torch.manual_seed(0) + + def values(*shape): + return torch.rand(shape, dtype=torch.bfloat16, device=device) - 0.5 + + topk_indices = torch.full((num_tokens, topk), -1, dtype=torch.int32, device=device) + for tok, count in enumerate(counts): + topk_indices[tok, :count] = torch.randperm( + num_blocks * block_size, device=device + )[:count].to(torch.int32) + + return SimpleNamespace( + kv_lora_rank=kv_lora_rank, + rope_dim=rope_dim, + block_size=block_size, + topk=topk, + kv_cache=values(num_blocks, block_size, kv_lora_rank + rope_dim), + topk_indices=topk_indices, + valid_counts=torch.tensor(counts, dtype=torch.int32, device=device), + ql_nope=values(num_tokens, num_heads, kv_lora_rank), + q_pe=values(num_tokens, num_heads, rope_dim), + scale=(kv_lora_rank + rope_dim) ** -0.5, + metadata=SimpleNamespace( + req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32, device=device), + block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view( + 1, num_blocks + ), + block_size=block_size, + num_decode_tokens=num_tokens, + cp_kv_cache_interleave_size=1, + ), + ) + + +def _record_fa4_kernels(monkeypatch, num_heads, *, device): + """Replace both lanes with recorders whose outputs identify the lane.""" + import flashinfer.decode + + fa4_calls: list[dict] = [] + trtllm_calls: list[dict] = [] + + def lse(rows): + return torch.arange(rows * num_heads, dtype=torch.float32, device=device).view( + rows, num_heads + ) + + def fake_fa4(**kwargs): + fa4_calls.append(kwargs) + rows = kwargs["q"].shape[0] + out = torch.ones(rows, num_heads, 512, dtype=torch.bfloat16, device=device) + return (out, lse(rows)) if kwargs.get("return_softmax_lse") else out + + def fake_trtllm(**kwargs): + trtllm_calls.append(kwargs) + rows = kwargs["query"].shape[0] + out = torch.full( + (rows, 1, num_heads, 512), 2.0, dtype=torch.bfloat16, device=device + ) + return (out, lse(rows)) if kwargs.get("return_lse") else out + + monkeypatch.setattr(flashattn_sparse_module, "flash_attn_varlen_func", fake_fa4) + monkeypatch.setattr( + flashinfer.decode, "trtllm_batch_decode_with_kv_cache_mla", fake_trtllm + ) + return SimpleNamespace(fa4=fa4_calls, trtllm=trtllm_calls, lse=lse) + + +@pytest.mark.parametrize( + "num_heads,return_lse", [(128, False), (16, True)], ids=["heads128", "heads16_lse"] +) +def test_fa4_sparse_decode_kernel_correctness(num_heads, return_lse): + if not current_platform.is_device_capability_family(100): + pytest.skip("FA4 sparse MLA requires SM 10.x") + if (reason := _fa4_cute_mla_available()) is not None: + pytest.skip(reason) + counts = [0, 1, 127, 128, 129, 2047, 2048] + inputs = _fa4_inputs(counts, num_heads=num_heads, topk=2048, num_blocks=64) + impl = _fa4_impl( + topk_indices_buffer=inputs.topk_indices, + scale=inputs.scale, + need_to_return_lse_for_decode=return_lse, + ) + q = (inputs.ql_nope, inputs.q_pe) + + with torch.inference_mode(): + out, lse = impl.forward_mqa(q, inputs.kv_cache, inputs.metadata, None) + if not return_lse: + assert lse is None + else: + out_fused, lse_fused = impl.forward_mqa( + torch.cat(q, dim=-1), inputs.kv_cache, inputs.metadata, None + ) + assert lse.shape == (len(counts), num_heads) + assert torch.equal(out, out_fused) and torch.equal(lse, lse_fused) + + kv_flat = inputs.kv_cache.view(-1, inputs.kv_lora_rank + inputs.rope_dim) + keys = kv_flat[inputs.topk_indices.clamp(min=0).long()].float() + scores = torch.einsum("thd,tkd->thk", torch.cat(q, dim=-1).float(), keys) + scores = (scores * inputs.scale).masked_fill( + (inputs.topk_indices < 0)[:, None, :], float("-inf") + ) + # The all-sentinel row softmaxes to NaN here; the kernel writes (0, -inf). + probs = torch.softmax(scores, dim=-1).nan_to_num() + torch.testing.assert_close( + out.float(), + torch.einsum("thk,tkd->thd", probs, keys[..., : inputs.kv_lora_rank]), + rtol=0.01, + atol=0.01, + ) + if return_lse: + torch.testing.assert_close( + lse, torch.logsumexp(scores, dim=-1), rtol=0, atol=5e-3 + ) + + +@pytest.mark.parametrize("is_decode", [True, False], ids=["decode", "prefill"]) +def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode): + """The hook runs a decode batch on FA4 and every other batch on trtllm-gen.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla + + counts = [1, 2, 3, 4, 5, 6, 7] + num_tokens, dcp = len(counts), 2 + # Under DCP the hook gets the heads of every DCP rank, not just its own. + kernel_heads = 16 * dcp + inputs = _fa4_inputs(counts, num_heads=kernel_heads, device="cpu") + calls = _record_fa4_kernels(monkeypatch, kernel_heads, device="cpu") + # DCP without PCP (the gate rejects the pair) always fuses the query. + q = torch.cat((inputs.ql_nope, inputs.q_pe), dim=-1) + impl = _fa4_impl( + # One row longer than the batch: the varlen-scalar slices become observable. + topk_indices_buffer=inputs.topk_indices.new_zeros(num_tokens + 1, inputs.topk), + dcp_world_size=dcp, + need_to_return_lse_for_decode=True, + ) + impl._convert_logical_to_physical_topk = MagicMock() + dcp_filter = MagicMock(return_value=(inputs.topk_indices, inputs.valid_counts)) + monkeypatch.setattr(sparse_mla, "triton_filter_and_convert_dcp_index", dcp_filter) + inputs.metadata.num_decode_tokens = num_tokens if is_decode else 1 + + with torch.inference_mode(): + out, lse = impl.forward_mqa(q, inputs.kv_cache, inputs.metadata, None) + + dcp_filter.assert_called_once() + impl._convert_logical_to_physical_topk.assert_not_called() + assert out.shape == (num_tokens, kernel_heads, 512) + if is_decode: + (fa4,) = calls.fa4 + assert calls.trtllm == [] + assert fa4["return_softmax_lse"] is True + assert fa4["cu_seqlens_q"].tolist() == list(range(num_tokens + 1)) + assert (out == 1.0).all() + # FA4's LSE is already natural-log. + torch.testing.assert_close(lse, calls.lse(num_tokens), rtol=0, atol=0) + else: + (trtllm,) = calls.trtllm + assert calls.fa4 == [] + assert trtllm["return_lse"] is True + assert (out == 2.0).all() + # trtllm-gen's LSE is base 2; the DCP merge expects natural log. + expected = calls.lse(num_tokens) * math.log(2.0) + torch.testing.assert_close(lse, expected, rtol=0, atol=1e-5) + + +@pytest.mark.parametrize( + "config,expected", + [ + ({}, None), + ({"heads": 128}, None), + ({"heads": 12}, "heads"), + ({"dcp": 2}, None), + ({"dcp": 8}, "heads"), + ({"index_topk": 192}, "index_topk"), + ({"rope": 32}, "(kv_lora_rank, qk_rope_head_dim)"), + ({"flashinfer": False}, "FlashInfer"), + ], +) +def test_fa4_sparse_gate(monkeypatch, config, expected): + """``validate_configuration`` on a fixed SM100: one row per gate reason.""" + has_flashinfer = config.get("flashinfer", True) + monkeypatch.setattr("vllm.utils.flashinfer.has_flashinfer", lambda: has_flashinfer) + monkeypatch.setattr( + flashattn_sparse_module, "_fa4_cute_mla_available", lambda: None + ) + dcp = config.get("dcp", 1) + vllm_config = _build_sparse_dcp_vllm_config(config.get("heads", 16), dcp) + hf_text_config = vllm_config.model_config.hf_text_config + hf_text_config.index_topk = config.get("index_topk", 128) + hf_text_config.qk_rope_head_dim = config.get("rope", 64) + + with set_current_vllm_config(vllm_config): + reasons = FlashAttnMLASparseFA4Backend.validate_configuration( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="auto", + block_size=64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=DeviceCapability(10, 0), + attn_type="decoder", + use_dcp=dcp > 1, + ) + + if expected is None: + assert reasons == [] + else: + (reason,) = reasons + assert expected in reason + + def test_flashmla_cache_dtype_aliases_use_ds_layout(): # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer # stores kv_cache_dtype, so they cannot bypass the gate. diff --git a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py index 45ba1ac0dc55..b322acf0d5ee 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py @@ -13,15 +13,12 @@ from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache from vllm.platforms import current_platform from vllm.v1.attention.backend import MultipleOf -from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( - FlashInferMLASparseTRTLLMBackend, -) -from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend from vllm.v1.attention.backends.mla.sparse_utils import ( _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL, flat_kv_row_view, ) +from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs @@ -49,17 +46,27 @@ @pytest.mark.parametrize( "backend,cache_dtype,sm100,stride_alignment", [ - (FlashInferMLASparseTRTLLMBackend, "auto", True, 32 * 1152), - (FlashInferMLASparseTRTLLMBackend, "fp8", True, 32 * 576), - (FlashMLASparseBackend, "auto", True, 1152), - (FlashMLASparseBackend, "fp8_ds_mla", False, MultipleOf(64)), - (FlashMLASparseBackend, "fp8_ds_mla", True, 656), - (FlashMLASparseBackend, "nvfp4_ds_mla", True, 352), + (AttentionBackendEnum.FLASHINFER_MLA_SPARSE, "auto", True, 32 * 1152), + (AttentionBackendEnum.FLASHINFER_MLA_SPARSE, "fp8", True, 32 * 576), + pytest.param( + AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE_FA4, + "auto", + True, + 64 * 1152, + marks=pytest.mark.skipif( + not current_platform.is_cuda(), reason="CUDA only" + ), + ), + (AttentionBackendEnum.FLASHMLA_SPARSE, "auto", True, 1152), + (AttentionBackendEnum.FLASHMLA_SPARSE, "fp8_ds_mla", False, MultipleOf(64)), + (AttentionBackendEnum.FLASHMLA_SPARSE, "fp8_ds_mla", True, 656), + (AttentionBackendEnum.FLASHMLA_SPARSE, "nvfp4_ds_mla", True, 352), ], ) def test_allocation_and_warmup_follow_addressing_mode( monkeypatch, layout, backend, cache_dtype, sm100, stride_alignment, device ): + backend = backend.get_class() monkeypatch.setattr( current_platform, "is_device_capability_family", lambda family: sm100 ) diff --git a/tests/v1/e2e/general/test_hisparse.py b/tests/v1/e2e/general/test_hisparse.py index 2409b5cb690b..d9e1fabf03f3 100644 --- a/tests/v1/e2e/general/test_hisparse.py +++ b/tests/v1/e2e/general/test_hisparse.py @@ -83,13 +83,20 @@ def _get_hisparse_worker(runner: VllmRunner) -> HiSparseConnectorWorker: not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA" ) @pytest.mark.parametrize( - "with_offloading", [False, True], ids=["standalone", "offload"] + "with_offloading,attention_backend", + [ + pytest.param(False, None, id="standalone"), + pytest.param(True, None, id="offload"), + # Offloading does not touch the attention kernels: FA4 runs standalone. + pytest.param(False, "FLASH_ATTN_MLA_SPARSE_FA4", id="standalone-fa4"), + ], ) @fork_new_process_for_each_test def test_hisparse_spill_and_prefix_restore( monkeypatch: pytest.MonkeyPatch, vllm_runner: type[VllmRunner], with_offloading: bool, + attention_backend: str | None, ): """Spilled prefixes restore and FULL-graph decode writes reach host KV. @@ -101,6 +108,8 @@ def test_hisparse_spill_and_prefix_restore( capability = current_platform.get_device_capability() if capability is None or capability.major < 9: pytest.skip("Sparse MLA requires Hopper or newer") + if attention_backend and not current_platform.is_device_capability_family(100): + pytest.skip("FA4 sparse MLA requires SM 10.x") monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip") @@ -142,6 +151,7 @@ def test_hisparse_spill_and_prefix_restore( ) ), kv_transfer_config=kv_transfer_config, + attention_backend=attention_backend, block_size=64, max_model_len=320, max_num_batched_tokens=1024, diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index db9e8dbc88ca..806faba50774 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -661,6 +661,7 @@ def __init__( "FLASHMLA_SPARSE", "FLASHINFER_MLA_SPARSE", "FLASHINFER_MLA_SPARSE_SM120", + "FLASH_ATTN_MLA_SPARSE_FA4", "DEEPSEEK_V32_INDEXER", ): from vllm.v1.attention.backends.mla.compressor_utils import ( @@ -751,7 +752,7 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: def _uses_flat_kv_cache(self) -> bool: backend = self.attn_backend.get_name() - return backend == "FLASHINFER_MLA_SPARSE" or ( + return backend in ("FLASHINFER_MLA_SPARSE", "FLASH_ATTN_MLA_SPARSE_FA4") or ( backend == "FLASHMLA_SPARSE" and self.kv_cache_dtype not in ("fp8_ds_mla", "nvfp4_ds_mla") ) @@ -1857,6 +1858,9 @@ def _use_masked_mha( seq_len: int, has_context: bool, ) -> bool: + if backend_name == "FLASH_ATTN_MLA_SPARSE_FA4": + # Its prefill batches run on FlashInfer's kernel, so it shares those rows. + backend_name = "FLASHINFER_MLA_SPARSE" thresholds = ( _MASKED_MHA_THRESHOLDS.get((qk_head_dim, v_head_dim), {}) .get(backend_name, {}) diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 37883593421a..a8fa74f91f62 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -36,12 +36,18 @@ is_quantized_kv_cache, np_to_pinned_tensor, ) -from vllm.v1.attention.backend import AttentionMetadataBuilder +from vllm.v1.attention.backend import AttentionLayer, AttentionMetadataBuilder from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from vllm.v1.attention.backends.mla.index_group import ( + HiSparseMLAIndexGroup, SparseMLAIndexGroup, SparseMLAIndexGroupBuilder, ) +from vllm.v1.attention.backends.mla.sparse_utils import ( + flat_kv_row_view, + triton_convert_req_index_to_global_index, + triton_filter_and_convert_dcp_index, +) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.dcp import MLADCPManager from vllm.v1.attention.ops.merge_attn_states import merge_attn_states @@ -630,6 +636,9 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic """Sparse MLA base with dense and masked-MHA prefill paths.""" is_sparse = True + # Reuse the index group's top-k conversion only for uniform decode batches + # of multi-layer groups; other batches convert per layer. + group_conversion_decode_only: ClassVar[bool] = False def __init__( self, @@ -728,6 +737,154 @@ def _convert_logical_to_physical_topk( return_valid_counts=return_valid_counts, ) + def forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: T, + layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + assert self.topk_indices_buffer is not None + num_tokens = q[0].shape[0] if isinstance(q, tuple) else q.shape[0] + topk = self.topk_indices_buffer[:num_tokens] + num_decode_tokens = attn_metadata.num_decode_tokens + block_size = attn_metadata.block_size + index_group = self.index_group + if not isinstance(index_group, HiSparseMLAIndexGroup): + _, block_stride_rows = flat_kv_row_view(kv_c_and_k_pe_cache, block_size) + if self.dcp_world_size > 1: + physical_topk, valid_counts = triton_filter_and_convert_dcp_index( + attn_metadata.req_id_per_token[:num_tokens], + attn_metadata.block_table, + topk, + dcp_size=self.dcp_world_size, + dcp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=topk.shape[1], + return_valid_counts=True, + ) + elif self.group_conversion_decode_only and not ( + index_group is not None + and index_group.num_layers > 1 + and num_decode_tokens >= num_tokens + ): + physical_topk, valid_counts = triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[:num_tokens], + attn_metadata.block_table, + topk, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=topk.shape[1], + return_valid_counts=True, + ) + else: + physical_topk, valid_counts = self._convert_logical_to_physical_topk( + topk, + attn_metadata, + block_stride_rows=block_stride_rows, + return_valid_counts=True, + ) + return self._forward_mqa_kernel( + q, + kv_c_and_k_pe_cache, + physical_topk, + valid_counts, + layer=layer, + block_size=block_size, + is_decode=num_decode_tokens >= num_tokens, + ) + + assert self.dcp_world_size == 1 + layer_index = self.index_group_index + hot_cache = index_group.physical_kv_cache(layer_index).view( + kv_c_and_k_pe_cache.dtype + ) + assert hot_cache.is_contiguous(), "HiSparse hot buffer must be contiguous" + decode_out = decode_lse = None + if num_decode_tokens: + physical_topk, valid_counts = index_group.convert_logical_to_physical_topk( + layer_index, + topk[:num_decode_tokens], + attn_metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + q_decode = ( + (q[0][:num_decode_tokens], q[1][:num_decode_tokens]) + if isinstance(q, tuple) + else q[:num_decode_tokens] + ) + decode_out, decode_lse = self._forward_mqa_kernel( + q_decode, + hot_cache, + physical_topk, + valid_counts, + layer=layer, + block_size=block_size, + is_decode=True, + ) + if num_decode_tokens >= num_tokens: + return decode_out, decode_lse + + cache = index_group.cache(layer_index) + if num_decode_tokens == 0 and cache.all_context_pages_resident: + physical_topk, valid_counts = index_group.convert_logical_to_physical_topk( + layer_index, + topk, + attn_metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + prefill_cache = hot_cache + else: + prefill_cache, block_table, req_ids = index_group.stage_prefill_rows( + layer_index, kv_c_and_k_pe_cache, attn_metadata + ) + physical_topk, valid_counts = triton_convert_req_index_to_global_index( + req_ids, + block_table, + topk[num_decode_tokens:], + BLOCK_SIZE=block_size, + NUM_TOPK_TOKENS=topk.shape[1], + return_valid_counts=True, + ) + q_prefill = ( + (q[0][num_decode_tokens:], q[1][num_decode_tokens:]) + if isinstance(q, tuple) + else q[num_decode_tokens:] + ) + prefill_out, prefill_lse = self._forward_mqa_kernel( + q_prefill, + prefill_cache, + physical_topk, + valid_counts, + layer=layer, + block_size=block_size, + is_decode=False, + ) + if decode_out is None: + return prefill_out, prefill_lse + out = torch.cat((decode_out, prefill_out)) + if decode_lse is None: + return out, None + assert prefill_lse is not None + return out, torch.cat((decode_lse, prefill_lse)) + + def _forward_mqa_kernel( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + valid_counts: torch.Tensor, + *, + layer: AttentionLayer, + block_size: int, + is_decode: bool, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + raise NotImplementedError + @staticmethod def masked_mha_workspace_fits(prefill: MLACommonPrefillMetadata) -> bool: """Whether this prefill batch's top-k masks fit the workspace.""" diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 401a98f66feb..7a28dcbc96d7 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -105,9 +105,10 @@ def _get_backend_priorities( ] else: # BF16 KV Cache - # Prefer FlashInfer at low head counts (FlashMLA uses padding) + # Prefer FA4, then FlashInfer at low head counts (FlashMLA uses padding) if num_heads is not None and num_heads <= 16: sparse_backends = [ + AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE_FA4, AttentionBackendEnum.FLASHINFER_MLA_SPARSE, AttentionBackendEnum.FLASHMLA_SPARSE, ] diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 0016eafedc8a..47347c8fe84e 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math from dataclasses import dataclass +from functools import cache from typing import Any, ClassVar import torch @@ -22,11 +24,14 @@ MultipleOf, ) from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla -from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + FlashInferMLASparseImpl, + FlashInferMLASparseTRTLLMBackend, + FlashInferMLASparseTRTLLMMetadataBuilder, +) from vllm.v1.attention.backends.mla.sparse_utils import ( align_blocks_to_rows, flat_kv_row_view, - triton_convert_req_index_to_global_index, ) from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.kv_cache_interface import AttentionSpec @@ -153,6 +158,8 @@ def update_draft_decode_metadata(self, metadata: SparseMLACommonMetadata) -> Non class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): + group_conversion_decode_only: ClassVar[bool] = True + def __init__( self, num_heads: int, @@ -205,144 +212,23 @@ def __init__( ) self.supports_quant_query_input = False - def forward_mqa( + def _forward_mqa_kernel( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - kv_c_and_k_pe_cache: torch.Tensor, - attn_metadata: FlashAttnMLASparseMetadata, + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + valid_counts: torch.Tensor, + *, layer: AttentionLayer, + block_size: int, + is_decode: bool, ) -> tuple[torch.Tensor, torch.Tensor | None]: if not isinstance(q, tuple): raise NotImplementedError( "FlashAttnMLASparseImpl expects split (q_nope, q_rope) input." ) q_nope, q_rope = q - num_actual_toks = q_rope.shape[0] - - assert self.topk_indices_buffer is not None - topk_indices = self.topk_indices_buffer[:num_actual_toks] - index_group = self.index_group - if isinstance(index_group, HiSparseMLAIndexGroup): - num_decode_tokens = attn_metadata.num_decode_tokens - outputs = [] - if num_decode_tokens: - physical_topk, valid_counts = ( - index_group.convert_logical_to_physical_topk( - self.index_group_index, - topk_indices[:num_decode_tokens], - attn_metadata, - block_stride_rows=None, - return_valid_counts=True, - ) - ) - outputs.append( - self._run_mqa_kernel( - q_nope[:num_decode_tokens], - q_rope[:num_decode_tokens], - index_group.physical_kv_cache(self.index_group_index).view( - kv_c_and_k_pe_cache.dtype - ), - physical_topk, - valid_counts, - attn_metadata.block_size, - ) - ) - if num_decode_tokens < num_actual_toks: - cache = index_group.cache(self.index_group_index) - if num_decode_tokens == 0 and cache.all_context_pages_resident: - physical_topk, valid_counts = ( - index_group.convert_logical_to_physical_topk( - self.index_group_index, - topk_indices, - attn_metadata, - block_stride_rows=None, - return_valid_counts=True, - ) - ) - prefill_cache = index_group.physical_kv_cache( - self.index_group_index - ).view(kv_c_and_k_pe_cache.dtype) - else: - prefill_cache, block_table, req_ids = ( - index_group.stage_prefill_rows( - self.index_group_index, - kv_c_and_k_pe_cache, - attn_metadata, - ) - ) - physical_topk, valid_counts = ( - triton_convert_req_index_to_global_index( - req_ids, - block_table, - topk_indices[num_decode_tokens:], - BLOCK_SIZE=attn_metadata.block_size, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) - ) - outputs.append( - self._run_mqa_kernel( - q_nope[num_decode_tokens:], - q_rope[num_decode_tokens:], - prefill_cache, - physical_topk, - valid_counts, - attn_metadata.block_size, - ) - ) - return torch.cat(outputs) if len(outputs) > 1 else outputs[0], None - - kv_rows, block_stride_rows = flat_kv_row_view( - kv_c_and_k_pe_cache, attn_metadata.block_size - ) - if ( - index_group is not None - and index_group.num_layers > 1 - and attn_metadata.num_decode_tokens == num_actual_toks - ): - topk_indices, valid_counts = self._convert_logical_to_physical_topk( - topk_indices, - attn_metadata, - block_stride_rows=block_stride_rows, - return_valid_counts=True, - ) - else: - topk_indices, valid_counts = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[:num_actual_toks], - attn_metadata.block_table, - topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) - return ( - self._run_mqa_kernel( - q_nope, - q_rope, - kv_rows, - topk_indices, - valid_counts, - attn_metadata.block_size, - cache_is_flat=True, - ), - None, - ) - - def _run_mqa_kernel( - self, - q_nope: torch.Tensor, - q_rope: torch.Tensor, - kv_cache: torch.Tensor, - topk_indices: torch.Tensor, - valid_counts: torch.Tensor, - block_size: int, - *, - cache_is_flat: bool = False, - ) -> torch.Tensor: - kv_rows = ( - kv_cache if cache_is_flat else flat_kv_row_view(kv_cache, block_size)[0] - ) + kv_rows, _ = flat_kv_row_view(kv_cache, block_size) cu_seqlens_q = self.cu_seqlens_q_buffer[: q_rope.shape[0] + 1] v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1) @@ -370,4 +256,167 @@ def _run_mqa_kernel( causal=True, fa_version=3, ) - return out + return out, None + + +def _fa4_cute_mla_available() -> str | None: + """Reason the FA4 cute-DSL MLA entry point is unusable, or None.""" + from vllm.vllm_flash_attn import flash_attn_interface as fa + + if not fa.is_fa_version_supported(4): + return f"FA4 unavailable: {fa.fa_version_unsupported_reason(4)}" + try: + from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd # noqa: F401 + except Exception as e: # cute-DSL / cutlass import chain + return f"FA4 cute-DSL entry point failed to import: {e!r}" + return None + + +class FlashAttnMLASparseFA4Backend(FlashInferMLASparseTRTLLMBackend): + """FA4 for is_decode kernel calls; trtllm-gen for the rest.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + # The qv kernel asserts every descale is None: BF16 KV cache only. + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"] + + @staticmethod + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: + return [MultipleOf(64)] + + @staticmethod + def get_name() -> str: + return "FLASH_ATTN_MLA_SPARSE_FA4" + + @staticmethod + def get_builder_cls() -> type["FlashAttnMLASparseFA4MetadataBuilder"]: + return FlashAttnMLASparseFA4MetadataBuilder + + @staticmethod + def get_impl_cls() -> type[MLAAttentionImpl[Any]]: + return FlashAttnMLASparseFA4Impl + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + # 512 latent + 64 RoPE; split back apart for the qv kernel. + return [576] + + @classmethod + def supports_combination(cls, *args: Any, **kwargs: Any) -> str | None: + from vllm.config import get_current_vllm_config_or_none + from vllm.utils.flashinfer import has_flashinfer + + if not has_flashinfer(): + return "FA4 sparse MLA runs its prefill batches on FlashInfer's kernel" + if (reason := _fa4_cute_mla_available()) is not None: + return reason + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return None + if vllm_config.model_config is not None: + hf_config = vllm_config.model_config.hf_text_config + topk = getattr(hf_config, "index_topk", None) + # gather_kv_indices' last dim must be a whole number of n-tiles. + if topk is None or topk % 128 != 0: + return f"FA4 sparse MLA requires index_topk % 128 == 0, got {topk}" + dims = ( + getattr(hf_config, "kv_lora_rank", None), + getattr(hf_config, "qk_rope_head_dim", None), + ) + if dims != (512, 64): + return ( + "FA4 sparse MLA requires (kv_lora_rank, qk_rope_head_dim) " + f"== (512, 64), got {dims}" + ) + # Under DCP the query is all-gathered: the kernel sees heads * dcp_size. + dcp_size = vllm_config.parallel_config.decode_context_parallel_size + num_heads = vllm_config.model_config.get_num_attention_heads( + vllm_config.parallel_config + ) + # 128 gathered heads take FA's prefill kernel, unvalidated under DCP. + head_counts = (8, 16, 32, 64) if dcp_size > 1 else (8, 16, 32, 64, 128) + if num_heads * dcp_size not in head_counts: + return ( + f"FA4 sparse MLA requires {head_counts} gathered query heads, " + f"got num_heads={num_heads} * dcp_size={dcp_size}" + ) + return super().supports_combination(*args, **kwargs) + + +class FlashAttnMLASparseFA4MetadataBuilder(FlashInferMLASparseTRTLLMMetadataBuilder): + """The trtllm-gen builder, for the prefill lane's DCP workspace pre-size. + + Its UNIFORM_BATCH support and varlen decode bound also fit FA4, which runs + one query row per token and bakes the row count into its varlen scalars. + """ + + +@cache +def _fa4_varlen_scalars( + num_tokens: int, num_kv_rows: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Per-token varlen scalars; never evict, CUDA graphs bake in the pointers.""" + cu_seqlens_q = torch.arange(num_tokens + 1, dtype=torch.int32, device=device) + cu_seqlens_k = torch.zeros(num_tokens + 1, dtype=torch.int32, device=device) + seqused_k = torch.full((num_tokens,), num_kv_rows, dtype=torch.int32, device=device) + return cu_seqlens_q, cu_seqlens_k, seqused_k + + +class FlashAttnMLASparseFA4Impl(FlashInferMLASparseImpl): + # FA4 emits a natural-log LSE; prefill converts trtllm-gen's log2 LSE. + lse_base_on_e: bool = True + + def _forward_mqa_kernel( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + valid_counts: torch.Tensor, + *, + layer: AttentionLayer, + block_size: int, + is_decode: bool, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if not is_decode: + out, lse = super()._forward_mqa_kernel( + q, + kv_cache, + topk_indices, + valid_counts, + layer=layer, + block_size=block_size, + is_decode=is_decode, + ) + return out, None if lse is None else lse * math.log(2.0) + if isinstance(q, tuple): + ql_nope, q_pe = q + else: + # 16B-aligned halves of the fused query; FA4 reads them in place. + ql_nope, q_pe = q[..., : self.kv_lora_rank], q[..., self.kv_lora_rank :] + kv_rows, _ = flat_kv_row_view(kv_cache, block_size) + num_tokens = q_pe.shape[0] + num_kv_rows = kv_rows.shape[0] + assert self.topk_indices_buffer is not None + cu_seqlens_q, cu_seqlens_k, seqused_k = _fa4_varlen_scalars( + self.topk_indices_buffer.shape[0], num_kv_rows, kv_rows.device + ) + # An all-sentinel row yields (0, -inf), the DCP merge identity; no masking. + kernel_out = flash_attn_varlen_func( + q=q_pe, + k=kv_rows[:, self.kv_lora_rank :].unsqueeze(1), + v=kv_rows[:, : self.kv_lora_rank].unsqueeze(1), + q_v=ql_nope, + max_seqlen_q=1, + cu_seqlens_q=cu_seqlens_q[: num_tokens + 1], + max_seqlen_k=num_kv_rows, + cu_seqlens_k=cu_seqlens_k[: num_tokens + 1], + seqused_k=seqused_k[:num_tokens], + gather_kv_indices=topk_indices, + # Part of FA4's compile key: pass it unconditionally. + gather_kv_valid_length=valid_counts, + softmax_scale=self.scale, + # Causality is in the top-k list; causal=True would clamp flat rows. + causal=False, + fa_version=4, + return_softmax_lse=self.need_to_return_lse_for_decode, + ) + return kernel_out if self.need_to_return_lse_for_decode else (kernel_out, None) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 21eb66aae490..d673868975ab 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -34,8 +34,6 @@ align_blocks_to_rows, flat_kv_row_view, prepare_sparse_mla_safe_lengths, - triton_convert_req_index_to_global_index, - triton_filter_and_convert_dcp_index, ) from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec @@ -129,14 +127,12 @@ def supports_combination( and parallel_config.decode_context_parallel_size > 1 ): return ( - "FLASHINFER_MLA_SPARSE does not support combined PCP+DCP; " + f"{cls.get_name()} does not support combined PCP+DCP; " "use FLASHMLA_SPARSE, which gathers each DCP KV shard before " "running the rank-local PCP prefill queries" ) if kv_cache_dtype == "fp8_ds_mla": - return ( - "FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype" - ) + return f"{cls.get_name()} SM10 does not support fp8_ds_mla kv-cache dtype" # FlashInfer MLA sparse SM10 kernel requires qk_nope_head_dim in [128, 192]. if vllm_config.model_config is not None: @@ -492,135 +488,6 @@ def __init__( # for query and kv_cache (mixed bf16+fp8 is not supported). self.supports_quant_query_input = True - def forward_mqa( - self, - q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - kv_c_and_k_pe_cache: torch.Tensor, - attn_metadata: FlashInferMLASparseMetadata, - layer: AttentionLayer, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - if isinstance(q, tuple): - ql_nope, q_pe = q - if q_pe.shape[-1] == 0 and ql_nope.is_contiguous(): - q = ql_nope - else: - q = torch.cat(q, dim=-1) - - num_actual_toks = q.shape[0] - - assert self.topk_indices_buffer is not None - topk_indices = self.topk_indices_buffer[:num_actual_toks] - - self._prepare_mqa_kernel(layer, q.device) - - index_group = self.index_group - if isinstance(index_group, HiSparseMLAIndexGroup): - num_decode_tokens = attn_metadata.num_decode_tokens - decode_out: torch.Tensor | None = None - decode_lse: torch.Tensor | None = None - if num_decode_tokens > 0: - physical_topk, valid_counts = ( - index_group.convert_logical_to_physical_topk( - self.index_group_index, - topk_indices[:num_decode_tokens], - attn_metadata, - block_stride_rows=None, - return_valid_counts=True, - ) - ) - decode_out, decode_lse = self._run_mqa_kernel( - q[:num_decode_tokens], - index_group.physical_kv_cache(self.index_group_index).view( - kv_c_and_k_pe_cache.dtype - ), - physical_topk, - valid_counts, - ) - if num_decode_tokens == num_actual_toks: - return decode_out, decode_lse - - cache = index_group.cache(self.index_group_index) - if num_decode_tokens == 0 and cache.all_context_pages_resident: - physical_topk, valid_counts = ( - index_group.convert_logical_to_physical_topk( - self.index_group_index, - topk_indices, - attn_metadata, - block_stride_rows=None, - return_valid_counts=True, - ) - ) - return self._run_mqa_kernel( - q, - index_group.physical_kv_cache(self.index_group_index).view( - kv_c_and_k_pe_cache.dtype - ), - physical_topk, - valid_counts, - ) - - prefill_cache, block_table, req_ids = index_group.stage_prefill_rows( - self.index_group_index, kv_c_and_k_pe_cache, attn_metadata - ) - prefill_indices, prefill_lens = triton_convert_req_index_to_global_index( - req_ids, - block_table, - topk_indices[num_decode_tokens:], - BLOCK_SIZE=attn_metadata.block_size, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) - prefill_out, prefill_lse = self._run_mqa_kernel( - q[num_decode_tokens:], - prefill_cache, - prefill_indices, - prefill_lens, - ) - if decode_out is None: - return prefill_out, prefill_lse - output = torch.cat((decode_out, prefill_out)) - if decode_lse is None: - return output, None - assert prefill_lse is not None - return output, torch.cat((decode_lse, prefill_lse)) - - kv_rows, block_stride_rows = flat_kv_row_view( - kv_c_and_k_pe_cache, attn_metadata.block_size - ) - - if self.dcp_world_size > 1: - topk_indices_physical, seq_lens = triton_filter_and_convert_dcp_index( - attn_metadata.req_id_per_token[:num_actual_toks], - attn_metadata.block_table, - topk_indices, - dcp_size=self.dcp_world_size, - dcp_rank=self.dcp_rank, - cp_kv_cache_interleave_size=(attn_metadata.cp_kv_cache_interleave_size), - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) - else: - topk_indices_physical, seq_lens = self._convert_logical_to_physical_topk( - topk_indices, - attn_metadata, - block_stride_rows=block_stride_rows, - return_valid_counts=True, - ) - - return self._run_mqa_kernel( - q, - # Block sizes and strides have been aligned to 32 rows. - kv_rows.view( - -1, - math.gcd(block_stride_rows, attn_metadata.block_size, 64), - kv_rows.shape[-1], - ), - topk_indices_physical, - seq_lens, - ) - def _prepare_mqa_kernel( self, layer: AttentionLayer, @@ -648,7 +515,6 @@ def autotune_hisparse_decode(self, layer: AttentionLayer) -> None: kv_cache = runtime.hot.attention_cache num_tokens = runtime.max_num_reqs topk_tokens = self.topk_indices_buffer.shape[1] - self._prepare_mqa_kernel(layer, kv_cache.device) q_dtype = ( current_platform.fp8_dtype() @@ -679,22 +545,48 @@ def autotune_hisparse_decode(self, layer: AttentionLayer) -> None: dtype=torch.int32, device=kv_cache.device, ) - self._run_mqa_kernel(q, kv_cache, topk_indices, seq_lens) + self._forward_mqa_kernel( + q, + kv_cache, + topk_indices, + seq_lens, + layer=layer, + block_size=runtime.hot.block_size, + is_decode=True, + ) - def _run_mqa_kernel( + def _forward_mqa_kernel( self, - q: torch.Tensor, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], kv_cache: torch.Tensor, topk_indices: torch.Tensor, seq_lens: torch.Tensor, + *, + layer: AttentionLayer, + block_size: int, + is_decode: bool, ) -> tuple[torch.Tensor, torch.Tensor | None]: + if isinstance(q, tuple): + ql_nope, q_pe = q + if q_pe.shape[-1] == 0 and ql_nope.is_contiguous(): + q = ql_nope + else: + q = torch.cat(q, dim=-1) + + self._prepare_mqa_kernel(layer, q.device) assert self._workspace_buffer is not None assert self.bmm1_scale is not None assert self.bmm2_scale is not None from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla - kv_cache = kv_cache.view(q.dtype) + kv_rows, block_stride_rows = flat_kv_row_view( + kv_cache.view(q.dtype), block_size + ) + # Block sizes and strides have been aligned to 32 rows. + kv_cache = kv_rows.view( + -1, math.gcd(block_stride_rows, block_size, 64), kv_rows.shape[-1] + ) # Single-token sparse decode. trtllm-gen requires the q_len_per_request # dim, but the sparse attention mask is fully per-token (each query token diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 207c112a384c..e1c4c8d06741 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -125,6 +125,10 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): FLASH_ATTN_MLA_SPARSE = ( "vllm.v1.attention.backends.mla.flashattn_mla_sparse.FlashAttnMLASparseBackend" ) + FLASH_ATTN_MLA_SPARSE_FA4 = ( + "vllm.v1.attention.backends.mla.flashattn_mla_sparse." + "FlashAttnMLASparseFA4Backend" + ) MINIMAX_M3_SPARSE = ( "vllm.models.minimax_m3.common.sparse_attention.MiniMaxM3SparseBackend" ) diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index 3c4c0d8948be..c746dd0a54a8 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -214,6 +214,8 @@ def flash_attn_varlen_func( aux_tensors=None, aux_tensor_leading_dims=None, dynamic_causal: "torch.Tensor | None" = None, + gather_kv_indices=None, + gather_kv_valid_length=None, ): """dropout_p should be set to 0.0 during evaluation Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads @@ -262,9 +264,10 @@ def flash_attn_varlen_func( testing only. The returned probabilities are not guaranteed to be correct (they might not have the right scaling). seqused_k: (batch_size,), dtype torch.int32. Actual number of key tokens per - sequence. Mutually exclusive with cu_seqlens_k, and required when + sequence. Mutually exclusive with cu_seqlens_k except on FA4, which reads + cu_seqlens_k as the KV offset and seqused_k as the length. Required when block_table is given. - q_v: optional value tensor fused with q. FA3 only. + q_v: optional value tensor fused with q. FA3 and FA4 only. block_table: (batch_size, max_blocks_per_seq), dtype torch.int32. Paged KV cache block table. Requires seqused_k. return_softmax_lse: bool. Whether to also return softmax_lse. @@ -289,18 +292,21 @@ def flash_attn_varlen_func( aux_tensors: auxiliary tensors consumed by mask_mod. FA4 only. aux_tensor_leading_dims: leading dimensions of each entry in aux_tensors. dynamic_causal: optional per-sequence causal offsets. FA4 only. + gather_kv_indices: (total_q, k) int32 KV rows per query token, `-1` padded, + k a multiple of 128. Requires q_v, excludes block_table. FA4 only. + gather_kv_valid_length: (total_q,) int32, real entries per row. FA4 only. Return: out: (total, nheads, headdim). - softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The - logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax - normalization factor). + softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen), or + (total_q, nheads) on FA4's q_v path. The logsumexp of each row of the matrix + QK^T * scaling (e.g., log of the softmax normalization factor). """ assert cu_seqlens_k is not None or seqused_k is not None, ( "cu_seqlens_k or seqused_k must be provided" ) - assert cu_seqlens_k is None or seqused_k is None, ( + assert fa_version == 4 or cu_seqlens_k is None or seqused_k is None, ( "cu_seqlens_k and seqused_k cannot be provided at the same time" ) assert block_table is None or seqused_k is not None, ( @@ -311,6 +317,9 @@ def flash_attn_varlen_func( f"Fused FP8 output (output_scale) is only supported by FA4, " f"got fa_version={fa_version}" ) + assert fa_version == 4 or ( + gather_kv_indices is None and gather_kv_valid_length is None + ), f"gather_kv_* is only supported by FA4, got fa_version={fa_version}" if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) @@ -446,12 +455,15 @@ def flash_attn_varlen_func( q, k, v, + qv=q_v, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, seqused_k=seqused_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, page_table=block_table, + gather_kv_indices=gather_kv_indices, + gather_kv_valid_length=gather_kv_valid_length, softmax_scale=softmax_scale, causal=causal, dynamic_causal=dynamic_causal,