From a09db3f14e6d015170407f7ca5c4aa7377c11ef8 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Mon, 21 Sep 2026 06:46:51 +0000 Subject: [PATCH 1/9] [Attention] Share sparse MLA MQA cache handling FA3 and FlashInfer each carry their own copy of the sparse MLA index conversion and of the HiSparse cache resolution and mixed-batch split. Move both into SparseMLACommonImpl.forward_mqa. A backend now implements only _forward_mqa_kernel(q, kv_cache, topk_indices, valid_counts, *, layer, block_size, is_decode), including its query formatting and kernel setup. The arguments reaching the FA3 and the trtllm-gen kernels are unchanged on every path (plain, DCP, HiSparse decode-only, mixed, resident and staged prefill). FA3 keeps #59464's rule through a class flag, group_conversion_decode_only: its plain path reuses the index group's conversion only for uniform decode batches of multi-layer groups and converts per layer otherwise. FlashInfer keeps converting through the group for every batch. Also guards the HiSparse hot buffer layout with a contiguity assert and names the gated backend in the TRTLLM backend's reasons through cls.get_name(). Tests cover the hook contract, the plain and DCP index wiring, the four HiSparse lanes, that decode indices are consumed before prefill can overwrite them, and the FA3 plain path for leader and follower layers. Co-authored-by: Lucas Wilkinson Co-authored-by: Claude Signed-off-by: Simon Veitner --- .../attention/test_flashinfer_mla_decode.py | 10 +- .../v1/attention/test_sparse_mla_backends.py | 403 +++++++++++++++++- .../layers/attention/sparse_mla_attention.py | 164 ++++++- .../backends/mla/flashattn_mla_sparse.py | 147 +------ .../backends/mla/flashinfer_mla_sparse.py | 160 ++----- 5 files changed, 596 insertions(+), 288 deletions(-) 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_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 4d33075c9303..710fc2740c57 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, @@ -65,9 +68,6 @@ ) from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import current_stream -from vllm.v1.attention.backends.mla import ( - flashattn_mla_sparse as flashattn_sparse_module, -) from vllm.v1.attention.backends.mla import index_group as index_group_module from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( FlashAttnMLASparseImpl, @@ -141,6 +141,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=1, + num_decode_tokens=1, ) recorded_scale = None @@ -3369,15 +3370,14 @@ def test_flashinfer_hisparse_decode_runs_batched_attention(): physical_topk = topk + 10 valid_counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=device) kernel_shapes: list[torch.Size] = [] + kernel_kwargs: list[dict] = [] 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 + def run_kernel(self, q, cache, indices, counts, **kwargs): # noqa: ARG001 kernel_shapes.append(q.shape) + kernel_kwargs.append(kwargs) return q[..., :1], None cache_handle = SimpleNamespace( @@ -3392,21 +3392,23 @@ def run_kernel(self, q, cache, indices, counts): # noqa: ARG001 ) impl = object.__new__(FlashInferMLASparseImpl) impl.topk_indices_buffer = topk + impl.dcp_world_size = 1 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) + impl._forward_mqa_kernel = MethodType(run_kernel, impl) + metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64) + layer = SimpleNamespace(layer_name="mqa") output, lse = FlashInferMLASparseImpl.forward_mqa( impl, q, torch.empty(1, device=device), metadata, - SimpleNamespace(), + layer, ) assert kernel_shapes == [q.shape] + assert kernel_kwargs == [{"layer": layer, "block_size": 64, "is_decode": True}] assert output.shape == (num_tokens, 2, 1) assert lse is None @@ -3482,7 +3484,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 ) @@ -3527,8 +3529,14 @@ 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() @@ -3550,9 +3558,10 @@ def test_flashattn_hisparse_decode_uses_index_group(): ) impl = object.__new__(FlashAttnMLASparseImpl) impl.topk_indices_buffer = topk + impl.dcp_world_size = 1 impl.index_group = index_group impl.index_group_index = 0 - impl._run_mqa_kernel = MagicMock(return_value=q_nope[..., :1]) + impl._forward_mqa_kernel = MagicMock(return_value=(q_nope[..., :1], None)) metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64) output, lse = FlashAttnMLASparseImpl.forward_mqa( @@ -3566,7 +3575,7 @@ def test_flashattn_hisparse_decode_uses_index_group(): 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() + impl._forward_mqa_kernel.assert_called_once() def test_flashinfer_sm120_hisparse_decode_uses_index_group(): @@ -3603,6 +3612,368 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group(): impl._run_mqa_kernel.assert_called_once() +@pytest.mark.parametrize("backend", ["fa3", "flashinfer"]) +def test_sparse_mqa_consumes_decode_before_prefill_reuses_indices(monkeypatch, backend): + """Preparing a later portion must not overwrite unconsumed decode indices.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla + + num_tokens, topk_tokens = 4, 128 + q_nope = torch.zeros(num_tokens, 16, 512) + q_rope = torch.zeros(num_tokens, 16, 64) + kv_cache = torch.zeros(1, 1, 576) + topk_indices = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32) + metadata = SimpleNamespace( + num_decode_tokens=2, + block_size=1, + block_table=torch.zeros((num_tokens, 1), dtype=torch.int32), + ) + physical = torch.full((2, topk_tokens), 3, dtype=torch.int32) + counts = torch.full((2,), topk_tokens, dtype=torch.int32) + group = object.__new__(HiSparseMLAIndexGroup) + group.physical_kv_cache = MagicMock(return_value=kv_cache) + group.convert_logical_to_physical_topk = MagicMock(return_value=(physical, counts)) + group.cache = MagicMock( + return_value=SimpleNamespace(all_context_pages_resident=False) + ) + group.stage_prefill_rows = MagicMock( + return_value=(kv_cache, metadata.block_table, torch.zeros(2)) + ) + + def convert_prefill(*args, **kwargs): + physical.fill_(7) + return physical, counts + + monkeypatch.setattr( + sparse_mla, "triton_convert_req_index_to_global_index", convert_prefill + ) + + lanes: list[bool] = [] + + def run_kernel(q, cache, indices, counts, *, is_decode, **kwargs): # noqa: ARG001 + lanes.append(is_decode) + out = indices[:, :1].clone().view(-1, 1, 1).expand(-1, 16, 512) + return out, None + + impl_cls = FlashAttnMLASparseImpl if backend == "fa3" else FlashInferMLASparseImpl + impl = object.__new__(impl_cls) + impl.topk_indices_buffer = topk_indices + impl.index_group = group + impl.index_group_index = 0 + impl.dcp_world_size = 1 + impl._forward_mqa_kernel = run_kernel + + out, lse = impl.forward_mqa((q_nope, q_rope), kv_cache, metadata, None) + assert lse is None + assert lanes == [True, False] + assert out.shape == (num_tokens, 16, 512) + assert (out[:2] == 3).all() + assert (out[2:] == 7).all() + + +@pytest.mark.parametrize("layer_index", [0, 1]) +def test_flashattn_plain_path_converts_through_index_group(monkeypatch, layer_index): + """FA3's plain path converts once per group; a follower reuses the buffer.""" + import vllm.v1.attention.backends.mla.flashattn_mla_sparse as fa3_module + + device = torch.device(DEVICE_TYPE) + num_tokens, topk_tokens, lora, rope, block_size = 3, 8, 4, 2, 4 + kv_cache = torch.arange( + 2 * block_size * (lora + rope), dtype=torch.bfloat16, device=device + ).view(2, block_size, lora + rope) + logical = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32, device=device) + group = index_group_module.SparseMLAIndexGroup( + logical_topk_indices=logical, + physical_topk_indices=torch.full_like(logical, 9), + valid_topk_counts=torch.ones(num_tokens, dtype=torch.int32, device=device), + row_indices=torch.arange(num_tokens, dtype=torch.int32, device=device), + request_ids=torch.zeros(num_tokens, dtype=torch.int32, device=device), + side_stream=torch.Stream(device=device), + logical_topk_ready=torch.Event(), + physical_topk_ready=torch.Event(), + has_indexer=False, + num_layers=2, + ) + converts: list[int] = [] + + def convert(req_ids, block_table, topk, **kwargs): # noqa: ARG001 + converts.append(kwargs["BLOCK_STRIDE_ROWS"]) + kwargs["out"].fill_(5) + kwargs["valid_counts_out"].fill_(topk_tokens) + + monkeypatch.setattr( + index_group_module, "triton_convert_req_index_to_global_index", convert + ) + recorded: dict = {} + + def fake_kernel(**kwargs): + recorded.update(kwargs) + return torch.zeros(num_tokens, 2, lora, device=device) + + monkeypatch.setattr(fa3_module, "flash_attn_varlen_func", fake_kernel) + + impl = object.__new__(FlashAttnMLASparseImpl) + impl.topk_indices_buffer = logical + impl.index_group = group + impl.index_group_index = layer_index + impl.dcp_world_size = 1 + impl.kv_lora_rank = lora + impl.qk_rope_head_dim = rope + impl.cu_seqlens_q_buffer = torch.arange( + num_tokens + 1, dtype=torch.int32, device=device + ) + impl.scale = 0.25 + q_nope = torch.zeros(num_tokens, 2, lora, device=device) + q_rope = torch.zeros(num_tokens, 2, rope, device=device) + metadata = SimpleNamespace( + num_decode_tokens=num_tokens, + block_size=block_size, + block_table=torch.zeros((num_tokens, 1), dtype=torch.int32, device=device), + req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32, device=device), + ) + + out, lse = impl.forward_mqa((q_nope, q_rope), kv_cache, metadata, None) + + # The leading layer converts with the cache's block stride into the group's + # buffer; a follower reuses that buffer without converting again. + assert converts == ([block_size] if layer_index == 0 else []) + assert (recorded["block_table"] == (5 if layer_index == 0 else 9)).all() + assert (recorded["seqused_k"] == (topk_tokens if layer_index == 0 else 1)).all() + assert recorded["block_table"].data_ptr() == group.physical_topk_indices.data_ptr() + assert recorded["seqused_k"].data_ptr() == group.valid_topk_counts.data_ptr() + kv_rows = kv_cache.view(-1, lora + rope) + assert recorded["q"] is q_rope and recorded["q_v"] is q_nope + torch.testing.assert_close(recorded["k"][:, 0, 0], kv_rows[:, lora:]) + torch.testing.assert_close(recorded["v"][:, 0, 0], kv_rows[:, :lora]) + torch.testing.assert_close( + recorded["cu_seqlens_q"], + torch.arange(num_tokens + 1, dtype=torch.int32, device=device), + ) + assert ( + recorded["max_seqlen_q"], + recorded["max_seqlen_k"], + recorded["softmax_scale"], + recorded["causal"], + recorded["fa_version"], + ) == (1, topk_tokens, impl.scale, True, 3) + assert lse is None and out.shape == (num_tokens, 2, lora) + + +@pytest.mark.parametrize("dcp", [1, 2], ids=["native", "dcp2"]) +@pytest.mark.parametrize( + "num_decode_tokens,is_decode", [(3, True), (1, False)], ids=["decode", "mixed"] +) +def test_sparse_mqa_plain_lane_wiring(monkeypatch, dcp, num_decode_tokens, is_decode): + """The plain lane converts once and hands the hook one whole-batch call.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla + + num_tokens, topk_tokens, block_size, head_dim = 3, 8, 4, 6 + # A page whose flat row stride (8) differs from its block size (4): a + # converter handed the block size instead would misaddress the flat rows. + kv_cache = torch.zeros(2, 8, head_dim)[:, :block_size] + # One row longer than the batch, so slicing to num_tokens is observable. + topk_buffer = torch.zeros(num_tokens + 1, topk_tokens, dtype=torch.int32) + req_id_per_token = torch.zeros(num_tokens + 1, dtype=torch.int32) + q = torch.zeros(num_tokens, 2, head_dim) + physical = torch.full((num_tokens, topk_tokens), 5, dtype=torch.int32) + counts = torch.full((num_tokens,), topk_tokens, dtype=torch.int32) + layer = object() + metadata = SimpleNamespace( + num_decode_tokens=num_decode_tokens, + block_size=block_size, + block_table=torch.zeros((num_tokens, 1), dtype=torch.int32), + req_id_per_token=req_id_per_token, + cp_kv_cache_interleave_size=2, + ) + index_filter = MagicMock(return_value=(physical, counts)) + monkeypatch.setattr(sparse_mla, "triton_filter_and_convert_dcp_index", index_filter) + + impl = object.__new__(FlashInferMLASparseImpl) + impl.topk_indices_buffer = topk_buffer + impl.index_group = MagicMock() + impl.index_group_index = 0 + impl.dcp_world_size = dcp + impl.dcp_rank = dcp - 1 + impl._convert_logical_to_physical_topk = MagicMock(return_value=(physical, counts)) + impl._forward_mqa_kernel = MagicMock(return_value=(q, None)) + + out, lse = impl.forward_mqa(q, kv_cache, metadata, layer) + + assert out is q and lse is None + if dcp > 1: + impl._convert_logical_to_physical_topk.assert_not_called() + args, kwargs = index_filter.call_args + assert args[0].data_ptr() == req_id_per_token.data_ptr() + assert args[0].shape == (num_tokens,) + assert args[1] is metadata.block_table + assert args[2].data_ptr() == topk_buffer.data_ptr() + assert args[2].shape == (num_tokens, topk_tokens) + assert kwargs == { + "dcp_size": dcp, + "dcp_rank": dcp - 1, + "cp_kv_cache_interleave_size": 2, + "BLOCK_SIZE": block_size, + "BLOCK_STRIDE_ROWS": 8, + "NUM_TOPK_TOKENS": topk_tokens, + "return_valid_counts": True, + } + else: + index_filter.assert_not_called() + args, kwargs = impl._convert_logical_to_physical_topk.call_args + assert args[0].data_ptr() == topk_buffer.data_ptr() + assert args[0].shape == (num_tokens, topk_tokens) + assert args[1] is metadata + assert kwargs == {"block_stride_rows": 8, "return_valid_counts": True} + + # Identity, not equality: a copy would break the group's buffer stability. + hook_args, hook_kwargs = impl._forward_mqa_kernel.call_args + assert hook_args[0] is q + assert hook_args[1] is kv_cache + assert hook_args[2] is physical + assert hook_args[3] is counts + assert hook_kwargs == { + "layer": layer, + "block_size": block_size, + "is_decode": is_decode, + } + + +@pytest.mark.parametrize( + "num_decode_tokens,all_resident,lane", + [ + (5, True, "decode"), + (0, True, "resident"), + (0, False, "staged"), + # Fewer prefill rows than decode rows, so the row counts alone cannot + # tell the prefill call from the decode call. + (2, True, "mixed"), + ], +) +def test_sparse_mqa_hisparse_lanes(monkeypatch, num_decode_tokens, all_resident, lane): + """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) + topk_buffer = torch.zeros(num_tokens + 1, topk_tokens, dtype=torch.int32) + q = torch.arange(num_tokens, dtype=torch.float32).view(num_tokens, 1, 1) + layer = object() + 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), + ) + + decode_ret = converted(num_decode_tokens, 1) + resident_ret = converted(num_tokens, 2) + staged_ret = converted(num_prefill, 3) + req_ids = torch.zeros(num_prefill, dtype=torch.int32) + staged_block_table = torch.zeros(1, staged.shape[0], 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) + ) + # Decode rows and a decode-free resident prefill share the group's converter. + index_group.convert_logical_to_physical_topk = MagicMock( + return_value=decode_ret if num_decode_tokens else resident_ret + ) + index_group.stage_prefill_rows = MagicMock( + return_value=(staged, staged_block_table, req_ids) + ) + convert = MagicMock(return_value=staged_ret) + monkeypatch.setattr(sparse_mla, "triton_convert_req_index_to_global_index", convert) + + calls: list[tuple] = [] + + def hook(rows, cache, indices, counts, **kwargs): + assert kwargs["layer"] is layer and kwargs["block_size"] == block_size + is_decode = kwargs["is_decode"] + calls.append((rows, cache.data_ptr(), indices, counts, is_decode)) + 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 = topk_buffer + 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, host, metadata, layer) + + assert lse is None + # A lane handed the whole batch would concatenate to more rows than this. + assert out.shape == (num_tokens, 1, 1) + if num_decode_tokens: + rows, kv_ptr, indices, counts, is_decode = calls[0] + assert is_decode is True + assert kv_ptr == hot.data_ptr() + # Identity, not equality: a sub-slice would break pointer stability. + assert indices is decode_ret[0] and counts is decode_ret[1] + assert rows.tolist() == q[:num_decode_tokens].tolist() + args, kwargs = index_group.convert_logical_to_physical_topk.call_args_list[0] + assert args[0] == 0 + assert args[1].data_ptr() == topk_buffer.data_ptr() + assert args[1].shape == (num_decode_tokens, topk_tokens) + assert args[2] is metadata + assert kwargs == {"block_stride_rows": None, "return_valid_counts": True} + + if lane == "decode": + assert [call[-1] for call in calls] == [True] + assert (out == 1.0).all() + return + + rows, kv_ptr, indices, counts, is_decode = calls[-1] + assert is_decode is False + assert rows.tolist() == q[num_decode_tokens:].tolist() + if lane == "resident": + index_group.stage_prefill_rows.assert_not_called() + convert.assert_not_called() + assert kv_ptr == hot.data_ptr() + assert indices is resident_ret[0] and counts is resident_ret[1] + args, kwargs = index_group.convert_logical_to_physical_topk.call_args + assert args[0] == 0 + assert args[1].data_ptr() == topk_buffer.data_ptr() + assert args[1].shape == (num_tokens, topk_tokens) + assert args[2] is metadata + # The resident cache carries its own stride, so the group resolves it. + assert kwargs == {"block_stride_rows": None, "return_valid_counts": True} + else: + # Staged and mixed both stage; the resident lane is decode-free only. + assert index_group.convert_logical_to_physical_topk.call_count == min( + num_decode_tokens, 1 + ) + assert index_group.stage_prefill_rows.call_args.args == (0, host, metadata) + assert kv_ptr == staged.data_ptr() + assert indices is staged_ret[0] and counts is staged_ret[1] + args, kwargs = convert.call_args + assert args[0] is req_ids and args[1] is staged_block_table + assert args[2].data_ptr() == topk_buffer[num_decode_tokens:].data_ptr() + assert args[2].shape == (num_prefill, topk_tokens) + # The staging buffer is densely packed: no block stride to pass on. + assert kwargs == { + "BLOCK_SIZE": block_size, + "NUM_TOPK_TOKENS": topk_tokens, + "return_valid_counts": True, + } + + if lane == "mixed": + assert [call[-1] for call in calls] == [True, False] + assert (out[:num_decode_tokens] == 1.0).all() + assert (out[num_decode_tokens:] == 2.0).all() + else: + assert [call[-1] for call in calls] == [False] + assert (out == 2.0).all() + + def test_hisparse_resident_prefill_uses_attention_block_stride(): expected = torch.tensor([[19]], dtype=torch.int32) cache_handle = SimpleNamespace( diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 37883593421a..e10edba484f3 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,159 @@ 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 + ) + # HiSparse bakes the hot buffer's block stride into the indices it + # produces, which is only correct while that buffer is densely packed. + assert hot_cache.is_contiguous(), ( + "HiSparse hot buffer must be contiguous; a strided buffer would " + "misaddress the flat rows the top-k indices point at" + ) + 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/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index aabd68224093..2c3de4aa4250 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -22,11 +22,7 @@ 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.sparse_utils import ( - flat_kv_row_view, - triton_convert_req_index_to_global_index, -) +from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.kv_cache_interface import AttentionSpec from vllm.vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func @@ -148,6 +144,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, @@ -200,144 +198,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) @@ -365,4 +242,4 @@ def _run_mqa_kernel( causal=True, fa_version=3, ) - return out + return 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 4ebda50178cb..b1aed5974ef2 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -30,10 +30,7 @@ ) from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( - 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 @@ -120,14 +117,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: @@ -483,130 +478,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)) - - _, 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, - kv_c_and_k_pe_cache, - topk_indices_physical, - seq_lens, - ) - def _prepare_mqa_kernel( self, layer: AttentionLayer, @@ -634,7 +505,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() @@ -665,15 +535,35 @@ 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 b8945ec26a1dd566dfb954af1e635a7432faa776 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Tue, 29 Sep 2026 13:48:23 +0000 Subject: [PATCH 2/9] [Attention] Add the FA4 sparse MQA decode backend for SM100 Adds FLASH_ATTN_MLA_SPARSE_FA4 on top of the shared sparse MLA MQA hook. The FA4 CuTe DSL kernel (vllm-project/flash-attention#195) serves decode-only batches and the decode rows of mixed HiSparse batches; every other batch runs on the trtllm-gen lane, as in FLASHINFER_MLA_SPARSE. BF16 KV cache only. - flash_attn_varlen_func forwards gather_kv_indices, gather_kv_valid_length and q_v to FA4, and lets FA4 take cu_seqlens_k (KV offset) and seqused_k (length) together. - The impl subclasses FlashInferMLASparseImpl and overrides _forward_mqa_kernel: the FA4 decode kernel when is_decode, otherwise FlashInfer's trtllm-gen call with the LSE converted to base e. - The metadata builder subclasses the FlashInfer TRTLLM builder, which brings the DCP workspace pre-size and its UNIFORM_BATCH CUDA graph support with the varlen decode bound. The backend class subclasses the FlashInfer TRTLLM backend for its gates (PCP+DCP reject, head dims) and adds the FA4 gates, including the gathered head count under DCP. - HiSparse: decode reads the hot buffer through the shared lanes; the inherited autotune_hisparse_decode compiles the FA4 kernel before graph capture. - Priority on SM100 with a BF16 KV cache and at most 16 heads: FA4, FlashInfer, FlashMLA, with and without HiSparse, as the attention backends doc footnote now says. On GLM-5.2 NVFP4-FP8 (B300, TP4) FA4 leads FlashInfer by 4-5% ITL without HiSparse and by 2-7% with DCP2, and is on par under HiSparse. - FA4 shares FlashInfer's masked-MHA thresholds. - FA4 reads the same flat KV cache as FlashInfer: it joins `_uses_flat_kv_cache` (packed block-stride alignment and convert-kernel JIT warmup) and the compressor/indexer JIT warmup list. - Bumps vllm-flash-attn to vllm-project/flash-attention 343f2f11 (#195). Co-authored-by: Claude Signed-off-by: Simon Veitner --- cmake/external_projects/vllm_flash_attn.cmake | 2 +- docs/design/attention_backends.md | 2 +- ...st_flashmla_nope_sm90_backend_selection.py | 9 + .../v1/attention/test_sparse_mla_backends.py | 365 +++++++++++++++++- .../test_sparse_mla_kv_cache_layout.py | 4 + tests/v1/e2e/general/test_hisparse.py | 16 + .../layers/attention/mla_attention.py | 6 +- vllm/platforms/cuda.py | 3 +- .../backends/mla/flashattn_mla_sparse.py | 186 +++++++++ vllm/v1/attention/backends/registry.py | 5 + vllm/vllm_flash_attn/flash_attn_interface.py | 28 +- 11 files changed, 614 insertions(+), 12 deletions(-) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 2a26023f740b..47c7c865ce1b 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 343f2f112b87a6c83647fb980d86ab619bae401d 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/v1/attention/test_flashmla_nope_sm90_backend_selection.py b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py index 652358876949..74fbbb43e979 100644 --- a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py +++ b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py @@ -141,6 +141,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", @@ -179,3 +180,11 @@ 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 + + +def test_sm100_576_bf16_low_heads_prefers_fa4(): + assert _sparse_order("auto", 576, num_heads=16)[:3] == [ + "FLASH_ATTN_MLA_SPARSE_FA4", + "FLASHINFER_MLA_SPARSE", + "FLASHMLA_SPARSE", + ] diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 710fc2740c57..bafb721522d9 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -49,6 +49,7 @@ ) from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability # TODO: Integrate ROCMAiterMLASparseBackend for ROCm. # The ROCm sparse MLA backend (rocm_aiter_mla_sparse.py) has a compatible @@ -70,8 +71,11 @@ from vllm.utils.torch_utils import current_stream from vllm.v1.attention.backends.mla import index_group as index_group_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, @@ -404,8 +408,12 @@ def unpack_e2m1(packed: torch.Tensor) -> torch.Tensor: @pytest.mark.parametrize( "backend_cls", - [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend], - ids=["FlashMLA", "FlashInferTRTLLM"], + [ + FlashMLASparseBackend, + FlashInferMLASparseTRTLLMBackend, + FlashAttnMLASparseFA4Backend, + ], + ids=["FlashMLA", "FlashInferTRTLLM", "FlashAttnFA4"], ) @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys())) @pytest.mark.parametrize( @@ -495,6 +503,8 @@ def test_sparse_backend_decode_correctness( device_capability ): pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability") + elif backend_cls == FlashAttnMLASparseFA4Backend: + _require_fa4() batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name] use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla" @@ -1252,6 +1262,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( @@ -4042,6 +4053,356 @@ def test_hisparse_mixed_single_token_batch_is_not_decode_only(): assert not cache.decode_batch +def _require_fa4(): + """Skip unless this GPU can run the FA4 cute-DSL MLA kernel.""" + 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) + + +def _fa4_impl(**fields): + """Skip constructor dependencies while retaining the implementation methods.""" + buffer = fields.pop("topk_indices_buffer") + # DeepSeek-V3.2 shapes, and the DCP fields AttentionImpl.__new__ would fill. + fields.setdefault("num_heads", 16) + fields.setdefault("kv_lora_rank", 512) + fields.setdefault("qk_nope_head_dim", 128) + fields.setdefault("qk_rope_head_dim", 64) + fields.setdefault("scale", 1.0) + fields.setdefault("kv_cache_dtype", "auto") + fields.setdefault("dcp_world_size", 1) + fields.setdefault("dcp_rank", 0) + fields.setdefault("need_to_return_lse_for_decode", False) + fields.setdefault("index_group_index", 0) + # FlashInferMLASparseImpl.__init__ state; the kernel scales resolve lazily. + fields.setdefault("is_nope_mla", False) + fields.setdefault("_workspace_buffer", None) + fields.setdefault("bmm1_scale", None) + fields.setdefault("bmm2_scale", None) + if "index_group" not in fields: + # A non-index-producing layer does not wait on the indexer's ready event. + fields["index_group"] = SparseMLAIndexGroupBuilder(buffer).register_layer( + False + )[0] + stub = object.__new__(FlashAttnMLASparseFA4Impl) + stub.__dict__.update(fields) + # topk_indices_buffer is a data descriptor, so it needs its own setter. + stub.topk_indices_buffer = buffer + return stub + + +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, + # index_kpool can widen the buffer beyond metadata.topk_tokens. + topk_tokens=topk // 2, + 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 + + import vllm.v1.attention.backends.mla.flashattn_mla_sparse as fa4_sparse + + 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(fa4_sparse, "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): + _require_fa4() + 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("dcp", [1, 2], ids=["native", "dcp2"]) +@pytest.mark.parametrize("is_decode", [True, False]) +def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode, dcp): + """The hook runs a decode batch on FA4 and every other batch on trtllm-gen.""" + counts = [1, 2, 3, 4, 5, 6, 7] + num_tokens = len(counts) + # The query reaches the hook all-gathered: more heads than the impl's own. + kernel_heads = 16 * dcp + inputs = _fa4_inputs(counts, num_heads=kernel_heads, device="cpu") + calls = _record_fa4_kernels(monkeypatch, kernel_heads, device="cpu") + valid_counts = inputs.valid_counts.clone() + q = (inputs.ql_nope, inputs.q_pe) + if dcp > 1: + # This rank's share after filtering: token 0 owns none of its top-k, so + # the filter leaves its whole index row at the -1 sentinel. + valid_counts[0] = 0 + inputs.topk_indices[0] = -1 + # mla_attention fuses the query whenever DCP is on without PCP, and the + # gate rejects PCP+DCP, so the hook only ever sees a fused DCP query. + q = torch.cat(q, dim=-1) + impl = _fa4_impl( + topk_indices_buffer=inputs.topk_indices, + dcp_world_size=dcp, + need_to_return_lse_for_decode=dcp > 1, + ) + + with torch.inference_mode(): + out, lse = impl._forward_mqa_kernel( + q, + inputs.kv_cache, + inputs.topk_indices, + valid_counts, + layer=None, + block_size=inputs.block_size, + is_decode=is_decode, + ) + + # Only DCP asks for an LSE, and both lanes then return it as (tokens, heads). + assert lse is None if dcp == 1 else lse.shape == (num_tokens, kernel_heads) + + if is_decode: + (fa4,) = calls.fa4 + assert calls.trtllm == [] + assert fa4["return_softmax_lse"] is (dcp > 1) + assert fa4["gather_kv_indices"] is inputs.topk_indices + assert fa4["gather_kv_valid_length"] is valid_counts + # Every token is its own one-row FA4 sequence, spec-decode rows included. + assert fa4["max_seqlen_q"] == 1 + assert fa4["cu_seqlens_q"].tolist() == list(range(num_tokens + 1)) + # The flat KV rows of the whole page pool, addressed by the top-k list. + assert fa4["v"].data_ptr() == inputs.kv_cache.data_ptr() + assert fa4["seqused_k"].tolist() == [4 * inputs.block_size] * num_tokens + # The qv kernel takes the two halves, fused query or not. + assert fa4["q"].shape[-1] == inputs.rope_dim + assert fa4["q_v"].shape[-1] == inputs.kv_lora_rank + assert (out == 1.0).all() + if dcp > 1: + torch.testing.assert_close(lse, calls.lse(num_tokens), rtol=0, atol=0) + return + + (trtllm,) = calls.trtllm + assert calls.fa4 == [] + assert trtllm["return_lse"] is (dcp > 1) + assert trtllm["query"].shape[0] == num_tokens + assert trtllm["seq_lens"] is valid_counts + assert trtllm["sparse_mla_top_k"] == inputs.topk + if dcp == 1: + assert (out == 2.0).all() + return + expected = calls.lse(num_tokens) * math.log(2.0) + torch.testing.assert_close(lse[1:], expected[1:], rtol=0, atol=1e-5) + assert (out[1:] == 2.0).all() + # Token 0 owns no slot on this rank: the DCP merge identity. + assert (out[0] == 0).all() + assert torch.isneginf(lse[0]).all() + + +def test_fa4_sparse_autotune_hisparse_decode_runs_decode_lane(monkeypatch): + """The warmup dispatch reaches FA4's hook, which runs the FA4 decode lane.""" + from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( + autotune_hisparse_flashinfer_attention, + ) + + block_size, topk, max_num_reqs = 64, 128, 3 + hot = torch.zeros(5, block_size, 576, dtype=torch.bfloat16) + calls = _record_fa4_kernels(monkeypatch, 16, device="cpu") + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [ + SimpleNamespace( + runtime=SimpleNamespace( + hot=SimpleNamespace(attention_cache=hot, block_size=block_size), + max_num_reqs=max_num_reqs, + ) + ) + ] + + def fa4_layer(num_topk, warm=True): + buffer = torch.zeros(max_num_reqs, num_topk, dtype=torch.int32) + return SimpleNamespace( + impl=_fa4_impl(topk_indices_buffer=buffer, index_group=index_group), + hisparse_cache=object() if warm else None, + ) + + layers = { + "a": fa4_layer(topk), + # The same warmup key, so one compile covers both layers. + "b": fa4_layer(topk), + # A distinct key, but no HiSparse cache to warm over. + "cold": fa4_layer(topk * 2, warm=False), + # An impl from another backend, which has no warmup hook at all. + "other": SimpleNamespace(impl=SimpleNamespace(), hisparse_cache=object()), + } + runner = SimpleNamespace( + vllm_config=SimpleNamespace( + compilation_config=SimpleNamespace(static_forward_context=layers) + ) + ) + + with torch.inference_mode(): + autotune_hisparse_flashinfer_attention(runner) + + # The warmed lane is FA4's, and only the first layer of the key warms it. + (fa4,) = calls.fa4 + assert calls.trtllm == [] + assert fa4["v"].data_ptr() == hot.data_ptr() + # The decode lane's cached varlen scalars, keyed on the hot row count. + assert fa4["seqused_k"].tolist() == [hot.shape[0] * block_size] * max_num_reqs + assert fa4["gather_kv_indices"].shape == (max_num_reqs, topk) + assert fa4["gather_kv_valid_length"].tolist() == [topk] * max_num_reqs + + +@pytest.mark.parametrize( + "config,expected", + [ + ({}, None), + ({"heads": 128}, None), + ({"heads": 12}, "heads"), + # Under DCP the kernel sees num_heads * dcp_size gathered heads. + ({"dcp": 2}, None), + ({"dcp": 8}, "heads"), + # gather_kv_indices' last dim must be a whole number of n-tiles. + ({"index_topk": None}, "index_topk"), + ({"index_topk": 192}, "index_topk"), + ({"pcp": 2}, None), + # The trtllm lane cannot serve PCP queries over a DCP-sharded cache; the + # inherited reason names this backend, not the FlashInfer one. + ({"dcp": 2, "pcp": 2}, "FLASH_ATTN_MLA_SPARSE_FA4 does not support"), + ({"hisparse": True}, None), + # The qv kernel asserts every descale is None: BF16 KV cache only. + ({"kv_dtype": "fp8"}, "kv_cache_dtype"), + ({"flashinfer": False}, "FlashInfer"), + ], +) +def test_fa4_sparse_gate(monkeypatch, config, expected): + """``validate_configuration`` on a fixed SM100: one row per gate reason.""" + import vllm.v1.attention.backends.mla.flashattn_mla_sparse as fa4_sparse + + has_flashinfer = config.get("flashinfer", True) + monkeypatch.setattr("vllm.utils.flashinfer.has_flashinfer", lambda: has_flashinfer) + monkeypatch.setattr(fa4_sparse, "_fa4_cute_mla_available", lambda: None) + dcp = config.get("dcp", 1) + vllm_config = _build_sparse_dcp_vllm_config(config.get("heads", 16), dcp) + vllm_config.parallel_config.prefill_context_parallel_size = config.get("pcp", 1) + vllm_config.model_config.hf_text_config.index_topk = config.get("index_topk", 128) + if config.get("hisparse"): + vllm_config.attention_config.hisparse_config = HiSparseConfig() + + with set_current_vllm_config(vllm_config): + reasons = FlashAttnMLASparseFA4Backend.validate_configuration( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype=config.get("kv_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 cd3c9b218147..5abc6d3b1312 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py @@ -12,6 +12,9 @@ from vllm.model_executor.layers.attention.mla_attention import MLAAttention from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache from vllm.platforms import current_platform +from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( + FlashAttnMLASparseFA4Backend, +) from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseTRTLLMBackend, ) @@ -50,6 +53,7 @@ [ (FlashInferMLASparseTRTLLMBackend, "auto", True, 1152), (FlashInferMLASparseTRTLLMBackend, "fp8", True, 576), + (FlashAttnMLASparseFA4Backend, "auto", True, 1152), (FlashMLASparseBackend, "auto", True, 1152), (FlashMLASparseBackend, "fp8_ds_mla", False, None), (FlashMLASparseBackend, "fp8_ds_mla", True, 656), diff --git a/tests/v1/e2e/general/test_hisparse.py b/tests/v1/e2e/general/test_hisparse.py index 2409b5cb690b..45d3f48990df 100644 --- a/tests/v1/e2e/general/test_hisparse.py +++ b/tests/v1/e2e/general/test_hisparse.py @@ -85,11 +85,17 @@ def _get_hisparse_worker(runner: VllmRunner) -> HiSparseConnectorWorker: @pytest.mark.parametrize( "with_offloading", [False, True], ids=["standalone", "offload"] ) +@pytest.mark.parametrize( + "attention_backend", + ["FLASHINFER_MLA_SPARSE", "FLASH_ATTN_MLA_SPARSE_FA4"], + ids=["flashinfer", "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, ): """Spilled prefixes restore and FULL-graph decode writes reach host KV. @@ -101,6 +107,12 @@ 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") + forced_backend: str | None = attention_backend + if not current_platform.is_device_capability_family(100): + if attention_backend == "FLASH_ATTN_MLA_SPARSE_FA4": + pytest.skip("FA4 sparse MLA requires SM 10.x") + # FLASHINFER_MLA_SPARSE is SM 10.x only; elsewhere use the default. + forced_backend = None monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip") @@ -142,6 +154,7 @@ def test_hisparse_spill_and_prefix_restore( ) ), kv_transfer_config=kv_transfer_config, + attention_backend=forced_backend, block_size=64, max_model_len=320, max_num_batched_tokens=1024, @@ -159,6 +172,9 @@ def test_hisparse_spill_and_prefix_restore( worker = _get_hisparse_worker(runner) engine_core = runner.llm.llm_engine.engine_core.engine_core model_runner = engine_core.model_executor.driver_worker.worker.model_runner + if forced_backend is not None: + selected = model_runner.attn_groups[0][0].backend.get_name() + assert selected == forced_backend full_graph_calls = 0 full_replay_pending = False original_run_fullgraph = model_runner.cudagraph_manager.run_fullgraph diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 2a3048096c98..ea430491589d 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") ) @@ -1865,6 +1866,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/platforms/cuda.py b/vllm/platforms/cuda.py index 68c3aa7851cd..ff4ba9817dfb 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 2c3de4aa4250..5d75a11ba9c3 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,6 +24,11 @@ MultipleOf, ) from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + FlashInferMLASparseImpl, + FlashInferMLASparseTRTLLMBackend, + FlashInferMLASparseTRTLLMMetadataBuilder, +) from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.kv_cache_interface import AttentionSpec @@ -243,3 +250,182 @@ def _forward_mqa_kernel( fa_version=3, ) 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): + """Decode-only batches on FA4; every other batch on the trtllm-gen kernel.""" + + 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 [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: + # No engine context, so only the static gates above can be applied. + 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}" + ) + # The prefill lane's own gates: PCP+DCP, qk_nope_head_dim, index_topk. + # Those reasons name ``cls``, so they name this backend. Pinned by test. + 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 for decode rows; FlashInfer's trtllm-gen kernel for prefill rows.""" + + # 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. + assert q.shape[-1] == self.kv_lora_rank + self.qk_rope_head_dim + ql_nope, q_pe = q[..., : self.kv_lora_rank], q[..., self.kv_lora_rank :] + kv_rows, _ = flat_kv_row_view(kv_cache, block_size) + return self._decode(ql_nope, q_pe, kv_rows, topk_indices, valid_counts) + + def _decode( + self, + ql_nope: torch.Tensor, + q_pe: torch.Tensor, + kv_rows: torch.Tensor, + topk_indices: torch.Tensor, + valid_counts: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + 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/registry.py b/vllm/v1/attention/backends/registry.py index 207c112a384c..24f621177bb2 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -125,6 +125,11 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): FLASH_ATTN_MLA_SPARSE = ( "vllm.v1.attention.backends.mla.flashattn_mla_sparse.FlashAttnMLASparseBackend" ) + # Blackwell sibling on FA4's cute-DSL MLA gather kernel. + 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..ed78e4b9916d 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,25 @@ 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, gather_kv_length) int32, one KV row list per query + token, with `-1` as the "no token" sentinel. Requires q_v, excludes + block_table, and gather_kv_length must be a multiple of 128. FA4 only. + gather_kv_valid_length: (total_q,) int32, per query token, how many leading + entries of its gather_kv_indices row are real. The kernel attends + `round_up(length, 128)` entries, so everything from `length` on must + still be a sentinel. FA4 only, requires gather_kv_indices. 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 +321,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 +459,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, From f89c3de33884630c05a1a60cd50dc79be23fd380 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Mon, 21 Sep 2026 09:47:29 +0000 Subject: [PATCH 3/9] cleanup: drop redundant sparse MLA tests and shorten comments - Delete tests subsumed by test_sparse_mqa_hisparse_lanes and test_fa4_sparse_hook_routes_lanes; fold their unique asserts in. - Slim the FA4 autotune and FA3 plain-path tests. - Shorten comments and docstrings; no executable source change. Co-authored-by: Claude Signed-off-by: Simon Veitner --- .../v1/attention/test_sparse_mla_backends.py | 189 ++++++------------ .../layers/attention/sparse_mla_attention.py | 2 - .../backends/mla/flashattn_mla_sparse.py | 5 +- vllm/vllm_flash_attn/flash_attn_interface.py | 7 +- 4 files changed, 61 insertions(+), 142 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index bafb721522d9..9c00be03d451 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -80,6 +80,7 @@ from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseImpl, FlashInferMLASparseTRTLLMBackend, + _get_workspace_buffer, ) from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import ( FlashInferMLASparseSM120Impl, @@ -3381,14 +3382,12 @@ def test_flashinfer_hisparse_decode_runs_batched_attention(): physical_topk = topk + 10 valid_counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=device) kernel_shapes: list[torch.Size] = [] - kernel_kwargs: list[dict] = [] def convert_decode(self, *args, **kwargs): # noqa: ARG001 return physical_topk, valid_counts def run_kernel(self, q, cache, indices, counts, **kwargs): # noqa: ARG001 kernel_shapes.append(q.shape) - kernel_kwargs.append(kwargs) return q[..., :1], None cache_handle = SimpleNamespace( @@ -3408,18 +3407,16 @@ def run_kernel(self, q, cache, indices, counts, **kwargs): # noqa: ARG001 impl.index_group_index = 0 impl._forward_mqa_kernel = MethodType(run_kernel, impl) metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64) - layer = SimpleNamespace(layer_name="mqa") output, lse = FlashInferMLASparseImpl.forward_mqa( impl, q, torch.empty(1, device=device), metadata, - layer, + SimpleNamespace(), ) assert kernel_shapes == [q.shape] - assert kernel_kwargs == [{"layer": layer, "block_size": 64, "is_decode": True}] assert output.shape == (num_tokens, 2, 1) assert lse is None @@ -3623,64 +3620,6 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group(): impl._run_mqa_kernel.assert_called_once() -@pytest.mark.parametrize("backend", ["fa3", "flashinfer"]) -def test_sparse_mqa_consumes_decode_before_prefill_reuses_indices(monkeypatch, backend): - """Preparing a later portion must not overwrite unconsumed decode indices.""" - import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla - - num_tokens, topk_tokens = 4, 128 - q_nope = torch.zeros(num_tokens, 16, 512) - q_rope = torch.zeros(num_tokens, 16, 64) - kv_cache = torch.zeros(1, 1, 576) - topk_indices = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32) - metadata = SimpleNamespace( - num_decode_tokens=2, - block_size=1, - block_table=torch.zeros((num_tokens, 1), dtype=torch.int32), - ) - physical = torch.full((2, topk_tokens), 3, dtype=torch.int32) - counts = torch.full((2,), topk_tokens, dtype=torch.int32) - group = object.__new__(HiSparseMLAIndexGroup) - group.physical_kv_cache = MagicMock(return_value=kv_cache) - group.convert_logical_to_physical_topk = MagicMock(return_value=(physical, counts)) - group.cache = MagicMock( - return_value=SimpleNamespace(all_context_pages_resident=False) - ) - group.stage_prefill_rows = MagicMock( - return_value=(kv_cache, metadata.block_table, torch.zeros(2)) - ) - - def convert_prefill(*args, **kwargs): - physical.fill_(7) - return physical, counts - - monkeypatch.setattr( - sparse_mla, "triton_convert_req_index_to_global_index", convert_prefill - ) - - lanes: list[bool] = [] - - def run_kernel(q, cache, indices, counts, *, is_decode, **kwargs): # noqa: ARG001 - lanes.append(is_decode) - out = indices[:, :1].clone().view(-1, 1, 1).expand(-1, 16, 512) - return out, None - - impl_cls = FlashAttnMLASparseImpl if backend == "fa3" else FlashInferMLASparseImpl - impl = object.__new__(impl_cls) - impl.topk_indices_buffer = topk_indices - impl.index_group = group - impl.index_group_index = 0 - impl.dcp_world_size = 1 - impl._forward_mqa_kernel = run_kernel - - out, lse = impl.forward_mqa((q_nope, q_rope), kv_cache, metadata, None) - assert lse is None - assert lanes == [True, False] - assert out.shape == (num_tokens, 16, 512) - assert (out[:2] == 3).all() - assert (out[2:] == 7).all() - - @pytest.mark.parametrize("layer_index", [0, 1]) def test_flashattn_plain_path_converts_through_index_group(monkeypatch, layer_index): """FA3's plain path converts once per group; a follower reuses the buffer.""" @@ -3744,8 +3683,6 @@ def fake_kernel(**kwargs): out, lse = impl.forward_mqa((q_nope, q_rope), kv_cache, metadata, None) - # The leading layer converts with the cache's block stride into the group's - # buffer; a follower reuses that buffer without converting again. assert converts == ([block_size] if layer_index == 0 else []) assert (recorded["block_table"] == (5 if layer_index == 0 else 9)).all() assert (recorded["seqused_k"] == (topk_tokens if layer_index == 0 else 1)).all() @@ -3755,17 +3692,6 @@ def fake_kernel(**kwargs): assert recorded["q"] is q_rope and recorded["q_v"] is q_nope torch.testing.assert_close(recorded["k"][:, 0, 0], kv_rows[:, lora:]) torch.testing.assert_close(recorded["v"][:, 0, 0], kv_rows[:, :lora]) - torch.testing.assert_close( - recorded["cu_seqlens_q"], - torch.arange(num_tokens + 1, dtype=torch.int32, device=device), - ) - assert ( - recorded["max_seqlen_q"], - recorded["max_seqlen_k"], - recorded["softmax_scale"], - recorded["causal"], - recorded["fa_version"], - ) == (1, topk_tokens, impl.scale, True, 3) assert lse is None and out.shape == (num_tokens, 2, lora) @@ -3778,8 +3704,7 @@ def test_sparse_mqa_plain_lane_wiring(monkeypatch, dcp, num_decode_tokens, is_de import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla num_tokens, topk_tokens, block_size, head_dim = 3, 8, 4, 6 - # A page whose flat row stride (8) differs from its block size (4): a - # converter handed the block size instead would misaddress the flat rows. + # Flat row stride (8) differs from block size (4): the wrong one misaddresses. kv_cache = torch.zeros(2, 8, head_dim)[:, :block_size] # One row longer than the batch, so slicing to num_tokens is observable. topk_buffer = torch.zeros(num_tokens + 1, topk_tokens, dtype=torch.int32) @@ -3848,18 +3773,20 @@ def test_sparse_mqa_plain_lane_wiring(monkeypatch, dcp, num_decode_tokens, is_de } +@pytest.mark.parametrize("split_q", [False, True], ids=["fused", "split"]) @pytest.mark.parametrize( "num_decode_tokens,all_resident,lane", [ (5, True, "decode"), (0, True, "resident"), (0, False, "staged"), - # Fewer prefill rows than decode rows, so the row counts alone cannot - # tell the prefill call from the decode call. + # Fewer prefill than decode rows: row counts cannot tell the calls apart. (2, True, "mixed"), ], ) -def test_sparse_mqa_hisparse_lanes(monkeypatch, num_decode_tokens, all_resident, lane): +def test_sparse_mqa_hisparse_lanes( + monkeypatch, num_decode_tokens, all_resident, lane, 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 @@ -3900,12 +3827,21 @@ def converted(rows, fill): index_group.stage_prefill_rows = MagicMock( return_value=(staged, staged_block_table, req_ids) ) - convert = MagicMock(return_value=staged_ret) - monkeypatch.setattr(sparse_mla, "triton_convert_req_index_to_global_index", convert) - 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 staged_ret + + convert = MagicMock(side_effect=convert_staged) + monkeypatch.setattr(sparse_mla, "triton_convert_req_index_to_global_index", convert) + def hook(rows, cache, indices, counts, **kwargs): + if split_q: + assert rows[0].tolist() == rows[1].tolist() + rows = rows[0] + assert kwargs.keys() == {"layer", "block_size", "is_decode"} assert kwargs["layer"] is layer and kwargs["block_size"] == block_size is_decode = kwargs["is_decode"] calls.append((rows, cache.data_ptr(), indices, counts, is_decode)) @@ -3918,7 +3854,7 @@ def hook(rows, cache, indices, counts, **kwargs): impl.dcp_world_size = 1 impl._forward_mqa_kernel = hook - out, lse = impl.forward_mqa(q, host, metadata, layer) + out, lse = impl.forward_mqa((q, q) if split_q else q, host, metadata, layer) assert lse is None # A lane handed the whole batch would concatenate to more rows than this. @@ -3931,7 +3867,6 @@ def hook(rows, cache, indices, counts, **kwargs): assert indices is decode_ret[0] and counts is decode_ret[1] assert rows.tolist() == q[:num_decode_tokens].tolist() args, kwargs = index_group.convert_logical_to_physical_topk.call_args_list[0] - assert args[0] == 0 assert args[1].data_ptr() == topk_buffer.data_ptr() assert args[1].shape == (num_decode_tokens, topk_tokens) assert args[2] is metadata @@ -3951,7 +3886,6 @@ def hook(rows, cache, indices, counts, **kwargs): assert kv_ptr == hot.data_ptr() assert indices is resident_ret[0] and counts is resident_ret[1] args, kwargs = index_group.convert_logical_to_physical_topk.call_args - assert args[0] == 0 assert args[1].data_ptr() == topk_buffer.data_ptr() assert args[1].shape == (num_tokens, topk_tokens) assert args[2] is metadata @@ -4126,8 +4060,6 @@ def values(*shape): ), block_size=block_size, num_decode_tokens=num_tokens, - # index_kpool can widen the buffer beyond metadata.topk_tokens. - topk_tokens=topk // 2, cp_kv_cache_interleave_size=1, ), ) @@ -4217,6 +4149,8 @@ def test_fa4_sparse_decode_kernel_correctness(num_heads, return_lse): @pytest.mark.parametrize("is_decode", [True, False]) def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode, dcp): """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 = len(counts) # The query reaches the hook all-gathered: more heads than the impl's own. @@ -4226,32 +4160,31 @@ def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode, dcp): valid_counts = inputs.valid_counts.clone() q = (inputs.ql_nope, inputs.q_pe) if dcp > 1: - # This rank's share after filtering: token 0 owns none of its top-k, so - # the filter leaves its whole index row at the -1 sentinel. + # This rank's share: token 0 owns none of its top-k, so its row stays all -1. valid_counts[0] = 0 inputs.topk_indices[0] = -1 - # mla_attention fuses the query whenever DCP is on without PCP, and the - # gate rejects PCP+DCP, so the hook only ever sees a fused DCP query. + # DCP without PCP (the gate rejects the pair) always fuses the query. q = torch.cat(q, dim=-1) impl = _fa4_impl( - topk_indices_buffer=inputs.topk_indices, + # 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=dcp > 1, ) + # Entered through forward_mqa on a plain (non-HiSparse) index group, whose + # converter is stubbed: the lane choice is the decode/prefill token split. + inputs.metadata.num_decode_tokens = num_tokens if is_decode else 0 + convert = MagicMock(return_value=(inputs.topk_indices, valid_counts)) + impl._convert_logical_to_physical_topk = convert + monkeypatch.setattr(sparse_mla, "triton_filter_and_convert_dcp_index", convert) + with torch.inference_mode(): - out, lse = impl._forward_mqa_kernel( - q, - inputs.kv_cache, - inputs.topk_indices, - valid_counts, - layer=None, - block_size=inputs.block_size, - is_decode=is_decode, - ) + out, lse = impl.forward_mqa(q, inputs.kv_cache, inputs.metadata, None) # Only DCP asks for an LSE, and both lanes then return it as (tokens, heads). assert lse is None if dcp == 1 else lse.shape == (num_tokens, kernel_heads) + assert out.shape == (num_tokens, kernel_heads, 512) if is_decode: (fa4,) = calls.fa4 @@ -4262,6 +4195,8 @@ def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode, dcp): # Every token is its own one-row FA4 sequence, spec-decode rows included. assert fa4["max_seqlen_q"] == 1 assert fa4["cu_seqlens_q"].tolist() == list(range(num_tokens + 1)) + # Gathered rows are addressed by the top-k list, so the K offsets stay 0. + assert fa4["cu_seqlens_k"].tolist() == [0] * (num_tokens + 1) # The flat KV rows of the whole page pool, addressed by the top-k list. assert fa4["v"].data_ptr() == inputs.kv_cache.data_ptr() assert fa4["seqused_k"].tolist() == [4 * inputs.block_size] * num_tokens @@ -4278,7 +4213,10 @@ def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode, dcp): assert trtllm["return_lse"] is (dcp > 1) assert trtllm["query"].shape[0] == num_tokens assert trtllm["seq_lens"] is valid_counts + assert trtllm["block_tables"].shape == (num_tokens, 1, inputs.topk) assert trtllm["sparse_mla_top_k"] == inputs.topk + # The lane runs on the shared buffer the builder pre-sizes under DCP. + assert trtllm["workspace_buffer"] is _get_workspace_buffer(torch.device("cpu")) if dcp == 1: assert (out == 2.0).all() return @@ -4291,7 +4229,7 @@ def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode, dcp): def test_fa4_sparse_autotune_hisparse_decode_runs_decode_lane(monkeypatch): - """The warmup dispatch reaches FA4's hook, which runs the FA4 decode lane.""" + """The HiSparse warmup reaches FA4's hook once per key, on the decode lane.""" from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( autotune_hisparse_flashinfer_attention, ) @@ -4299,31 +4237,23 @@ def test_fa4_sparse_autotune_hisparse_decode_runs_decode_lane(monkeypatch): block_size, topk, max_num_reqs = 64, 128, 3 hot = torch.zeros(5, block_size, 576, dtype=torch.bfloat16) calls = _record_fa4_kernels(monkeypatch, 16, device="cpu") + runtime = SimpleNamespace( + hot=SimpleNamespace(attention_cache=hot, block_size=block_size), + max_num_reqs=max_num_reqs, + ) index_group = object.__new__(HiSparseMLAIndexGroup) - index_group.caches = [ - SimpleNamespace( - runtime=SimpleNamespace( - hot=SimpleNamespace(attention_cache=hot, block_size=block_size), - max_num_reqs=max_num_reqs, - ) - ) - ] + index_group.caches = [SimpleNamespace(runtime=runtime)] - def fa4_layer(num_topk, warm=True): - buffer = torch.zeros(max_num_reqs, num_topk, dtype=torch.int32) - return SimpleNamespace( - impl=_fa4_impl(topk_indices_buffer=buffer, index_group=index_group), - hisparse_cache=object() if warm else None, - ) + def fa4_layer(): + buffer = torch.zeros(max_num_reqs, topk, dtype=torch.int32) + impl = _fa4_impl(topk_indices_buffer=buffer, index_group=index_group) + return SimpleNamespace(impl=impl, hisparse_cache=object()) + # Two layers share one warmup key; another backend's impl is skipped. layers = { - "a": fa4_layer(topk), - # The same warmup key, so one compile covers both layers. - "b": fa4_layer(topk), - # A distinct key, but no HiSparse cache to warm over. - "cold": fa4_layer(topk * 2, warm=False), - # An impl from another backend, which has no warmup hook at all. - "other": SimpleNamespace(impl=SimpleNamespace(), hisparse_cache=object()), + "a": fa4_layer(), + "b": fa4_layer(), + "other": SimpleNamespace(impl=object(), hisparse_cache=object()), } runner = SimpleNamespace( vllm_config=SimpleNamespace( @@ -4334,14 +4264,12 @@ def fa4_layer(num_topk, warm=True): with torch.inference_mode(): autotune_hisparse_flashinfer_attention(runner) - # The warmed lane is FA4's, and only the first layer of the key warms it. (fa4,) = calls.fa4 assert calls.trtllm == [] assert fa4["v"].data_ptr() == hot.data_ptr() # The decode lane's cached varlen scalars, keyed on the hot row count. assert fa4["seqused_k"].tolist() == [hot.shape[0] * block_size] * max_num_reqs assert fa4["gather_kv_indices"].shape == (max_num_reqs, topk) - assert fa4["gather_kv_valid_length"].tolist() == [topk] * max_num_reqs @pytest.mark.parametrize( @@ -4350,18 +4278,13 @@ def fa4_layer(num_topk, warm=True): ({}, None), ({"heads": 128}, None), ({"heads": 12}, "heads"), - # Under DCP the kernel sees num_heads * dcp_size gathered heads. ({"dcp": 2}, None), ({"dcp": 8}, "heads"), - # gather_kv_indices' last dim must be a whole number of n-tiles. ({"index_topk": None}, "index_topk"), ({"index_topk": 192}, "index_topk"), - ({"pcp": 2}, None), - # The trtllm lane cannot serve PCP queries over a DCP-sharded cache; the - # inherited reason names this backend, not the FlashInfer one. + # PCP over a DCP-sharded cache: the inherited reason names this backend. ({"dcp": 2, "pcp": 2}, "FLASH_ATTN_MLA_SPARSE_FA4 does not support"), ({"hisparse": True}, None), - # The qv kernel asserts every descale is None: BF16 KV cache only. ({"kv_dtype": "fp8"}, "kv_cache_dtype"), ({"flashinfer": False}, "FlashInfer"), ], diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index e10edba484f3..504124176a2f 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -801,8 +801,6 @@ def forward_mqa( hot_cache = index_group.physical_kv_cache(layer_index).view( kv_c_and_k_pe_cache.dtype ) - # HiSparse bakes the hot buffer's block stride into the indices it - # produces, which is only correct while that buffer is densely packed. assert hot_cache.is_contiguous(), ( "HiSparse hot buffer must be contiguous; a strided buffer would " "misaddress the flat rows the top-k indices point at" diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 5d75a11ba9c3..e43ebfa6c870 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -304,7 +304,7 @@ def supports_combination(cls, *args: Any, **kwargs: Any) -> str | None: return reason vllm_config = get_current_vllm_config_or_none() if vllm_config is None: - # No engine context, so only the static gates above can be applied. + # No engine context: only the static gates above apply. return None if vllm_config.model_config is not None: hf_config = vllm_config.model_config.hf_text_config @@ -333,8 +333,7 @@ def supports_combination(cls, *args: Any, **kwargs: Any) -> str | None: f"FA4 sparse MLA requires {head_counts} gathered query heads, " f"got num_heads={num_heads} * dcp_size={dcp_size}" ) - # The prefill lane's own gates: PCP+DCP, qk_nope_head_dim, index_topk. - # Those reasons name ``cls``, so they name this backend. Pinned by test. + # super()'s remaining gates; their reasons name ``cls``, hence this backend. return super().supports_combination(*args, **kwargs) diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index ed78e4b9916d..8ec4128f8be6 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -295,10 +295,9 @@ def flash_attn_varlen_func( gather_kv_indices: (total_q, gather_kv_length) int32, one KV row list per query token, with `-1` as the "no token" sentinel. Requires q_v, excludes block_table, and gather_kv_length must be a multiple of 128. FA4 only. - gather_kv_valid_length: (total_q,) int32, per query token, how many leading - entries of its gather_kv_indices row are real. The kernel attends - `round_up(length, 128)` entries, so everything from `length` on must - still be a sentinel. FA4 only, requires gather_kv_indices. + gather_kv_valid_length: (total_q,) int32, leading real entries per + gather_kv_indices row. The kernel attends `round_up(length, 128)` + entries, so the rest must be `-1`. Needs gather_kv_indices. FA4 only. Return: out: (total, nheads, headdim). From 0140d70253e0dfe96220888cbdfdb5c5898493db Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Wed, 30 Sep 2026 10:16:37 +0000 Subject: [PATCH 4/9] [Build] Bump vllm-flash-attn to vllm-project/flash-attention b273a026 Picks up the tuned sparse decode split rule for DCP shapes on top of #195 (vllm-project/flash-attention#200), and the SM90 hdim 512 forward fix (vllm-project/flash-attention#203), the current head of main. Signed-off-by: Simon Veitner --- cmake/external_projects/vllm_flash_attn.cmake | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 47c7c865ce1b..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 343f2f112b87a6c83647fb980d86ab619bae401d + 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 From 0850ebcf6274d413ab43f7f031668099feef108d Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Wed, 30 Sep 2026 18:24:20 +0000 Subject: [PATCH 5/9] remove some slop Signed-off-by: Simon Veitner --- .../v1/attention/test_sparse_mla_backends.py | 158 ++++-------------- tests/v1/e2e/general/test_hisparse.py | 3 - .../layers/attention/sparse_mla_attention.py | 5 +- .../backends/mla/flashattn_mla_sparse.py | 17 +- vllm/v1/attention/backends/registry.py | 1 - vllm/vllm_flash_attn/flash_attn_interface.py | 9 +- 6 files changed, 38 insertions(+), 155 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 9c00be03d451..c380c5d6d033 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -80,7 +80,6 @@ from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseImpl, FlashInferMLASparseTRTLLMBackend, - _get_workspace_buffer, ) from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import ( FlashInferMLASparseSM120Impl, @@ -409,12 +408,8 @@ def unpack_e2m1(packed: torch.Tensor) -> torch.Tensor: @pytest.mark.parametrize( "backend_cls", - [ - FlashMLASparseBackend, - FlashInferMLASparseTRTLLMBackend, - FlashAttnMLASparseFA4Backend, - ], - ids=["FlashMLA", "FlashInferTRTLLM", "FlashAttnFA4"], + [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend], + ids=["FlashMLA", "FlashInferTRTLLM"], ) @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys())) @pytest.mark.parametrize( @@ -504,8 +499,6 @@ def test_sparse_backend_decode_correctness( device_capability ): pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability") - elif backend_cls == FlashAttnMLASparseFA4Backend: - _require_fa4() batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name] use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla" @@ -3620,9 +3613,8 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group(): impl._run_mqa_kernel.assert_called_once() -@pytest.mark.parametrize("layer_index", [0, 1]) -def test_flashattn_plain_path_converts_through_index_group(monkeypatch, layer_index): - """FA3's plain path converts once per group; a follower reuses the buffer.""" +def test_flashattn_plain_path_passes_converted_topk(monkeypatch): + """FA3's plain path hands the converted top-k straight to the kernel.""" import vllm.v1.attention.backends.mla.flashattn_mla_sparse as fa3_module device = torch.device(DEVICE_TYPE) @@ -3631,28 +3623,8 @@ def test_flashattn_plain_path_converts_through_index_group(monkeypatch, layer_in 2 * block_size * (lora + rope), dtype=torch.bfloat16, device=device ).view(2, block_size, lora + rope) logical = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32, device=device) - group = index_group_module.SparseMLAIndexGroup( - logical_topk_indices=logical, - physical_topk_indices=torch.full_like(logical, 9), - valid_topk_counts=torch.ones(num_tokens, dtype=torch.int32, device=device), - row_indices=torch.arange(num_tokens, dtype=torch.int32, device=device), - request_ids=torch.zeros(num_tokens, dtype=torch.int32, device=device), - side_stream=torch.Stream(device=device), - logical_topk_ready=torch.Event(), - physical_topk_ready=torch.Event(), - has_indexer=False, - num_layers=2, - ) - converts: list[int] = [] - - def convert(req_ids, block_table, topk, **kwargs): # noqa: ARG001 - converts.append(kwargs["BLOCK_STRIDE_ROWS"]) - kwargs["out"].fill_(5) - kwargs["valid_counts_out"].fill_(topk_tokens) - - monkeypatch.setattr( - index_group_module, "triton_convert_req_index_to_global_index", convert - ) + physical = torch.full_like(logical, 5) + counts = torch.full((num_tokens,), topk_tokens, dtype=torch.int32, device=device) recorded: dict = {} def fake_kernel(**kwargs): @@ -3663,9 +3635,9 @@ def fake_kernel(**kwargs): impl = object.__new__(FlashAttnMLASparseImpl) impl.topk_indices_buffer = logical - impl.index_group = group - impl.index_group_index = layer_index + impl.index_group = SimpleNamespace(num_layers=2) impl.dcp_world_size = 1 + impl._convert_logical_to_physical_topk = MagicMock(return_value=(physical, counts)) impl.kv_lora_rank = lora impl.qk_rope_head_dim = rope impl.cu_seqlens_q_buffer = torch.arange( @@ -3674,20 +3646,12 @@ def fake_kernel(**kwargs): impl.scale = 0.25 q_nope = torch.zeros(num_tokens, 2, lora, device=device) q_rope = torch.zeros(num_tokens, 2, rope, device=device) - metadata = SimpleNamespace( - num_decode_tokens=num_tokens, - block_size=block_size, - block_table=torch.zeros((num_tokens, 1), dtype=torch.int32, device=device), - req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32, device=device), - ) + metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=block_size) out, lse = impl.forward_mqa((q_nope, q_rope), kv_cache, metadata, None) - assert converts == ([block_size] if layer_index == 0 else []) - assert (recorded["block_table"] == (5 if layer_index == 0 else 9)).all() - assert (recorded["seqused_k"] == (topk_tokens if layer_index == 0 else 1)).all() - assert recorded["block_table"].data_ptr() == group.physical_topk_indices.data_ptr() - assert recorded["seqused_k"].data_ptr() == group.valid_topk_counts.data_ptr() + assert recorded["block_table"] is physical + assert recorded["seqused_k"] is counts kv_rows = kv_cache.view(-1, lora + rope) assert recorded["q"] is q_rope and recorded["q_v"] is q_nope torch.testing.assert_close(recorded["k"][:, 0, 0], kv_rows[:, lora:]) @@ -3696,10 +3660,7 @@ def fake_kernel(**kwargs): @pytest.mark.parametrize("dcp", [1, 2], ids=["native", "dcp2"]) -@pytest.mark.parametrize( - "num_decode_tokens,is_decode", [(3, True), (1, False)], ids=["decode", "mixed"] -) -def test_sparse_mqa_plain_lane_wiring(monkeypatch, dcp, num_decode_tokens, is_decode): +def test_sparse_mqa_plain_lane_wiring(monkeypatch, dcp): """The plain lane converts once and hands the hook one whole-batch call.""" import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla @@ -3714,7 +3675,7 @@ def test_sparse_mqa_plain_lane_wiring(monkeypatch, dcp, num_decode_tokens, is_de counts = torch.full((num_tokens,), topk_tokens, dtype=torch.int32) layer = object() metadata = SimpleNamespace( - num_decode_tokens=num_decode_tokens, + num_decode_tokens=num_tokens, block_size=block_size, block_table=torch.zeros((num_tokens, 1), dtype=torch.int32), req_id_per_token=req_id_per_token, @@ -3743,15 +3704,7 @@ def test_sparse_mqa_plain_lane_wiring(monkeypatch, dcp, num_decode_tokens, is_de assert args[1] is metadata.block_table assert args[2].data_ptr() == topk_buffer.data_ptr() assert args[2].shape == (num_tokens, topk_tokens) - assert kwargs == { - "dcp_size": dcp, - "dcp_rank": dcp - 1, - "cp_kv_cache_interleave_size": 2, - "BLOCK_SIZE": block_size, - "BLOCK_STRIDE_ROWS": 8, - "NUM_TOPK_TOKENS": topk_tokens, - "return_valid_counts": True, - } + assert kwargs["BLOCK_STRIDE_ROWS"] == 8 else: index_filter.assert_not_called() args, kwargs = impl._convert_logical_to_physical_topk.call_args @@ -3759,25 +3712,21 @@ def test_sparse_mqa_plain_lane_wiring(monkeypatch, dcp, num_decode_tokens, is_de assert args[0].shape == (num_tokens, topk_tokens) assert args[1] is metadata assert kwargs == {"block_stride_rows": 8, "return_valid_counts": True} - - # Identity, not equality: a copy would break the group's buffer stability. - hook_args, hook_kwargs = impl._forward_mqa_kernel.call_args - assert hook_args[0] is q - assert hook_args[1] is kv_cache - assert hook_args[2] is physical - assert hook_args[3] is counts - assert hook_kwargs == { - "layer": layer, - "block_size": block_size, - "is_decode": is_decode, - } + impl._forward_mqa_kernel.assert_called_once_with( + q, + kv_cache, + physical, + counts, + layer=layer, + block_size=block_size, + is_decode=True, + ) @pytest.mark.parametrize("split_q", [False, True], ids=["fused", "split"]) @pytest.mark.parametrize( "num_decode_tokens,all_resident,lane", [ - (5, True, "decode"), (0, True, "resident"), (0, False, "staged"), # Fewer prefill than decode rows: row counts cannot tell the calls apart. @@ -3820,7 +3769,6 @@ def converted(rows, fill): index_group.cache = MagicMock( return_value=SimpleNamespace(all_context_pages_resident=all_resident) ) - # Decode rows and a decode-free resident prefill share the group's converter. index_group.convert_logical_to_physical_topk = MagicMock( return_value=decode_ret if num_decode_tokens else resident_ret ) @@ -3841,7 +3789,6 @@ def hook(rows, cache, indices, counts, **kwargs): if split_q: assert rows[0].tolist() == rows[1].tolist() rows = rows[0] - assert kwargs.keys() == {"layer", "block_size", "is_decode"} assert kwargs["layer"] is layer and kwargs["block_size"] == block_size is_decode = kwargs["is_decode"] calls.append((rows, cache.data_ptr(), indices, counts, is_decode)) @@ -3872,11 +3819,6 @@ def hook(rows, cache, indices, counts, **kwargs): assert args[2] is metadata assert kwargs == {"block_stride_rows": None, "return_valid_counts": True} - if lane == "decode": - assert [call[-1] for call in calls] == [True] - assert (out == 1.0).all() - return - rows, kv_ptr, indices, counts, is_decode = calls[-1] assert is_decode is False assert rows.tolist() == q[num_decode_tokens:].tolist() @@ -3892,7 +3834,6 @@ def hook(rows, cache, indices, counts, **kwargs): # The resident cache carries its own stride, so the group resolves it. assert kwargs == {"block_stride_rows": None, "return_valid_counts": True} else: - # Staged and mixed both stage; the resident lane is decode-free only. assert index_group.convert_logical_to_physical_topk.call_count == min( num_decode_tokens, 1 ) @@ -3987,18 +3928,8 @@ def test_hisparse_mixed_single_token_batch_is_not_decode_only(): assert not cache.decode_batch -def _require_fa4(): - """Skip unless this GPU can run the FA4 cute-DSL MLA kernel.""" - 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) - - def _fa4_impl(**fields): - """Skip constructor dependencies while retaining the implementation methods.""" buffer = fields.pop("topk_indices_buffer") - # DeepSeek-V3.2 shapes, and the DCP fields AttentionImpl.__new__ would fill. fields.setdefault("num_heads", 16) fields.setdefault("kv_lora_rank", 512) fields.setdefault("qk_nope_head_dim", 128) @@ -4009,7 +3940,6 @@ def _fa4_impl(**fields): fields.setdefault("dcp_rank", 0) fields.setdefault("need_to_return_lse_for_decode", False) fields.setdefault("index_group_index", 0) - # FlashInferMLASparseImpl.__init__ state; the kernel scales resolve lazily. fields.setdefault("is_nope_mla", False) fields.setdefault("_workspace_buffer", None) fields.setdefault("bmm1_scale", None) @@ -4104,7 +4034,10 @@ def fake_trtllm(**kwargs): "num_heads,return_lse", [(128, False), (16, True)], ids=["heads128", "heads16_lse"] ) def test_fa4_sparse_decode_kernel_correctness(num_heads, return_lse): - _require_fa4() + 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( @@ -4160,9 +4093,6 @@ def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode, dcp): valid_counts = inputs.valid_counts.clone() q = (inputs.ql_nope, inputs.q_pe) if dcp > 1: - # This rank's share: token 0 owns none of its top-k, so its row stays all -1. - valid_counts[0] = 0 - inputs.topk_indices[0] = -1 # DCP without PCP (the gate rejects the pair) always fuses the query. q = torch.cat(q, dim=-1) impl = _fa4_impl( @@ -4172,9 +4102,7 @@ def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode, dcp): need_to_return_lse_for_decode=dcp > 1, ) - # Entered through forward_mqa on a plain (non-HiSparse) index group, whose - # converter is stubbed: the lane choice is the decode/prefill token split. - inputs.metadata.num_decode_tokens = num_tokens if is_decode else 0 + inputs.metadata.num_decode_tokens = num_tokens if is_decode else 1 convert = MagicMock(return_value=(inputs.topk_indices, valid_counts)) impl._convert_logical_to_physical_topk = convert monkeypatch.setattr(sparse_mla, "triton_filter_and_convert_dcp_index", convert) @@ -4182,7 +4110,6 @@ def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode, dcp): with torch.inference_mode(): out, lse = impl.forward_mqa(q, inputs.kv_cache, inputs.metadata, None) - # Only DCP asks for an LSE, and both lanes then return it as (tokens, heads). assert lse is None if dcp == 1 else lse.shape == (num_tokens, kernel_heads) assert out.shape == (num_tokens, kernel_heads, 512) @@ -4200,7 +4127,6 @@ def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode, dcp): # The flat KV rows of the whole page pool, addressed by the top-k list. assert fa4["v"].data_ptr() == inputs.kv_cache.data_ptr() assert fa4["seqused_k"].tolist() == [4 * inputs.block_size] * num_tokens - # The qv kernel takes the two halves, fused query or not. assert fa4["q"].shape[-1] == inputs.rope_dim assert fa4["q_v"].shape[-1] == inputs.kv_lora_rank assert (out == 1.0).all() @@ -4213,19 +4139,10 @@ def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode, dcp): assert trtllm["return_lse"] is (dcp > 1) assert trtllm["query"].shape[0] == num_tokens assert trtllm["seq_lens"] is valid_counts - assert trtllm["block_tables"].shape == (num_tokens, 1, inputs.topk) - assert trtllm["sparse_mla_top_k"] == inputs.topk - # The lane runs on the shared buffer the builder pre-sizes under DCP. - assert trtllm["workspace_buffer"] is _get_workspace_buffer(torch.device("cpu")) - if dcp == 1: - assert (out == 2.0).all() - return - expected = calls.lse(num_tokens) * math.log(2.0) - torch.testing.assert_close(lse[1:], expected[1:], rtol=0, atol=1e-5) - assert (out[1:] == 2.0).all() - # Token 0 owns no slot on this rank: the DCP merge identity. - assert (out[0] == 0).all() - assert torch.isneginf(lse[0]).all() + assert (out == 2.0).all() + if dcp > 1: + expected = calls.lse(num_tokens) * math.log(2.0) + torch.testing.assert_close(lse, expected, rtol=0, atol=1e-5) def test_fa4_sparse_autotune_hisparse_decode_runs_decode_lane(monkeypatch): @@ -4249,12 +4166,7 @@ def fa4_layer(): impl = _fa4_impl(topk_indices_buffer=buffer, index_group=index_group) return SimpleNamespace(impl=impl, hisparse_cache=object()) - # Two layers share one warmup key; another backend's impl is skipped. - layers = { - "a": fa4_layer(), - "b": fa4_layer(), - "other": SimpleNamespace(impl=object(), hisparse_cache=object()), - } + layers = {"a": fa4_layer()} runner = SimpleNamespace( vllm_config=SimpleNamespace( compilation_config=SimpleNamespace(static_forward_context=layers) @@ -4267,7 +4179,6 @@ def fa4_layer(): (fa4,) = calls.fa4 assert calls.trtllm == [] assert fa4["v"].data_ptr() == hot.data_ptr() - # The decode lane's cached varlen scalars, keyed on the hot row count. assert fa4["seqused_k"].tolist() == [hot.shape[0] * block_size] * max_num_reqs assert fa4["gather_kv_indices"].shape == (max_num_reqs, topk) @@ -4284,7 +4195,6 @@ def fa4_layer(): ({"index_topk": 192}, "index_topk"), # PCP over a DCP-sharded cache: the inherited reason names this backend. ({"dcp": 2, "pcp": 2}, "FLASH_ATTN_MLA_SPARSE_FA4 does not support"), - ({"hisparse": True}, None), ({"kv_dtype": "fp8"}, "kv_cache_dtype"), ({"flashinfer": False}, "FlashInfer"), ], @@ -4300,8 +4210,6 @@ def test_fa4_sparse_gate(monkeypatch, config, expected): vllm_config = _build_sparse_dcp_vllm_config(config.get("heads", 16), dcp) vllm_config.parallel_config.prefill_context_parallel_size = config.get("pcp", 1) vllm_config.model_config.hf_text_config.index_topk = config.get("index_topk", 128) - if config.get("hisparse"): - vllm_config.attention_config.hisparse_config = HiSparseConfig() with set_current_vllm_config(vllm_config): reasons = FlashAttnMLASparseFA4Backend.validate_configuration( diff --git a/tests/v1/e2e/general/test_hisparse.py b/tests/v1/e2e/general/test_hisparse.py index 45d3f48990df..b122c6990404 100644 --- a/tests/v1/e2e/general/test_hisparse.py +++ b/tests/v1/e2e/general/test_hisparse.py @@ -172,9 +172,6 @@ def test_hisparse_spill_and_prefix_restore( worker = _get_hisparse_worker(runner) engine_core = runner.llm.llm_engine.engine_core.engine_core model_runner = engine_core.model_executor.driver_worker.worker.model_runner - if forced_backend is not None: - selected = model_runner.attn_groups[0][0].backend.get_name() - assert selected == forced_backend full_graph_calls = 0 full_replay_pending = False original_run_fullgraph = model_runner.cudagraph_manager.run_fullgraph diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 504124176a2f..a8fa74f91f62 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -801,10 +801,7 @@ def forward_mqa( 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; a strided buffer would " - "misaddress the flat rows the top-k indices point at" - ) + 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( diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index e43ebfa6c870..724839881253 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -266,7 +266,7 @@ def _fa4_cute_mla_available() -> str | None: class FlashAttnMLASparseFA4Backend(FlashInferMLASparseTRTLLMBackend): - """Decode-only batches on FA4; every other batch on the trtllm-gen kernel.""" + """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. @@ -304,7 +304,6 @@ def supports_combination(cls, *args: Any, **kwargs: Any) -> str | None: return reason vllm_config = get_current_vllm_config_or_none() if vllm_config is None: - # No engine context: only the static gates above apply. return None if vllm_config.model_config is not None: hf_config = vllm_config.model_config.hf_text_config @@ -333,7 +332,6 @@ def supports_combination(cls, *args: Any, **kwargs: Any) -> str | None: f"FA4 sparse MLA requires {head_counts} gathered query heads, " f"got num_heads={num_heads} * dcp_size={dcp_size}" ) - # super()'s remaining gates; their reasons name ``cls``, hence this backend. return super().supports_combination(*args, **kwargs) @@ -357,8 +355,6 @@ def _fa4_varlen_scalars( class FlashAttnMLASparseFA4Impl(FlashInferMLASparseImpl): - """FA4 for decode rows; FlashInfer's trtllm-gen kernel for prefill rows.""" - # FA4 emits a natural-log LSE; prefill converts trtllm-gen's log2 LSE. lse_base_on_e: bool = True @@ -388,19 +384,8 @@ def _forward_mqa_kernel( ql_nope, q_pe = q else: # 16B-aligned halves of the fused query; FA4 reads them in place. - assert q.shape[-1] == self.kv_lora_rank + self.qk_rope_head_dim ql_nope, q_pe = q[..., : self.kv_lora_rank], q[..., self.kv_lora_rank :] kv_rows, _ = flat_kv_row_view(kv_cache, block_size) - return self._decode(ql_nope, q_pe, kv_rows, topk_indices, valid_counts) - - def _decode( - self, - ql_nope: torch.Tensor, - q_pe: torch.Tensor, - kv_rows: torch.Tensor, - topk_indices: torch.Tensor, - valid_counts: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor | None]: num_tokens = q_pe.shape[0] num_kv_rows = kv_rows.shape[0] assert self.topk_indices_buffer is not None diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 24f621177bb2..e1c4c8d06741 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -125,7 +125,6 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): FLASH_ATTN_MLA_SPARSE = ( "vllm.v1.attention.backends.mla.flashattn_mla_sparse.FlashAttnMLASparseBackend" ) - # Blackwell sibling on FA4's cute-DSL MLA gather kernel. FLASH_ATTN_MLA_SPARSE_FA4 = ( "vllm.v1.attention.backends.mla.flashattn_mla_sparse." "FlashAttnMLASparseFA4Backend" diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index 8ec4128f8be6..c746dd0a54a8 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -292,12 +292,9 @@ 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, gather_kv_length) int32, one KV row list per query - token, with `-1` as the "no token" sentinel. Requires q_v, excludes - block_table, and gather_kv_length must be a multiple of 128. FA4 only. - gather_kv_valid_length: (total_q,) int32, leading real entries per - gather_kv_indices row. The kernel attends `round_up(length, 128)` - entries, so the rest must be `-1`. Needs gather_kv_indices. 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). From 50d451294c5160d0a2626025d17ebba894f69819 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Fri, 9 Oct 2026 07:37:04 +0000 Subject: [PATCH 6/9] [Tests] Trim sparse MLA wiring tests Drop mock-only tests that pin internals and duplicate coverage, fold the HiSparse decode-only tests into the lanes test, split the shared converter mock in the FA4 routing test, and replace the SM100 priority restatement with a fallback test. Signed-off-by: Simon Veitner Co-Authored-By: Claude --- ...st_flashmla_nope_sm90_backend_selection.py | 36 +- .../v1/attention/test_sparse_mla_backends.py | 457 ++++-------------- tests/v1/e2e/general/test_hisparse.py | 25 +- 3 files changed, 137 insertions(+), 381 deletions(-) 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 74fbbb43e979..f60b1a5ed29e 100644 --- a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py +++ b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py @@ -182,9 +182,33 @@ def test_sm90_nope_fp8_ds_mla_resolves_to_flashmla(): assert not backend_cls.supports_kv_cache_dtype("fp8_ds_mla"), name -def test_sm100_576_bf16_low_heads_prefers_fa4(): - assert _sparse_order("auto", 576, num_heads=16)[:3] == [ - "FLASH_ATTN_MLA_SPARSE_FA4", - "FLASHINFER_MLA_SPARSE", - "FLASHMLA_SPARSE", - ] +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 c380c5d6d033..c8bc5ee7cdc9 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -3367,53 +3367,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 run_kernel(self, q, cache, indices, counts, **kwargs): # 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.dcp_world_size = 1 - impl.index_group = index_group - impl.index_group_index = 0 - impl._forward_mqa_kernel = MethodType(run_kernel, impl) - metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64) - - 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", @@ -3543,42 +3496,6 @@ def run(): 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.dcp_world_size = 1 - impl.index_group = index_group - impl.index_group_index = 0 - impl._forward_mqa_kernel = MagicMock(return_value=(q_nope[..., :1], None)) - 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._forward_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) @@ -3613,128 +3530,51 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group(): impl._run_mqa_kernel.assert_called_once() -def test_flashattn_plain_path_passes_converted_topk(monkeypatch): - """FA3's plain path hands the converted top-k straight to the kernel.""" - import vllm.v1.attention.backends.mla.flashattn_mla_sparse as fa3_module - - device = torch.device(DEVICE_TYPE) - num_tokens, topk_tokens, lora, rope, block_size = 3, 8, 4, 2, 4 - kv_cache = torch.arange( - 2 * block_size * (lora + rope), dtype=torch.bfloat16, device=device - ).view(2, block_size, lora + rope) - logical = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32, device=device) - physical = torch.full_like(logical, 5) - counts = torch.full((num_tokens,), topk_tokens, dtype=torch.int32, device=device) - recorded: dict = {} - - def fake_kernel(**kwargs): - recorded.update(kwargs) - return torch.zeros(num_tokens, 2, lora, device=device) - - monkeypatch.setattr(fa3_module, "flash_attn_varlen_func", fake_kernel) - - impl = object.__new__(FlashAttnMLASparseImpl) - impl.topk_indices_buffer = logical - impl.index_group = SimpleNamespace(num_layers=2) - impl.dcp_world_size = 1 - impl._convert_logical_to_physical_topk = MagicMock(return_value=(physical, counts)) - impl.kv_lora_rank = lora - impl.qk_rope_head_dim = rope - impl.cu_seqlens_q_buffer = torch.arange( - num_tokens + 1, dtype=torch.int32, device=device - ) - impl.scale = 0.25 - q_nope = torch.zeros(num_tokens, 2, lora, device=device) - q_rope = torch.zeros(num_tokens, 2, rope, device=device) - metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=block_size) - - out, lse = impl.forward_mqa((q_nope, q_rope), kv_cache, metadata, None) - - assert recorded["block_table"] is physical - assert recorded["seqused_k"] is counts - kv_rows = kv_cache.view(-1, lora + rope) - assert recorded["q"] is q_rope and recorded["q_v"] is q_nope - torch.testing.assert_close(recorded["k"][:, 0, 0], kv_rows[:, lora:]) - torch.testing.assert_close(recorded["v"][:, 0, 0], kv_rows[:, :lora]) - assert lse is None and out.shape == (num_tokens, 2, lora) - - -@pytest.mark.parametrize("dcp", [1, 2], ids=["native", "dcp2"]) -def test_sparse_mqa_plain_lane_wiring(monkeypatch, dcp): - """The plain lane converts once and hands the hook one whole-batch call.""" +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] - # One row longer than the batch, so slicing to num_tokens is observable. - topk_buffer = torch.zeros(num_tokens + 1, topk_tokens, dtype=torch.int32) - req_id_per_token = torch.zeros(num_tokens + 1, dtype=torch.int32) q = torch.zeros(num_tokens, 2, head_dim) - physical = torch.full((num_tokens, topk_tokens), 5, dtype=torch.int32) - counts = torch.full((num_tokens,), topk_tokens, dtype=torch.int32) - layer = object() 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=req_id_per_token, + req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32), cp_kv_cache_interleave_size=2, ) - index_filter = MagicMock(return_value=(physical, counts)) + 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 = topk_buffer + 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 = dcp - impl.dcp_rank = dcp - 1 - impl._convert_logical_to_physical_topk = MagicMock(return_value=(physical, counts)) + 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)) - out, lse = impl.forward_mqa(q, kv_cache, metadata, layer) - - assert out is q and lse is None - if dcp > 1: - impl._convert_logical_to_physical_topk.assert_not_called() - args, kwargs = index_filter.call_args - assert args[0].data_ptr() == req_id_per_token.data_ptr() - assert args[0].shape == (num_tokens,) - assert args[1] is metadata.block_table - assert args[2].data_ptr() == topk_buffer.data_ptr() - assert args[2].shape == (num_tokens, topk_tokens) - assert kwargs["BLOCK_STRIDE_ROWS"] == 8 - else: - index_filter.assert_not_called() - args, kwargs = impl._convert_logical_to_physical_topk.call_args - assert args[0].data_ptr() == topk_buffer.data_ptr() - assert args[0].shape == (num_tokens, topk_tokens) - assert args[1] is metadata - assert kwargs == {"block_stride_rows": 8, "return_valid_counts": True} - impl._forward_mqa_kernel.assert_called_once_with( - q, - kv_cache, - physical, - counts, - layer=layer, - block_size=block_size, - is_decode=True, - ) + 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("split_q", [False, True], ids=["fused", "split"]) @pytest.mark.parametrize( - "num_decode_tokens,all_resident,lane", + "num_decode_tokens,all_resident,split_q", [ - (0, True, "resident"), - (0, False, "staged"), - # Fewer prefill than decode rows: row counts cannot tell the calls apart. - (2, True, "mixed"), + 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, lane, split_q + 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 @@ -3745,9 +3585,7 @@ def test_sparse_mqa_hisparse_lanes( 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) - topk_buffer = torch.zeros(num_tokens + 1, topk_tokens, dtype=torch.int32) q = torch.arange(num_tokens, dtype=torch.float32).view(num_tokens, 1, 1) - layer = object() metadata = SimpleNamespace( num_decode_tokens=num_decode_tokens, block_size=block_size ) @@ -3758,106 +3596,63 @@ def converted(rows, fill): torch.full((rows,), rows, dtype=torch.int32), ) - decode_ret = converted(num_decode_tokens, 1) - resident_ret = converted(num_tokens, 2) - staged_ret = converted(num_prefill, 3) - req_ids = torch.zeros(num_prefill, dtype=torch.int32) - staged_block_table = torch.zeros(1, staged.shape[0], 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=decode_ret if num_decode_tokens else resident_ret - ) - index_group.stage_prefill_rows = MagicMock( - return_value=(staged, staged_block_table, req_ids) + 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 staged_ret + return converted(num_prefill, 3) - convert = MagicMock(side_effect=convert_staged) - monkeypatch.setattr(sparse_mla, "triton_convert_req_index_to_global_index", convert) + monkeypatch.setattr( + sparse_mla, + "triton_convert_req_index_to_global_index", + MagicMock(side_effect=convert_staged), + ) - def hook(rows, cache, indices, counts, **kwargs): + def hook(rows, cache, indices, counts, *, is_decode, **kwargs): if split_q: assert rows[0].tolist() == rows[1].tolist() rows = rows[0] - assert kwargs["layer"] is layer and kwargs["block_size"] == block_size - is_decode = kwargs["is_decode"] - calls.append((rows, cache.data_ptr(), indices, counts, is_decode)) + 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 = topk_buffer + 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, layer) + 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.shape == (num_tokens, 1, 1) - if num_decode_tokens: - rows, kv_ptr, indices, counts, is_decode = calls[0] - assert is_decode is True - assert kv_ptr == hot.data_ptr() - # Identity, not equality: a sub-slice would break pointer stability. - assert indices is decode_ret[0] and counts is decode_ret[1] - assert rows.tolist() == q[:num_decode_tokens].tolist() - args, kwargs = index_group.convert_logical_to_physical_topk.call_args_list[0] - assert args[1].data_ptr() == topk_buffer.data_ptr() - assert args[1].shape == (num_decode_tokens, topk_tokens) - assert args[2] is metadata - assert kwargs == {"block_stride_rows": None, "return_valid_counts": True} - - rows, kv_ptr, indices, counts, is_decode = calls[-1] - assert is_decode is False - assert rows.tolist() == q[num_decode_tokens:].tolist() - if lane == "resident": - index_group.stage_prefill_rows.assert_not_called() - convert.assert_not_called() - assert kv_ptr == hot.data_ptr() - assert indices is resident_ret[0] and counts is resident_ret[1] - args, kwargs = index_group.convert_logical_to_physical_topk.call_args - assert args[1].data_ptr() == topk_buffer.data_ptr() - assert args[1].shape == (num_tokens, topk_tokens) - assert args[2] is metadata - # The resident cache carries its own stride, so the group resolves it. - assert kwargs == {"block_stride_rows": None, "return_valid_counts": True} - else: - assert index_group.convert_logical_to_physical_topk.call_count == min( - num_decode_tokens, 1 - ) - assert index_group.stage_prefill_rows.call_args.args == (0, host, metadata) - assert kv_ptr == staged.data_ptr() - assert indices is staged_ret[0] and counts is staged_ret[1] - args, kwargs = convert.call_args - assert args[0] is req_ids and args[1] is staged_block_table - assert args[2].data_ptr() == topk_buffer[num_decode_tokens:].data_ptr() - assert args[2].shape == (num_prefill, topk_tokens) - # The staging buffer is densely packed: no block stride to pass on. - assert kwargs == { - "BLOCK_SIZE": block_size, - "NUM_TOPK_TOKENS": topk_tokens, - "return_valid_counts": True, - } - - if lane == "mixed": - assert [call[-1] for call in calls] == [True, False] - assert (out[:num_decode_tokens] == 1.0).all() - assert (out[num_decode_tokens:] == 2.0).all() - else: - assert [call[-1] for call in calls] == [False] - assert (out == 2.0).all() + assert out.flatten().tolist() == [1.0] * num_decode_tokens + [2.0] * num_prefill def test_hisparse_resident_prefill_uses_attention_block_stride(): @@ -3928,32 +3723,31 @@ def test_hisparse_mixed_single_token_batch_is_not_decode_only(): assert not cache.decode_batch -def _fa4_impl(**fields): - buffer = fields.pop("topk_indices_buffer") - fields.setdefault("num_heads", 16) - fields.setdefault("kv_lora_rank", 512) - fields.setdefault("qk_nope_head_dim", 128) - fields.setdefault("qk_rope_head_dim", 64) - fields.setdefault("scale", 1.0) - fields.setdefault("kv_cache_dtype", "auto") - fields.setdefault("dcp_world_size", 1) - fields.setdefault("dcp_rank", 0) - fields.setdefault("need_to_return_lse_for_decode", False) - fields.setdefault("index_group_index", 0) - fields.setdefault("is_nope_mla", False) - fields.setdefault("_workspace_buffer", None) - fields.setdefault("bmm1_scale", None) - fields.setdefault("bmm2_scale", None) - if "index_group" not in fields: +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. - fields["index_group"] = SparseMLAIndexGroupBuilder(buffer).register_layer( + index_group=SparseMLAIndexGroupBuilder(topk_indices_buffer).register_layer( False - )[0] - stub = object.__new__(FlashAttnMLASparseFA4Impl) - stub.__dict__.update(fields) + )[0], + index_group_index=0, + ) + impl.__dict__.update(fields) # topk_indices_buffer is a data descriptor, so it needs its own setter. - stub.topk_indices_buffer = buffer - return stub + impl.topk_indices_buffer = topk_indices_buffer + return impl def _fa4_inputs(counts, *, num_heads=16, topk=128, num_blocks=4, device=DEVICE_TYPE): @@ -4078,111 +3872,54 @@ def test_fa4_sparse_decode_kernel_correctness(num_heads, return_lse): ) -@pytest.mark.parametrize("dcp", [1, 2], ids=["native", "dcp2"]) -@pytest.mark.parametrize("is_decode", [True, False]) -def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode, dcp): +@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 = len(counts) - # The query reaches the hook all-gathered: more heads than the impl's own. + 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") - valid_counts = inputs.valid_counts.clone() - q = (inputs.ql_nope, inputs.q_pe) - if dcp > 1: - # DCP without PCP (the gate rejects the pair) always fuses the query. - q = torch.cat(q, dim=-1) + # 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=dcp > 1, + 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 - convert = MagicMock(return_value=(inputs.topk_indices, valid_counts)) - impl._convert_logical_to_physical_topk = convert - monkeypatch.setattr(sparse_mla, "triton_filter_and_convert_dcp_index", convert) with torch.inference_mode(): out, lse = impl.forward_mqa(q, inputs.kv_cache, inputs.metadata, None) - assert lse is None if dcp == 1 else lse.shape == (num_tokens, kernel_heads) + 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 (dcp > 1) - assert fa4["gather_kv_indices"] is inputs.topk_indices - assert fa4["gather_kv_valid_length"] is valid_counts - # Every token is its own one-row FA4 sequence, spec-decode rows included. - assert fa4["max_seqlen_q"] == 1 + assert fa4["return_softmax_lse"] is True assert fa4["cu_seqlens_q"].tolist() == list(range(num_tokens + 1)) - # Gathered rows are addressed by the top-k list, so the K offsets stay 0. - assert fa4["cu_seqlens_k"].tolist() == [0] * (num_tokens + 1) - # The flat KV rows of the whole page pool, addressed by the top-k list. - assert fa4["v"].data_ptr() == inputs.kv_cache.data_ptr() - assert fa4["seqused_k"].tolist() == [4 * inputs.block_size] * num_tokens - assert fa4["q"].shape[-1] == inputs.rope_dim - assert fa4["q_v"].shape[-1] == inputs.kv_lora_rank assert (out == 1.0).all() - if dcp > 1: - torch.testing.assert_close(lse, calls.lse(num_tokens), rtol=0, atol=0) - return - - (trtllm,) = calls.trtllm - assert calls.fa4 == [] - assert trtllm["return_lse"] is (dcp > 1) - assert trtllm["query"].shape[0] == num_tokens - assert trtllm["seq_lens"] is valid_counts - assert (out == 2.0).all() - if dcp > 1: + # 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) -def test_fa4_sparse_autotune_hisparse_decode_runs_decode_lane(monkeypatch): - """The HiSparse warmup reaches FA4's hook once per key, on the decode lane.""" - from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( - autotune_hisparse_flashinfer_attention, - ) - - block_size, topk, max_num_reqs = 64, 128, 3 - hot = torch.zeros(5, block_size, 576, dtype=torch.bfloat16) - calls = _record_fa4_kernels(monkeypatch, 16, device="cpu") - runtime = SimpleNamespace( - hot=SimpleNamespace(attention_cache=hot, block_size=block_size), - max_num_reqs=max_num_reqs, - ) - index_group = object.__new__(HiSparseMLAIndexGroup) - index_group.caches = [SimpleNamespace(runtime=runtime)] - - def fa4_layer(): - buffer = torch.zeros(max_num_reqs, topk, dtype=torch.int32) - impl = _fa4_impl(topk_indices_buffer=buffer, index_group=index_group) - return SimpleNamespace(impl=impl, hisparse_cache=object()) - - layers = {"a": fa4_layer()} - runner = SimpleNamespace( - vllm_config=SimpleNamespace( - compilation_config=SimpleNamespace(static_forward_context=layers) - ) - ) - - with torch.inference_mode(): - autotune_hisparse_flashinfer_attention(runner) - - (fa4,) = calls.fa4 - assert calls.trtllm == [] - assert fa4["v"].data_ptr() == hot.data_ptr() - assert fa4["seqused_k"].tolist() == [hot.shape[0] * block_size] * max_num_reqs - assert fa4["gather_kv_indices"].shape == (max_num_reqs, topk) - - @pytest.mark.parametrize( "config,expected", [ @@ -4191,11 +3928,8 @@ def fa4_layer(): ({"heads": 12}, "heads"), ({"dcp": 2}, None), ({"dcp": 8}, "heads"), - ({"index_topk": None}, "index_topk"), ({"index_topk": 192}, "index_topk"), - # PCP over a DCP-sharded cache: the inherited reason names this backend. - ({"dcp": 2, "pcp": 2}, "FLASH_ATTN_MLA_SPARSE_FA4 does not support"), - ({"kv_dtype": "fp8"}, "kv_cache_dtype"), + ({"rope": 32}, "(kv_lora_rank, qk_rope_head_dim)"), ({"flashinfer": False}, "FlashInfer"), ], ) @@ -4208,14 +3942,15 @@ def test_fa4_sparse_gate(monkeypatch, config, expected): monkeypatch.setattr(fa4_sparse, "_fa4_cute_mla_available", lambda: None) dcp = config.get("dcp", 1) vllm_config = _build_sparse_dcp_vllm_config(config.get("heads", 16), dcp) - vllm_config.parallel_config.prefill_context_parallel_size = config.get("pcp", 1) - vllm_config.model_config.hf_text_config.index_topk = config.get("index_topk", 128) + 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=config.get("kv_dtype", "auto"), + kv_cache_dtype="auto", block_size=64, use_mla=True, has_sink=False, diff --git a/tests/v1/e2e/general/test_hisparse.py b/tests/v1/e2e/general/test_hisparse.py index b122c6990404..d9e1fabf03f3 100644 --- a/tests/v1/e2e/general/test_hisparse.py +++ b/tests/v1/e2e/general/test_hisparse.py @@ -83,19 +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"] -) -@pytest.mark.parametrize( - "attention_backend", - ["FLASHINFER_MLA_SPARSE", "FLASH_ATTN_MLA_SPARSE_FA4"], - ids=["flashinfer", "fa4"], + "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, + attention_backend: str | None, ): """Spilled prefixes restore and FULL-graph decode writes reach host KV. @@ -107,12 +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") - forced_backend: str | None = attention_backend - if not current_platform.is_device_capability_family(100): - if attention_backend == "FLASH_ATTN_MLA_SPARSE_FA4": - pytest.skip("FA4 sparse MLA requires SM 10.x") - # FLASHINFER_MLA_SPARSE is SM 10.x only; elsewhere use the default. - forced_backend = None + 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") @@ -154,7 +151,7 @@ def test_hisparse_spill_and_prefix_restore( ) ), kv_transfer_config=kv_transfer_config, - attention_backend=forced_backend, + attention_backend=attention_backend, block_size=64, max_model_len=320, max_num_batched_tokens=1024, From 99e55b6538046245b589ccb5670d55183a8fa649 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Fri, 9 Oct 2026 17:07:54 +0000 Subject: [PATCH 7/9] [CI] Resolve sparse MLA backends lazily in kv cache layout test for ROCm The FA4 sparse MLA backend imports vllm.vllm_flash_attn, which is not built on ROCm, so importing it at module level broke test collection there. Parametrize over AttentionBackendEnum and resolve the backend class inside the test, and mark the FA4 case CUDA only, so the CPU cases for the other backends keep running on ROCm. Signed-off-by: Simon Veitner --- .../test_sparse_mla_kv_cache_layout.py | 31 ++++++++++--------- 1 file changed, 17 insertions(+), 14 deletions(-) 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 5abc6d3b1312..2d9b35ddeec9 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py @@ -12,18 +12,12 @@ from vllm.model_executor.layers.attention.mla_attention import MLAAttention from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache from vllm.platforms import current_platform -from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( - FlashAttnMLASparseFA4Backend, -) -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_utils import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs @@ -51,18 +45,27 @@ @pytest.mark.parametrize( "backend,cache_dtype,sm100,stride_alignment", [ - (FlashInferMLASparseTRTLLMBackend, "auto", True, 1152), - (FlashInferMLASparseTRTLLMBackend, "fp8", True, 576), - (FlashAttnMLASparseFA4Backend, "auto", True, 1152), - (FlashMLASparseBackend, "auto", True, 1152), - (FlashMLASparseBackend, "fp8_ds_mla", False, None), - (FlashMLASparseBackend, "fp8_ds_mla", True, 656), - (FlashMLASparseBackend, "nvfp4_ds_mla", True, 352), + (AttentionBackendEnum.FLASHINFER_MLA_SPARSE, "auto", True, 1152), + (AttentionBackendEnum.FLASHINFER_MLA_SPARSE, "fp8", True, 576), + pytest.param( + AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE_FA4, + "auto", + True, + 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, None), + (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 ) From 7fdd3ad28091e0428ae6612b47f195b37ebc7d63 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Sun, 11 Oct 2026 21:32:29 +0000 Subject: [PATCH 8/9] [CI] Skip FA4 sparse MLA fallback selection test on non-CUDA platforms The test imports vllm.v1.attention.backends.mla.flashattn_mla_sparse, which imports vllm.vllm_flash_attn; that extension is not built on ROCm, so the test failed with ImportError on AMD CI. Mark it CUDA only, matching the neighbouring selection test. Signed-off-by: Simon Veitner --- tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py | 1 + 1 file changed, 1 insertion(+) 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 bf534dd03a17..8f31bd488aa4 100644 --- a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py +++ b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py @@ -279,6 +279,7 @@ def test_sm90_nope_fp8_ds_mla_resolves_to_flashmla(): 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 4131823817088fab11ac142f2ea9b826f177764e Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Mon, 12 Oct 2026 19:56:27 +0000 Subject: [PATCH 9/9] [Misc] Use the CUDA-guarded flashattn_mla_sparse alias in FA4 tests Restore main's `flashattn_sparse_module` alias for vllm.v1.attention.backends.mla.flashattn_mla_sparse, imported under the `if current_platform.is_cuda():` guard, which was dropped during the merge resolution. Use it in `_record_fa4_kernels` and `test_fa4_sparse_gate` instead of their per-test imports. No behaviour change. Co-authored-by: Claude Signed-off-by: Simon Veitner --- tests/v1/attention/test_sparse_mla_backends.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 57b671539d3b..bb7b38acd910 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -71,6 +71,9 @@ from vllm.v1.attention.backends.mla import index_group as index_group_module if current_platform.is_cuda(): + from vllm.v1.attention.backends.mla import ( + flashattn_mla_sparse as flashattn_sparse_module, + ) from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( FlashAttnMLASparseFA4Backend, FlashAttnMLASparseFA4Impl, @@ -4009,8 +4012,6 @@ def _record_fa4_kernels(monkeypatch, num_heads, *, device): """Replace both lanes with recorders whose outputs identify the lane.""" import flashinfer.decode - import vllm.v1.attention.backends.mla.flashattn_mla_sparse as fa4_sparse - fa4_calls: list[dict] = [] trtllm_calls: list[dict] = [] @@ -4033,7 +4034,7 @@ def fake_trtllm(**kwargs): ) return (out, lse(rows)) if kwargs.get("return_lse") else out - monkeypatch.setattr(fa4_sparse, "flash_attn_varlen_func", fake_fa4) + monkeypatch.setattr(flashattn_sparse_module, "flash_attn_varlen_func", fake_fa4) monkeypatch.setattr( flashinfer.decode, "trtllm_batch_decode_with_kv_cache_mla", fake_trtllm ) @@ -4151,11 +4152,11 @@ def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode): ) def test_fa4_sparse_gate(monkeypatch, config, expected): """``validate_configuration`` on a fixed SM100: one row per gate reason.""" - import vllm.v1.attention.backends.mla.flashattn_mla_sparse as fa4_sparse - has_flashinfer = config.get("flashinfer", True) monkeypatch.setattr("vllm.utils.flashinfer.has_flashinfer", lambda: has_flashinfer) - monkeypatch.setattr(fa4_sparse, "_fa4_cute_mla_available", lambda: None) + 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