@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b38e111d3e

[Attention][MLA] Per-request scheduling for MLA chunked context (#50613)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Matthew Bonanni committed Aug 6, 2026 at 12:45 UTC b38e111d3e4806a553ec2798e2b075da7a8b03d3
15 files changed +940 -552
tests/kernels/attention/test_merge_attn_states.py
+1 -31
@@ -11,9 +11,6 @@ from vllm._custom_ops import (
11 scaled_fp8_quant,
12 )
13 from vllm.platforms import current_platform
14 -from vllm.v1.attention.ops.triton_merge_attn_states import (
15 - mask_empty_context,
16 -)
14 from vllm.v1.attention.ops.triton_merge_attn_states import (
15 merge_attn_states as merge_attn_states_triton,
16 )
@@ -76,32 +73,6 @@ DTYPES = [torch.float32, torch.half, torch.bfloat16]
73 all_case_info: list[tuple] = []
74
75
79 -def test_mask_empty_context() -> None:
80 - query_lens = torch.tensor([2] + [1] * 31 + [131, 1], dtype=torch.int32)
81 - query_start_loc = torch.cat(
82 - (torch.zeros(1, dtype=torch.int32), query_lens.cumsum(0))
83 - ).cuda()
84 - context_lens = torch.tensor([4] * 32 + [0, 3], dtype=torch.int32)
85 - context_start_loc = torch.cat(
86 - (torch.zeros(1, dtype=torch.int32), context_lens.cumsum(0))
87 - ).cuda()
88 - num_heads, num_tokens, head_dim = 4, 165, 16
89 - lse = torch.randn(num_heads, num_tokens, device="cuda")
90 - output = torch.randn(num_tokens, num_heads, head_dim, device="cuda")
91 - # Empty-context rows carry undefined (possibly non-finite) attention output.
92 - output[33:164] = float("nan")
93 -
94 - expected_lse = lse.clone()
95 - expected_lse[:, 33:164] = float("-inf")
96 - expected_output = output.clone()
97 - expected_output[33:164] = 0.0
98 -
99 - mask_empty_context(lse, output, query_start_loc, context_start_loc)
100 -
101 - torch.testing.assert_close(lse, expected_lse)
102 - torch.testing.assert_close(output, expected_output)
103 -
104 -
76 @pytest.mark.parametrize("merge_fn", [merge_attn_states_cuda, merge_attn_states_triton])
77 @pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16])
78 def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:
@@ -117,8 +88,7 @@ def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:
88 )
89 suffix_lse = torch.randn(num_heads, num_tokens, device="cuda")
90
120 - # Tokens 2 and 3 are empty on both sides (mask_empty_context already zeroed
121 - # their outputs and set both LSEs to -inf).
91 + # Tokens 2 and 3 are empty on both sides.
92 empty = slice(2, 4)
93 prefix_lse[:, empty] = float("-inf")
94 suffix_lse[:, empty] = float("-inf")
tests/v1/attention/test_mla_backends.py
+102 -26
@@ -280,6 +280,10 @@ BATCH_SPECS = {
280 "spec_decode_medium": BatchSpec(
281 seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[8, 8, 8, 8, 8, 8]
282 ),
283 + "chunked_context_prefill": BatchSpec(
284 + seq_lens=[1568, 520, 8, 80, 96, 8],
285 + query_lens=[32, 8, 8, 16, 16, 8],
286 + ),
287 }
288
289
@@ -1049,6 +1053,7 @@ def run_attention_backend(
1053 k_scale: float,
1054 kv_cache_dtype: str = "auto",
1055 prefill_backend: MLAPrefillBackendEnum | None = None,
1056 + chunked_prefill_workspace_size: int | None = None,
1057 ) -> torch.Tensor:
1058 """Run attention computation using the specified backend's AttentionImpl."""
1059
@@ -1137,6 +1142,13 @@ def run_attention_backend(
1142
1143 # Build metadata
1144 builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
1145 + if chunked_prefill_workspace_size is not None:
1146 + builder.chunked_prefill_workspace_size = chunked_prefill_workspace_size
1147 + builder.chunked_prefill_workspace = (
1148 + builder.chunked_prefill_workspace.new_empty(
1149 + chunked_prefill_workspace_size, head_size
1150 + )
1151 + )
1152 attn_metadata = builder.build(
1153 common_prefix_len=0,
1154 common_attn_metadata=common_attn_metadata,
@@ -1156,32 +1168,7 @@ def run_attention_backend(
1168 return output
1169
1170
1159 -@pytest.mark.parametrize(
1160 - "batch_spec_name",
1161 - [
1162 - "small_decode",
1163 - "small_prefill",
1164 - "mixed_small",
1165 - "medium_decode",
1166 - "medium_prefill",
1167 - "mixed_medium",
1168 - "large_decode",
1169 - "large_prefill",
1170 - "single_decode",
1171 - "single_prefill",
1172 - "spec_decode_small",
1173 - "spec_decode_medium",
1174 - ],
1175 -)
1176 -@pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-R1"])
1177 -@pytest.mark.parametrize("tensor_parallel_size", [1, 4, 8, 16])
1178 -@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
1179 -@pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])
1180 -@pytest.mark.parametrize(
1181 - ("prefill_backend", "qk_nope_head_dim", "v_head_dim"),
1182 - _prefill_backend_dimension_params(),
1183 -)
1184 -def test_backend_correctness(
1171 +def _run_backend_correctness(
1172 default_vllm_config,
1173 dist_init,
1174 workspace_init,
@@ -1194,6 +1181,7 @@ def test_backend_correctness(
1181 prefill_backend: MLAPrefillBackendEnum,
1182 qk_nope_head_dim: int,
1183 v_head_dim: int,
1184 + chunked_prefill_workspace_size: int | None = None,
1185 ):
1186 """
1187 Test that all backends produce similar outputs to a reference implementation
@@ -1592,6 +1580,7 @@ def test_backend_correctness(
1580 q_scale=q_scale,
1581 k_scale=k_scale,
1582 kv_cache_dtype=kv_cache_dtype,
1583 + chunked_prefill_workspace_size=chunked_prefill_workspace_size,
1584 )
1585
1586 # Use backend_idx to get the correct SDPA output for this backend
@@ -1643,3 +1632,90 @@ def test_backend_correctness(
1632 summary = f"{len(failures)} backend(s) failed: {', '.join(backend_names)}"
1633 detailed_msg = "\n".join(failures)
1634 pytest.fail(f"{summary}\n{detailed_msg}")
1635 +
1636 +
1637 +@pytest.mark.parametrize(
1638 + "batch_spec_name",
1639 + [
1640 + "small_decode",
1641 + "small_prefill",
1642 + "mixed_small",
1643 + "medium_decode",
1644 + "medium_prefill",
1645 + "mixed_medium",
1646 + "large_decode",
1647 + "large_prefill",
1648 + "single_decode",
1649 + "single_prefill",
1650 + "spec_decode_small",
1651 + "spec_decode_medium",
1652 + ],
1653 +)
1654 +@pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-R1"])
1655 +@pytest.mark.parametrize("tensor_parallel_size", [1, 4, 8, 16])
1656 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
1657 +@pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])
1658 +@pytest.mark.parametrize(
1659 + ("prefill_backend", "qk_nope_head_dim", "v_head_dim"),
1660 + _prefill_backend_dimension_params(),
1661 +)
1662 +def test_backend_correctness(
1663 + default_vllm_config,
1664 + dist_init,
1665 + workspace_init,
1666 + batch_spec_name: str,
1667 + model: str,
1668 + tensor_parallel_size: int,
1669 + kv_cache_dtype: str,
1670 + q_scale: float,
1671 + k_scale: float,
1672 + prefill_backend: MLAPrefillBackendEnum,
1673 + qk_nope_head_dim: int,
1674 + v_head_dim: int,
1675 +):
1676 + _run_backend_correctness(
1677 + default_vllm_config,
1678 + dist_init,
1679 + workspace_init,
1680 + batch_spec_name,
1681 + model,
1682 + tensor_parallel_size,
1683 + kv_cache_dtype,
1684 + q_scale,
1685 + k_scale,
1686 + prefill_backend,
1687 + qk_nope_head_dim,
1688 + v_head_dim,
1689 + )
1690 +
1691 +
1692 +@pytest.mark.parametrize(
1693 + ("prefill_backend", "qk_nope_head_dim", "v_head_dim"),
1694 + _prefill_backend_dimension_params(),
1695 +)
1696 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
1697 +def test_chunked_context_backend_correctness(
1698 + default_vllm_config,
1699 + dist_init,
1700 + workspace_init,
1701 + prefill_backend: MLAPrefillBackendEnum,
1702 + qk_nope_head_dim: int,
1703 + v_head_dim: int,
1704 + kv_cache_dtype: str,
1705 +):
1706 + """Split, packed, and context-free requests match the SDPA reference."""
1707 + _run_backend_correctness(
1708 + default_vllm_config,
1709 + dist_init,
1710 + workspace_init,
1711 + batch_spec_name="chunked_context_prefill",
1712 + model="deepseek-ai/DeepSeek-R1",
1713 + tensor_parallel_size=16,
1714 + kv_cache_dtype=kv_cache_dtype,
1715 + q_scale=1.0,
1716 + k_scale=1.0,
1717 + prefill_backend=prefill_backend,
1718 + qk_nope_head_dim=qk_nope_head_dim,
1719 + v_head_dim=v_head_dim,
1720 + chunked_prefill_workspace_size=1024,
1721 + )
tests/v1/attention/test_mla_context_chunks.py new
+288
@@ -0,0 +1,288 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Per-request scheduling of MLA chunked-context prefill.
4 +
5 +These cover the metadata contract `_compute_prefill_context` relies on: every
6 +context row is gathered exactly once, no chunk exceeds the workspace, and a
7 +chunk never covers a prefill without context (which is why the partial no
8 +longer needs an empty-span masking pass).
9 +"""
10 +
11 +import pytest
12 +import torch
13 +
14 +from vllm.model_executor.layers.attention.mla_attention import (
15 + build_mla_chunked_context_metadata,
16 + init_mla_context_partial,
17 + reorg_kvcache,
18 +)
19 +
20 +BLOCK_SIZE = 16
21 +
22 +
23 +def build_chunked_context(
24 + context_lens: list[int],
25 + query_lens: list[int],
26 + workspace_size: int,
27 + block_size: int = BLOCK_SIZE,
28 + dcp_world_size: int = 1,
29 + dcp_local_block_size: int = 1,
30 +):
31 + query_start_loc = torch.zeros(len(query_lens) + 1, dtype=torch.int32)
32 + query_start_loc[1:] = torch.tensor(query_lens, dtype=torch.int32).cumsum(0)
33 + workspace_rows = workspace_size + workspace_size // dcp_world_size
34 + return build_mla_chunked_context_metadata(
35 + context_lens_cpu=torch.tensor(context_lens, dtype=torch.int32),
36 + prefill_query_start_loc_cpu=query_start_loc,
37 + chunked_prefill_workspace=torch.empty((workspace_rows, 1)),
38 + chunked_prefill_workspace_size=workspace_size,
39 + block_size=block_size,
40 + align_chunk_to_block=True,
41 + device=torch.device("cpu"),
42 + dcp_world_size=dcp_world_size,
43 + dcp_local_block_size=dcp_local_block_size,
44 + dcp_virtual_block_size=dcp_local_block_size * dcp_world_size,
45 + )
46 +
47 +
48 +@pytest.mark.parametrize(
49 + "context_lens,workspace_size",
50 + [
51 + # Heterogeneous batch: the long request must not shrink the chunk the
52 + # short ones share.
53 + ([960, 16, 16, 16, 16, 16, 16, 16], 1024),
54 + # A single request that does not fit is split on its own.
55 + ([2048, 320], 1024),
56 + # Every request needs its own chunk.
57 + ([1024, 1024, 1024], 1024),
58 + # Unaligned tails.
59 + ([37, 1000, 5], 1024),
60 + ([1], 1024),
61 + ],
62 +)
63 +def test_chunks_gather_every_context_row_exactly_once(context_lens, workspace_size):
64 + """Chunks tile the batch's context without gaps, overlap, or overflow.
65 +
66 + This is the invariant the whole schedule rests on: attention over a chunk
67 + only sees the rows that chunk gathered, so a row counted twice or dropped
68 + silently corrupts the context partial.
69 + """
70 + query_lens = [4] * len(context_lens)
71 + metadata = build_chunked_context(context_lens, query_lens, workspace_size)
72 + assert metadata is not None
73 +
74 + gathered: dict[int, list[tuple[int, int]]] = {}
75 + previous_request_start = -1
76 + for chunk in metadata.chunks:
77 + assert chunk.num_context_tokens <= workspace_size
78 + assert chunk.num_context_tokens == int(chunk.cu_seq_lens[-1])
79 + assert chunk.token_to_seq.shape[0] == chunk.num_context_tokens
80 + # Chunks are emitted in request order, which is what lets the partial
81 + # treat only a chunk's first request as a continuation.
82 + assert chunk.request_slice.start >= previous_request_start
83 + previous_request_start = chunk.request_slice.start
84 +
85 + starts = chunk.starts.tolist()
86 + seq_lens = chunk.seq_lens.tolist()
87 + assert len(starts) == len(seq_lens) == chunk.num_requests
88 + for offset, (start, length) in enumerate(zip(starts, seq_lens)):
89 + assert length > 0, "a chunk must not cover an empty context span"
90 + gathered.setdefault(chunk.request_slice.start + offset, []).append(
91 + (start, length)
92 + )
93 +
94 + expected = {i: length for i, length in enumerate(context_lens) if length > 0}
95 + assert gathered.keys() == expected.keys()
96 + for request, spans in gathered.items():
97 + cursor = 0
98 + for start, length in spans:
99 + assert start == cursor
100 + cursor += length
101 + assert cursor == expected[request]
102 +
103 +
104 +def test_continuation_is_confined_to_a_chunks_first_request():
105 + """Only a chunk's first request may continue an earlier chunk.
106 +
107 + `accumulate_mla_context_chunk` merges the continuation token slice and
108 + writes the rest, so a chunk that continues a request other than its first,
109 + or reports the wrong boundary, folds the partial into the wrong tokens.
110 + """
111 + context_lens = [2048, 32, 32]
112 + query_lens = [3, 5, 7]
113 + metadata = build_chunked_context(context_lens, query_lens, 1024)
114 + assert metadata is not None
115 + query_start_loc = [0, 3, 8, 15]
116 +
117 + assert [chunk.is_continuation for chunk in metadata.chunks] == [
118 + False,
119 + True,
120 + False,
121 + ]
122 + for chunk in metadata.chunks:
123 + starts = chunk.starts.tolist()
124 + assert chunk.is_continuation == (starts[0] > 0)
125 + # Requests after the first always start at the beginning of their
126 + # context, so they can only ever be initialized, never merged.
127 + assert all(start == 0 for start in starts[1:])
128 + request_start = chunk.request_slice.start
129 + request_end = chunk.request_slice.stop
130 + assert chunk.continuation_token_end == query_start_loc[request_start + 1]
131 + assert chunk.token_slice == slice(
132 + query_start_loc[request_start], query_start_loc[request_end]
133 + )
134 + assert chunk.query_start_loc.tolist() == [
135 + offset - chunk.token_slice.start
136 + for offset in query_start_loc[request_start : request_end + 1]
137 + ]
138 +
139 +
140 +def test_tail_splitting_minimizes_chunks():
141 + """A tail request may fill one chunk and continue at the next chunk's head.
142 +
143 + Without tail splitting these contexts need three chunks. Splitting on the
144 + block boundary reduces that to the workspace lower bound of two.
145 + """
146 + metadata = build_chunked_context([768, 512, 512], [3, 5, 7], 1024)
147 + assert metadata is not None
148 +
149 + assert len(metadata.chunks) == 2
150 + first, second = metadata.chunks
151 + assert first.starts.tolist() == [0, 0]
152 + assert first.seq_lens.tolist() == [768, 256]
153 + assert not first.is_continuation
154 +
155 + assert second.starts.tolist() == [256, 0]
156 + assert second.seq_lens.tolist() == [256, 512]
157 + assert second.is_continuation
158 +
159 +
160 +def test_prefills_without_context_are_skipped_and_neutralized():
161 + """Context-free prefills get neutral partials instead of chunk work.
162 +
163 + If the full-query partial leaves either an internal or trailing gap
164 + uninitialized, the final merge combines the suffix with undefined scratch.
165 + """
166 + metadata = build_chunked_context([64, 0, 64, 0], [4, 4, 4, 4], 1024)
167 + assert metadata is not None
168 +
169 + covered = {
170 + request
171 + for chunk in metadata.chunks
172 + for request in range(chunk.request_slice.start, chunk.request_slice.stop)
173 + }
174 + assert covered == {0, 2}
175 + assert metadata.empty_token_slices == [slice(4, 8), slice(12, 16)]
176 +
177 + output, output_lse = init_mla_context_partial(
178 + metadata,
179 + attn_output=torch.empty(4, 2, 3),
180 + attn_softmax_lse=torch.empty(2, 4),
181 + num_tokens=16,
182 + )
183 + assert output.shape == (16, 2, 3)
184 + assert output_lse.shape == (2, 16)
185 + for token_slice in metadata.empty_token_slices:
186 + assert not torch.count_nonzero(output[token_slice])
187 + assert torch.isneginf(output_lse[:, token_slice]).all()
188 +
189 +
190 +def test_no_context_needs_no_chunks():
191 + assert build_chunked_context([0, 0], [4, 4], 1024) is None
192 +
193 +
194 +def test_dcp_chunks_fit_the_per_rank_row_budget():
195 + """Under DCP each rank gathers its own shard, so the budget is 1/world.
196 +
197 + The local starts and lengths must tile the rank's padded context the same
198 + way the global ones tile the real context, since `reorg_kvcache` unpads
199 + using them.
200 + """
201 + dcp_world_size, interleave, workspace_size = 2, 64, 1024
202 + context_lens = [3000, 200, 200]
203 + metadata = build_chunked_context(
204 + context_lens,
205 + [4] * len(context_lens),
206 + workspace_size,
207 + block_size=128,
208 + dcp_world_size=dcp_world_size,
209 + dcp_local_block_size=interleave,
210 + )
211 + assert metadata is not None
212 + virtual_block_size = interleave * dcp_world_size
213 +
214 + local_cursor: dict[int, int] = {}
215 + for chunk in metadata.chunks:
216 + assert chunk.num_local_context_tokens <= workspace_size // dcp_world_size
217 + assert chunk.local_starts == chunk.starts.tolist()
218 + assert chunk.padded_local_cu_seq_lens.tolist() == [0] + list(
219 + torch.tensor(chunk.padded_local_seq_lens).cumsum(0)
220 + )
221 + assert (
222 + chunk.padded_local_token_to_seq.shape[0] == chunk.num_local_context_tokens
223 + )
224 + for offset, (start, length) in enumerate(
225 + zip(chunk.local_starts, chunk.padded_local_seq_lens)
226 + ):
227 + request = chunk.request_slice.start + offset
228 + assert start == local_cursor.get(request, 0)
229 + local_cursor[request] = start + length
230 +
231 + for request, length in enumerate(context_lens):
232 + padded_local = -(-length // virtual_block_size) * interleave
233 + assert local_cursor[request] == padded_local
234 +
235 +
236 +def test_dcp_reorg_uses_each_chunks_local_starts():
237 + """Reorganization drops DCP padding at each request's own chunk offset."""
238 + dcp_world_size, interleave, workspace_size = 2, 64, 1024
239 + metadata = build_chunked_context(
240 + [3000, 200, 200],
241 + [4, 4, 4],
242 + workspace_size,
243 + block_size=128,
244 + dcp_world_size=dcp_world_size,
245 + dcp_local_block_size=interleave,
246 + )
247 + assert metadata is not None
248 +
249 + for chunk in metadata.chunks:
250 + assert chunk.padded_local_seq_lens is not None
251 + assert chunk.local_context_lens_allranks is not None
252 + assert chunk.local_starts is not None
253 +
254 + toks = chunk.num_local_context_tokens
255 + rank_buffers = [torch.full((toks, 1, 1), -1) for _ in range(dcp_world_size)]
256 + expected = []
257 + src_token_idx = 0
258 + for request, (padded_len, local_lens, local_start) in enumerate(
259 + zip(
260 + chunk.padded_local_seq_lens,
261 + chunk.local_context_lens_allranks,
262 + chunk.local_starts,
263 + )
264 + ):
265 + for rank, local_len in enumerate(local_lens):
266 + actual_len = min(max(0, local_len - local_start), padded_len)
267 + values = (
268 + request * 100_000
269 + + rank * 10_000
270 + + torch.arange(local_start, local_start + actual_len)
271 + ).view(-1, 1, 1)
272 + rank_buffers[rank][src_token_idx : src_token_idx + actual_len] = values
273 + expected.append(values)
274 + src_token_idx += padded_len
275 +
276 + allgathered = torch.cat(rank_buffers)
277 + reorganized, _ = reorg_kvcache(
278 + allgathered,
279 + allgathered,
280 + padded_local_chunk_seq_lens_lst=chunk.padded_local_seq_lens,
281 + local_context_lens_allranks=chunk.local_context_lens_allranks,
282 + local_starts=chunk.local_starts,
283 + sum_seq_len=chunk.num_context_tokens,
284 + max_seq_len=chunk.max_seq_len,
285 + toks=toks,
286 + )
287 +
288 + torch.testing.assert_close(reorganized, torch.cat(expected))
tests/v1/attention/test_mla_prefill_registry.py
+2 -2
@@ -24,7 +24,7 @@ class CustomMLAPrefillBackend(MLAPrefillBackend):
24 def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
25 raise NotImplementedError
26
27 - def run_prefill_context_chunk(self, chunk_idx, q, k, v):
27 + def run_prefill_context_chunk(self, chunk, q, k, v):
28 raise NotImplementedError
29
30
@@ -112,7 +112,7 @@ def test_register_custom_backend_as_decorator():
112 def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
113 raise NotImplementedError
114
115 - def run_prefill_context_chunk(self, chunk_idx, q, k, v):
115 + def run_prefill_context_chunk(self, chunk, q, k, v):
116 raise NotImplementedError
117
118 assert MLAPrefillBackendEnum.CUSTOM.is_overridden()
tests/v1/attention/test_sparse_mla_backends.py
+3 -3
@@ -875,7 +875,7 @@ def test_masked_mha_workspace_fits_single_request_boundary(max_query_len, expect
875 _masked_mha_workspace_fits(
876 batch_size=1,
877 max_query_len=max_query_len,
878 - context_chunk_max_seq_lens=None,
878 + max_context_chunk_seq_len=0,
879 workspace_numel=GLOBAL_TOPK_MASK_MAX_BYTES // torch.int32.itemsize,
880 )
881 is expected
@@ -884,7 +884,7 @@ def test_masked_mha_workspace_fits_single_request_boundary(max_query_len, expect
884
885 def test_masked_mha_workspace_fits_accounts_for_batch_and_context():
886 """Request count and context chunk length are independent multipliers."""
887 - base = dict(batch_size=2, max_query_len=2048, context_chunk_max_seq_lens=[2048])
887 + base = dict(batch_size=2, max_query_len=2048, max_context_chunk_seq_len=2048)
888 exact = math.prod(_topk_mask_shape(2, 2048, 2048))
889
890 assert _masked_mha_workspace_fits(**base, workspace_numel=exact)
@@ -892,7 +892,7 @@ def test_masked_mha_workspace_fits_accounts_for_batch_and_context():
892 **{**base, "batch_size": 3}, workspace_numel=exact
893 )
894 assert not _masked_mha_workspace_fits(
895 - **{**base, "context_chunk_max_seq_lens": [4096]}, workspace_numel=exact
895 + **{**base, "max_context_chunk_seq_len": 4096}, workspace_numel=exact
896 )
897
898
vllm/model_executor/layers/attention/mla_attention.py
+454 -272
@@ -133,9 +133,19 @@ fixed workspace size.
133
134 The chunked prefill approach is as follows:
135
136 -MCC Max chunk of context to process per iter, computed dynamically,
136 +W Workspace rows, i.e. the context rows we may gather at once,
137 used to bound the memory usage
138
139 +The context is scheduled per request: requests are packed in order, splitting
140 +the next request on an aligned boundary when needed to fill `W`.
141 +So a chunk covers a contiguous run of prefills and only ever the tokens and
142 +context rows of those prefills — attention, up-projection and merging are
143 +charged only to the requests it covers, and no chunk contains an empty context
144 +span. See `plan_mla_context_chunks`.
145 +
146 +Only a chunk's first request can be a continuation, so accumulation performs at
147 +most one request-slice merge and one bulk write per chunk.
148 +
149 q_c = h_t @ W_DQ
150 q_nope = (q_c @ W_UQ).view(Sq, N, P)
151 q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
@@ -157,10 +167,12 @@ curr_o, curr_lse = scaled_dot_product_attention(
167 return_softmax_lse=True
168 )
169
160 -// Compute attention with the already existing context
161 -for chunk_idx in range(cdiv(C, MCC)):
162 - chunk_start = chunk_idx * MCC
163 - chunk_end = min(chunk_start + MCC, C)
170 +// Compute attention with the already existing context. Shown for a single
171 +// request; with a batch, a chunk covers a run of requests and both `q` and the
172 +// gathered context below are sliced to that run.
173 +for chunk_idx in range(cdiv(C, W)):
174 + chunk_start = chunk_idx * W
175 + chunk_end = min(chunk_start + W, C)
176 Sc = chunk_end - chunk_start
177 cache_kv_c_chunk = cache_kv_c[chunk_start:chunk_end]
178 cache_k_pe_chunk = cache_k_pe[chunk_start:chunk_end]
@@ -188,11 +200,15 @@ return curr_o @ W_O
200 """
201
202 import functools
203 +import itertools
204 +import math
205 from abc import abstractmethod
206 +from collections.abc import Callable
207 from dataclasses import dataclass
208 from enum import Enum
209 from typing import ClassVar, Generic, TypeVar, cast
210
211 +import numpy as np
212 import torch
213 import torch.nn as nn
214 from tqdm import tqdm
@@ -256,6 +272,7 @@ from vllm.utils.torch_utils import (
272 direct_register_custom_op,
273 is_quantized_kv_cache,
274 kv_cache_dtype_str_to_dtype,
275 + np_to_pinned_tensor,
276 )
277 from vllm.v1.attention.backend import (
278 AttentionBackend,
@@ -278,7 +295,6 @@ from vllm.v1.attention.backends.utils import (
295 from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs
296 from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
297 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
281 -from vllm.v1.attention.ops.triton_merge_attn_states import mask_empty_context
298 from vllm.v1.attention.selector import get_attn_backend
299 from vllm.v1.kv_cache_interface import (
300 AttentionSpec,
@@ -1375,28 +1391,42 @@ class MLACommonPrefillMetadata:
1391 """Prefill Specific Metadata"""
1392
1393 @dataclass
1378 - class ChunkedContextMetadata:
1379 - # New for MLA (compared to FlashAttention)
1380 - # For handling chunked prefill
1394 + class ContextChunk:
1395 + """One workspace-sized slice of paged context for a run of prefills."""
1396 +
1397 + index: int
1398 + request_slice: slice
1399 + token_slice: slice
1400 + continuation_token_end: int
1401 + is_continuation: bool
1402 + num_context_tokens: int
1403 + query_start_loc: torch.Tensor
1404 + max_query_len: int
1405 cu_seq_lens: torch.Tensor
1406 starts: torch.Tensor
1383 - seq_tot: list[int]
1384 - max_seq_lens: list[int]
1407 + max_seq_len: int
1408 seq_lens: torch.Tensor
1386 - context_lens: torch.Tensor
1387 - workspace: torch.Tensor
1409 token_to_seq: torch.Tensor
1389 - chunk_total_token: list[int]
1390 - has_empty_context: list[bool]
1410
1411 # for mla DCP
1393 - padded_local_chunk_seq_lens: list[list[int]] | None = None
1412 + padded_local_seq_lens: list[int] | None = None
1413 local_context_lens_allranks: list[list[int]] | None = None
1414 padded_local_cu_seq_lens: torch.Tensor | None = None
1415 padded_local_token_to_seq: torch.Tensor | None = None
1397 - cu_seq_lens_lst: list[list[int]] | None = None
1398 - chunk_size: int | None = None
1399 - prefill_tokens_with_context: int | None = None
1416 + num_local_context_tokens: int = 0
1417 + local_starts: list[int] | None = None
1418 +
1419 + @property
1420 + def num_requests(self) -> int:
1421 + return self.request_slice.stop - self.request_slice.start
1422 +
1423 + @dataclass
1424 + class ChunkedContextMetadata:
1425 + context_lens: torch.Tensor
1426 + workspace: torch.Tensor
1427 + chunks: "list[MLACommonPrefillMetadata.ContextChunk]"
1428 + context_lens_list: list[int]
1429 + empty_token_slices: list[slice]
1430
1431 block_table: torch.Tensor
1432 query_start_loc: torch.Tensor
@@ -1566,11 +1596,98 @@ def backend_supports_prefill_query_quantization() -> bool:
1596 )
1597
1598
1599 +@dataclass
1600 +class _ContextChunkPlan:
1601 + """Request-space layout of one context chunk, before tensors are built."""
1602 +
1603 + request_start: int
1604 + request_end: int
1605 + # Per-request context offset and row count, in request order.
1606 + starts: list[int]
1607 + seq_lens: list[int]
1608 + is_continuation: bool
1609 +
1610 +
1611 +def plan_mla_context_chunks(
1612 + context_lens: list[int],
1613 + row_budget: int,
1614 + max_context_chunk: int,
1615 + split_alignment: int,
1616 + padded_rows: Callable[[int], int],
1617 +) -> list[_ContextChunkPlan]:
1618 + """Pack per-request contexts into workspace-sized chunks."""
1619 + assert max_context_chunk > 0
1620 + assert split_alignment > 0
1621 +
1622 + def aligned_split_len(start: int, remaining: int, available: int) -> int:
1623 + max_split = min(
1624 + max_context_chunk,
1625 + round_down(remaining - 1, split_alignment),
1626 + )
1627 + low, high = 0, max_split // split_alignment
1628 + while low < high:
1629 + mid = (low + high + 1) // 2
1630 + length = mid * split_alignment
1631 + rows = padded_rows(start + length) - padded_rows(start)
1632 + if rows <= available:
1633 + low = mid
1634 + else:
1635 + high = mid - 1
1636 + return low * split_alignment
1637 +
1638 + plans: list[_ContextChunkPlan] = []
1639 + num_requests = len(context_lens)
1640 + request = 0
1641 + start = 0
1642 + while request < num_requests:
1643 + context_len = context_lens[request]
1644 + if context_len == 0:
1645 + assert start == 0
1646 + request += 1
1647 + continue
1648 +
1649 + request_start = request
1650 + starts: list[int] = []
1651 + seq_lens: list[int] = []
1652 + rows = 0
1653 + while request < num_requests and context_lens[request] > 0:
1654 + remaining = context_lens[request] - start
1655 + request_rows = padded_rows(start + remaining) - padded_rows(start)
1656 + if rows + request_rows > row_budget:
1657 + split_len = aligned_split_len(start, remaining, row_budget - rows)
1658 + if split_len > 0:
1659 + starts.append(start)
1660 + seq_lens.append(split_len)
1661 + rows += padded_rows(start + split_len) - padded_rows(start)
1662 + start += split_len
1663 + break
1664 + rows += request_rows
1665 + starts.append(start)
1666 + seq_lens.append(remaining)
1667 + request += 1
1668 + start = 0
1669 + assert seq_lens
1670 + plans.append(
1671 + _ContextChunkPlan(
1672 + request_start=request_start,
1673 + request_end=request_start + len(seq_lens),
1674 + starts=starts,
1675 + seq_lens=seq_lens,
1676 + is_continuation=starts[0] > 0,
1677 + )
1678 + )
1679 + return plans
1680 +
1681 +
1682 +def _flat_int32(values: list[int] | np.ndarray) -> torch.Tensor:
1683 + """Pinned int32 CPU tensor backing one concatenated per-chunk field."""
1684 + return np_to_pinned_tensor(np.asarray(values, dtype=np.int32))
1685 +
1686 +
1687 def build_mla_chunked_context_metadata(
1688 *,
1689 context_lens_cpu: torch.Tensor,
1690 prefill_query_start_loc_cpu: torch.Tensor,
1573 - num_prefills: int,
1691 chunked_prefill_workspace: torch.Tensor,
1692 chunked_prefill_workspace_size: int,
1693 block_size: int,
@@ -1582,14 +1699,13 @@ def build_mla_chunked_context_metadata(
1699 ) -> "MLACommonPrefillMetadata.ChunkedContextMetadata | None":
1700 """Build chunked-context metadata for an MLA prefill.
1701
1585 - Shared by dense and sparse builders. Splits each prefill's context
1586 - into workspace-sized chunks and, under DCP, plans the per-rank interleaved
1587 - local chunks the all-gather reduction consumes.
1702 + Shared by dense and sparse builders. Packs the prefill contexts into a flat
1703 + list of workspace-sized per-request chunks and, under DCP, plans the
1704 + per-rank interleaved local chunks the all-gather reduction consumes.
1705
1706 Args:
1707 context_lens_cpu: Per-prefill context length (seq_len - query_len).
1708 prefill_query_start_loc_cpu: Prefill query cumulative offsets (0-based).
1592 - num_prefills: Number of prefill requests.
1709 chunked_prefill_workspace: Scratch buffer the context gather writes to.
1710 chunked_prefill_workspace_size: Row capacity of the workspace.
1711 block_size: KV cache page size for chunk-start alignment.
@@ -1604,147 +1720,189 @@ def build_mla_chunked_context_metadata(
1720 """
1721 # NOTE: it is recommended you read the `Chunked Prefill` section in the
1722 # comment at the top of the file before trying to understand this code.
1607 - max_context_len = context_lens_cpu.max().item()
1608 - if max_context_len <= 0:
1723 + context_lens = context_lens_cpu.tolist()
1724 + if max(context_lens, default=0) <= 0:
1725 return None
1610 - num_prefills_with_context = int((context_lens_cpu > 0).sum().item())
1611 -
1612 - # Currently we allocate an equal amount of workspace for each prefill with
1613 - # context; we could probably use a more advanced algorithm here and allocate
1614 - # more workspace to prefills with longer context lengths.
1615 - max_context_chunk = chunked_prefill_workspace_size // num_prefills_with_context
1616 - if align_chunk_to_block:
1617 - # The `gather_and_maybe_dequant_cache` kernel cannot handle chunk
1618 - # starts that are not aligned to block_size, so round down.
1619 - max_context_chunk = round_down(max_context_chunk, block_size)
1620 - assert max_context_chunk > 0
1726
1622 - num_chunks = cdiv(max_context_len, max_context_chunk)
1623 - # e.g. max_context_chunk=256, num_chunks=3, num_prefills=4 ->
1624 - # [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]
1625 - # Note(simon): this is done on CPU because of downstream's use of `to_list`.
1626 - chunk_starts = torch.empty(
1627 - num_chunks, num_prefills, dtype=torch.int32, pin_memory=True
1628 - ).copy_(
1629 - torch.arange(num_chunks, dtype=torch.int32)
1630 - .multiply_(max_context_chunk)
1631 - .unsqueeze(1)
1632 - )
1633 - chunk_ends = torch.min(
1634 - context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk
1635 - )
1636 - chunk_seq_lens = chunk_ends - chunk_starts
1637 - chunk_seq_lens.clamp_(min=0)
1638 - has_empty_context = torch.any(chunk_seq_lens == 0, dim=1).tolist()
1727 + # The `gather_and_maybe_dequant_cache` kernel cannot handle chunk starts
1728 + # that are not aligned to block_size, so a split request advances in
1729 + # block-aligned steps.
1730 + chunk_alignment = block_size if align_chunk_to_block else 1
1731 + if dcp_world_size > 1:
1732 + # Each rank gathers only its own shard, so a chunk's workspace cost is
1733 + # 1/dcp_world_size of its context rows, rounded up to the interleave
1734 + # block size. Sub-chunks are additionally aligned to the virtual block
1735 + # size so that every rank's local start stays interleave-aligned.
1736 + row_budget = chunked_prefill_workspace_size // dcp_world_size
1737 + chunk_alignment = math.lcm(chunk_alignment, dcp_virtual_block_size)
1738 +
1739 + def padded_rows(rows: int) -> int:
1740 + return cdiv(rows, dcp_virtual_block_size) * dcp_local_block_size
1741 + else:
1742 + row_budget = chunked_prefill_workspace_size
1743 +
1744 + def padded_rows(rows: int) -> int:
1745 + return rows
1746
1640 - cu_seq_lens_cpu = torch.zeros(
1641 - num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
1747 + max_context_chunk = round_down(chunked_prefill_workspace_size, chunk_alignment)
1748 + assert max_context_chunk > 0, (
1749 + f"chunked prefill workspace ({chunked_prefill_workspace_size} rows) is "
1750 + f"smaller than the context chunk alignment ({chunk_alignment} rows)"
1751 )
1643 - torch.cumsum(chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32)
1644 - chunk_total_token = cu_seq_lens_cpu[:, -1]
1752
1646 - max_tokens_over_chunk = chunk_total_token.max().item()
1647 - token_to_seq_cpu = torch.zeros(
1648 - (num_chunks, max_tokens_over_chunk), dtype=torch.int32, pin_memory=True
1753 + plans = plan_mla_context_chunks(
1754 + context_lens,
1755 + row_budget,
1756 + max_context_chunk,
1757 + chunk_alignment,
1758 + padded_rows,
1759 )
1650 - req_indices = torch.arange(num_prefills, dtype=torch.int32)
1651 - for i in range(num_chunks):
1652 - token_to_seq = torch.repeat_interleave(req_indices, chunk_seq_lens[i])
1653 - token_to_seq_cpu[i, : token_to_seq.shape[0]] = token_to_seq
1760
1655 - prefill_tokens_with_context = prefill_query_start_loc_cpu[
1656 - num_prefills_with_context
1657 - ].item()
1761 + query_start_loc = prefill_query_start_loc_cpu.tolist()
1762 + empty_token_slices = [
1763 + slice(query_start_loc[i], query_start_loc[i + 1])
1764 + for i, context_len in enumerate(context_lens)
1765 + if context_len == 0
1766 + ]
1767
1659 - metadata_cls = MLACommonPrefillMetadata.ChunkedContextMetadata
1660 - if dcp_world_size > 1:
1661 - local_context_lens_allranks = get_dcp_local_seq_lens(
1768 + use_dcp = dcp_world_size > 1
1769 + local_context_lens_allranks = (
1770 + get_dcp_local_seq_lens(
1771 context_lens_cpu, dcp_world_size, None, dcp_local_block_size
1772 + ).tolist()
1773 + if use_dcp
1774 + else None
1775 + )
1776 +
1777 + # Concatenate the per-chunk fields so each one costs a single
1778 + # host-to-device transfer, and remember where every chunk's slice lands.
1779 + starts_flat: list[int] = []
1780 + seq_lens_flat: list[int] = []
1781 + cu_seq_lens_flat: list[int] = []
1782 + cu_seqlens_q_flat: list[int] = []
1783 + token_to_seq_parts: list[np.ndarray] = []
1784 + padded_local_cu_seq_lens_flat: list[int] = []
1785 + padded_local_token_to_seq_parts: list[np.ndarray] = []
1786 + local_starts_per_chunk: list[list[int]] = []
1787 + local_seq_lens_per_chunk: list[list[int]] = []
1788 + layouts: list[tuple[slice, slice, slice, slice]] = []
1789 +
1790 + request_offset = boundary_offset = token_offset = local_token_offset = 0
1791 + for plan in plans:
1792 + num_requests = len(plan.seq_lens)
1793 + num_tokens = sum(plan.seq_lens)
1794 + num_local_tokens = num_tokens
1795 +
1796 + seq_lens_flat.extend(plan.seq_lens)
1797 + cu_seq_lens_flat.extend(itertools.accumulate(plan.seq_lens, initial=0))
1798 + query_base = query_start_loc[plan.request_start]
1799 + cu_seqlens_q_flat.extend(
1800 + query_start_loc[request] - query_base
1801 + for request in range(plan.request_start, plan.request_end + 1)
1802 )
1664 - # Note(qcs): Per-rank local context lengths, padded to
1665 - # `dcp_local_block_size`.
1666 - padded_local_context_lens_cpu: torch.Tensor = (
1667 - cdiv(context_lens_cpu, dcp_virtual_block_size) * dcp_local_block_size
1668 - )
1669 - # Note(hc): The above max_context_chunk already enforces block_size
1670 - # alignment; DCP only requires block_size be divisible by dcp_world_size,
1671 - # because DCP uses cp_gather_cache, which does not require chunk starts
1672 - # aligned to block_size.
1673 - assert max_context_chunk % dcp_world_size == 0
1674 - padded_local_max_context_chunk = (
1675 - cdiv(max_context_chunk, dcp_virtual_block_size) * dcp_local_block_size
1676 - )
1677 - local_chunk_starts = torch.empty(
1678 - num_chunks, num_prefills, dtype=torch.int32, pin_memory=True
1679 - ).copy_(
1680 - torch.arange(num_chunks, dtype=torch.int32)
1681 - .multiply_(padded_local_max_context_chunk)
1682 - .unsqueeze(1)
1683 - )
1684 - local_chunk_ends = torch.min(
1685 - padded_local_context_lens_cpu.unsqueeze(0),
1686 - local_chunk_starts + padded_local_max_context_chunk,
1803 + token_to_seq_parts.append(
1804 + np.repeat(np.arange(num_requests, dtype=np.int32), plan.seq_lens)
1805 )
1688 - padded_local_chunk_seq_lens = local_chunk_ends - local_chunk_starts
1689 - padded_local_chunk_seq_lens.clamp_(min=0)
1806
1691 - padded_local_cu_seq_lens_cpu = torch.zeros(
1692 - num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
1693 - )
1694 - torch.cumsum(
1695 - padded_local_chunk_seq_lens,
1696 - dim=1,
1697 - out=padded_local_cu_seq_lens_cpu[:, 1:],
1698 - dtype=torch.int32,
1807 + if use_dcp:
1808 + # A request's local rows are its context rows sharded across ranks
1809 + # and rounded up to the interleave block size, so the local start of
1810 + # a continuation is the padded row count of the context before it.
1811 + local_starts = [padded_rows(start) for start in plan.starts]
1812 + local_seq_lens = [
1813 + padded_rows(start + length) - local_start
1814 + for start, length, local_start in zip(
1815 + plan.starts, plan.seq_lens, local_starts
1816 + )
1817 + ]
1818 + local_starts_per_chunk.append(local_starts)
1819 + local_seq_lens_per_chunk.append(local_seq_lens)
1820 + padded_local_cu_seq_lens_flat.extend(
1821 + itertools.accumulate(local_seq_lens, initial=0)
1822 + )
1823 + padded_local_token_to_seq_parts.append(
1824 + np.repeat(np.arange(num_requests, dtype=np.int32), local_seq_lens)
1825 + )
1826 + num_local_tokens = sum(local_seq_lens)
1827 + # The gather takes per-rank local offsets under DCP.
1828 + starts_flat.extend(local_starts)
1829 + else:
1830 + starts_flat.extend(plan.starts)
1831 + assert num_local_tokens <= row_budget
1832 +
1833 + layouts.append(
1834 + (
1835 + slice(request_offset, request_offset + num_requests),
1836 + slice(boundary_offset, boundary_offset + num_requests + 1),
1837 + slice(token_offset, token_offset + num_tokens),
1838 + slice(local_token_offset, local_token_offset + num_local_tokens),
1839 + )
1840 )
1700 - max_padded_local_tokens = padded_local_cu_seq_lens_cpu[:, -1].max().item()
1701 - padded_local_token_to_seq_cpu = torch.zeros(
1702 - (num_chunks, max_padded_local_tokens), dtype=torch.int32
1841 + request_offset += num_requests
1842 + boundary_offset += num_requests + 1
1843 + token_offset += num_tokens
1844 + local_token_offset += num_local_tokens
1845 +
1846 + seq_lens_cpu = _flat_int32(seq_lens_flat)
1847 + starts = _flat_int32(starts_flat).to(device, non_blocking=True)
1848 + cu_seq_lens = _flat_int32(cu_seq_lens_flat).to(device, non_blocking=True)
1849 + cu_seqlens_q = _flat_int32(cu_seqlens_q_flat).to(device, non_blocking=True)
1850 + token_to_seq = _flat_int32(np.concatenate(token_to_seq_parts)).to(
1851 + device, non_blocking=True
1852 + )
1853 + if use_dcp:
1854 + padded_local_cu_seq_lens = _flat_int32(padded_local_cu_seq_lens_flat).to(
1855 + device, non_blocking=True
1856 )
1704 - for i in range(num_chunks):
1705 - tts = torch.repeat_interleave(req_indices, padded_local_chunk_seq_lens[i])
1706 - padded_local_token_to_seq_cpu[i, : tts.shape[0]] = tts
1707 -
1708 - chunked_context_metadata = metadata_cls(
1709 - cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),
1710 - starts=local_chunk_starts.to(device, non_blocking=True),
1711 - seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(),
1712 - max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
1713 - seq_lens=chunk_seq_lens,
1714 - context_lens=context_lens_cpu.to(device, non_blocking=True),
1715 - token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
1716 - chunk_total_token=chunk_total_token.tolist(),
1717 - workspace=chunked_prefill_workspace,
1718 - has_empty_context=has_empty_context,
1719 - prefill_tokens_with_context=prefill_tokens_with_context,
1720 - padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),
1721 - local_context_lens_allranks=local_context_lens_allranks.tolist(),
1722 - padded_local_cu_seq_lens=padded_local_cu_seq_lens_cpu.to(
1723 - device, non_blocking=True
1857 + padded_local_token_to_seq = _flat_int32(
1858 + np.concatenate(padded_local_token_to_seq_parts)
1859 + ).to(device, non_blocking=True)
1860 +
1861 + chunks: list[MLACommonPrefillMetadata.ContextChunk] = []
1862 + for index, (plan, layout) in enumerate(zip(plans, layouts)):
1863 + request_slice, boundary_slice, token_slice, local_token_slice = layout
1864 + query_lens = [
1865 + query_start_loc[request + 1] - query_start_loc[request]
1866 + for request in range(plan.request_start, plan.request_end)
1867 + ]
1868 + chunk = MLACommonPrefillMetadata.ContextChunk(
1869 + index=index,
1870 + request_slice=slice(plan.request_start, plan.request_end),
1871 + token_slice=slice(
1872 + query_start_loc[plan.request_start], query_start_loc[plan.request_end]
1873 ),
1725 - padded_local_token_to_seq=padded_local_token_to_seq_cpu.to(
1726 - device, non_blocking=True
1727 - ),
1728 - cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),
1729 - chunk_size=padded_local_max_context_chunk,
1730 - )
1731 - else:
1732 - chunked_context_metadata = metadata_cls(
1733 - cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),
1734 - starts=chunk_starts.to(device, non_blocking=True),
1735 - seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
1736 - max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
1737 - seq_lens=chunk_seq_lens,
1738 - context_lens=context_lens_cpu.to(device, non_blocking=True),
1739 - token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
1740 - chunk_total_token=chunk_total_token,
1741 - workspace=chunked_prefill_workspace,
1742 - has_empty_context=has_empty_context,
1743 - prefill_tokens_with_context=prefill_tokens_with_context,
1874 + continuation_token_end=query_start_loc[plan.request_start + 1],
1875 + is_continuation=plan.is_continuation,
1876 + num_context_tokens=token_slice.stop - token_slice.start,
1877 + query_start_loc=cu_seqlens_q[boundary_slice],
1878 + max_query_len=max(query_lens),
1879 + cu_seq_lens=cu_seq_lens[boundary_slice],
1880 + starts=starts[request_slice],
1881 + max_seq_len=max(plan.seq_lens),
1882 + seq_lens=seq_lens_cpu[request_slice],
1883 + token_to_seq=token_to_seq[token_slice],
1884 + num_local_context_tokens=local_token_slice.stop - local_token_slice.start,
1885 )
1745 -
1746 - assert max(chunked_context_metadata.max_seq_lens) <= chunked_prefill_workspace_size
1747 - return chunked_context_metadata
1886 + if use_dcp:
1887 + assert local_context_lens_allranks is not None
1888 + chunk.padded_local_seq_lens = local_seq_lens_per_chunk[index]
1889 + chunk.local_context_lens_allranks = local_context_lens_allranks[
1890 + plan.request_start : plan.request_end
1891 + ]
1892 + chunk.padded_local_cu_seq_lens = padded_local_cu_seq_lens[boundary_slice]
1893 + chunk.padded_local_token_to_seq = padded_local_token_to_seq[
1894 + local_token_slice
1895 + ]
1896 + chunk.local_starts = local_starts_per_chunk[index]
1897 + chunks.append(chunk)
1898 +
1899 + return MLACommonPrefillMetadata.ChunkedContextMetadata(
1900 + context_lens=context_lens_cpu.to(device, non_blocking=True),
1901 + workspace=chunked_prefill_workspace,
1902 + chunks=chunks,
1903 + context_lens_list=context_lens,
1904 + empty_token_slices=empty_token_slices,
1905 + )
1906
1907
1908 class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
@@ -1796,13 +1954,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1954 64 * 1024,
1955 )
1956
1799 - # Enforce that we enough for at least 1 page per request
1800 - chunked_prefill_workspace_size = max(
1801 - chunked_prefill_workspace_size,
1802 - scheduler_config.max_num_seqs * cache_config.block_size,
1803 - )
1804 -
1805 - return chunked_prefill_workspace_size
1957 + return max(chunked_prefill_workspace_size, cache_config.block_size)
1958
1959 @staticmethod
1960 def determine_prefill_query_data_type(
@@ -2063,7 +2215,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2215 chunked_context_metadata = build_mla_chunked_context_metadata(
2216 context_lens_cpu=context_lens_cpu,
2217 prefill_query_start_loc_cpu=prefill_query_start_loc_cpu,
2066 - num_prefills=num_prefills,
2218 chunked_prefill_workspace=self.chunked_prefill_workspace,
2219 chunked_prefill_workspace_size=self.chunked_prefill_workspace_size,
2220 block_size=self.page_size,
@@ -2138,10 +2289,9 @@ def reorg_kvcache(
2289 allgatered_k_pe: torch.Tensor,
2290 padded_local_chunk_seq_lens_lst: list[int],
2291 local_context_lens_allranks: list[list[int]],
2292 + local_starts: list[int],
2293 sum_seq_len: int,
2294 max_seq_len: int,
2143 - chunk_size: int,
2144 - chunk_idx: int,
2295 toks: int,
2296 ) -> tuple[torch.Tensor, torch.Tensor]:
2297 """
@@ -2155,19 +2305,18 @@ def reorg_kvcache(
2305 padded_local_chunk_seq_lens_lst: local chunk context lengths
2306 under current CP rank.
2307 local_context_lens_allranks: local context lengths on each CP rank.
2308 + local_starts: per-request local offset into the context this chunk
2309 + starts at.
2310 sum_seq_len: the sum of cp_chunk_seq_lens_lst.
2311 max_seq_len: the max value of cp_chunk_seq_lens_lst.
2160 - chunk_size: the local padded max context chunk from
2161 - chunked_context_metadata building.
2162 - chunk_idx: chunk idx of chunked_prefill.
2312 toks: the number of tokens for local gather cache.
2313 """
2314 kv_c_segments = []
2315 k_pe_segments = []
2316 src_token_idx = 0
2317 max_seq_len_check = 0
2169 - for padded_local_chunk_seq_len, local_context_lens in zip(
2170 - padded_local_chunk_seq_lens_lst, local_context_lens_allranks
2318 + for padded_local_chunk_seq_len, local_context_lens, local_start in zip(
2319 + padded_local_chunk_seq_lens_lst, local_context_lens_allranks, local_starts
2320 ):
2321 cur_seq_len = 0
2322 for rank, local_context_len in enumerate(local_context_lens):
@@ -2179,7 +2328,7 @@ def reorg_kvcache(
2328 # local_chunk_len in dcp1: |-----|-----|--|
2329 # so we need update the last chunk length in dcp1.
2330 local_chunk_len = min(
2182 - max(0, local_context_len - chunk_idx * chunk_size),
2331 + max(0, local_context_len - local_start),
2332 padded_local_chunk_seq_len,
2333 )
2334 if local_chunk_len != 0:
@@ -2206,6 +2355,66 @@ def reorg_kvcache(
2355 return reorganized_kv_c_normed, reorganized_k_pe
2356
2357
2358 +def init_mla_context_partial(
2359 + chunked_context: "MLACommonPrefillMetadata.ChunkedContextMetadata",
2360 + attn_output: torch.Tensor,
2361 + attn_softmax_lse: torch.Tensor,
2362 + num_tokens: int,
2363 +) -> tuple[torch.Tensor, torch.Tensor]:
2364 + """Allocate the running context partial over all prefill tokens.
2365 +
2366 + Laid out like the chunk partials so the final whole-batch merge against the
2367 + suffix partial sees matching head strides.
2368 + """
2369 + output = torch.empty(
2370 + (num_tokens, *attn_output.shape[1:]),
2371 + dtype=attn_output.dtype,
2372 + device=attn_output.device,
2373 + )
2374 + output_lse = torch.empty(
2375 + (attn_softmax_lse.shape[0], num_tokens),
2376 + dtype=attn_softmax_lse.dtype,
2377 + device=attn_softmax_lse.device,
2378 + )
2379 + # No chunk covers a prefill without context, so neutralize its partial.
2380 + for token_slice in chunked_context.empty_token_slices:
2381 + output[token_slice].zero_()
2382 + output_lse[:, token_slice].fill_(float("-inf"))
2383 + return output, output_lse
2384 +
2385 +
2386 +def accumulate_mla_context_chunk(
2387 + chunk: "MLACommonPrefillMetadata.ContextChunk",
2388 + attn_output: torch.Tensor,
2389 + attn_softmax_lse: torch.Tensor,
2390 + output: torch.Tensor,
2391 + output_lse: torch.Tensor,
2392 +) -> None:
2393 + """Fold one chunk's partial into the running context partial.
2394 +
2395 + Only the first request may be a continuation; its tokens are merged and the
2396 + remaining token range is initialized.
2397 + """
2398 + token_start = chunk.token_slice.start
2399 + token_end = chunk.token_slice.stop
2400 + init_start = token_start
2401 + if chunk.is_continuation:
2402 + init_start = chunk.continuation_token_end
2403 + num_merged = init_start - token_start
2404 + merge_attn_states(
2405 + output=output[token_start:init_start],
2406 + output_lse=output_lse[:, token_start:init_start],
2407 + prefix_output=output[token_start:init_start],
2408 + prefix_lse=output_lse[:, token_start:init_start],
2409 + suffix_output=attn_output[:num_merged],
2410 + suffix_lse=attn_softmax_lse[:, :num_merged],
2411 + )
2412 + if init_start < token_end:
2413 + written = init_start - token_start
2414 + output[init_start:token_end].copy_(attn_output[written:])
2415 + output_lse[:, init_start:token_end].copy_(attn_softmax_lse[:, written:])
2416 +
2417 +
2418 class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2419 """
2420 Shared MLA base providing dense-MHA prefill (via the selected
@@ -2282,53 +2491,53 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2491 assert attn_metadata.prefill is not None
2492 prefill_metadata = attn_metadata.prefill
2493 assert prefill_metadata.prefill_backend is not None
2285 - assert prefill_metadata.chunked_context is not None
2494 + chunked_context = prefill_metadata.chunked_context
2495 + assert chunked_context is not None
2496
2497 use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
2498 kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(
2499 self.kv_b_proj, use_fp8_prefill
2500 )
2291 -
2292 - output = None
2293 - merge_output = None
2294 - iters = len(prefill_metadata.chunked_context.seq_tot)
2295 - workspace = prefill_metadata.chunked_context.workspace
2501 + workspace = chunked_context.workspace
2502
2503 if use_fp8_prefill:
2504 q = q.to(prefill_metadata.q_data_type)
2505
2300 - for i in range(iters):
2301 - toks = prefill_metadata.chunked_context.seq_tot[i]
2506 + output = None
2507 + output_lse = None
2508 + for chunk in chunked_context.chunks:
2509 + toks = chunk.num_context_tokens
2510 + block_table = prefill_metadata.block_table[chunk.request_slice]
2511 if self.kv_cache_dtype == "fp8_ds_mla":
2512 ops.cp_gather_and_upconvert_fp8_kv_cache(
2513 src_cache=kv_c_and_k_pe_cache,
2514 dst=workspace[:toks],
2306 - block_table=prefill_metadata.block_table,
2307 - workspace_starts=prefill_metadata.chunked_context.cu_seq_lens[i],
2308 - batch_size=attn_metadata.num_prefills,
2309 - seq_starts=prefill_metadata.chunked_context.starts[i],
2515 + block_table=block_table,
2516 + workspace_starts=chunk.cu_seq_lens,
2517 + batch_size=chunk.num_requests,
2518 + seq_starts=chunk.starts,
2519 )
2520 elif not use_fp8_prefill:
2521 ops.gather_and_maybe_dequant_cache(
2522 src_cache=kv_c_and_k_pe_cache,
2523 dst=workspace,
2315 - block_table=prefill_metadata.block_table,
2316 - cu_seq_lens=prefill_metadata.chunked_context.cu_seq_lens[i],
2317 - token_to_seq=prefill_metadata.chunked_context.token_to_seq[i],
2318 - num_tokens=prefill_metadata.chunked_context.chunk_total_token[i],
2524 + block_table=block_table,
2525 + cu_seq_lens=chunk.cu_seq_lens,
2526 + token_to_seq=chunk.token_to_seq,
2527 + num_tokens=toks,
2528 kv_cache_dtype=self.kv_cache_dtype,
2529 scale=k_scale,
2321 - seq_starts=prefill_metadata.chunked_context.starts[i],
2530 + seq_starts=chunk.starts,
2531 )
2532 else:
2533 # FP8 path: gather cache without dequantization
2534 ops.cp_gather_cache(
2535 src_cache=kv_c_and_k_pe_cache,
2536 dst=workspace,
2328 - block_table=prefill_metadata.block_table,
2329 - cu_seq_lens=prefill_metadata.chunked_context.cu_seq_lens[i],
2330 - batch_size=attn_metadata.num_prefills,
2331 - seq_starts=prefill_metadata.chunked_context.starts[i],
2537 + block_table=block_table,
2538 + cu_seq_lens=chunk.cu_seq_lens,
2539 + batch_size=chunk.num_requests,
2540 + seq_starts=chunk.starts,
2541 )
2542
2543 # Extract kv_c_normed from workspace
@@ -2351,37 +2560,28 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2560
2561 attn_output, attn_softmax_lse = (
2562 prefill_metadata.prefill_backend.run_prefill_context_chunk(
2354 - chunk_idx=i,
2355 - q=q,
2563 + chunk=chunk,
2564 + q=q[chunk.token_slice],
2565 k=k,
2566 v=v,
2567 )
2568 )
2360 - if prefill_metadata.chunked_context.has_empty_context[i]:
2361 - mask_empty_context(
2362 - attn_softmax_lse,
2363 - attn_output,
2364 - prefill_metadata.query_start_loc,
2365 - prefill_metadata.chunked_context.cu_seq_lens[i],
2366 - )
2569
2570 if output is None:
2369 - output = attn_output
2370 - output_lse = attn_softmax_lse
2371 - else:
2372 - if merge_output is None:
2373 - merge_output = torch.empty_like(output)
2374 - merge_output_lse = torch.empty_like(output_lse)
2375 - merge_attn_states(
2376 - output=merge_output,
2377 - output_lse=merge_output_lse,
2378 - prefix_output=output,
2379 - prefix_lse=output_lse,
2380 - suffix_output=attn_output,
2381 - suffix_lse=attn_softmax_lse,
2571 + if (
2572 + len(chunked_context.chunks) == 1
2573 + and not chunked_context.empty_token_slices
2574 + ):
2575 + return attn_output, attn_softmax_lse
2576 + output, output_lse = init_mla_context_partial(
2577 + chunked_context,
2578 + attn_output,
2579 + attn_softmax_lse,
2580 + num_tokens=q.shape[0],
2581 )
2383 - output, merge_output = merge_output, output
2384 - output_lse, merge_output_lse = merge_output_lse, output_lse
2582 + accumulate_mla_context_chunk(
2583 + chunk, attn_output, attn_softmax_lse, output, output_lse
2584 + )
2585
2586 return output, output_lse
2587
@@ -2396,62 +2596,57 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2596 assert attn_metadata.prefill is not None
2597 prefill_metadata = attn_metadata.prefill
2598 assert prefill_metadata.prefill_backend is not None
2399 - assert prefill_metadata.chunked_context is not None
2400 - assert prefill_metadata.chunked_context.padded_local_chunk_seq_lens is not None
2401 - assert prefill_metadata.chunked_context.local_context_lens_allranks is not None
2402 - assert prefill_metadata.chunked_context.padded_local_cu_seq_lens is not None
2403 - assert prefill_metadata.chunked_context.padded_local_token_to_seq is not None
2404 - assert prefill_metadata.chunked_context.cu_seq_lens_lst is not None
2405 - assert prefill_metadata.chunked_context.chunk_size is not None
2599 + chunked_context = prefill_metadata.chunked_context
2600 + assert chunked_context is not None
2601
2602 use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
2603 kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(
2604 self.kv_b_proj, use_fp8_prefill
2605 )
2606 output = None
2412 - merge_output = None
2413 - iters = len(prefill_metadata.chunked_context.seq_tot)
2414 - workspace = prefill_metadata.chunked_context.workspace
2415 -
2416 - for i in range(iters):
2417 - toks = prefill_metadata.chunked_context.seq_tot[i]
2418 - if toks == 0:
2419 - continue
2420 - padded_local_cu_seq_lens = (
2421 - prefill_metadata.chunked_context.padded_local_cu_seq_lens[i]
2422 - )
2607 + output_lse = None
2608 + workspace = chunked_context.workspace
2609 +
2610 + for chunk in chunked_context.chunks:
2611 + assert chunk.padded_local_seq_lens is not None
2612 + assert chunk.local_context_lens_allranks is not None
2613 + assert chunk.padded_local_cu_seq_lens is not None
2614 + assert chunk.padded_local_token_to_seq is not None
2615 + assert chunk.local_starts is not None
2616 +
2617 + toks = chunk.num_local_context_tokens
2618 + padded_local_cu_seq_lens = chunk.padded_local_cu_seq_lens
2619 + block_table = prefill_metadata.block_table[chunk.request_slice]
2620 if self.kv_cache_dtype == "fp8_ds_mla":
2621 ops.cp_gather_and_upconvert_fp8_kv_cache(
2622 src_cache=kv_c_and_k_pe_cache,
2623 dst=workspace[:toks],
2427 - block_table=prefill_metadata.block_table,
2624 + block_table=block_table,
2625 workspace_starts=padded_local_cu_seq_lens,
2429 - batch_size=attn_metadata.num_prefills,
2430 - seq_starts=prefill_metadata.chunked_context.starts[i],
2626 + batch_size=chunk.num_requests,
2627 + seq_starts=chunk.starts,
2628 )
2629 elif is_quantized_kv_cache(self.kv_cache_dtype):
2630 assert k_scale is not None
2631 ops.gather_and_maybe_dequant_cache(
2632 src_cache=kv_c_and_k_pe_cache,
2633 dst=workspace,
2437 - block_table=prefill_metadata.block_table,
2634 + block_table=block_table,
2635 cu_seq_lens=padded_local_cu_seq_lens,
2439 - token_to_seq=prefill_metadata.chunked_context.padded_local_token_to_seq[
2440 - i
2441 - ],
2636 + token_to_seq=chunk.padded_local_token_to_seq,
2637 num_tokens=toks,
2638 kv_cache_dtype=self.kv_cache_dtype,
2639 scale=k_scale,
2445 - seq_starts=prefill_metadata.chunked_context.starts[i],
2640 + seq_starts=chunk.starts,
2641 )
2642 else:
2643 ops.cp_gather_cache(
2644 src_cache=kv_c_and_k_pe_cache,
2645 dst=workspace,
2451 - block_table=prefill_metadata.block_table,
2646 + block_table=block_table,
2647 cu_seq_lens=padded_local_cu_seq_lens,
2453 - batch_size=attn_metadata.num_prefills,
2454 - seq_starts=prefill_metadata.chunked_context.starts[i],
2648 + batch_size=chunk.num_requests,
2649 + seq_starts=chunk.starts,
2650 )
2651 # workspace
2652 # |------- N tokens --------|--------- N*dcp_size tokens ----------|
@@ -2479,14 +2674,11 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2674 kv_c_normed, k_pe = reorg_kvcache(
2675 allgatered_kv_c_normed,
2676 allgatered_k_pe,
2482 - padded_local_chunk_seq_lens_lst=prefill_metadata.chunked_context.padded_local_chunk_seq_lens[
2483 - i
2484 - ],
2485 - local_context_lens_allranks=prefill_metadata.chunked_context.local_context_lens_allranks,
2486 - sum_seq_len=prefill_metadata.chunked_context.cu_seq_lens_lst[i][-1],
2487 - max_seq_len=prefill_metadata.chunked_context.max_seq_lens[i],
2488 - chunk_size=prefill_metadata.chunked_context.chunk_size,
2489 - chunk_idx=i,
2677 + padded_local_chunk_seq_lens_lst=chunk.padded_local_seq_lens,
2678 + local_context_lens_allranks=chunk.local_context_lens_allranks,
2679 + local_starts=chunk.local_starts,
2680 + sum_seq_len=chunk.num_context_tokens,
2681 + max_seq_len=chunk.max_seq_len,
2682 toks=toks,
2683 )
2684 if kv_b_proj_input_dtype is not None:
@@ -2503,37 +2695,28 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2695
2696 attn_output, attn_softmax_lse = (
2697 prefill_metadata.prefill_backend.run_prefill_context_chunk(
2506 - chunk_idx=i,
2507 - q=q,
2698 + chunk=chunk,
2699 + q=q[chunk.token_slice],
2700 k=k,
2701 v=v,
2702 )
2703 )
2512 - if prefill_metadata.chunked_context.has_empty_context[i]:
2513 - mask_empty_context(
2514 - attn_softmax_lse,
2515 - attn_output,
2516 - prefill_metadata.query_start_loc,
2517 - prefill_metadata.chunked_context.cu_seq_lens[i],
2518 - )
2704
2705 if output is None:
2521 - output = attn_output
2522 - output_lse = attn_softmax_lse
2523 - else:
2524 - if merge_output is None:
2525 - merge_output = torch.empty_like(output)
2526 - merge_output_lse = torch.empty_like(output_lse)
2527 - merge_attn_states(
2528 - output=merge_output,
2529 - output_lse=merge_output_lse,
2530 - prefix_output=output,
2531 - prefix_lse=output_lse,
2532 - suffix_output=attn_output,
2533 - suffix_lse=attn_softmax_lse,
2706 + if (
2707 + len(chunked_context.chunks) == 1
2708 + and not chunked_context.empty_token_slices
2709 + ):
2710 + return attn_output, attn_softmax_lse
2711 + output, output_lse = init_mla_context_partial(
2712 + chunked_context,
2713 + attn_output,
2714 + attn_softmax_lse,
2715 + num_tokens=q.shape[0],
2716 )
2535 - output, merge_output = merge_output, output
2536 - output_lse, merge_output_lse = merge_output_lse, output_lse
2717 + accumulate_mla_context_chunk(
2718 + chunk, attn_output, attn_softmax_lse, output, output_lse
2719 + )
2720
2721 return output, output_lse
2722
@@ -2615,7 +2798,6 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2798 prefix_lse=context_lse,
2799 suffix_output=suffix_output,
2800 suffix_lse=suffix_lse,
2618 - prefill_tokens_with_context=prefill_metadata.chunked_context.prefill_tokens_with_context,
2801 )
2802 elif output_scale is None:
2803 # With output_scale set, backend already wrote into `output` in place.
vllm/model_executor/layers/attention/sparse_mla_attention.py
+50 -54
@@ -19,8 +19,10 @@ from vllm.model_executor.layers.attention.mla_attention import (
19 MLACommonBaseImpl,
20 MLACommonMetadata,
21 MLACommonPrefillMetadata,
22 + accumulate_mla_context_chunk,
23 build_mla_chunked_context_metadata,
24 get_mla_dims,
25 + init_mla_context_partial,
26 )
27 from vllm.platforms import current_platform
28 from vllm.triton_utils import tl, triton
@@ -60,17 +62,14 @@ def _topk_mask_shape(
62 def _masked_mha_workspace_fits(
63 batch_size: int,
64 max_query_len: int,
63 - context_chunk_max_seq_lens: list[int] | None,
65 + max_context_chunk_seq_len: int,
66 workspace_numel: int,
67 ) -> bool:
68 """Return whether the suffix and per-context-chunk masks fit the workspace.
69
70 The global mask is excluded: it always needs more, and has its own check.
71 """
70 - max_key_len = max(
71 - max_query_len,
72 - max(context_chunk_max_seq_lens or (), default=0),
73 - )
72 + max_key_len = max(max_query_len, max_context_chunk_seq_len)
73 needed = math.prod(_topk_mask_shape(batch_size, max_query_len, max_key_len))
74 return needed <= workspace_numel
75
@@ -185,10 +184,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
184 64 * 1024,
185 scheduler_config.max_num_seqs * topk_tokens,
186 )
188 - return max(
189 - workspace_size,
190 - scheduler_config.max_num_seqs * cache_config.block_size,
191 - )
187 + return max(workspace_size, cache_config.block_size)
188
189 def _build_req_id_per_token(
190 self,
@@ -228,7 +224,6 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
224 return build_mla_chunked_context_metadata(
225 context_lens_cpu=context_lens_cpu,
226 prefill_query_start_loc_cpu=prefill_query_start_loc_cpu,
231 - num_prefills=num_prefills,
227 chunked_prefill_workspace=self.chunked_prefill_workspace,
228 chunked_prefill_workspace_size=self.chunked_prefill_workspace_size,
229 block_size=self.kv_cache_spec.block_size,
@@ -542,14 +537,15 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):
537 workspace = prefill.topk_mask_workspace
538 if workspace is None or prefill.query_lens_cpu is None:
539 return False
540 + max_context_chunk_seq_len = 0
541 + if prefill.chunked_context is not None:
542 + max_context_chunk_seq_len = max(
543 + chunk.max_seq_len for chunk in prefill.chunked_context.chunks
544 + )
545 fits = _masked_mha_workspace_fits(
546 batch_size=len(prefill.query_lens_cpu),
547 max_query_len=prefill.max_query_len,
548 - context_chunk_max_seq_lens=(
549 - None
550 - if prefill.chunked_context is None
551 - else prefill.chunked_context.max_seq_lens
552 - ),
548 + max_context_chunk_seq_len=max_context_chunk_seq_len,
549 workspace_numel=workspace.numel(),
550 )
551 if not fits:
@@ -716,69 +712,70 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):
712
713 chunked_context = prefill_metadata.chunked_context
714 assert chunked_context is not None
719 - use_global_mask = dense_mask is not None
715 output: torch.Tensor | None = None
716 output_lse: torch.Tensor | None = None
717 workspace = chunked_context.workspace
718
724 - for i, toks in enumerate(chunked_context.seq_tot):
725 - if toks == 0:
726 - continue
719 + for chunk in chunked_context.chunks:
720 + toks = chunk.num_context_tokens
721 + requests = chunk.request_slice
722 ops.gather_and_maybe_dequant_cache(
723 src_cache=kv_c_and_k_pe_cache,
724 dst=workspace,
730 - block_table=prefill_metadata.block_table,
731 - cu_seq_lens=chunked_context.cu_seq_lens[i],
732 - token_to_seq=chunked_context.token_to_seq[i],
733 - num_tokens=chunked_context.chunk_total_token[i],
725 + block_table=prefill_metadata.block_table[requests],
726 + cu_seq_lens=chunk.cu_seq_lens,
727 + token_to_seq=chunk.token_to_seq,
728 + num_tokens=toks,
729 kv_cache_dtype=self.kv_cache_dtype,
730 scale=k_scale,
736 - seq_starts=chunked_context.starts[i],
731 + seq_starts=chunk.starts,
732 )
733
734 chunk_kv_c = workspace[:toks, : self.kv_lora_rank]
735 chunk_k_pe = workspace[:toks, self.kv_lora_rank :].unsqueeze(1)
736 k, v = self._project_kv(chunk_kv_c, chunk_k_pe)
742 - chunk_lens = chunked_context.seq_lens[i].tolist()
743 - chunk_topk = (
744 - topk_per_req
745 - if use_global_mask
746 - else self._remap_topk_to_ranges(
747 - topk_per_req,
748 - chunked_context.starts[i],
749 - chunk_lens,
737 + if dense_mask is not None:
738 + chunk_mask: torch.Tensor | None = dense_mask[requests]
739 + chunk_topk = topk_per_req[requests]
740 + key_starts: torch.Tensor | None = chunk.starts
741 + else:
742 + chunk_mask = None
743 + chunk_topk = self._remap_topk_to_ranges(
744 + topk_per_req[requests],
745 + chunk.starts,
746 + chunk.seq_lens.tolist(),
747 )
751 - )
748 + key_starts = None
749 attn_out, lse = self._run_masked_mha(
753 - q=q,
750 + q=q[chunk.token_slice],
751 k=k,
752 v=v,
756 - cu_seqlens_q=prefill_metadata.query_start_loc,
757 - cu_seqlens_k=chunked_context.cu_seq_lens[i],
758 - max_seqlen_q=prefill_metadata.max_query_len,
759 - max_seqlen_k=chunked_context.max_seq_lens[i],
753 + cu_seqlens_q=chunk.query_start_loc,
754 + cu_seqlens_k=chunk.cu_seq_lens,
755 + max_seqlen_q=chunk.max_query_len,
756 + max_seqlen_k=chunk.max_seq_len,
757 topk_per_req=chunk_topk,
761 - q_lens=q_lens,
758 + q_lens=q_lens[requests],
759 causal=False,
760 return_softmax_lse=True,
764 - dense_mask=dense_mask,
765 - key_starts=(chunked_context.starts[i] if use_global_mask else None),
761 + dense_mask=chunk_mask,
762 + key_starts=key_starts,
763 topk_mask_workspace=prefill_metadata.topk_mask_workspace,
764 )
765
766 if output is None:
770 - output = attn_out
771 - output_lse = lse
772 - else:
773 - assert output_lse is not None
774 - merge_attn_states(
775 - output=output,
776 - output_lse=output_lse,
777 - prefix_output=output,
778 - prefix_lse=output_lse,
779 - suffix_output=attn_out,
780 - suffix_lse=lse,
767 + if (
768 + len(chunked_context.chunks) == 1
769 + and not chunked_context.empty_token_slices
770 + ):
771 + return attn_out, lse
772 + output, output_lse = init_mla_context_partial(
773 + chunked_context,
774 + attn_out,
775 + lse,
776 + num_tokens=q.shape[0],
777 )
778 + accumulate_mla_context_chunk(chunk, attn_out, lse, output, output_lse)
779
780 assert output is not None and output_lse is not None
781 return output, output_lse
@@ -844,7 +841,7 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):
841 output.copy_(attn_out[..., : self.v_head_dim].flatten(start_dim=-2))
842 return
843
847 - context_lens = chunked_context.seq_lens.sum(dim=0).tolist()
844 + context_lens = chunked_context.context_lens_list
845 dense_mask = self._try_build_global_mask(
846 topk_per_req,
847 q_lens,
@@ -889,5 +886,4 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):
886 prefix_lse=context_lse,
887 suffix_output=suffix_output[..., : self.v_head_dim],
888 suffix_lse=suffix_lse,
892 - prefill_tokens_with_context=chunked_context.prefill_tokens_with_context,
889 )
vllm/models/kimi_k3/nvidia/mla.py
-1
@@ -763,7 +763,6 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
763 prefix_lse=context_lse,
764 suffix_output=suffix_output[..., : self.v_head_dim],
765 suffix_lse=suffix_lse,
766 - prefill_tokens_with_context=prefill.chunked_context.prefill_tokens_with_context,
766 )
767 elif not writes_out:
768 out.copy_(output_prefill[..., : self.v_head_dim].flatten(start_dim=-2))
vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py
+8 -7
@@ -17,6 +17,9 @@ from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
17
18 if TYPE_CHECKING:
19 from vllm.config import VllmConfig
20 + from vllm.model_executor.layers.attention.mla_attention import (
21 + MLACommonPrefillMetadata,
22 + )
23 from vllm.platforms.interface import DeviceCapability
24
25
@@ -99,21 +102,19 @@ class AiterFlashAttnPrefillBackend(MLAPrefillBackend):
102
103 def run_prefill_context_chunk(
104 self,
102 - chunk_idx: int,
105 + chunk: "MLACommonPrefillMetadata.ContextChunk",
106 q: torch.Tensor,
107 k: torch.Tensor,
108 v: torch.Tensor,
109 ) -> tuple[torch.Tensor, torch.Tensor]:
107 - assert self._prefill_metadata.chunked_context is not None
108 - chunked = self._prefill_metadata.chunked_context
110 out, lse = self.flash_attn_varlen_func(
111 q=q,
112 k=k,
113 v=v,
113 - cu_seqlens_q=self._prefill_metadata.query_start_loc,
114 - cu_seqlens_k=chunked.cu_seq_lens[chunk_idx],
115 - max_seqlen_q=self._prefill_metadata.max_query_len,
116 - max_seqlen_k=chunked.max_seq_lens[chunk_idx],
114 + cu_seqlens_q=chunk.query_start_loc,
115 + cu_seqlens_k=chunk.cu_seq_lens,
116 + max_seqlen_q=chunk.max_query_len,
117 + max_seqlen_k=chunk.max_seq_len,
118 softmax_scale=self.scale,
119 causal=False,
120 return_lse=True,
vllm/v1/attention/backends/mla/prefill/base.py
+1 -1
@@ -175,7 +175,7 @@ class MLAPrefillBackend(ABC):
175 @abstractmethod
176 def run_prefill_context_chunk(
177 self,
178 - chunk_idx: int,
178 + chunk: "MLACommonPrefillMetadata.ContextChunk",
179 q: torch.Tensor,
180 k: torch.Tensor,
181 v: torch.Tensor,
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+9 -7
@@ -27,7 +27,10 @@ from vllm.v1.attention.backends.mla.prefill.base import (
27
28 if TYPE_CHECKING:
29 from vllm.config import VllmConfig
30 - from vllm.model_executor.layers.attention.mla_attention import MLADims
30 + from vllm.model_executor.layers.attention.mla_attention import (
31 + MLACommonPrefillMetadata,
32 + MLADims,
33 + )
34 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
35
36 if is_flash_attn_varlen_func_available():
@@ -450,20 +453,19 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
453
454 def run_prefill_context_chunk(
455 self,
453 - chunk_idx: int,
456 + chunk: "MLACommonPrefillMetadata.ContextChunk",
457 q: torch.Tensor,
458 k: torch.Tensor,
459 v: torch.Tensor,
460 ) -> tuple[torch.Tensor, torch.Tensor]:
458 - assert self._prefill_metadata.chunked_context is not None
461 return self._flash_attn_varlen_diff_headdims(
462 q=q,
463 k=k,
464 v=v,
463 - cu_seqlens_q=self._prefill_metadata.query_start_loc,
464 - cu_seqlens_k=self._prefill_metadata.chunked_context.cu_seq_lens[chunk_idx],
465 - max_seqlen_q=self._prefill_metadata.max_query_len,
466 - max_seqlen_k=self._prefill_metadata.chunked_context.max_seq_lens[chunk_idx],
465 + cu_seqlens_q=chunk.query_start_loc,
466 + cu_seqlens_k=chunk.cu_seq_lens,
467 + max_seqlen_q=chunk.max_query_len,
468 + max_seqlen_k=chunk.max_seq_len,
469 softmax_scale=self.scale,
470 causal=False, # Context is unmasked
471 return_softmax_lse=True,
vllm/v1/attention/backends/mla/prefill/flashinfer.py
+7 -10
@@ -150,8 +150,7 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
150 if has_context:
151 chunked_context = prefill_metadata.chunked_context
152 assert chunked_context is not None
153 - num_chunks = chunked_context.cu_seq_lens.shape[0]
154 - self._ensure_chunks(num_chunks, self._workspace_buffer)
153 + self._ensure_chunks(len(chunked_context.chunks), self._workspace_buffer)
154
155 num_qo_heads = self.num_heads
156 num_kv_heads = num_qo_heads
@@ -179,12 +178,10 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
178 if has_context:
179 chunked_context = prefill_metadata.chunked_context
180 assert chunked_context is not None
182 - for i in range(num_chunks):
183 - kv_indptr_chunk = chunked_context.cu_seq_lens[i]
184 -
185 - self._prefill_chunks[i].plan(
186 - qo_indptr=qo_indptr,
187 - kv_indptr=kv_indptr_chunk,
181 + for chunk in chunked_context.chunks:
182 + self._prefill_chunks[chunk.index].plan(
183 + qo_indptr=chunk.query_start_loc,
184 + kv_indptr=chunk.cu_seq_lens,
185 num_qo_heads=num_qo_heads,
186 num_kv_heads=num_kv_heads,
187 head_dim_qk=head_dim_qk,
@@ -227,12 +224,12 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
224
225 def run_prefill_context_chunk(
226 self,
230 - chunk_idx: int,
227 + chunk: "MLACommonPrefillMetadata.ContextChunk",
228 q: torch.Tensor,
229 k: torch.Tensor,
230 v: torch.Tensor,
231 ) -> tuple[torch.Tensor, torch.Tensor]:
235 - attn_out, lse = self._prefill_chunks[chunk_idx].run(
232 + attn_out, lse = self._prefill_chunks[chunk.index].run(
233 q=q,
234 k=k,
235 v=v,
vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py
+7 -10
@@ -162,16 +162,13 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
162
163 def run_prefill_context_chunk(
164 self,
165 - chunk_idx: int,
165 + chunk: "MLACommonPrefillMetadata.ContextChunk",
166 q: torch.Tensor,
167 k: torch.Tensor,
168 v: torch.Tensor,
169 ) -> tuple[torch.Tensor, torch.Tensor]:
170 from tokenspeed_mla import tokenspeed_mla_prefill
171
172 - assert self._prefill_metadata.chunked_context is not None
173 - chunked = self._prefill_metadata.chunked_context
174 -
172 # See note in run_prefill_new_tokens — `v` is a split-view of `kv_nope`
173 # in `_compute_prefill_context` and arrives non-contiguous.
174 v = v.contiguous()
@@ -180,15 +177,15 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
177 query=q,
178 key=k,
179 value=v,
183 - seq_lens=chunked.seq_lens[chunk_idx],
184 - cum_seq_lens=chunked.cu_seq_lens[chunk_idx],
185 - max_seq_len=chunked.max_seq_lens[chunk_idx],
186 - batch_size=chunked.seq_lens[chunk_idx].shape[0],
180 + seq_lens=chunk.seq_lens,
181 + cum_seq_lens=chunk.cu_seq_lens,
182 + max_seq_len=chunk.max_seq_len,
183 + batch_size=chunk.num_requests,
184 softmax_scale=self.scale,
185 is_causal=False,
186 return_lse=True,
190 - cum_seq_lens_q=self._prefill_metadata.query_start_loc,
191 - max_seq_len_q=self._prefill_metadata.max_query_len,
187 + cum_seq_lens_q=chunk.query_start_loc,
188 + max_seq_len_q=chunk.max_query_len,
189 enable_pdl=False,
190 )
191
vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py
+7 -14
@@ -143,16 +143,13 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
143
144 def run_prefill_context_chunk(
145 self,
146 - chunk_idx: int,
146 + chunk: "MLACommonPrefillMetadata.ContextChunk",
147 q: torch.Tensor,
148 k: torch.Tensor,
149 v: torch.Tensor,
150 ) -> tuple[torch.Tensor, torch.Tensor]:
151 from flashinfer.prefill import trtllm_ragged_attention_deepseek
152
153 - assert self._prefill_metadata.chunked_context is not None
154 - assert self._prefill_metadata.chunked_context.seq_lens[chunk_idx] is not None
155 -
153 out = torch.empty(
154 q.shape[0],
155 q.shape[1],
@@ -166,20 +163,16 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
163 key=k,
164 value=v,
165 workspace_buffer=self._workspace_buffer,
169 - seq_lens=self._prefill_metadata.chunked_context.seq_lens[chunk_idx],
170 - max_q_len=self._prefill_metadata.max_query_len,
171 - max_kv_len=self._prefill_metadata.chunked_context.max_seq_lens[chunk_idx],
166 + seq_lens=chunk.seq_lens,
167 + max_q_len=chunk.max_query_len,
168 + max_kv_len=chunk.max_seq_len,
169 bmm1_scale=self.scale,
170 bmm2_scale=1.0,
171 o_sf_scale=1.0,
175 - batch_size=self._prefill_metadata.chunked_context.seq_lens[chunk_idx].shape[
176 - 0
177 - ],
172 + batch_size=chunk.num_requests,
173 window_left=-1,
179 - cum_seq_lens_q=self._prefill_metadata.query_start_loc,
180 - cum_seq_lens_kv=self._prefill_metadata.chunked_context.cu_seq_lens[
181 - chunk_idx
182 - ],
174 + cum_seq_lens_q=chunk.query_start_loc,
175 + cum_seq_lens_kv=chunk.cu_seq_lens,
176 enable_pdl=False,
177 is_causal=False,
178 return_lse=True,
vllm/v1/attention/ops/triton_merge_attn_states.py
+1 -114
@@ -9,118 +9,6 @@ from vllm.triton_utils import tl, triton
9 float8_info = torch.finfo(current_platform.fp8_dtype())
10
11
12 -def mask_empty_context(
13 - lse: torch.Tensor,
14 - output: torch.Tensor,
15 - query_start_loc: torch.Tensor,
16 - context_start_loc: torch.Tensor,
17 -) -> None:
18 - """Neutralize context chunks that cover no keys before merging.
19 -
20 - A prefill query whose context chunk is empty attended to no keys, so its
21 - partial attention is undefined: the backend leaves the output rows as
22 - uninitialized scratch (which may hold NaN/Inf) even when it reports an LSE
23 - of -inf. Sanitize both here so ``merge_attn_states`` can stay generic:
24 - force the LSE to -inf (zero softmax weight) and zero the undefined output
25 - rows (so a zero weight cannot combine with NaN/Inf). Emptiness is derived
26 - from the context offsets, not from the -inf LSE, so no merge kernel has to
27 - reason about undefined partials.
28 -
29 - Args:
30 - lse: Chunk log-sum-exp, shape [num_heads, num_tokens].
31 - output: Chunk attention output, shape [num_tokens, num_heads, ...].
32 - query_start_loc: Prefill query cumulative offsets, shape [num_reqs + 1].
33 - context_start_loc: Chunk context cumulative offsets,
34 - shape [num_reqs + 1]; an empty chunk has a zero-length span.
35 - """
36 - num_heads, num_tokens = lse.shape
37 - num_reqs = query_start_loc.shape[0] - 1
38 - block_size = 128
39 - # Reserve the worst-case number of request-local blocks.
40 - num_query_blocks = num_tokens // block_size + num_reqs
41 - is_empty = torch.zeros(num_tokens, dtype=torch.bool, device=lse.device)
42 - mask_empty_context_kernel[(num_query_blocks,)](
43 - lse,
44 - is_empty,
45 - query_start_loc,
46 - context_start_loc,
47 - lse.stride(0),
48 - lse.stride(1),
49 - num_reqs,
50 - NUM_HEADS=num_heads,
51 - BLOCK_SIZE=block_size,
52 - BLOCK_HEADS=8,
53 - num_warps=8,
54 - )
55 - output.masked_fill_(is_empty[:, None, None], 0.0)
56 -
57 -
58 -@triton.jit
59 -def mask_empty_context_kernel(
60 - lse,
61 - is_empty,
62 - query_start_loc,
63 - context_start_loc,
64 - lse_head_stride,
65 - lse_token_stride,
66 - num_reqs,
67 - NUM_HEADS: tl.constexpr,
68 - BLOCK_SIZE: tl.constexpr,
69 - BLOCK_HEADS: tl.constexpr,
70 -):
71 - query_block_idx = tl.program_id(0)
72 -
73 - lanes = tl.arange(0, 32)
74 - chunk_start = 0
75 - req_idx = 0
76 - req_idx_found = False
77 - while (chunk_start < num_reqs) & (not req_idx_found):
78 - req_offsets = chunk_start + lanes
79 - req_mask = req_offsets < num_reqs
80 - query_starts = tl.load(query_start_loc + req_offsets, mask=req_mask)
81 - # Assume the worst-case number of blocks for each request.
82 - req_block_starts = query_starts // BLOCK_SIZE + req_offsets
83 - matched_idx = tl.sum(
84 - (req_mask & (req_block_starts <= query_block_idx)).to(tl.int32)
85 - )
86 - # matched_idx == 32 means the match is past this warp chunk.
87 - req_idx = chunk_start + matched_idx - 1
88 - req_idx_found = matched_idx < 32
89 - chunk_start += 32
90 -
91 - query_start = tl.load(query_start_loc + req_idx)
92 - query_end = tl.load(query_start_loc + req_idx + 1)
93 - query_len = query_end - query_start
94 - req_first_block = query_start // BLOCK_SIZE + req_idx
95 - block_in_req = query_block_idx - req_first_block
96 - token_offset = block_in_req * BLOCK_SIZE
97 - if token_offset >= query_len:
98 - return
99 -
100 - context_start = tl.load(context_start_loc + req_idx)
101 - context_end = tl.load(context_start_loc + req_idx + 1)
102 - if context_start != context_end:
103 - return
104 -
105 - token_offsets = token_offset + tl.arange(0, BLOCK_SIZE)
106 - token_indices = query_start + token_offsets
107 - token_lse_offsets = token_indices * lse_token_stride
108 - valid_tokens = token_offsets < query_len
109 - tl.store(is_empty + token_indices, True, mask=valid_tokens)
110 - head_offsets = tl.arange(0, BLOCK_HEADS)
111 - for head_start in range(0, NUM_HEADS, BLOCK_HEADS):
112 - head_indices = head_start + head_offsets
113 - lse_ptrs = (
114 - lse + head_indices[:, None] * lse_head_stride + token_lse_offsets[None, :]
115 - )
116 - valid_heads = head_indices < NUM_HEADS
117 - tl.store(
118 - lse_ptrs,
119 - float("-inf"),
120 - mask=valid_heads[:, None] & valid_tokens[None, :],
121 - )
122 -
123 -
12 # Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
13 # can be used to combine partial attention results (in the split-KV case)
14 def merge_attn_states(
@@ -317,8 +205,7 @@ def merge_attn_states_kernel(
205 s_scale = s_se / out_se
206 out = p_out * p_scale + s_out * s_scale
207 # If both sides are empty (max_lse == -inf) the scales are 0/0 = NaN; emit
320 - # zeros rather than NaN. Callers with empty chunks (see mask_empty_context)
321 - # zero those inputs, so this only guards the fully-undefined corner.
208 + # zeros rather than NaN.
209 out = tl.where(max_lse == float("-inf"), 0.0, out)
210
211 if USE_FP8: