diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 00411eaa4ee..243db16a2bc 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -47,6 +47,19 @@ from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite +@functools.cache +def _flash_attention_kv_cache_layout() -> str: + """Return the installed vLLM backend's K/V packing contract.""" + cache_shape = FlashAttentionBackend.get_kv_cache_shape(3, 16, 1, 16) + if cache_shape == (2, 3, 16, 1, 16): + return "kv-first" + if cache_shape == (3, 2, 16, 1, 16): + return "blocks-first" + if cache_shape == (3, 1, 16, 32): + return "packed" + raise RuntimeError(f"Unsupported vLLM FlashAttention KV cache shape {cache_shape}") + + def _target_sparse_ratio_for_phase(target_sparse_ratio, phase: str) -> float: """Return target sparsity for a phase, defaulting old checkpoint metadata.""" if isinstance(target_sparse_ratio, float | int): @@ -514,7 +527,13 @@ def native_forward(): if resolved is None: return native_forward() - key_cache, value_cache = kv_cache.unbind(0) + cache_layout = _flash_attention_kv_cache_layout() + if cache_layout == "kv-first": + key_cache, value_cache = kv_cache.unbind(0) + elif cache_layout == "blocks-first": + key_cache, value_cache = kv_cache.unbind(1) + else: + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) is_decode_only = attn_metadata.max_query_len <= 1 common_kw = { "layer": layer, diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 05b3c6ace8d..578922db077 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -28,7 +28,7 @@ import vllm from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends import flashinfer as flashinfer_backend -from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend, FlashAttentionImpl from vllm.v1.attention.backends.flashinfer import ( FlashInferBackend, FlashInferImpl, @@ -541,6 +541,76 @@ def _make_flash_attention_impl(*, sparse=False, quantized=False): return impl +def _flash_attention_kv_cache(num_blocks, page_size, num_kv_heads, head_size): + layout = vllm_plugin._flash_attention_kv_cache_layout() + if layout == "packed": + shape = [num_blocks, num_kv_heads, page_size, 2 * head_size] + else: + shape = [num_blocks, page_size, num_kv_heads, head_size] + shape.insert(0 if layout == "kv-first" else 1, 2) + return torch.zeros(shape, dtype=torch.float16) + + +@pytest.mark.parametrize( + ("layout", "backend_shape"), + [ + ("kv-first", (2, 3, 16, 1, 16)), + ("blocks-first", (3, 2, 16, 1, 16)), + ("packed", (3, 1, 16, 32)), + ], +) +def test_flash_attention_forward_follows_backend_kv_cache_layout( + monkeypatch, layout, backend_shape +): + impl = _make_flash_attention_impl(sparse=True) + if layout == "packed": + shape = [3, impl.num_kv_heads, 16, 2 * impl.head_size] + else: + shape = [3, 16, impl.num_kv_heads, impl.head_size] + shape.insert(0 if layout == "kv-first" else 1, 2) + monkeypatch.setattr( + FlashAttentionBackend, "get_kv_cache_shape", staticmethod(lambda *_args: backend_shape) + ) + vllm_plugin._flash_attention_kv_cache_layout.cache_clear() + kv_cache = torch.zeros(shape, dtype=torch.float16) + query = torch.zeros(4, impl.num_heads, impl.head_size, dtype=torch.float16) + metadata = _flash_attention_metadata(query.shape[0], 16) + captured = {} + + def fake_attention(query, **kwargs): + captured.update(kwargs) + return torch.zeros_like(query) + + monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + + try: + impl.forward( + layer=None, + query=query, + key=query, + value=query, + kv_cache=kv_cache, + attn_metadata=metadata, + output=torch.empty_like(query), + ) + finally: + vllm_plugin._flash_attention_kv_cache_layout.cache_clear() + + if layout == "packed": + expected_key_cache, expected_value_cache = kv_cache.transpose(1, 2).split( + impl.head_size, dim=-1 + ) + else: + expected_key_cache, expected_value_cache = kv_cache.unbind(0 if layout == "kv-first" else 1) + assert captured["k_cache"].shape == expected_key_cache.shape + assert captured["v_cache"].shape == expected_value_cache.shape + assert captured["k_cache"].stride() == expected_key_cache.stride() + assert captured["v_cache"].stride() == expected_value_cache.stride() + assert captured["k_cache"].data_ptr() == expected_key_cache.data_ptr() + assert captured["v_cache"].data_ptr() == expected_value_cache.data_ptr() + assert captured["page_size"] == 16 + + def _flash_attention_mixed_metadata(decode_len=1, prefill_len=17): query_lens = (decode_len, prefill_len) seq_lens = (16, 34) @@ -568,7 +638,7 @@ def test_flash_attention_mixed_batch_splits_decode_and_prefill(monkeypatch, quan prefill_tokens = 17 impl = _make_flash_attention_impl(sparse=True, quantized=quantized) query = torch.zeros(1 + prefill_tokens, 2, 64, dtype=torch.float16) - kv_cache = torch.zeros(2, 4, 16, 2, 64, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(4, 16, 2, 64) metadata = _flash_attention_mixed_metadata(decode_len=1, prefill_len=prefill_tokens) layer = SimpleNamespace( _query_quant_in_kernel=quantized, @@ -761,7 +831,7 @@ def test_forward_delegates_cascade_metadata_to_vllm(monkeypatch): """Cascade/prefix-cache metadata should use vLLM's native implementation.""" impl = _clone_sparse_impl(_make_old_impl()) q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size) output = torch.empty_like(q) attn_metadata = type("AttnMetadata", (), {"use_cascade": True})() called = {} @@ -838,9 +908,7 @@ def test_forward_delegates_launches_without_effective_sparse_work( impl = _clone_sparse_impl(_make_old_impl()) impl.sparse_kw = sparse_kw q = torch.zeros(max_query_len, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros( - 2, 1, max_seq_len, impl.num_kv_heads, impl.head_size, dtype=torch.float16 - ) + kv_cache = _flash_attention_kv_cache(1, max_seq_len, impl.num_kv_heads, impl.head_size) output = torch.empty_like(q) attn_metadata = _flash_attention_metadata(max_query_len, max_seq_len) called = {} @@ -903,7 +971,7 @@ def test_forward_resolves_calibrated_skip_softmax_threshold(monkeypatch): "target_sparse_ratio": {"prefill": 0.4, "decode": 0.6}, } q = torch.zeros(max_query_len, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, seq_len, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(1, seq_len, impl.num_kv_heads, impl.head_size) attn_metadata = _flash_attention_metadata(max_query_len, seq_len) captured = {} @@ -980,7 +1048,7 @@ def quantize_q(query): } q = torch.full((4, impl.num_heads, impl.head_size), 2.0, dtype=torch.float16) q[2:] = 10_000 - kv_cache = torch.zeros(2, 4, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(4, 16, impl.num_kv_heads, impl.head_size) metadata = SimpleNamespace( num_actual_tokens=q.shape[0], max_query_len=1, @@ -1023,8 +1091,15 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): "v_qdq_scale": 1.0, } key_cache, value_cache, block_table, seq_lens, decode_kw = calls["decode"] - assert key_cache.data_ptr() == kv_cache[0].data_ptr() - assert value_cache.data_ptr() == kv_cache[1].data_ptr() + layout = vllm_plugin._flash_attention_kv_cache_layout() + if layout == "packed": + expected_key_cache, expected_value_cache = kv_cache.transpose(1, 2).split( + impl.head_size, dim=-1 + ) + else: + expected_key_cache, expected_value_cache = kv_cache.unbind(0 if layout == "kv-first" else 1) + assert key_cache.data_ptr() == expected_key_cache.data_ptr() + assert value_cache.data_ptr() == expected_value_cache.data_ptr() assert block_table is metadata.block_table assert seq_lens is metadata.seq_lens assert calls["query"].shape[0] == metadata.seq_lens.shape[0] @@ -1048,7 +1123,7 @@ def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch): } impl.sparse_kw = {"skip_softmax_threshold": 0.001} q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size) metadata = _flash_attention_metadata(1, 16) captured = {} @@ -1140,7 +1215,7 @@ def test_forward_allows_chunked_prefill_metadata(monkeypatch): q_len = 4 kv_len = 10 q = torch.zeros(q_len, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size) attn_metadata = _flash_attention_metadata(q_len, kv_len) captured = {} diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py index fa11b144354..d08a4072b4d 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py @@ -20,7 +20,7 @@ * ``query_start_loc`` -> ``b_start_loc`` / ``b_seq_len`` * ``seq_lens`` -> ``b_seq_len_k`` -* ``kv_cache.unbind(0)`` -> key_cache / value_cache (axis order) +* backend-declared K/V axis -> key_cache / value_cache * ``k_cache.shape[1]`` -> ``page_size`` Asserted against a contiguous reference call to the underlying Triton kernel. @@ -33,7 +33,7 @@ from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE -from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ModelOptSparseAttentionImpl +from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as vllm_plugin if TRITON_KERNEL_AVAILABLE: from modelopt.torch.kernels.common.attention import attention as triton_attention @@ -46,11 +46,19 @@ } +def _make_backend_paged_cache(k_cache, v_cache): + layout = vllm_plugin._flash_attention_kv_cache_layout() + if layout == "kv-first": + return torch.stack([k_cache, v_cache], dim=0) + if layout == "blocks-first": + return torch.stack([k_cache, v_cache], dim=1) + return torch.cat([k_cache, v_cache], dim=-1).transpose(1, 2) + + def _make_paged_cache(k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page_size): - """Scatter contiguous K/V into a paged KV cache stacked as [2, ...]. + """Scatter contiguous K/V into the installed vLLM paged-cache layout. - Returns a single ``kv_cache`` tensor (matching vLLM's layout that - ``ModelOptSparseAttentionImpl`` consumes via ``kv_cache.unbind(0)``). + Returns a single ``kv_cache`` tensor with the backend-declared K/V axis. """ batch = b_seq_len.shape[0] device, dtype = k.device, k.dtype @@ -76,14 +84,13 @@ def _make_paged_cache(k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page v_cache[g, :n] = v[start + ts : start + te] g += 1 - # Stack on a new leading axis so kv_cache.unbind(0) recovers (k_cache, v_cache). - kv_cache = torch.stack([k_cache, v_cache], dim=0) + kv_cache = _make_backend_paged_cache(k_cache, v_cache) return kv_cache, block_table def _make_impl(num_heads, head_dim, num_kv_heads): """Construct ModelOptSparseAttentionImpl with minimal valid kwargs.""" - return ModelOptSparseAttentionImpl( + return vllm_plugin.ModelOptSparseAttentionImpl( num_heads=num_heads, head_size=head_dim, scale=1.0 / (head_dim**0.5), @@ -132,7 +139,7 @@ def test_prefill_matches_contiguous(self): **_ACTIVE_PREFILL_SPARSE_KW, ) - # Build paged kv_cache shaped [2, num_blocks, page_size, num_kv_heads, head_dim]. + # Build the paged cache using the installed backend's K/V axis. kv_cache, block_table = _make_paged_cache( k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page_size ) @@ -174,7 +181,8 @@ def test_chunked_prefill_is_forwarded_to_kernel(self): block_table=torch.zeros(1, 1, device="cuda", dtype=torch.int32), ) q = torch.zeros(4, 2, 64, device="cuda", dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, 2, 64, device="cuda", dtype=torch.float16) + k_cache = torch.zeros(1, 16, 2, 64, device="cuda", dtype=torch.float16) + kv_cache = _make_backend_paged_cache(k_cache, torch.zeros_like(k_cache)) out = impl.forward( layer=None, query=q, @@ -346,8 +354,7 @@ def test_page_size_inferred_from_k_cache(self): kv_cache, block_table = _make_paged_cache( k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page_size ) - # Sanity: kv_cache axis 1 is page_size. - assert kv_cache.shape == (2, seq_len // page_size, page_size, num_kv_heads, head_dim) + assert kv_cache.shape[2] == page_size attn_metadata = SimpleNamespace( num_actual_tokens=seq_len,