A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 15:15 UTC
[Update] Use FlashInfer fast_decode_plan directly instead of replication (#34687)
Signed-off-by: Andrii <askliar@nvidia.com> Co-authored-by: Andrii <askliar@nvidia.com>
Andrii Skliar committed
Feb 27, 2026 at 01:31 UTC
56a6371706bc32c9e2e996ec29911c087a547ac0
2 files changed
+286
-131
tests/kernels/attention/test_flashinfer.py
+203
@@ -84,6 +84,209 @@ def ref_paged_attn(
84
return torch.cat(outputs, dim=0)
85
86
87
+def _make_paged_kv_metadata(
88
+ kv_lens: list[int],
89
+ block_size: int,
90
+ num_blocks: int,
91
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
92
+ """Build paged-KV metadata tensors for fast_plan_decode tests.
93
+
94
+ Returns:
95
+ kv_indptr – CPU int32, shape [num_seqs + 1]
96
+ kv_indices – CUDA int32, shape [total_blocks]
97
+ kv_last_page_lens – CPU int32, shape [num_seqs]
98
+ block_tables – CUDA int32, shape [num_seqs, max_blocks_per_seq]
99
+ """
100
+ num_seqs = len(kv_lens)
101
+ max_blocks = (max(kv_lens) + block_size - 1) // block_size
102
+ block_tables = torch.randint(
103
+ 0, num_blocks, (num_seqs, max_blocks), dtype=torch.int32, device="cuda"
104
+ )
105
+
106
+ indptr_list = [0]
107
+ indices_list: list[int] = []
108
+ last_lens_list: list[int] = []
109
+ for i, seq_len in enumerate(kv_lens):
110
+ n = (seq_len + block_size - 1) // block_size
111
+ indices_list.extend(block_tables[i, :n].cpu().tolist())
112
+ indptr_list.append(indptr_list[-1] + n)
113
+ last_lens_list.append(seq_len % block_size or block_size)
114
+
115
+ return (
116
+ torch.tensor(indptr_list, dtype=torch.int32, device="cpu"),
117
+ torch.tensor(indices_list, dtype=torch.int32, device="cuda"),
118
+ torch.tensor(last_lens_list, dtype=torch.int32, device="cpu"),
119
+ block_tables,
120
+ )
121
+
122
+
123
+def _make_cg_decode_wrapper(
124
+ num_seqs: int,
125
+ kv_indices_buffer: torch.Tensor,
126
+ workspace_buffer: torch.Tensor,
127
+ use_tensor_cores: bool = True,
128
+) -> "flashinfer.BatchDecodeWithPagedKVCacheWrapper":
129
+ """Create a cudagraph-enabled BatchDecodeWithPagedKVCacheWrapper.
130
+
131
+ *kv_indices_buffer* is shared with the caller so that fast_plan_decode
132
+ can avoid the device-to-device index copy on subsequent (cudagraph) calls.
133
+ """
134
+ return flashinfer.BatchDecodeWithPagedKVCacheWrapper(
135
+ workspace_buffer,
136
+ "NHD",
137
+ use_cuda_graph=True,
138
+ paged_kv_indptr_buffer=torch.zeros(
139
+ num_seqs + 1, dtype=torch.int32, device="cuda"
140
+ ),
141
+ paged_kv_indices_buffer=kv_indices_buffer,
142
+ paged_kv_last_page_len_buffer=torch.zeros(
143
+ num_seqs, dtype=torch.int32, device="cuda"
144
+ ),
145
+ use_tensor_cores=use_tensor_cores,
146
+ )
147
+
148
+
149
+def test_fast_decode_plan_importable() -> None:
150
+ """fast_decode_plan must be importable from flashinfer.decode.
151
+
152
+ This is a forward-compatibility smoke test: if FlashInfer reorganises its
153
+ public API the import will fail before any other test does.
154
+ """
155
+ from flashinfer.decode import fast_decode_plan # noqa: F401
156
+
157
+ assert callable(fast_decode_plan)
158
+
159
+
160
+@pytest.mark.parametrize("dtype", DTYPES)
161
+@torch.inference_mode
162
+def test_fast_plan_decode_warmup_uses_full_plan(dtype: torch.dtype) -> None:
163
+ """On the first call fast_plan_decode must route through self.plan() and
164
+ flip vllm_first_call to False on the wrapper object."""
165
+ from unittest.mock import patch
166
+
167
+ from vllm.v1.attention.backends.flashinfer import fast_plan_decode
168
+
169
+ torch.set_default_device("cuda")
170
+ set_random_seed(0)
171
+
172
+ kv_lens = [128, 64]
173
+ block_size = 16
174
+ num_seqs = len(kv_lens)
175
+ num_query_heads, num_kv_heads = 8, 2
176
+ head_size = 128
177
+
178
+ kv_indptr, kv_indices, kv_last_page_lens, _ = _make_paged_kv_metadata(
179
+ kv_lens, block_size, NUM_BLOCKS
180
+ )
181
+
182
+ workspace = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
183
+ wrapper = _make_cg_decode_wrapper(num_seqs, kv_indices.clone(), workspace)
184
+
185
+ assert getattr(wrapper, "vllm_first_call", True) is True
186
+
187
+ with patch.object(wrapper, "plan", wraps=wrapper.plan) as mock_plan:
188
+ fast_plan_decode(
189
+ wrapper,
190
+ indptr_cpu=kv_indptr,
191
+ indices=kv_indices,
192
+ last_page_len_cpu=kv_last_page_lens,
193
+ num_qo_heads=num_query_heads,
194
+ num_kv_heads=num_kv_heads,
195
+ head_dim=head_size,
196
+ page_size=block_size,
197
+ q_data_type=dtype,
198
+ kv_data_type=dtype,
199
+ )
200
+ mock_plan.assert_called_once()
201
+
202
+ assert wrapper.vllm_first_call is False, (
203
+ "vllm_first_call should be False after the first fast_plan_decode call"
204
+ )
205
+
206
+
207
+@pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]])
208
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
209
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
210
+@pytest.mark.parametrize("block_size", BLOCK_SIZES)
211
+@pytest.mark.parametrize("dtype", DTYPES)
212
+@torch.inference_mode
213
+def test_fast_plan_decode_matches_full_plan(
214
+ kv_lens: list[int],
215
+ num_heads: tuple[int, int],
216
+ head_size: int,
217
+ block_size: int,
218
+ dtype: torch.dtype,
219
+) -> None:
220
+ """fast_plan_decode's cudagraph path (delegating to FlashInfer's
221
+ fast_decode_plan) must produce attention output numerically identical to
222
+ a standard plan() call.
223
+
224
+ Both the warmup call (self.plan) and the subsequent fast call
225
+ (fast_decode_plan) are verified against the same reference.
226
+ """
227
+ from vllm.v1.attention.backends.flashinfer import fast_plan_decode
228
+
229
+ torch.set_default_device("cuda")
230
+ set_random_seed(0)
231
+ num_seqs = len(kv_lens)
232
+ num_query_heads, num_kv_heads = num_heads
233
+
234
+ query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dtype)
235
+ key_value_cache = torch.randn(
236
+ NUM_BLOCKS, 2, block_size, num_kv_heads, head_size, dtype=dtype
237
+ )
238
+
239
+ kv_indptr, kv_indices, kv_last_page_lens, _ = _make_paged_kv_metadata(
240
+ kv_lens, block_size, NUM_BLOCKS
241
+ )
242
+
243
+ # Reference output via the standard plan()
244
+ workspace_ref = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
245
+ ref_wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
246
+ workspace_ref, "NHD", use_tensor_cores=True
247
+ )
248
+ ref_wrapper.plan(
249
+ kv_indptr,
250
+ kv_indices,
251
+ kv_last_page_lens,
252
+ num_query_heads,
253
+ num_kv_heads,
254
+ head_size,
255
+ block_size,
256
+ "NONE",
257
+ q_data_type=dtype,
258
+ kv_data_type=dtype,
259
+ )
260
+ ref_output = ref_wrapper.run(query, key_value_cache)
261
+
262
+ # CUDAGraph wrapper exercised through fast_plan_decode
263
+ kv_indices_buf = kv_indices.clone()
264
+ workspace_cg = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
265
+ cg_wrapper = _make_cg_decode_wrapper(num_seqs, kv_indices_buf, workspace_cg)
266
+
267
+ plan_kwargs: dict = dict(
268
+ indptr_cpu=kv_indptr,
269
+ indices=kv_indices_buf,
270
+ last_page_len_cpu=kv_last_page_lens,
271
+ num_qo_heads=num_query_heads,
272
+ num_kv_heads=num_kv_heads,
273
+ head_dim=head_size,
274
+ page_size=block_size,
275
+ q_data_type=dtype,
276
+ kv_data_type=dtype,
277
+ )
278
+
279
+ # First call – warmup path (routes through self.plan)
280
+ fast_plan_decode(cg_wrapper, **plan_kwargs)
281
+ warmup_output = cg_wrapper.run(query, key_value_cache)
282
+ torch.testing.assert_close(warmup_output, ref_output, atol=1e-2, rtol=1e-2)
283
+
284
+ # Second call – fast path (routes through fast_decode_plan from FlashInfer)
285
+ fast_plan_decode(cg_wrapper, **plan_kwargs)
286
+ fast_output = cg_wrapper.run(query, key_value_cache)
287
+ torch.testing.assert_close(fast_output, ref_output, atol=1e-2, rtol=1e-2)
288
+
289
+
290
@pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]])
291
@pytest.mark.parametrize("num_heads", NUM_HEADS)
292
@pytest.mark.parametrize("head_size", HEAD_SIZES)
vllm/v1/attention/backends/flashinfer.py
+83
-131
@@ -13,7 +13,7 @@ from flashinfer import (
13
BatchPrefillWithRaggedKVCacheWrapper,
14
MultiLevelCascadeAttentionWrapper,
15
)
16
-from flashinfer.decode import _get_range_buf, trtllm_batch_decode_with_kv_cache
16
+from flashinfer.decode import fast_decode_plan, trtllm_batch_decode_with_kv_cache
17
from flashinfer.prefill import trtllm_batch_context_with_kv_cache
18
from flashinfer.utils import FP4Tensor
19
from typing_extensions import override
@@ -199,14 +199,14 @@ class BatchDCPPrefillWrapper:
199
):
200
"""Plan the prefill operation with given parameters."""
201
self._context.plan(
202
- qo_indptr_cpu,
203
- paged_kv_indptr_cpu,
204
- paged_kv_indices,
205
- paged_kv_last_page_len_cpu,
206
- num_qo_heads * dcp_world_size,
207
- num_kv_heads,
208
- head_dim,
209
- page_size,
202
+ qo_indptr=qo_indptr_cpu,
203
+ paged_kv_indptr=paged_kv_indptr_cpu,
204
+ paged_kv_indices=paged_kv_indices,
205
+ paged_kv_last_page_len=paged_kv_last_page_len_cpu,
206
+ num_qo_heads=num_qo_heads * dcp_world_size,
207
+ num_kv_heads=num_kv_heads,
208
+ head_dim_qk=head_dim,
209
+ page_size=page_size,
210
causal=False, # This is context run
211
sm_scale=sm_scale,
212
window_left=window_left,
@@ -818,6 +818,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
818
page_size,
819
paged_kv_last_page_len_np,
820
)
821
+ self.paged_kv_last_page_len.gpu[:num_reqs].copy_(
822
+ self.paged_kv_last_page_len.cpu[:num_reqs], non_blocking=True
823
+ )
824
return paged_kv_indices
825
826
def build(
@@ -999,14 +1002,17 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1002
1003
attn_metadata.cascade_wrapper = self._get_cascade_wrapper()
1004
attn_metadata.cascade_wrapper.plan(
1002
- [shared_qo_indptr_cpu, qo_indptr_cpu],
1003
- [shared_kv_page_indptr_cpu, paged_kv_indptr_cpu],
1004
- [shared_kv_page_indices_cpu, paged_kv_indices],
1005
- [shared_kv_last_page_len_cpu, paged_kv_last_page_len_cpu],
1006
- self.num_qo_heads,
1007
- self.num_kv_heads,
1008
- self.head_dim,
1009
- self.page_size,
1005
+ qo_indptr_arr=[shared_qo_indptr_cpu, qo_indptr_cpu],
1006
+ paged_kv_indptr_arr=[shared_kv_page_indptr_cpu, paged_kv_indptr_cpu],
1007
+ paged_kv_indices_arr=[shared_kv_page_indices_cpu, paged_kv_indices],
1008
+ paged_kv_last_page_len=[
1009
+ shared_kv_last_page_len_cpu,
1010
+ paged_kv_last_page_len_cpu,
1011
+ ],
1012
+ num_qo_heads=self.num_qo_heads,
1013
+ num_kv_heads=self.num_kv_heads,
1014
+ head_dim=self.head_dim,
1015
+ page_size=self.page_size,
1016
causal=True,
1017
sm_scale=self.sm_scale,
1018
window_left=self.window_left,
@@ -1084,14 +1090,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1090
BatchPrefillWithPagedKVCacheWrapper,
1091
)
1092
prefill_wrapper.plan(
1087
- qo_indptr_prefill_cpu,
1088
- paged_kv_indptr_prefill_cpu,
1089
- paged_kv_indices,
1090
- paged_kv_last_page_len_prefill_cpu,
1091
- self.num_qo_heads,
1092
- self.num_kv_heads,
1093
- self.head_dim,
1094
- self.page_size,
1093
+ qo_indptr=qo_indptr_prefill_cpu,
1094
+ paged_kv_indptr=paged_kv_indptr_prefill_cpu,
1095
+ paged_kv_indices=paged_kv_indices,
1096
+ paged_kv_last_page_len=paged_kv_last_page_len_prefill_cpu,
1097
+ num_qo_heads=self.num_qo_heads,
1098
+ num_kv_heads=self.num_kv_heads,
1099
+ head_dim_qk=self.head_dim,
1100
+ page_size=self.page_size,
1101
causal=True,
1102
sm_scale=self.sm_scale,
1103
window_left=self.window_left,
@@ -1132,14 +1138,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1138
# in atten_metadata when using cudagraph.
1139
fast_plan_decode(
1140
decode_wrapper,
1135
- self.paged_kv_indptr.cpu[: num_input_tokens + 1],
1136
- paged_kv_indices,
1137
- self.paged_kv_last_page_len.cpu[:num_input_tokens],
1138
- seq_lens_cpu[:num_input_tokens],
1139
- self.num_qo_heads * self.dcp_world_size,
1140
- self.num_kv_heads,
1141
- self.head_dim,
1142
- self.page_size,
1141
+ indptr_cpu=self.paged_kv_indptr.cpu[: num_input_tokens + 1],
1142
+ indices=paged_kv_indices,
1143
+ last_page_len_cpu=self.paged_kv_last_page_len.cpu[
1144
+ :num_input_tokens
1145
+ ],
1146
+ num_qo_heads=self.num_qo_heads * self.dcp_world_size,
1147
+ num_kv_heads=self.num_kv_heads,
1148
+ head_dim=self.head_dim,
1149
+ page_size=self.page_size,
1150
# Disable flashinfer's pos encoding and use vllm's rope.
1151
pos_encoding_mode="NONE",
1152
sm_scale=self.sm_scale,
@@ -1617,7 +1624,6 @@ def fast_plan_decode(
1624
indptr_cpu: torch.Tensor,
1625
indices: torch.Tensor,
1626
last_page_len_cpu: torch.Tensor,
1620
- seq_lens_cpu: torch.Tensor,
1627
num_qo_heads: int,
1628
num_kv_heads: int,
1629
head_dim: int,
@@ -1654,111 +1660,57 @@ def fast_plan_decode(
1660
# this warm up is to generate the _cached_module for the decode wrapper.
1661
if not self.is_cuda_graph_enabled or getattr(self, "vllm_first_call", True):
1662
self.plan(
1657
- indptr_cpu,
1658
- indices,
1659
- last_page_len_cpu,
1660
- num_qo_heads,
1661
- num_kv_heads,
1662
- head_dim,
1663
- page_size,
1664
- pos_encoding_mode,
1665
- window_left,
1666
- logits_soft_cap,
1667
- q_data_type,
1668
- kv_data_type,
1669
- o_data_type,
1670
- data_type,
1671
- sm_scale,
1672
- rope_scale,
1673
- rope_theta,
1674
- non_blocking,
1675
- None, # block_tables
1676
- None, # seq_lens
1677
- fixed_split_size,
1678
- disable_split_kv,
1663
+ indptr=indptr_cpu,
1664
+ indices=indices,
1665
+ last_page_len=last_page_len_cpu,
1666
+ num_qo_heads=num_qo_heads,
1667
+ num_kv_heads=num_kv_heads,
1668
+ head_dim=head_dim,
1669
+ page_size=page_size,
1670
+ pos_encoding_mode=pos_encoding_mode,
1671
+ window_left=window_left,
1672
+ logits_soft_cap=logits_soft_cap,
1673
+ q_data_type=q_data_type,
1674
+ kv_data_type=kv_data_type,
1675
+ o_data_type=o_data_type,
1676
+ data_type=data_type,
1677
+ sm_scale=sm_scale,
1678
+ rope_scale=rope_scale,
1679
+ rope_theta=rope_theta,
1680
+ non_blocking=non_blocking,
1681
+ block_tables=None,
1682
+ seq_lens=None,
1683
+ fixed_split_size=fixed_split_size,
1684
+ disable_split_kv=disable_split_kv,
1685
)
1686
self.vllm_first_call = False
1687
return
1688
1689
assert self.is_cuda_graph_enabled, "Should be cudagraph only here"
1690
1685
- batch_size = len(last_page_len_cpu)
1686
- if logits_soft_cap is None:
1687
- logits_soft_cap = 0.0
1688
-
1689
- # Handle data types consistently
1690
- if data_type is not None:
1691
- if q_data_type is None:
1692
- q_data_type = data_type
1693
- if kv_data_type is None:
1694
- kv_data_type = data_type
1695
- elif q_data_type is None:
1696
- q_data_type = "float16"
1697
-
1698
- if kv_data_type is None:
1699
- kv_data_type = q_data_type
1700
- q_data_type = (
1701
- getattr(torch, q_data_type) if isinstance(q_data_type, str) else q_data_type
1702
- )
1703
- kv_data_type = (
1704
- getattr(torch, kv_data_type) if isinstance(kv_data_type, str) else kv_data_type
1691
+ fast_decode_plan(
1692
+ self,
1693
+ indptr=indptr_cpu,
1694
+ indices=indices,
1695
+ last_page_len=last_page_len_cpu,
1696
+ num_qo_heads=num_qo_heads,
1697
+ num_kv_heads=num_kv_heads,
1698
+ head_dim=head_dim,
1699
+ page_size=page_size,
1700
+ pos_encoding_mode=pos_encoding_mode,
1701
+ window_left=window_left,
1702
+ logits_soft_cap=logits_soft_cap,
1703
+ q_data_type=q_data_type,
1704
+ kv_data_type=kv_data_type,
1705
+ data_type=data_type,
1706
+ sm_scale=sm_scale,
1707
+ rope_scale=rope_scale,
1708
+ rope_theta=rope_theta,
1709
+ non_blocking=non_blocking,
1710
+ fixed_split_size=fixed_split_size,
1711
+ disable_split_kv=disable_split_kv,
1712
)
1713
1707
- if batch_size != self._fixed_batch_size:
1708
- raise ValueError(
1709
- "The batch size should be fixed in cudagraph mode, the runtime "
1710
- "batch size {} mismatches the batch size set during "
1711
- "initialization {}".format(batch_size, self._fixed_batch_size)
1712
- )
1713
- if len(indices) > len(self._paged_kv_indices_buf):
1714
- raise ValueError(
1715
- "The size of indices should be less than or equal to the allocated buffer"
1716
- )
1717
-
1718
- # host-to-device copy for the indptr buffer
1719
- self._paged_kv_indptr_buf.copy_(indptr_cpu, non_blocking=True)
1720
- # host-to-device copy for the last_page_len buffer
1721
- self._paged_kv_last_page_len_buf.copy_(last_page_len_cpu, non_blocking=True)
1722
-
1723
- qo_indptr_host = _get_range_buf(batch_size + 1, "cpu")
1724
-
1725
- try:
1726
- # Make sure we pass exactly 19 arguments for tensor core version
1727
- args = [
1728
- self._float_workspace_buffer,
1729
- self._int_workspace_buffer,
1730
- self._pin_memory_int_workspace_buffer,
1731
- qo_indptr_host,
1732
- indptr_cpu,
1733
- seq_lens_cpu,
1734
- batch_size, # total_num_rows
1735
- batch_size,
1736
- num_qo_heads,
1737
- num_kv_heads,
1738
- page_size,
1739
- self.is_cuda_graph_enabled,
1740
- head_dim,
1741
- head_dim,
1742
- False, # causal
1743
- window_left,
1744
- ]
1745
- if self._backend == "fa2":
1746
- args.append(fixed_split_size)
1747
- args.append(disable_split_kv)
1748
- args.append(0) # num_colocated_ctas
1749
- self._plan_info = self._cached_module.plan(
1750
- *args,
1751
- )
1752
- except Exception as e:
1753
- raise RuntimeError(f"Error in tensor core plan: {e}") from e
1754
-
1755
- self._pos_encoding_mode = pos_encoding_mode
1756
- self._window_left = window_left
1757
- self._logits_soft_cap = logits_soft_cap
1758
- self._sm_scale = sm_scale
1759
- self._rope_scale = rope_scale
1760
- self._rope_theta = rope_theta
1761
-
1714
1715
@triton.jit
1716
def _copy_page_indices_kernel(