| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the multi-token decode flatten in ``TritonMLAImpl.forward_mqa``. |
| 4 | |
| 5 | ``decode_attention_fwd`` has no causal flag: it launches one program per row of |
| 6 | ``q`` and reads that row's KV extent from ``B_Seqlen``. So a ``query_len``-token |
| 7 | decode block is flattened to one row per query token, and intra-block causality |
| 8 | has to be expressed through the per-row sequence lengths -- row ``t`` gets |
| 9 | ``seq_len - (query_len - 1) + t``, i.e. the committed prefix plus block tokens |
| 10 | ``0..t``. A non-causal draft block instead gives every row the same full extent. |
| 11 | |
| 12 | The kernel is replaced by a spy, so these tests assert on the ``block_table`` |
| 13 | and ``seq_lens`` the backend actually submits and need no GPU or Triton runtime. |
| 14 | """ |
| 15 | |
| 16 | from types import SimpleNamespace |
| 17 | from unittest.mock import patch |
| 18 | |
| 19 | import pytest |
| 20 | import torch |
| 21 | |
| 22 | from vllm.model_executor.layers.attention.mla_attention import QueryLenSupport |
| 23 | from vllm.v1.attention.backend import AttentionCGSupport |
| 24 | from vllm.v1.attention.backends.mla.triton_mla import ( |
| 25 | TritonMLAImpl, |
| 26 | TritonMLAMetadataBuilder, |
| 27 | ) |
| 28 | |
| 29 | NUM_HEADS = 16 |
| 30 | KV_LORA_RANK = 512 |
| 31 | QK_ROPE_HEAD_DIM = 64 |
| 32 | HEAD_SIZE = KV_LORA_RANK + QK_ROPE_HEAD_DIM |
| 33 | PAGE_SIZE = 16 |
| 34 | |
| 35 | # Committed context per request, including a fresh request at zero context |
| 36 | # where the whole KV extent is the block itself. A decode block's seq_len |
| 37 | # already counts the block, so seq_len = context + query_len. |
| 38 | CONTEXT_LENS = [0, 1, 37, 512] |
| 39 | |
| 40 | |
| 41 | def _seq_lens(query_len: int) -> list[int]: |
| 42 | return [c + query_len for c in CONTEXT_LENS] |
| 43 | |
| 44 | |
| 45 | def _make_impl(dcp_world_size: int = 1) -> TritonMLAImpl: |
| 46 | """Only the attributes ``forward_mqa`` reads; __init__ wants a full config.""" |
| 47 | impl = object.__new__(TritonMLAImpl) |
| 48 | impl.kv_lora_rank = KV_LORA_RANK |
| 49 | impl.scale = HEAD_SIZE**-0.5 |
| 50 | impl._sm_count = 304 |
| 51 | impl.dcp_world_size = dcp_world_size |
| 52 | return impl |
| 53 | |
| 54 | |
| 55 | def _make_metadata(seq_lens: list[int], query_len: int, causal: bool): |
| 56 | num_decodes = len(seq_lens) |
| 57 | return SimpleNamespace( |
| 58 | num_decodes=num_decodes, |
| 59 | num_decode_tokens=num_decodes * query_len, |
| 60 | max_seq_len=max(seq_lens), |
| 61 | causal=causal, |
| 62 | decode=SimpleNamespace( |
| 63 | # arange page ids so each request's window is an identifiable slice. |
| 64 | block_table=torch.arange(num_decodes * 64, dtype=torch.int32).reshape( |
| 65 | num_decodes, 64 |
| 66 | ), |
| 67 | seq_lens=torch.tensor(seq_lens, dtype=torch.int32), |
| 68 | ), |
| 69 | ) |
| 70 | |
| 71 | |
| 72 | def _run_forward_mqa( |
| 73 | seq_lens: list[int], |
| 74 | query_len: int, |
| 75 | causal: bool, |
| 76 | q_rows=None, |
| 77 | dcp_world_size: int = 1, |
| 78 | ): |
| 79 | """Drive forward_mqa with the decode kernel spied out. |
| 80 | |
| 81 | Returns the ``(block_table, seq_lens)`` handed to ``decode_attention_fwd``. |
| 82 | """ |
| 83 | captured: dict = {} |
| 84 | |
| 85 | def spy(q, kv_cache, kv_c_cache, o, lse, block_table, b_seq_len, *args, **kwargs): |
| 86 | captured["block_table"] = block_table.detach().clone() |
| 87 | captured["seq_lens"] = b_seq_len.detach().clone() |
| 88 | |
| 89 | metadata = _make_metadata(seq_lens, query_len, causal) |
| 90 | rows = metadata.num_decode_tokens if q_rows is None else q_rows |
| 91 | q = torch.zeros(rows, NUM_HEADS, HEAD_SIZE, dtype=torch.bfloat16) |
| 92 | kv_cache = torch.zeros( |
| 93 | len(seq_lens) * 64, PAGE_SIZE, HEAD_SIZE, dtype=torch.bfloat16 |
| 94 | ) |
| 95 | layer = SimpleNamespace(_k_scale=torch.ones(1), layer_name="test") |
| 96 | |
| 97 | with patch("vllm.v1.attention.backends.mla.triton_mla.decode_attention_fwd", spy): |
| 98 | _make_impl(dcp_world_size).forward_mqa(q, kv_cache, metadata, layer) |
| 99 | |
| 100 | assert captured, "forward_mqa did not reach the decode kernel" |
| 101 | return captured |
| 102 | |
| 103 | |
| 104 | def test_multi_token_decode_flags_are_declared_together(): |
| 105 | """Capture probes max_query_len > 1, which only clears the assert in |
| 106 | build_for_cudagraph_capture once query_len_support raised the threshold.""" |
| 107 | assert ( |
| 108 | TritonMLAMetadataBuilder._cudagraph_support == AttentionCGSupport.UNIFORM_BATCH |
| 109 | ) |
| 110 | assert TritonMLAMetadataBuilder.query_len_support == QueryLenSupport.UNIFORM |
| 111 | |
| 112 | |
| 113 | @pytest.mark.parametrize("query_len", [2, 3, 4, 5, 8]) |
| 114 | def test_causal_block_rows_see_prefix_plus_own_position(query_len): |
| 115 | """Causal verify: row t sees seq_len - (query_len - 1) + t KV entries. |
| 116 | |
| 117 | Regression guard: the causal branch used to be skipped entirely, so a verify |
| 118 | token could attend the draft siblings it is supposed to be checking. Fails |
| 119 | on unmodified upstream, where no flatten happens at all for causal blocks. |
| 120 | """ |
| 121 | seq_lens = _seq_lens(query_len) |
| 122 | captured = _run_forward_mqa(seq_lens, query_len, causal=True) |
| 123 | |
| 124 | got = captured["seq_lens"].tolist() |
| 125 | want = [s - (query_len - 1) + t for s in seq_lens for t in range(query_len)] |
| 126 | assert got == want, ( |
| 127 | "causal flatten is not causal: verify row r*query_len+t must get " |
| 128 | f"seq_len_r - {query_len - 1} + t entries.\n" |
| 129 | f" seq_lens {seq_lens}\n got {got}\n expected {want}" |
| 130 | ) |
| 131 | |
| 132 | # Per-request invariants the arithmetic must hold regardless of query_len. |
| 133 | for r, seq_len in enumerate(seq_lens): |
| 134 | rows = got[r * query_len : (r + 1) * query_len] |
| 135 | assert rows == sorted(rows) and len(set(rows)) == query_len, ( |
| 136 | f"request {r} rows must strictly increase, got {rows}" |
| 137 | ) |
| 138 | assert rows[-1] == seq_len, "last row must see the full sequence" |
| 139 | assert rows[0] == seq_len - (query_len - 1) |
| 140 | assert all(n > 0 for n in rows), f"non-positive KV extent in {rows}" |
| 141 | |
| 142 | |
| 143 | def test_cudagraph_padding_rows_present_no_kv_extent(): |
| 144 | """Causal flattening pins cudagraph padding rows (seq_len 0) to exactly |
| 145 | zero KV extent, keeping them in the kernel's skipped branch.""" |
| 146 | query_len = 4 |
| 147 | seq_lens = _seq_lens(query_len) + [0] |
| 148 | captured = _run_forward_mqa(seq_lens, query_len, causal=True) |
| 149 | |
| 150 | padding = captured["seq_lens"].tolist()[-query_len:] |
| 151 | assert padding == [0] * query_len, ( |
| 152 | f"padding rows must present exactly zero KV extent, got {padding}" |
| 153 | ) |
| 154 | |
| 155 | |
| 156 | @pytest.mark.parametrize("query_len", [2, 4]) |
| 157 | def test_non_causal_block_rows_all_see_the_full_prefix(query_len): |
| 158 | """The DSpark draft block is generated in one pass, so no row is masked.""" |
| 159 | seq_lens = _seq_lens(query_len) |
| 160 | captured = _run_forward_mqa(seq_lens, query_len, causal=False) |
| 161 | |
| 162 | want = [s for s in seq_lens for _ in range(query_len)] |
| 163 | assert captured["seq_lens"].tolist() == want |
| 164 | |
| 165 | |
| 166 | @pytest.mark.parametrize("causal", [True, False]) |
| 167 | def test_block_table_is_expanded_to_match_q_rows(causal): |
| 168 | """The kernel indexes block_table by program_id, so it needs one row per |
| 169 | query token -- a short block_table is read out of bounds.""" |
| 170 | query_len = 4 |
| 171 | captured = _run_forward_mqa(_seq_lens(query_len), query_len, causal=causal) |
| 172 | |
| 173 | num_rows = len(CONTEXT_LENS) * query_len |
| 174 | assert captured["block_table"].shape[0] == num_rows |
| 175 | assert captured["seq_lens"].shape[0] == num_rows |
| 176 | # repeat_interleave, not repeat: a request's rows must stay adjacent. |
| 177 | for r in range(len(CONTEXT_LENS)): |
| 178 | block = captured["block_table"][r * query_len : (r + 1) * query_len] |
| 179 | assert torch.equal(block, block[:1].expand_as(block)) |
| 180 | |
| 181 | |
| 182 | @pytest.mark.parametrize("causal", [True, False]) |
| 183 | def test_single_token_decode_is_untouched(causal): |
| 184 | """query_len == 1 must short-circuit: ordinary decode sees no expansion.""" |
| 185 | seq_lens = _seq_lens(1) |
| 186 | captured = _run_forward_mqa(seq_lens, query_len=1, causal=causal) |
| 187 | |
| 188 | assert captured["seq_lens"].tolist() == seq_lens |
| 189 | assert captured["block_table"].shape[0] == len(seq_lens) |
| 190 | |
| 191 | |
| 192 | def test_expansion_factor_follows_the_q_row_count(): |
| 193 | """The kernel indexes block_table and seq_lens by program_id, so the |
| 194 | expansion must follow q rows, not the token/request counts.""" |
| 195 | num_rows = len(CONTEXT_LENS) * 4 |
| 196 | captured = _run_forward_mqa(_seq_lens(1), query_len=1, causal=True, q_rows=num_rows) |
| 197 | |
| 198 | assert captured["block_table"].shape[0] == num_rows |
| 199 | assert captured["seq_lens"].shape[0] == num_rows |
| 200 | |
| 201 | |
| 202 | def test_non_uniform_block_is_rejected(): |
| 203 | """Q rows must divide evenly across the decode requests.""" |
| 204 | with pytest.raises(AssertionError, match="non-uniform decode block"): |
| 205 | _run_forward_mqa(_seq_lens(1), query_len=1, causal=True, q_rows=6) |
| 206 | |
| 207 | |
| 208 | def test_causal_multi_token_decode_is_rejected_under_dcp(): |
| 209 | """Per-row extents offset the global seq_len, but DCP passes the rank-local |
| 210 | slice, so the causal arithmetic would silently address the wrong KV.""" |
| 211 | with pytest.raises(AssertionError, match="not supported with DCP"): |
| 212 | _run_forward_mqa(_seq_lens(4), query_len=4, causal=True, dcp_world_size=2) |
| 213 | |
| 214 | # Non-causal rows all take the same extent, which stays correct when local. |
| 215 | _run_forward_mqa(_seq_lens(4), query_len=4, causal=False, dcp_world_size=2) |