flashinfer.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 2,813 lines 118 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	"""Attention layer with FlashInfer."""
4
5	from dataclasses import dataclass, replace
6	from enum import Enum
7	from functools import partial
8	from typing import ClassVar
9
10	import numpy as np
11	import torch
12	from flashinfer import (
13	BatchAttentionWithAttentionSinkWrapper,
14	BatchDecodeWithPagedKVCacheWrapper,
15	BatchPrefillWithPagedKVCacheWrapper,
16	BatchPrefillWithRaggedKVCacheWrapper,
17	MultiLevelCascadeAttentionWrapper,
18	get_seq_lens,
19	)
20	from flashinfer.decode import fast_decode_plan, trtllm_batch_decode_with_kv_cache
21	from flashinfer.prefill import trtllm_batch_context_with_kv_cache
22	from flashinfer.utils import FP4Tensor
23	from typing_extensions import override
24
25	from vllm import _custom_ops as custom_ops
26	from vllm import envs
27	from vllm.config import (
28	CUDAGraphMode,
29	VllmConfig,
30	get_current_vllm_config_or_none,
31	)
32	from vllm.config.cache import CacheDType
33	from vllm.distributed.parallel_state import (
34	get_dcp_group,
35	get_dcp_world_size_and_rank,
36	)
37	from vllm.logger import init_logger
38	from vllm.model_executor.layers.quantization.utils.quant_utils import (
39	QuantKey,
40	kFp8StaticTensorSym,
41	kNvfp4Dynamic,
42	)
43	from vllm.platforms import current_platform
44	from vllm.platforms.interface import DeviceCapability
45	from vllm.triton_utils import tl, triton
46	from vllm.utils.flashinfer import (
47	can_use_trtllm_attention,
48	flashinfer_xqa_batch_decode_with_kv_cache,
49	force_use_trtllm_attention,
50	pin_host_range_buf,
51	supports_trtllm_attention,
52	use_trtllm_attention,
53	)
54	from vllm.utils.gpu_sync_debug import gpu_sync_allowed
55	from vllm.utils.math_utils import cdiv
56	from vllm.utils.torch_utils import (
57	PIN_MEMORY,
58	canonicalize_singleton_dim_strides,
59	get_dtype_size,
60	is_quantized_kv_cache,
61	is_strictly_contiguous,
62	nvfp4_kv_cache_full_dim,
63	nvfp4_split_data_scale,
64	)
65	from vllm.v1.attention.backend import (
66	AttentionBackend,
67	AttentionCGSupport,
68	AttentionImpl,
69	AttentionMetadataBuilder,
70	AttentionType,
71	CommonAttentionMetadata,
72	MultipleOf,
73	max_decode_query_len,
74	)
75	from vllm.v1.attention.backends.utils import (
76	get_dcp_local_seq_lens,
77	get_flashinfer_layout_string,
78	get_num_attention_heads_from_layers,
79	get_per_layer_parameters,
80	infer_global_hyperparameters,
81	log2_lse_to_ln,
82	split_decodes_and_prefills,
83	)
84	from vllm.v1.attention.ops.dcp import (
85	cp_lse_ag_out_rs,
86	dcp_a2a_lse_reduce,
87	)
88	from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
89	from vllm.v1.kv_cache_interface import (
90	AttentionSpec,
91	KVCacheLayout,
92	KVCacheSpec,
93	KVQuantMode,
94	iter_layer_specs,
95	)
96	from vllm.v1.utils import CpuGpuBuffer
97
98	FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT = 2048 * 1024 * 1024
99	FLASHINFER_PREFILL_WORKSPACE_BYTES_PER_ELEM = 16
100
101	FP8_DTYPE = current_platform.fp8_dtype()
102	FP4_DTYPE = torch.uint8
103
104	logger = init_logger(__name__)
105
106	trtllm_workspace_buffer = None
107
108
109	def _get_trtllm_workspace_buffer():
110	global trtllm_workspace_buffer
111	if trtllm_workspace_buffer is None:
112	trtllm_workspace_buffer = torch.zeros(
113	envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8, device="cuda"
114	)
115	return trtllm_workspace_buffer
116
117
118	def _pack_draft_block_bool_mask(
119	bool_mask: torch.Tensor, num_packed: int
120	) -> torch.Tensor:
121	"""Pack a boolean draft mask into XQA's uint16 layout."""
122	num_rows = bool_mask.shape[0]
123	bits = 1 << torch.arange(32, device=bool_mask.device, dtype=torch.int64)
124	mask_u32 = (
125	(bool_mask.view(num_rows, num_packed, 32).to(torch.int64) * bits)
126	.sum(dim=-1)
127	.to(torch.uint32)
128	)
129	return mask_u32.view(torch.uint16).reshape(num_rows, num_packed * 2)
130
131
132	def _make_xqa_draft_block_mask(
133	q_len: int, causal: bool, device: torch.device
134	) -> torch.Tensor:
135	"""Build a packed XQA draft mask for a uniform request."""
136	num_packed = (q_len + 31) // 32
137	padded = num_packed * 32
138	q_idx = torch.arange(q_len, device=device).unsqueeze(1)
139	kv_idx = torch.arange(padded, device=device).unsqueeze(0)
140	bool_mask = kv_idx <= q_idx if causal else (kv_idx < q_len).expand(q_len, padded)
141	return _pack_draft_block_bool_mask(bool_mask, num_packed)
142
143
144	def _make_xqa_ragged_draft_block_mask(
145	q_lens: list[int], max_q_len: int, causal: bool, device: torch.device
146	) -> torch.Tensor:
147	"""Build a packed XQA draft mask for a ragged batch."""
148	num_packed = (max_q_len + 31) // 32
149	padded = num_packed * 32
150	q_lens_t = torch.tensor(q_lens, device=device)
151	row_lens = torch.repeat_interleave(q_lens_t, q_lens_t)
152	request_starts = torch.cumsum(q_lens_t, dim=0) - q_lens_t
153	row_starts = torch.repeat_interleave(request_starts, q_lens_t)
154	q_idx = torch.arange(sum(q_lens), device=device) - row_starts
155	kv_idx = torch.arange(padded, device=device).unsqueeze(0)
156	bool_mask = kv_idx < row_lens.unsqueeze(1)
157	if causal:
158	bool_mask &= kv_idx <= q_idx.unsqueeze(1)
159	return _pack_draft_block_bool_mask(bool_mask, num_packed)
160
161
162	@triton.jit
163	def _trtllm_prefill_attn_kvfp8_dequant(
164	kv_cache_ptr,
165	block_tables_prefill_ptr,
166	block_table_stride,
167	mock_kv_cache_ptr,
168	k_scale_ptr,
169	v_scale_ptr,
170	src_stride_page,
171	src_stride_kv,
172	src_stride_head,
173	src_stride_block,
174	src_stride_head_size,
175	DST_K_CACHE_STRIDE: tl.constexpr,
176	DST_KV_CACHE_STRIDE: tl.constexpr,
177	HEAD_STRIDE: tl.constexpr,
178	HEAD_SIZE: tl.constexpr,
179	NUM_KV_HEADS: tl.constexpr,
180	):
181	batch_idx = tl.program_id(0).to(tl.int64)
182	mock_block_table_idx = tl.program_id(1).to(tl.int64)
183	orig_page_num = tl.load(
184	block_tables_prefill_ptr + batch_idx * block_table_stride + mock_block_table_idx
185	).to(tl.int64)
186	if orig_page_num <= 0:
187	return
188	dequant_dtype = mock_kv_cache_ptr.dtype.element_ty
189
190	k_scale_val = tl.load(k_scale_ptr)
191	v_scale_val = tl.load(v_scale_ptr)
192
193	mock_page_idx = batch_idx * block_table_stride + mock_block_table_idx + 1
194	logical_offsets = tl.arange(0, HEAD_STRIDE)
195	block_offsets = logical_offsets // HEAD_SIZE
196	head_size_offsets = logical_offsets % HEAD_SIZE
197
198	for h in range(NUM_KV_HEADS):
199	h_off = tl.cast(h, tl.int64)
200
201	# Read K from source (supports non-contiguous page/kv/head strides)
202	src_k = (
203	orig_page_num * src_stride_page
204	+ h_off * src_stride_head
205	+ block_offsets * src_stride_block
206	+ head_size_offsets * src_stride_head_size
207	)
208	fp8_k = tl.load(kv_cache_ptr + src_k)
209	dequant_k = (fp8_k.to(tl.float32) * k_scale_val).to(dequant_dtype)
210
211	# Write K to contiguous mock cache
212	dst_k = mock_page_idx * DST_KV_CACHE_STRIDE + h * HEAD_STRIDE + logical_offsets
213	tl.store(mock_kv_cache_ptr + dst_k, dequant_k)
214
215	# Read V from source (offset by src_stride_kv for the V half)
216	src_v = (
217	orig_page_num * src_stride_page
218	+ src_stride_kv
219	+ h_off * src_stride_head
220	+ block_offsets * src_stride_block
221	+ head_size_offsets * src_stride_head_size
222	)
223	fp8_v = tl.load(kv_cache_ptr + src_v)
224	dequant_v = (fp8_v.to(tl.float32) * v_scale_val).to(dequant_dtype)
225
226	# Write V to contiguous mock cache
227	dst_v = (
228	mock_page_idx * DST_KV_CACHE_STRIDE
229	+ DST_K_CACHE_STRIDE
230	+ h * HEAD_STRIDE
231	+ logical_offsets
232	)
233	tl.store(mock_kv_cache_ptr + dst_v, dequant_v)
234
235
236	def trtllm_prefill_attn_kvfp8_dequant(
237	kv_cache: torch.Tensor,
238	block_tables_prefill: torch.Tensor,
239	k_scale: torch.Tensor,
240	v_scale: torch.Tensor,
241	dequant_dtype: torch.dtype,
242	) -> tuple[torch.Tensor, torch.Tensor]:
243	batch_size, num_of_page_per_token = block_tables_prefill.shape
244	s = kv_cache.shape
245	assert s[1] == 2
246	assert dequant_dtype in (torch.bfloat16, torch.float16)
247
248	# Logical source layout is (B, 2, H, N, D). The tensor may be a
249	# non-contiguous view, so the Triton kernel indexes it with actual strides.
250	strides = kv_cache.stride()
251	num_kv_heads, block_size, head_size = s[2], s[3], s[4]
252	head_stride = block_size * head_size
253	k_cache_stride = num_kv_heads * head_stride
254	kv_cache_stride = k_cache_stride * s[1]
255
256	new_s = (batch_size * num_of_page_per_token + 1, s[1], s[2], s[3], s[4])
257	# mock kv cache contains just the pages needed by this prefill
258	mock_kv_cache = torch.empty(new_s, dtype=dequant_dtype, device=kv_cache.device)
259	# we simply sequentially index the pages needed by this prefill
260	mock_block_table = torch.arange(
261	start=1,
262	end=batch_size * num_of_page_per_token + 1,
263	dtype=torch.int32,
264	device=block_tables_prefill.device,
265	).reshape(batch_size, num_of_page_per_token)
266	grid = (batch_size, num_of_page_per_token)
267	_trtllm_prefill_attn_kvfp8_dequant[grid](
268	kv_cache,
269	block_tables_prefill,
270	num_of_page_per_token,
271	mock_kv_cache,
272	k_scale,
273	v_scale,
274	strides[0],
275	strides[1],
276	strides[2],
277	strides[3],
278	strides[4],
279	k_cache_stride,
280	kv_cache_stride,
281	head_stride,
282	head_size,
283	num_kv_heads,
284	)
285	return mock_kv_cache, mock_block_table
286
287
288	class BatchDCPPrefillWrapper:
289	def __init__(
290	self,
291	kv_layout: str,
292	workspace_buffer: torch.Tensor \| None = None,
293	dcp_a2a: bool = False,
294	):
295	if dcp_a2a:
296	self._dcp_combine = partial(dcp_a2a_lse_reduce, is_lse_base_on_e=False)
297	else:
298	self._dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
299	self._context = BatchPrefillWithPagedKVCacheWrapper(workspace_buffer, kv_layout)
300	self._new_tokens = BatchPrefillWithRaggedKVCacheWrapper(workspace_buffer)
301
302	def plan(
303	self,
304	qo_indptr_cpu: torch.Tensor,
305	paged_kv_indptr_cpu: torch.Tensor,
306	paged_kv_indices: torch.Tensor,
307	paged_kv_last_page_len_cpu: torch.Tensor,
308	kv_lens_cpu: torch.Tensor,
309	page_size: int,
310	num_qo_heads: int,
311	dcp_world_size: int,
312	num_kv_heads: int,
313	head_dim: int,
314	sm_scale: float,
315	window_left: int,
316	logits_soft_cap: float \| None,
317	q_data_type: torch.dtype,
318	kv_cache_dtype: torch.dtype,
319	prefill_fixed_split_size: int,
320	disable_split_kv: bool,
321	):
322	"""Plan the prefill operation with given parameters."""
323	self._context.plan(
324	qo_indptr=qo_indptr_cpu,
325	paged_kv_indptr=paged_kv_indptr_cpu,
326	paged_kv_indices=paged_kv_indices,
327	paged_kv_last_page_len=paged_kv_last_page_len_cpu,
328	seq_lens=kv_lens_cpu,
329	num_qo_heads=num_qo_heads * dcp_world_size,
330	num_kv_heads=num_kv_heads,
331	head_dim_qk=head_dim,
332	page_size=page_size,
333	causal=False, # This is context run
334	sm_scale=sm_scale,
335	window_left=window_left,
336	logits_soft_cap=logits_soft_cap,
337	q_data_type=q_data_type,
338	kv_data_type=kv_cache_dtype,
339	fixed_split_size=prefill_fixed_split_size,
340	disable_split_kv=disable_split_kv,
341	)
342	self._new_tokens.plan(
343	qo_indptr=qo_indptr_cpu,
344	kv_indptr=qo_indptr_cpu,
345	num_qo_heads=num_qo_heads,
346	num_kv_heads=num_kv_heads,
347	head_dim_qk=head_dim,
348	head_dim_vo=head_dim,
349	causal=True, # This is newtokens run
350	sm_scale=sm_scale,
351	window_left=window_left,
352	logits_soft_cap=logits_soft_cap,
353	q_data_type=q_data_type,
354	)
355
356	def run(
357	self,
358	layer: torch.nn.Module,
359	prefill_query: torch.Tensor,
360	kv_cache_tuple: tuple[torch.Tensor, torch.Tensor],
361	key: torch.Tensor,
362	value: torch.Tensor,
363	out: torch.Tensor,
364	):
365	prefill_query_across_dcp = get_dcp_group().all_gather(
366	prefill_query.contiguous(), dim=1
367	)
368	output_context_tmp, lse_context_tmp = self._context.run(
369	prefill_query_across_dcp,
370	kv_cache_tuple,
371	k_scale=layer._k_scale_float,
372	v_scale=layer._v_scale_float,
373	return_lse=True,
374	)
375	output_context, lse_context = self._dcp_combine(
376	output_context_tmp,
377	lse_context_tmp,
378	get_dcp_group(),
379	return_lse=True,
380	)
381	lse_context = log2_lse_to_ln(lse_context.transpose(0, 1).contiguous())
382
383	output_query, lse_query = self._new_tokens.run(
384	prefill_query,
385	key,
386	value,
387	return_lse=True,
388	)
389	lse_query = log2_lse_to_ln(lse_query.transpose(0, 1).contiguous())
390
391	merge_attn_states(
392	out,
393	output_context,
394	lse_context,
395	output_query,
396	lse_query,
397	)
398	return out
399
400
401	class FlashInferBackend(AttentionBackend):
402	@classmethod
403	def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":
404	"""NVFP4 stores K and V as separate per-head slots of packed fp4 data
405	plus fp8 block scales."""
406	if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_nvfp4:
407	return spec
408	hs_k = nvfp4_kv_cache_full_dim(spec.head_size)
409	hs_v = nvfp4_kv_cache_full_dim(spec.head_size_v)
410	assert hs_k == hs_v, "nvfp4 with asymmetric K/V head sizes not yet supported"
411	return replace(
412	spec,
413	num_head_slots=2 * spec.num_kv_heads,
414	state_content_bytes=hs_k * get_dtype_size(spec.dtype),
415	)
416
417	supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
418	supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
419	"auto",
420	"float16",
421	"bfloat16",
422	"fp8",
423	"fp8_e4m3",
424	"fp8_e5m2",
425	"nvfp4",
426	"nvfp4_4over6",
427	]
428
429	@staticmethod
430	def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int \| MultipleOf]:
431	# Page sizes >= 128 only run on the trtllm-gen dynamic kernel (GQA/MQA
432	# on Blackwell); advertise them only when usable so selection never
433	# picks a large kernel block we cannot serve.
434	use_large_pages = False
435	vllm_config = get_current_vllm_config_or_none()
436	if vllm_config is not None and vllm_config.model_config is not None:
437	pc = vllm_config.parallel_config
438	mc = vllm_config.model_config
439	num_qo_heads = mc.get_num_attention_heads(pc)
440	num_kv_heads = mc.get_num_kv_heads(pc)
441	use_large_pages = (
442	num_kv_heads > 0
443	and num_qo_heads // num_kv_heads > 1
444	and current_platform.is_device_capability_family(100)
445	and can_use_trtllm_attention(num_qo_heads, num_kv_heads)
446	)
447	if not use_large_pages:
448	return [16, 32, 64]
449	return [16, 32, 64, 128, 256, 512, 1024]
450
451	@staticmethod
452	def get_name() -> str:
453	return "FLASHINFER"
454
455	@classmethod
456	def supports_non_causal(cls) -> bool:
457	return True
458
459	@classmethod
460	def supports_device_cpu_query_lens_mismatch(cls) -> bool:
461	# Only trtllm-gen decode (SM100 family) reads query offsets from the
462	# device tensor; XQA and wrapper decode plan from qo_indptr_cpu, and
463	# under DCP spec queries are routed to the prefill path.
464	if not current_platform.is_device_capability_family(100):
465	return False
466	vllm_config = get_current_vllm_config_or_none()
467	if vllm_config is None or vllm_config.model_config is None:
468	return True
469	pc = vllm_config.parallel_config
470	mc = vllm_config.model_config
471	return pc.decode_context_parallel_size == 1 and can_use_trtllm_attention(
472	mc.get_num_attention_heads(pc), mc.get_num_kv_heads(pc), is_prefill=False
473	)
474
475	@classmethod
476	def supports_sliding_window(cls) -> bool:
477	return True
478
479	@staticmethod
480	def get_impl_cls() -> type["FlashInferImpl"]:
481	return FlashInferImpl
482
483	@staticmethod
484	def get_builder_cls() -> type["FlashInferMetadataBuilder"]:
485	return FlashInferMetadataBuilder
486
487	@staticmethod
488	def get_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype:
489	if kv_cache_dtype in ("fp8", "fp8_e4m3"):
490	return torch.float8_e4m3fn
491	elif kv_cache_dtype == "fp8_e5m2":
492	return torch.float8_e5m2
493	elif kv_cache_dtype.startswith("nvfp4"):
494	return torch.uint8
495	else:
496	raise ValueError(f"Unrecognized dtype: {kv_cache_dtype}")
497
498	@classmethod
499	def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType \| None) -> bool:
500	if kv_cache_dtype is not None and kv_cache_dtype.startswith("nvfp4"):
501	return (
502	current_platform.is_device_capability_family(100)
503	and supports_trtllm_attention(is_prefill=True)
504	and supports_trtllm_attention(is_prefill=False)
505	)
506	return super().supports_kv_cache_dtype(kv_cache_dtype)
507
508	@classmethod
509	def get_supported_head_sizes(cls) -> list[int]:
510	# https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157
511	return [64, 128, 256, 512]
512
513	@classmethod
514	def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
515	# FlashInfer supports SM75+, but is currently broken on SM75 (Turing):
516	# https://github.com/flashinfer-ai/flashinfer/issues/3620 (fix:
517	# https://github.com/flashinfer-ai/flashinfer/pull/3621). Temporarily
518	# raise the floor to SM80 so it is not auto-selected on SM75 until
519	# that fix lands; revert to DeviceCapability(7, 5) once it does.
520	return capability >= DeviceCapability(8, 0) and capability <= DeviceCapability(
521	12, 1
522	)
523
524	@classmethod
525	def supports_sink(cls) -> bool:
526	"""Whether FlashInfer can serve attention sinks on this platform."""
527	from vllm.utils.flashinfer import (
528	force_use_trtllm_attention,
529	)
530
531	if force_use_trtllm_attention() is False:
532	return False
533
534	if current_platform.is_device_capability(
535	90
536	) or current_platform.is_device_capability_family(120):
537	return supports_trtllm_attention(is_prefill=False)
538
539	if not current_platform.is_device_capability_family(100):
540	return False
541
542	return supports_trtllm_attention(
543	is_prefill=False
544	) and supports_trtllm_attention(is_prefill=True)
545
546	@classmethod
547	def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] \| None:
548	capability = current_platform.get_device_capability()
549	if capability is not None and capability.major == 10:
550	# The trtllm-gen kernels consume head-major block interiors; the L/B
551	# nesting outside the block is immaterial to them.
552	return (KVCacheLayout.LBHNC, KVCacheLayout.BLHNC)
553	return super().supported_kv_cache_layouts()
554
555	forward_includes_kv_cache_update: bool = False
556
557
558	@dataclass
559	class FIPrefill:
560	"""Metadata for the native FlashInfer prefill pathway (non-TRTLLM)."""
561
562	wrapper: BatchPrefillWithPagedKVCacheWrapper \| BatchDCPPrefillWrapper
563
564
565	@dataclass
566	class FIDecode:
567	"""Metadata for the native FlashInfer decode pathway (non-TRTLLM)."""
568
569	wrapper: BatchDecodeWithPagedKVCacheWrapper
570
571
572	class FlashInferDecodeKernel(Enum):
573	"""Decode kernels selected inside the FlashInfer backend."""
574
575	XQA = "xqa"
576	TRTLLM_GEN = "trtllm-gen"
577
578
579	def _is_xqa_head_dim_supported(head_dim: int) -> bool:
580	return 16 <= head_dim <= 256 and head_dim % 16 == 0
581
582
583	@dataclass
584	class TRTLLMPrefill:
585	"""Metadata for the TRTLLM prefill pathway."""
586
587	block_tables: torch.Tensor
588	"""
589	The slice of the block table tensor corresponding only to prefill requests.
590	Shape: [num_prefills, max_num_blocks_per_seq]
591	"""
592
593	seq_lens: torch.Tensor
594	"""
595	The slice of the sequence lengths tensor corresponding only to prefill requests.
596	Shape: [num_prefills]
597	"""
598
599	cum_seq_lens_q: torch.Tensor
600	cum_seq_lens_kv: torch.Tensor
601
602	max_q_len: int
603	"""
604	The maximum query length among prefill requests.
605	"""
606
607	max_seq_len: int
608	"""The maximum sequence length for KV Cache."""
609
610
611	@dataclass
612	class FlashInferTrtllmAPIDecode:
613	"""Metadata for XQA and trtllm-gen decode."""
614
615	kernel: FlashInferDecodeKernel
616
617	block_tables: torch.Tensor
618	"""
619	The slice of the block table tensor corresponding only to decode requests.
620	Shape: [num_decodes, max_num_blocks_per_seq]
621	"""
622
623	seq_lens: torch.Tensor
624	"""
625	The slice of the sequence lengths tensor corresponding only to decode requests.
626	Shape: [num_decodes]
627	"""
628
629	max_seq_len: int
630	"""The maximum sequence length for KV Cache."""
631
632	q_len_per_req: int = 1
633	"""Query tokens per request for uniform speculative decode."""
634
635	q_cu_seq_lens: torch.Tensor \| None = None
636	"""Cumulative query lengths for ragged XQA decode."""
637
638	mask: torch.Tensor \| None = None
639	"""Packed XQA draft mask."""
640
641	dcp_query_start_loc: torch.Tensor \| None = None
642	"""Decode query boundaries for masking empty DCP shards during reduction."""
643
644
645	@dataclass
646	class FlashInferMetadata:
647	kv_cache_layout: KVCacheLayout
648	num_actual_tokens: int
649	"""Total number of tokens in the batch (excluding padding)."""
650
651	slot_mapping: torch.Tensor
652	"""Tensor for writing K/V to the cache. Shape: [num_actual_tokens]"""
653
654	# The data types of the query for prefill and decode.
655	# On SM90, these two data types may be different.
656	q_data_type_prefill: torch.dtype
657	q_data_type_decode: torch.dtype
658
659	num_decodes: int
660	num_decode_tokens: int
661	num_prefills: int
662	num_prefill_tokens: int
663	causal: bool
664
665	prefill: FIPrefill \| TRTLLMPrefill \| None
666	"""
667	Holds the metadata for the prefill portion of the batch.
668	Will be `None` if `num_prefill_tokens == 0`.
669	"""
670
671	decode: FIDecode \| FlashInferTrtllmAPIDecode \| None
672	"""
673	Holds the metadata for the decode portion of the batch.
674	Will be `None` if `num_decode_tokens == 0`.
675	"""
676
677	# --- Special Case: Cascade Attention ---
678
679	use_cascade: bool
680	"""
681	If True, the entire batch is a cascade attention call, and the
682	`prefill` and `decode` fields will both be None.
683	"""
684
685	cascade_wrapper: MultiLevelCascadeAttentionWrapper \| None
686
687
688	class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
689	kv_cache_spec: AttentionSpec
690	reorder_batch_threshold: int = 1
691
692	def __init__(
693	self,
694	kv_cache_spec: AttentionSpec,
695	layer_names: list[str],
696	vllm_config: VllmConfig,
697	device: torch.device,
698	):
699	super().__init__(kv_cache_spec, layer_names, vllm_config, device)
700	self.cache_config = vllm_config.cache_config
701	self.model_config = vllm_config.model_config
702	self.attention_config = vllm_config.attention_config
703	self._num_speculative_tokens = vllm_config.num_speculative_tokens
704	self._workspace_buffer = None
705	self._prefill_wrapper: (
706	BatchPrefillWithPagedKVCacheWrapper \| BatchDCPPrefillWrapper \| None
707	) = None # Wrapper for prefill/append
708	self._noncausal_prefill_wrapper: BatchPrefillWithPagedKVCacheWrapper \| None = (
709	None # Wrapper for non-causal prefill (DFlash)
710	)
711	self._decode_wrapper = None # Wrapper for decode (general shape)
712
713	if envs.VLLM_BATCH_INVARIANT:
714	self.decode_fixed_split_size = 2048
715	self.prefill_fixed_split_size = 4096
716	self.disable_split_kv = True
717	else:
718	self.decode_fixed_split_size = -1
719	self.prefill_fixed_split_size = -1
720	self.disable_split_kv = False
721
722	self.compilation_config = vllm_config.compilation_config
723	self.max_num_batched_tokens = (
724	vllm_config.scheduler_config.max_num_batched_tokens
725	)
726	max_num_pages_per_req = cdiv(
727	self.model_config.max_model_len, self.kv_cache_spec.block_size
728	)
729	max_num_reqs = vllm_config.scheduler_config.max_num_seqs
730	self.max_num_reqs = max_num_reqs
731	max_num_pages = max_num_reqs * max_num_pages_per_req
732	# Persistent uniform masks keep stable addresses for CUDA graphs.
733	self._decode_mask_cache: dict[tuple[int, bool], torch.Tensor] = {}
734	speculative_config = vllm_config.speculative_config
735	num_spec_tokens = (
736	speculative_config.num_speculative_tokens
737	if speculative_config is not None
738	else 0
739	)
740	self.enable_cuda_graph = (
741	self.compilation_config.cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
742	)
743	if self.enable_cuda_graph:
744	# For full cudagraph capture, one `decode_wrapper` for each batch
745	# size is needed for FlashInfer.
746	self._decode_wrappers_cudagraph: dict[
747	int, BatchDecodeWithPagedKVCacheWrapper
748	] = {}
749	self._decode_cudagraph_max_bs = (1 + num_spec_tokens) * max_num_reqs
750	if self.compilation_config.max_cudagraph_capture_size is not None:
751	self._decode_cudagraph_max_bs = min(
752	self._decode_cudagraph_max_bs,
753	self.compilation_config.max_cudagraph_capture_size,
754	)
755	self.dcp_world_size, self.dcp_rank = get_dcp_world_size_and_rank(
756	kv_cache_spec.dcp_sharded
757	)
758	self.dcp_kv_cache_interleave_size = (
759	vllm_config.parallel_config.cp_kv_cache_interleave_size
760	)
761	self.use_dcp = self.dcp_world_size > 1
762	self.dcp_a2a = (
763	self.use_dcp and vllm_config.parallel_config.dcp_comm_backend == "a2a"
764	)
765
766	# Compatible with models with non-uniform per-layer head counts.
767	self.num_qo_heads = get_num_attention_heads_from_layers(
768	vllm_config, layer_names
769	) or self.model_config.get_num_attention_heads(self.vllm_config.parallel_config)
770
771	self.num_kv_heads = self.kv_cache_spec.num_kv_heads
772	self.head_dim = self.kv_cache_spec.head_size
773	self.page_size = self.kv_cache_spec.block_size
774
775	if self.kv_cache_spec.kv_quant_mode != KVQuantMode.NONE:
776	self.cache_dtype = self.cache_config.cache_dtype
777	# Cannot use self.kv_cache_spec.dtype here because kv_cache_spec
778	# storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3)
779	self.is_kvcache_nvfp4 = self.cache_dtype.startswith("nvfp4")
780	if self.is_kvcache_nvfp4:
781	if (
782	force_use_trtllm_attention() is False
783	or not supports_trtllm_attention(is_prefill=True)
784	or not supports_trtllm_attention(is_prefill=False)
785	):
786	raise ValueError(
787	f"--kv-cache-dtype {self.cache_dtype} requires the "
788	"SM100 trtllm-gen "
789	"FlashInfer path."
790	)
791	# The scale search only affects the store kernel. FlashInfer
792	# reads both variants using the same NVFP4 layout.
793	self.kv_cache_dtype = "nvfp4"
794	else:
795	self.kv_cache_dtype = FlashInferBackend.get_dtype_for_flashinfer(
796	self.cache_dtype
797	)
798	else:
799	self.cache_dtype = "auto"
800	self.is_kvcache_nvfp4 = False
801	assert self.kv_cache_spec.dtype == self.model_config.dtype
802	self.kv_cache_dtype = self.kv_cache_spec.dtype
803
804	# Compute per-phase Q dtype. On SM90 (XQA decode), the prefill and
805	# decode phases require different Q dtypes when the KV cache is FP8
806	# (FP8-Q for the FI native prefill, BF16/FP16-Q for XQA decode),
807	# so both values must be tracked independently.
808	self.q_data_type_prefill = self.get_q_data_type(is_prefill=True)
809	self.q_data_type_decode = self.get_q_data_type(is_prefill=False)
810
811	# Prefer TRTLLM/XQA for decoding whenever supported. The decode kernel
812	# must be selected statically for FULL cudagraph capture.
813	can_use_xqa_or_trtllm_gen_decode = can_use_trtllm_attention(
814	self.num_qo_heads, self.num_kv_heads, is_prefill=False
815	)
816	# Page sizes >= 128 require the trtllm-gen GQA/MQA path (guaranteed by
817	# get_supported_kernel_block_sizes).
818	assert self.page_size <= 64 or (
819	current_platform.is_device_capability_family(100)
820	and can_use_xqa_or_trtllm_gen_decode
821	and self.num_qo_heads // self.num_kv_heads > 1
822	), f"Unexpected FlashInfer page size {self.page_size} without trtllm-gen GQA"
823	self.use_trtllm_decode_attention = can_use_xqa_or_trtllm_gen_decode
824	self.flashinfer_trtllm_api_decode_kernel: FlashInferDecodeKernel \| None = (
825	self._get_flashinfer_trtllm_api_decode_kernel()
826	if can_use_xqa_or_trtllm_gen_decode
827	else None
828	)
829	# The dedicated FlashInfer XQA API accepts head dimensions in
830	# [16, 256] that are divisible by 16. Some hybrid-attention models
831	# (for example Gemma 4) use XQA-compatible sliding-attention groups
832	# alongside global-attention groups with head_dim=512. Resolve the
833	# backend per KV-cache group so the eligible groups still use XQA and
834	# the wider groups fall back to native FlashInfer decode.
835	if (
836	self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
837	and not _is_xqa_head_dim_supported(self.head_dim)
838	):
839	logger.warning_once(
840	"FlashInfer XQA decode does not support head_dim=%d; "
841	"reverting this KV-cache group to native FlashInfer decode.",
842	self.head_dim,
843	)
844	self.use_trtllm_decode_attention = False
845	self.flashinfer_trtllm_api_decode_kernel = None
846	if (
847	self.use_dcp
848	and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
849	):
850	logger.warning_once(
851	"FlashInfer XQA decode does not support returning LSE and "
852	"therefore does not support DCP, reverting to native FlashInfer "
853	"decode."
854	)
855	self.use_trtllm_decode_attention = False
856	self.flashinfer_trtllm_api_decode_kernel = None
857	self.use_xqa = (
858	self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
859	)
860	self.use_trtllm_gen_varlen_decode = self._uses_trtllm_gen_varlen_decode(
861	vllm_config, self.flashinfer_trtllm_api_decode_kernel, self.use_dcp
862	)
863	# Fused draft steps advance seq_lens in place and keep block tables fixed,
864	# so one trtllm-gen decode build serves every step (not DCP-local seq_lens).
865	self.supports_draft_decode_metadata_update = (
866	self.flashinfer_trtllm_api_decode_kernel
867	== FlashInferDecodeKernel.TRTLLM_GEN
868	and not self.use_dcp
869	)
870	self._init_reorder_batch_threshold(
871	1,
872	supports_spec_as_decode=(
873	self.flashinfer_trtllm_api_decode_kernel is not None
874	),
875	# trtllm-gen decode receives no cp_rank/global-seq-len information,
876	# so its end-aligned causal mask is wrong for q_len > 1 over the
877	# DCP-interleaved local KV shard (spec token i misses up to
878	# (dcp_world_size - 1) * (q_len - 1 - i) KV entries, including its
879	# own). Keep the threshold at 1 under DCP so spec queries take the
880	# DCP-aware prefill path, until the kernel is CP-aware (compare
881	# flash_attn_varlen_func's cp_world_size/cp_rank/cp_tot_seqused_k).
882	supports_dcp_with_varlen=False,
883	)
884
885	self._cascade_wrapper = None # Wrapper for cascade attention
886
887	# Global hyperparameters shared by all attention layers
888	# TODO: discard this for trtllm-gen backend
889	per_layer_parameters = get_per_layer_parameters(
890	vllm_config, layer_names, FlashInferImpl
891	)
892	self.global_hyperparameters = infer_global_hyperparameters(per_layer_parameters)
893	self.sm_scale = self.global_hyperparameters.sm_scale
894	self.window_left = self.global_hyperparameters.window_left
895	self.logits_soft_cap = self.global_hyperparameters.logits_soft_cap
896	self.has_sinks = self.global_hyperparameters.has_sinks
897	if self.has_sinks and not FlashInferBackend.supports_sink():
898	raise NotImplementedError(
899	"FlashInfer backend currently does not support attention "
900	"sinks, please use trtllm on blackwell or flash attention on "
901	"earlier GPUs."
902	)
903	capability = current_platform.get_device_capability()
904	arch = f"sm{capability.major}{capability.minor}" if capability else "unknown"
905	decode_backend = (
906	self.flashinfer_trtllm_api_decode_kernel.value
907	if self.flashinfer_trtllm_api_decode_kernel is not None
908	else "flashinfer-native"
909	)
910	logger.info_once(
911	"FlashInfer resolved query dtypes: prefill=%s, decode=%s, "
912	"decode_backend=%s, kv_cache_dtype=%s, arch=%s",
913	self.q_data_type_prefill,
914	self.q_data_type_decode,
915	decode_backend,
916	self.kv_cache_dtype,
917	arch,
918	)
919	# Preparing persistent buffers
920	self.paged_kv_indptr = CpuGpuBuffer(
921	max_num_reqs + 1, dtype=torch.int32, device=self.device, pin_memory=False
922	)
923	self.paged_kv_indices = torch.zeros(
924	max_num_pages, dtype=torch.int32, device=self.device
925	)
926	self.paged_kv_last_page_len = CpuGpuBuffer(
927	max_num_reqs, dtype=torch.int32, device=self.device, pin_memory=False
928	)
929
930	@property
931	def kv_cache_layout(self) -> KVCacheLayout:
932	return self.cache_config.get_resolved_kv_cache_layout()
933
934	# Keep SM90 prefill/decode Q dtype selection in one place.
935	def get_q_data_type(self, is_prefill: bool) -> torch.dtype:
936	# The user sets --attention-config.disable_flashinfer_q_quantization
937	# to 1 explicitly, use model dtype for query.
938	if self.vllm_config.attention_config.disable_flashinfer_q_quantization:
939	return self.model_config.dtype
940
941	# self.cache_dtype is resolved per KV-cache group: it is "auto" when
942	# this group is unquantized (e.g. --kv-cache-dtype-skip-layers), even
943	# if cache_config requests a quantized dtype globally.
944	cache_dtype = self.cache_dtype
945
946	# On SM90/SM12x, XQA decode requires BF16/FP16-Q even with FP8 KV cache.
947	# FI native prefill on SM90 still uses FP8-Q in that case; SM12x prefill
948	# is fa2-only and keeps the model dtype (handled below).
949	if (
950	(
951	current_platform.is_device_capability(90)
952	or current_platform.is_device_capability_family(120)
953	)
954	and not is_prefill
955	and force_use_trtllm_attention() is not False
956	and cache_dtype.startswith("fp8")
957	):
958	return self.model_config.dtype
959
960	# Otherwise, match Q dtype to the KV cache dtype.
961	if cache_dtype.startswith("fp8"):
962	# FP8-Q requires an fp8 tensor-core attention path. SM90 uses
963	# native FA3 for prefill; XQA decode on SM90/SM12x is handled
964	# above with model-dtype Q, while SM100 uses trtllm-gen.
965	# Architectures with only fa2 (e.g. SM89, SM120) cannot
966	# consume FP8 queries, so keep the model dtype for Q there.
967	if current_platform.is_device_capability(
968	90
969	) or current_platform.is_device_capability_family(100):
970	return FlashInferBackend.get_dtype_for_flashinfer(cache_dtype)
971	return self.model_config.dtype
972	if cache_dtype.startswith("nvfp4"):
973	return FlashInferBackend.get_dtype_for_flashinfer("fp8_e4m3")
974	return self.kv_cache_spec.dtype
975
976	@override # type: ignore[misc]
977	@classmethod
978	def get_cudagraph_support(
979	cls: type["FlashInferMetadataBuilder"],
980	vllm_config: VllmConfig,
981	kv_cache_spec: KVCacheSpec,
982	) -> AttentionCGSupport:
983	"""Get the cudagraph support level for FlashInfer attention."""
984	# XQA lacks LSE for DCP; DCP also cannot graph variable-length trtllm-gen.
985	if vllm_config.parallel_config.decode_context_parallel_size > 1:
986	return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
987
988	kv_specs = iter_layer_specs(kv_cache_spec)
989	num_qo_heads = vllm_config.model_config.get_num_attention_heads(
990	vllm_config.parallel_config
991	)
992	is_xqa_arch = current_platform.is_device_capability(
993	90
994	) or current_platform.is_device_capability_family(120)
995	has_uniform_batch_support: bool = len(kv_specs) > 0
996	for spec in kv_specs:
997	if not isinstance(spec, AttentionSpec):
998	# FlashInfer only applies to attention, so we don't consider other types
999	# of KV spec (e.g. Mamba) here. This is mostly for type checking.
1000	continue
1001	if not can_use_trtllm_attention(
1002	num_qo_heads=num_qo_heads,
1003	num_kv_heads=spec.num_kv_heads,
1004	is_prefill=False,
1005	) or (is_xqa_arch and not _is_xqa_head_dim_supported(spec.head_size)):
1006	has_uniform_batch_support = False
1007	break
1008
1009	use_non_causal = vllm_config.attention_config.use_non_causal
1010	if has_uniform_batch_support and (not use_non_causal or is_xqa_arch):
1011	return AttentionCGSupport.UNIFORM_BATCH
1012	else:
1013	return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
1014
1015	@override # type: ignore[misc]
1016	@classmethod
1017	def get_varlen_cudagraph_max_query_len(
1018	cls: type["FlashInferMetadataBuilder"],
1019	vllm_config: VllmConfig,
1020	kv_cache_spec: KVCacheSpec,
1021	) -> int \| None:
1022	# The uniform gate already covers DCP, head geometry, the XQA head-dim
1023	# limit and causality, and yields None for subclasses that force NEVER.
1024	if (
1025	cls.get_cudagraph_support(vllm_config, kv_cache_spec)
1026	!= AttentionCGSupport.UNIFORM_BATCH
1027	):
1028	return None
1029	# The kernel selector asserts off the trtllm-capable architectures, so
1030	# call it only past the gate, as get_cudagraph_support does.
1031	if not cls._uses_trtllm_gen_varlen_decode(
1032	vllm_config,
1033	cls._get_flashinfer_trtllm_api_decode_kernel(),
1034	use_dcp=vllm_config.parallel_config.decode_context_parallel_size > 1,
1035	):
1036	return None
1037	# Longer requests are split off as prefills.
1038	return max_decode_query_len(vllm_config)
1039
1040	def _get_workspace_buffer(self):
1041	if self._workspace_buffer is None:
1042	buffer_size = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE
1043	if envs.VLLM_BATCH_INVARIANT:
1044	buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT
1045	else:
1046	# FlashInfer prefill temp buffers (batch_prefill_tmp_v, ...)
1047	# scale with the prefill chunk and query-head footprint, NOT
1048	# context length. The fixed ~394 MiB default is too small for
1049	# wide-head models at the default 8192-token chunk on some
1050	# archs (e.g. sm_120), where FlashInfer hard-errors instead of
1051	# growing. Size to the batch's head footprint; never shrink
1052	# below the configured default.
1053	est = (
1054	self.max_num_batched_tokens
1055	* self.num_qo_heads
1056	* self.head_dim
1057	* FLASHINFER_PREFILL_WORKSPACE_BYTES_PER_ELEM
1058	)
1059	buffer_size = max(buffer_size, est)
1060	self._workspace_buffer = torch.zeros(
1061	buffer_size, dtype=torch.uint8, device=self.device
1062	)
1063	return self._workspace_buffer
1064
1065	def set_workspace_buffer(self, workspace_buffer: torch.Tensor):
1066	self._workspace_buffer = workspace_buffer
1067
1068	@staticmethod
1069	def _uses_trtllm_gen_varlen_decode(
1070	vllm_config: VllmConfig,
1071	decode_kernel: FlashInferDecodeKernel \| None,
1072	use_dcp: bool,
1073	) -> bool:
1074	"""Whether decode reads per-request query lengths from the device
1075	qo_indptr. Adaptive verification trims drafts on device, and only
1076	trtllm-gen supports that (the selector already rejects the other
1077	configurations)."""
1078	speculative_config = vllm_config.speculative_config
1079	return (
1080	speculative_config is not None
1081	and speculative_config.enable_adaptive_verification
1082	and decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN
1083	and not use_dcp
1084	)
1085
1086	@staticmethod
1087	def _get_flashinfer_trtllm_api_decode_kernel() -> FlashInferDecodeKernel:
1088	if current_platform.is_device_capability(
1089	90
1090	) or current_platform.is_device_capability_family(120):
1091	return FlashInferDecodeKernel.XQA
1092	assert current_platform.is_device_capability_family(100)
1093	return FlashInferDecodeKernel.TRTLLM_GEN
1094
1095	def _compute_decode_query_lens(
1096	self,
1097	qo_indptr: torch.Tensor,
1098	qo_indptr_cpu: torch.Tensor,
1099	num_decodes: int,
1100	num_decode_tokens: int,
1101	) -> tuple[int, torch.Tensor \| None, list[int] \| None]:
1102	"""Return the query width, ragged offsets, and effective query lengths."""
1103	if num_decodes == 0 or num_decode_tokens == 0:
1104	return 1, None, None
1105
1106	if not self.use_xqa:
1107	decode_q_lens = (
1108	qo_indptr_cpu[1 : num_decodes + 1] - qo_indptr_cpu[:num_decodes]
1109	)
1110	q_len_per_req = num_decode_tokens // num_decodes
1111	if num_decode_tokens % num_decodes == 0 and bool(
1112	(decode_q_lens == q_len_per_req).all().item()
1113	):
1114	return q_len_per_req, None, None
1115
1116	# CUDA-graph padding can leave zero-length requests after a
1117	# uniform speculative decode batch. Keep the packed query and
1118	# describe it through the trtllm-gen varlen API.
1119	max_q_len = int(decode_q_lens.max().item())
1120	if max_q_len > 1:
1121	return max_q_len, qo_indptr[: num_decodes + 1], None
1122	return 1, None, None
1123
1124	decode_q_lens = qo_indptr_cpu[1 : num_decodes + 1] - qo_indptr_cpu[:num_decodes]
1125	nonzero = decode_q_lens[decode_q_lens > 0]
1126	q_len_per_req = int(nonzero.max().item()) if nonzero.numel() > 0 else 1
1127	uniform = nonzero.numel() <= 1 or bool((nonzero == nonzero[0]).all().item())
1128	if uniform and num_decode_tokens == num_decodes * q_len_per_req:
1129	return q_len_per_req, None, None
1130
1131	if q_len_per_req <= 1:
1132	return 1, None, None
1133
1134	q_lens = decode_q_lens.tolist()
1135	return (
1136	q_len_per_req,
1137	qo_indptr[: num_decodes + 1],
1138	q_lens,
1139	)
1140
1141	def _get_decode_mask(
1142	self,
1143	q_len_per_req: int,
1144	ragged_q_lens: list[int] \| None,
1145	num_decodes: int,
1146	causal: bool,
1147	) -> torch.Tensor \| None:
1148	"""Return the packed XQA draft mask for speculative decode."""
1149	if q_len_per_req <= 1:
1150	return None
1151
1152	if ragged_q_lens is None:
1153	key = (q_len_per_req, causal)
1154	buf = self._decode_mask_cache.get(key)
1155	if buf is None:
1156	per_req = _make_xqa_draft_block_mask(q_len_per_req, causal, self.device)
1157	buf = (
1158	per_req.unsqueeze(0).expand(self.max_num_reqs, -1, -1).contiguous()
1159	)
1160	self._decode_mask_cache[key] = buf
1161	return buf[:num_decodes]
1162
1163	return _make_xqa_ragged_draft_block_mask(
1164	ragged_q_lens, q_len_per_req, causal, self.device
1165	)
1166
1167	def _get_prefill_wrapper(
1168	self,
1169	causal: bool = True,
1170	) -> BatchPrefillWithPagedKVCacheWrapper \| BatchDCPPrefillWrapper:
1171	if not causal:
1172	if self.use_dcp:
1173	raise NotImplementedError(
1174	"FlashInfer non-causal prefill is not supported with DCP yet."
1175	)
1176	if self.is_kvcache_nvfp4:
1177	raise NotImplementedError(
1178	"FlashInfer non-causal attention is not supported with "
1179	"NVFP4 KV cache."
1180	)
1181	if self._noncausal_prefill_wrapper is None:
1182	if self.has_sinks and (
1183	current_platform.is_device_capability(90)
1184	or current_platform.is_device_capability_family(120)
1185	):
1186	self._noncausal_prefill_wrapper = (
1187	BatchAttentionWithAttentionSinkWrapper(
1188	self._get_workspace_buffer(),
1189	get_flashinfer_layout_string(self.kv_cache_layout),
1190	backend="auto",
1191	q_data_type=self.q_data_type_prefill,
1192	kv_data_type=self.kv_cache_dtype,
1193	head_dim_qk=self.head_dim,
1194	head_dim_vo=self.head_dim,
1195	window_left=self.window_left,
1196	)
1197	)
1198	else:
1199	self._noncausal_prefill_wrapper = (
1200	BatchPrefillWithPagedKVCacheWrapper(
1201	self._get_workspace_buffer(),
1202	get_flashinfer_layout_string(self.kv_cache_layout),
1203	backend="auto",
1204	)
1205	)
1206	return self._noncausal_prefill_wrapper
1207
1208	if self._prefill_wrapper is None:
1209	if self.use_dcp:
1210	self._prefill_wrapper = BatchDCPPrefillWrapper(
1211	kv_layout=get_flashinfer_layout_string(self.kv_cache_layout),
1212	workspace_buffer=self._get_workspace_buffer(),
1213	dcp_a2a=self.dcp_a2a,
1214	)
1215	else:
1216	if self.has_sinks and (
1217	current_platform.is_device_capability(90)
1218	or current_platform.is_device_capability_family(120)
1219	):
1220	assert not self.is_kvcache_nvfp4
1221	self._prefill_wrapper = BatchAttentionWithAttentionSinkWrapper(
1222	self._get_workspace_buffer(),
1223	get_flashinfer_layout_string(self.kv_cache_layout),
1224	backend="auto",
1225	q_data_type=self.q_data_type_prefill,
1226	kv_data_type=self.kv_cache_dtype,
1227	head_dim_qk=self.head_dim,
1228	head_dim_vo=self.head_dim,
1229	window_left=self.window_left,
1230	)
1231	else:
1232	# NVFP4 KV cache requires the trtllm-gen backend inside
1233	# the wrapper; fa2/fa3 do not support nvfp4.
1234	backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
1235	self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
1236	self._get_workspace_buffer(),
1237	get_flashinfer_layout_string(self.kv_cache_layout),
1238	backend=backend,
1239	)
1240	assert self._prefill_wrapper is not None
1241	return self._prefill_wrapper
1242
1243	def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False):
1244	if use_cudagraph:
1245	decode_wrapper = self._decode_wrappers_cudagraph.get(batch_size, None)
1246	else:
1247	decode_wrapper = self._decode_wrapper
1248
1249	if decode_wrapper is None:
1250	if use_cudagraph:
1251	paged_kv_indptr = self.paged_kv_indptr.gpu[: batch_size + 1]
1252	paged_kv_indices = self.paged_kv_indices
1253	paged_kv_last_page_len = self.paged_kv_last_page_len.gpu[:batch_size]
1254	else:
1255	paged_kv_indptr = None
1256	paged_kv_indices = None
1257	paged_kv_last_page_len = None
1258	# NVFP4 KV cache requires the trtllm-gen backend inside
1259	# the wrapper; fa2/fa3 do not support nvfp4.
1260	backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
1261	decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
1262	self._get_workspace_buffer(),
1263	get_flashinfer_layout_string(self.kv_cache_layout),
1264	use_cuda_graph=use_cudagraph,
1265	paged_kv_indptr_buffer=paged_kv_indptr,
1266	paged_kv_indices_buffer=paged_kv_indices,
1267	paged_kv_last_page_len_buffer=paged_kv_last_page_len,
1268	# Tensor cores are enabled by default because the perf would be
1269	# at least as good as cuda cores for all attention ops in latest
1270	# gpus.
1271	use_tensor_cores=True,
1272	backend=backend,
1273	)
1274
1275	# save the decode wrapper
1276	if use_cudagraph:
1277	self._decode_wrappers_cudagraph[batch_size] = decode_wrapper
1278	else:
1279	self._decode_wrapper = decode_wrapper
1280
1281	return decode_wrapper
1282
1283	def _get_cascade_wrapper(self):
1284	if self._cascade_wrapper is None:
1285	self._cascade_wrapper = MultiLevelCascadeAttentionWrapper(
1286	2,
1287	self._get_workspace_buffer(),
1288	get_flashinfer_layout_string(self.kv_cache_layout),
1289	)
1290	return self._cascade_wrapper
1291
1292	def _compute_flashinfer_kv_metadata(
1293	self,
1294	num_blocks_np: np.ndarray,
1295	seq_lens_np: np.ndarray,
1296	block_table_tensor: torch.Tensor,
1297	num_reqs: int,
1298	page_size: int,
1299	) -> torch.Tensor:
1300	"""Compute paged_kv_indptr, paged_kv_indices and paged_kv_last_page_len.
1301
1302	Results are stored in self.paged_kv_indptr,
1303	self.paged_kv_indices, self.paged_kv_last_page_len buffers.
1304
1305	Returns paged_kv_indices, a GPU tensor with shape [num_actual_pages].
1306	"""
1307	# write self.paged_kv_indptr_cpu inplace (0-index is always 0)
1308	np.cumsum(
1309	num_blocks_np, dtype=np.int32, out=self.paged_kv_indptr.np[1 : num_reqs + 1]
1310	)
1311	paged_kv_indptr = self.paged_kv_indptr.copy_to_gpu(num_reqs + 1)
1312
1313	# write self.paged_kv_indices inplace
1314	num_actual_pages = self.paged_kv_indptr.np[num_reqs]
1315	paged_kv_indices = self.paged_kv_indices[:num_actual_pages]
1316	_copy_page_indices_kernel[(num_reqs,)](
1317	paged_kv_indices,
1318	block_table_tensor,
1319	block_table_tensor.stride(0),
1320	paged_kv_indptr,
1321	BLOCK_SIZE=1024,
1322	)
1323
1324	# write self.paged_kv_last_page_len_cpu inplace
1325	paged_kv_last_page_len_np = seq_lens_np % page_size
1326	self.paged_kv_last_page_len.np[:num_reqs] = np.where(
1327	(paged_kv_last_page_len_np == 0) & (seq_lens_np != 0),
1328	page_size,
1329	paged_kv_last_page_len_np,
1330	)
1331	self.paged_kv_last_page_len.copy_to_gpu(num_reqs)
1332	return paged_kv_indices
1333
1334	def build(
1335	self,
1336	common_prefix_len: int,
1337	common_attn_metadata: CommonAttentionMetadata,
1338	fast_build: bool = False,
1339	) -> FlashInferMetadata:
1340	num_reqs = common_attn_metadata.num_reqs
1341	num_actual_tokens = common_attn_metadata.num_actual_tokens
1342	causal = common_attn_metadata.causal
1343	route_decode = causal or self.use_xqa
1344	if route_decode:
1345	num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
1346	split_decodes_and_prefills(
1347	common_attn_metadata,
1348	decode_threshold=self.reorder_batch_threshold,
1349	require_uniform=not (
1350	self.use_xqa or self.use_trtllm_gen_varlen_decode
1351	),
1352	)
1353	)
1354	else:
1355	num_decodes = 0
1356	num_prefills = num_reqs
1357	num_decode_tokens = 0
1358	num_prefill_tokens = num_actual_tokens
1359
1360	page_size = self.page_size
1361	max_seq_len = common_attn_metadata.max_seq_len
1362	seq_lens = common_attn_metadata.seq_lens
1363	block_table_tensor = common_attn_metadata.block_table_tensor
1364	qo_indptr = common_attn_metadata.query_start_loc
1365	qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu
1366
1367	# Step 1: Decide which dispatch modes to use:
1368	# - Cascade attention (distinct mode)
1369	# - Prefill (FI native or TRTLLM)
1370	# - Decode (FI native, XQA, or trtllm-gen)
1371	use_cascade = common_prefix_len > 0
1372	uses_spec_reorder = self.reorder_batch_threshold > 1
1373	# Page sizes >= 128 must use trtllm-gen; force it for prefill too.
1374	prefill_force_trtllm = (
1375	True if page_size >= 128 else self.attention_config.use_trtllm_attention
1376	)
1377	prefill_use_trtllm = causal and use_trtllm_attention(
1378	self.num_qo_heads,
1379	self.num_kv_heads,
1380	num_prefill_tokens,
1381	max_seq_len,
1382	self.dcp_world_size,
1383	self.cache_dtype,
1384	self.q_data_type_prefill,
1385	is_prefill=True,
1386	force_use_trtllm=prefill_force_trtllm,
1387	has_sinks=self.has_sinks,
1388	has_spec=uses_spec_reorder,
1389	)
1390	decode_with_flashinfer_trtllm_api = self.use_trtllm_decode_attention and (
1391	causal or self.use_xqa
1392	)
1393
1394	if not causal and self.use_dcp:
1395	raise NotImplementedError(
1396	"FlashInfer non-causal prefill is not supported with DCP yet."
1397	)
1398	if not causal and self.use_trtllm_decode_attention:
1399	logger.warning_once(
1400	"Using FlashInfer for draft model non-causal attention; TRTLLM "
1401	"can still be used for target model causal attention."
1402	)
1403	all_uses_trtllm = causal and (
1404	(num_prefills == 0 or prefill_use_trtllm)
1405	and (num_decodes == 0 or decode_with_flashinfer_trtllm_api)
1406	)
1407
1408	if not all_uses_trtllm:
1409	sinks_would_be_dropped = self.use_dcp or use_cascade
1410	if self.has_sinks and sinks_would_be_dropped:
1411	raise NotImplementedError(
1412	"FlashInfer BatchDCPPrefillWrapper and "
1413	"MultiLevelCascadeAttentionWrapper do not support attention sinks. "
1414	"Please use TRTLLM on Blackwell or FlashAttention on earlier GPUs."
1415	)
1416
1417	if not self.global_hyperparameters.has_same_window_lefts:
1418	raise ValueError(
1419	"Window left is not the same for all layers. "
1420	"One potential fix is to set disable_sliding_window=True"
1421	)
1422
1423	assert self.global_hyperparameters.has_same_all_params, (
1424	"FlashInfer backend currently only supports models in which "
1425	"all layers share the same values for the following "
1426	"hyperparameters: `window_left`, `logits_soft_cap`, "
1427	"`sm_scale`."
1428	)
1429
1430	# Step 2: Initialize the output metadata
1431	# Leave prefill/decode/cascade_wrapper empty, to be populated
1432	# case by case depending on the batch contents and backend selection.
1433	attn_metadata = FlashInferMetadata(
1434	kv_cache_layout=self.kv_cache_layout,
1435	num_actual_tokens=num_actual_tokens,
1436	slot_mapping=common_attn_metadata.slot_mapping,
1437	q_data_type_prefill=self.q_data_type_prefill,
1438	q_data_type_decode=self.q_data_type_decode,
1439	num_decodes=num_decodes,
1440	num_decode_tokens=num_decode_tokens,
1441	num_prefills=num_prefills,
1442	num_prefill_tokens=num_prefill_tokens,
1443	causal=causal,
1444	use_cascade=use_cascade,
1445	prefill=None,
1446	decode=None,
1447	cascade_wrapper=None,
1448	)
1449
1450	# Guard access to seq_lens_cpu, which may not always be needed
1451	# and can be expensive to retrieve in async mode.
1452	# When all attention (both prefill and decode) uses TRTLLM,
1453	# seq_lens_cpu is not needed since TRTLLM paths use GPU tensors
1454	# (block_tables, seq_lens) directly.
1455	needs_seq_lens_cpu = self.use_dcp or use_cascade or not all_uses_trtllm
1456	if needs_seq_lens_cpu:
1457	if (
1458	self._num_speculative_tokens == 0
1459	and common_attn_metadata.seq_lens_cpu_upper_bound is not None
1460	):
1461	# No speculative token accounting, so the upper bound is exact.
1462	seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
1463	else:
1464	with gpu_sync_allowed():
1465	seq_lens_cpu = common_attn_metadata.seq_lens.cpu()
1466	else:
1467	seq_lens_cpu = None
1468
1469	# Adjust seq_lens_cpu for DCP
1470	if self.use_dcp:
1471	assert seq_lens_cpu is not None
1472	if num_prefills > 0:
1473	# Other attention groups may reuse the same common metadata.
1474	seq_lens_cpu = seq_lens_cpu.clone()
1475	qo_indptr_prefill_cpu = (
1476	qo_indptr_cpu[num_decodes:] - qo_indptr_cpu[num_decodes]
1477	)
1478	query_lens_prefill_cpu = (
1479	qo_indptr_prefill_cpu[1:] - qo_indptr_prefill_cpu[:-1]
1480	)
1481	seq_lens_cpu[num_decodes:] = (
1482	seq_lens_cpu[num_decodes:] - query_lens_prefill_cpu
1483	)
1484
1485	seq_lens_cpu = get_dcp_local_seq_lens(
1486	seq_lens_cpu,
1487	self.dcp_world_size,
1488	self.dcp_rank,
1489	self.dcp_kv_cache_interleave_size,
1490	)
1491
1492	# Native paged attention consumes rank-local lengths (context only
1493	# for DCP prefills), so derive its page counts after the conversion.
1494	if seq_lens_cpu is not None:
1495	seq_lens_np = seq_lens_cpu.numpy()
1496	num_blocks_np = (seq_lens_np + (page_size - 1)) // page_size
1497	else:
1498	seq_lens_np = None
1499	num_blocks_np = None
1500
1501	# Adjust num_block_np for cascade attention
1502	if use_cascade:
1503	assert num_blocks_np is not None
1504	assert common_prefix_len % page_size == 0
1505	num_common_kv_blocks = common_prefix_len // page_size
1506	num_blocks_np -= num_common_kv_blocks
1507
1508	# Compute paged_kv_indices if necessary
1509	# paged_kv_indices is only needed for FlashInfer native paths;
1510	# XQA/trtllm-gen paths use block_tables directly on GPU.
1511	needs_native_paged_prefill = num_prefills > 0 and not prefill_use_trtllm
1512	needs_native_paged_decode = (
1513	num_decodes > 0 and not decode_with_flashinfer_trtllm_api
1514	)
1515	needs_paged_kv_indices = (
1516	use_cascade or needs_native_paged_prefill or needs_native_paged_decode
1517	)
1518	if needs_paged_kv_indices:
1519	assert num_blocks_np is not None
1520	assert seq_lens_np is not None
1521	paged_kv_indices = self._compute_flashinfer_kv_metadata(
1522	num_blocks_np,
1523	seq_lens_np,
1524	block_table_tensor,
1525	num_reqs,
1526	page_size,
1527	)
1528	else:
1529	paged_kv_indices = None
1530
1531	# Early-out for cascade attention
1532	if use_cascade:
1533	assert num_blocks_np is not None
1534	# Grab the blocks of the shared prefix from the first request.
1535	num_common_kv_blocks = common_prefix_len // page_size
1536
1537	# Create CPU versions directly for cascade (no GPU versions needed)
1538	shared_qo_indptr_cpu = torch.tensor(
1539	[0, num_actual_tokens], dtype=torch.int32, device="cpu"
1540	)
1541	shared_kv_page_indptr_cpu = torch.tensor(
1542	[0, num_common_kv_blocks], dtype=torch.int32, device="cpu"
1543	)
1544	shared_kv_page_indices_cpu = block_table_tensor[0, :num_common_kv_blocks]
1545	shared_kv_last_page_len_cpu = torch.tensor(
1546	[page_size], dtype=torch.int32, device="cpu"
1547	)
1548
1549	# Remove the blocks of the shared prefix from all requests.
1550	block_table_tensor = block_table_tensor[:, num_common_kv_blocks:]
1551	num_blocks_np -= num_common_kv_blocks
1552
1553	assert paged_kv_indices is not None
1554	paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: 1 + num_reqs]
1555	paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_reqs]
1556
1557	attn_metadata.cascade_wrapper = self._get_cascade_wrapper()
1558	# Cascade attention must use the same q dtype for prefill and decode
1559	# because it does not support FP8 kv-cache or FP8 query yet.
1560	assert self.q_data_type_prefill == self.q_data_type_decode
1561	attn_metadata.cascade_wrapper.plan(
1562	qo_indptr_arr=[shared_qo_indptr_cpu, qo_indptr_cpu],
1563	paged_kv_indptr_arr=[shared_kv_page_indptr_cpu, paged_kv_indptr_cpu],
1564	paged_kv_indices_arr=[shared_kv_page_indices_cpu, paged_kv_indices],
1565	paged_kv_last_page_len=[
1566	shared_kv_last_page_len_cpu,
1567	paged_kv_last_page_len_cpu,
1568	],
1569	num_qo_heads=self.num_qo_heads,
1570	num_kv_heads=self.num_kv_heads,
1571	head_dim=self.head_dim,
1572	page_size=self.page_size,
1573	causal=True,
1574	sm_scale=self.sm_scale,
1575	window_left=self.window_left,
1576	logits_soft_cap=self.logits_soft_cap,
1577	q_data_type=self.q_data_type_prefill,
1578	kv_data_type=self.kv_cache_dtype,
1579	)
1580	return attn_metadata
1581
1582	# Step 3: Handle prefill and decode pathways case by case
1583	## PREFILL PATHWAY
1584	if num_prefills > 0:
1585	# Slices for shared prefill metadata
1586	prefill_start = num_decodes
1587	qo_indptr_prefill_cpu = (
1588	qo_indptr_cpu[prefill_start:] - qo_indptr_cpu[prefill_start]
1589	)
1590	assert qo_indptr_prefill_cpu.shape[0] == num_prefills + 1
1591
1592	if prefill_use_trtllm:
1593	# TRTLLM prefill has no cross-rank combine for DCP-sharded KV;
1594	# use_trtllm_attention never selects it when DCP is enabled.
1595	assert not self.use_dcp
1596	# Create GPU versions
1597	qo_indptr_prefill_gpu = (
1598	qo_indptr[prefill_start:] - qo_indptr[prefill_start]
1599	)
1600	# Compute cum_seq_lens_kv on GPU to avoid CPU sync.
1601	# This is the cumulative sum of the number of KV cache
1602	# blocks per prefill request.
1603	prefill_seq_lens = seq_lens[prefill_start:]
1604	num_blocks_per_req = (prefill_seq_lens + page_size - 1) // page_size
1605	paged_kv_indptr_prefill_gpu = self.paged_kv_indptr.gpu[
1606	prefill_start : num_reqs + 1
1607	]
1608	# Assign to slice to avoid cpu sync.
1609	paged_kv_indptr_prefill_gpu[:1] = 0
1610	torch.cumsum(
1611	num_blocks_per_req,
1612	dim=0,
1613	out=paged_kv_indptr_prefill_gpu[1:],
1614	)
1615	# Compute max_q_len for prefill requests
1616	query_lens_prefill_cpu = (
1617	qo_indptr_prefill_cpu[1:] - qo_indptr_prefill_cpu[:-1]
1618	)
1619	max_q_len_prefill = int(query_lens_prefill_cpu.max().item())
1620	prefill_block_tables = block_table_tensor[prefill_start:]
1621	if (
1622	self.q_data_type_prefill != FP8_DTYPE
1623	and self.cache_dtype.startswith("fp8")
1624	):
1625	seq_lens_cpu_upper_bound = (
1626	common_attn_metadata.seq_lens_cpu_upper_bound
1627	)
1628	max_prefill_seq_len = max_seq_len
1629	if seq_lens_cpu_upper_bound is not None:
1630	max_prefill_seq_len = int(
1631	seq_lens_cpu_upper_bound[prefill_start:num_reqs]
1632	.max()
1633	.item()
1634	)
1635	# Dequantization allocates one page per table entry and
1636	# indexes rows by their logical width.
1637	prefill_block_tables = canonicalize_singleton_dim_strides(
1638	prefill_block_tables[
1639	:, : cdiv(max_prefill_seq_len, page_size)
1640	].contiguous()
1641	)
1642	attn_metadata.prefill = TRTLLMPrefill(
1643	block_tables=prefill_block_tables,
1644	seq_lens=prefill_seq_lens,
1645	cum_seq_lens_q=qo_indptr_prefill_gpu,
1646	cum_seq_lens_kv=paged_kv_indptr_prefill_gpu,
1647	max_q_len=max_q_len_prefill,
1648	max_seq_len=max_seq_len,
1649	)
1650	else:
1651	prefill_wrapper = self._get_prefill_wrapper(causal=attn_metadata.causal)
1652	# Slicing CPU buffers that are only needed for FI native prefills
1653	paged_kv_last_page_len_prefill_cpu = self.paged_kv_last_page_len.cpu[
1654	prefill_start:num_reqs
1655	]
1656	assert paged_kv_last_page_len_prefill_cpu.shape[0] == num_prefills
1657	paged_kv_indptr_prefill_cpu = self.paged_kv_indptr.cpu[
1658	prefill_start : num_reqs + 1
1659	]
1660	assert paged_kv_indptr_prefill_cpu.shape[0] == num_prefills + 1
1661	# plan() copies these to the GPU with non_blocking=True;
1662	# stage them in pinned memory so the copies stay async
1663	# (the reused buffers themselves are intentionally not pinned).
1664	if PIN_MEMORY:
1665	qo_indptr_prefill_cpu = qo_indptr_prefill_cpu.pin_memory()
1666	paged_kv_indptr_prefill_cpu = (
1667	paged_kv_indptr_prefill_cpu.pin_memory()
1668	)
1669	paged_kv_last_page_len_prefill_cpu = (
1670	paged_kv_last_page_len_prefill_cpu.pin_memory()
1671	)
1672	# plan() also copies kv_lens to the GPU with non_blocking=True,
1673	# deriving them via get_seq_lens() into unpinned memory when
1674	# seq_lens is not passed; pass an explicit pinned copy of the
1675	# same values so that copy stays async too.
1676	kv_lens_prefill_cpu = get_seq_lens(
1677	paged_kv_indptr_prefill_cpu,
1678	paged_kv_last_page_len_prefill_cpu,
1679	page_size,
1680	)
1681	if PIN_MEMORY:
1682	kv_lens_prefill_cpu = kv_lens_prefill_cpu.pin_memory()
1683	if self.use_dcp:
1684	assert isinstance(prefill_wrapper, BatchDCPPrefillWrapper)
1685	prefill_wrapper.plan(
1686	qo_indptr_cpu=qo_indptr_prefill_cpu,
1687	paged_kv_indptr_cpu=paged_kv_indptr_prefill_cpu,
1688	paged_kv_indices=paged_kv_indices,
1689	paged_kv_last_page_len_cpu=paged_kv_last_page_len_prefill_cpu,
1690	kv_lens_cpu=kv_lens_prefill_cpu,
1691	page_size=self.page_size,
1692	num_qo_heads=self.num_qo_heads,
1693	dcp_world_size=self.dcp_world_size,
1694	num_kv_heads=self.num_kv_heads,
1695	head_dim=self.head_dim,
1696	sm_scale=self.sm_scale,
1697	window_left=self.window_left,
1698	logits_soft_cap=self.logits_soft_cap,
1699	q_data_type=self.q_data_type_prefill,
1700	kv_cache_dtype=self.kv_cache_dtype,
1701	prefill_fixed_split_size=self.prefill_fixed_split_size,
1702	disable_split_kv=self.disable_split_kv,
1703	)
1704	else:
1705	assert isinstance(
1706	prefill_wrapper,
1707	BatchPrefillWithPagedKVCacheWrapper,
1708	)
1709	# NVFP4 trtllm kernel only supports FP8 output;
1710	# use FP8 o_data_type so the wrapper matches the
1711	# FP8 output buffer allocated in forward().
1712	o_dtype = (
1713	FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype
1714	)
1715	prefill_wrapper.plan(
1716	qo_indptr=qo_indptr_prefill_cpu,
1717	paged_kv_indptr=paged_kv_indptr_prefill_cpu,
1718	paged_kv_indices=paged_kv_indices,
1719	paged_kv_last_page_len=paged_kv_last_page_len_prefill_cpu,
1720	seq_lens=kv_lens_prefill_cpu,
1721	num_qo_heads=self.num_qo_heads,
1722	num_kv_heads=self.num_kv_heads,
1723	head_dim_qk=self.head_dim,
1724	page_size=self.page_size,
1725	causal=attn_metadata.causal,
1726	sm_scale=self.sm_scale,
1727	window_left=self.window_left,
1728	logits_soft_cap=self.logits_soft_cap,
1729	q_data_type=self.q_data_type_prefill,
1730	kv_data_type=self.kv_cache_dtype,
1731	o_data_type=o_dtype,
1732	fixed_split_size=self.prefill_fixed_split_size,
1733	disable_split_kv=self.disable_split_kv,
1734	)
1735	attn_metadata.prefill = FIPrefill(wrapper=prefill_wrapper)
1736
1737	## DECODE PATHWAY
1738	if num_decodes > 0:
1739	if decode_with_flashinfer_trtllm_api:
1740	assert self.flashinfer_trtllm_api_decode_kernel is not None
1741	seq_lens_decode = seq_lens[:num_decodes]
1742	if self.use_dcp:
1743	assert common_attn_metadata.dcp_local_seq_lens is not None
1744	seq_lens_decode = common_attn_metadata.dcp_local_seq_lens[
1745	:num_decodes
1746	]
1747	q_len_per_req, q_cu_seq_lens, ragged_q_lens = (
1748	self._compute_decode_query_lens(
1749	qo_indptr, qo_indptr_cpu, num_decodes, num_decode_tokens
1750	)
1751	)
1752	decode_mask = None
1753	if self.use_xqa:
1754	decode_mask = self._get_decode_mask(
1755	q_len_per_req,
1756	ragged_q_lens,
1757	num_decodes,
1758	bool(causal),
1759	)
1760	elif self.use_trtllm_gen_varlen_decode:
1761	# CPU placeholders may underestimate query lengths;
1762	# use a safe upper bound.
1763	q_len_per_req = min(
1764	common_attn_metadata.max_query_len, self.reorder_batch_threshold
1765	)
1766	# Always pass GPU offsets so a graph captured with uniform
1767	# query lengths can replay batches with variable query lengths.
1768	q_cu_seq_lens = qo_indptr[: num_decodes + 1]
1769	attn_metadata.decode = FlashInferTrtllmAPIDecode(
1770	kernel=self.flashinfer_trtllm_api_decode_kernel,
1771	block_tables=block_table_tensor[:num_decodes],
1772	seq_lens=seq_lens_decode,
1773	max_seq_len=max_seq_len,
1774	q_len_per_req=q_len_per_req,
1775	q_cu_seq_lens=q_cu_seq_lens,
1776	mask=decode_mask,
1777	dcp_query_start_loc=(
1778	qo_indptr[: num_decodes + 1] if self.use_dcp else None
1779	),
1780	)
1781	else:
1782	assert seq_lens_cpu is not None
1783	pure_decode = num_prefills == 0
1784	use_cudagraph = (
1785	self.enable_cuda_graph
1786	and pure_decode
1787	and num_decode_tokens <= self._decode_cudagraph_max_bs
1788	)
1789	num_input_tokens = num_decode_tokens
1790
1791	decode_wrapper = self._get_decode_wrapper(
1792	num_input_tokens, use_cudagraph
1793	)
1794	# Use the persistent buffer with padding length,
1795	# instead of the same address but chunked version
1796	# in atten_metadata when using cudagraph.
1797	# NVFP4 trtllm kernel only supports FP8 output;
1798	# use FP8 o_data_type so the wrapper matches the
1799	# FP8 output buffer allocated in forward().
1800	o_dtype = (
1801	FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype
1802	)
1803	paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: num_input_tokens + 1]
1804	paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[
1805	:num_input_tokens
1806	]
1807	# plan() copies these to the GPU with non_blocking=True;
1808	# stage them in pinned memory so the copies stay async
1809	# (the reused buffers themselves are intentionally not pinned).
1810	if PIN_MEMORY:
1811	paged_kv_indptr_cpu = paged_kv_indptr_cpu.pin_memory()
1812	paged_kv_last_page_len_cpu = paged_kv_last_page_len_cpu.pin_memory()
1813	# plan() also copies kv_lens to the GPU with non_blocking=True
1814	# (trtllm-gen/cute-dsl backends), deriving them via
1815	# get_seq_lens() into unpinned memory when seq_lens is not
1816	# passed; pass an explicit pinned copy of the same values
1817	# so that copy stays async too.
1818	kv_lens_decode_cpu = get_seq_lens(
1819	paged_kv_indptr_cpu, paged_kv_last_page_len_cpu, page_size
1820	)
1821	if PIN_MEMORY:
1822	kv_lens_decode_cpu = kv_lens_decode_cpu.pin_memory()
1823	fast_plan_decode(
1824	decode_wrapper,
1825	indptr_cpu=paged_kv_indptr_cpu,
1826	indices=paged_kv_indices,
1827	last_page_len_cpu=paged_kv_last_page_len_cpu,
1828	seq_lens_cpu=kv_lens_decode_cpu,
1829	num_qo_heads=self.num_qo_heads * self.dcp_world_size,
1830	num_kv_heads=self.num_kv_heads,
1831	head_dim=self.head_dim,
1832	page_size=self.page_size,
1833	# Disable flashinfer's pos encoding and use vllm's rope.
1834	pos_encoding_mode="NONE",
1835	sm_scale=self.sm_scale,
1836	window_left=self.window_left,
1837	logits_soft_cap=self.logits_soft_cap,
1838	q_data_type=self.q_data_type_decode,
1839	kv_data_type=self.kv_cache_dtype,
1840	o_data_type=o_dtype,
1841	fixed_split_size=self.decode_fixed_split_size,
1842	disable_split_kv=self.disable_split_kv,
1843	)
1844	attn_metadata.decode = FIDecode(wrapper=decode_wrapper)
1845	return attn_metadata
1846
1847	def update_draft_decode_metadata(self, _metadata: FlashInferMetadata) -> None:
1848	pass
1849
1850	def use_cascade_attention(self, args, *kwargs) -> bool:
1851	if self.kv_cache_spec.dtype != self.vllm_config.model_config.dtype:
1852	# TODO: The cascade wrapper currently does not support setting
1853	# kv cache dtype to something different from query dtype.
1854	return False
1855	# TODO: Cascade attention doesn't work, disable it for now
1856	# return use_cascade_attention(args, *kwargs)
1857	return False
1858
1859
1860	class FlashInferImpl(AttentionImpl):
1861	can_return_lse_for_decode: bool = True
1862	supports_dcp: bool = True
1863
1864	def __init__(
1865	self,
1866	num_heads: int,
1867	head_size: int,
1868	scale: float,
1869	num_kv_heads: int,
1870	alibi_slopes: list[float] \| None,
1871	sliding_window: int \| None,
1872	kv_cache_dtype: str,
1873	logits_soft_cap: float \| None = None,
1874	attn_type: AttentionType = AttentionType.DECODER,
1875	kv_sharing_target_layer_name: int \| None = None,
1876	sinks: torch.Tensor \| None = None,
1877	) -> None:
1878	self.num_heads = num_heads
1879	self.head_size = head_size
1880	self.scale = float(scale)
1881	self.num_kv_heads = num_kv_heads
1882	if alibi_slopes is not None:
1883	alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
1884	self.alibi_slopes = alibi_slopes
1885	if sliding_window is None:
1886	self.sliding_window = (-1, -1)
1887	else:
1888	self.sliding_window = (sliding_window - 1, 0)
1889	self.window_left = (
1890	self.sliding_window[0] if self.sliding_window is not None else -1
1891	)
1892	self.cache_dtype = kv_cache_dtype
1893	self.is_kvcache_nvfp4 = kv_cache_dtype.startswith("nvfp4")
1894	self.kv_cache_dtype = "nvfp4" if self.is_kvcache_nvfp4 else kv_cache_dtype
1895	self.fp4_data_dim = head_size // 2 if self.is_kvcache_nvfp4 else 0
1896	self.logits_soft_cap = logits_soft_cap
1897	self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
1898
1899	self.num_queries_per_kv = self.num_heads // self.num_kv_heads
1900
1901	if attn_type != AttentionType.DECODER:
1902	raise NotImplementedError(
1903	"Encoder self-attention and "
1904	"encoder/decoder cross-attention "
1905	"are not implemented for "
1906	"FlashInferImpl"
1907	)
1908
1909	self.sinks: torch.Tensor \| None = None
1910	# Keep the source so RL weight updates can refresh the runtime tensor.
1911	self._sinks_source = sinks
1912	if sinks is not None:
1913	if sinks.shape[0] != num_heads:
1914	raise ValueError(
1915	"Sinks must have the same number of heads as the number of "
1916	f"heads in the layer. Expected {num_heads}, but got "
1917	f"{sinks.shape[0]}."
1918	)
1919	self.sinks = sinks
1920
1921	self.supports_xqa_or_trtllm_gen_decode = can_use_trtllm_attention(
1922	num_heads, num_kv_heads, is_prefill=False
1923	)
1924	vllm_config = get_current_vllm_config_or_none()
1925	# Query pre-quantization needs a single dtype for the whole query tensor.
1926	# SM90 XQA needs BF16/FP16-Q for decode and FP8 for prefill,
1927	# so only enable this for SM100 trtllm-gen where both use FP8-Q.
1928	self.supports_quant_query_input = (
1929	self.supports_xqa_or_trtllm_gen_decode
1930	and is_quantized_kv_cache(self.kv_cache_dtype)
1931	and current_platform.is_device_capability_family(100)
1932	and vllm_config is not None
1933	and not vllm_config.attention_config.disable_flashinfer_q_quantization
1934	)
1935	self.bmm1_scale: float \| None = None
1936	self.bmm2_scale: float \| None = None
1937	self.o_sf_scale: float \| None = None
1938
1939	# Pre-allocated FP8 output buffer for NVFP4 without fused output quant.
1940	if self.is_kvcache_nvfp4 and vllm_config is not None:
1941	max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
1942	self._nvfp4_fp8_out = torch.empty(
1943	(max_num_tokens, num_heads, head_size),
1944	dtype=FP8_DTYPE,
1945	device="cuda",
1946	)
1947	else:
1948	self._nvfp4_fp8_out = None
1949
1950	dcp_a2a = (
1951	vllm_config is not None
1952	and vllm_config.parallel_config.decode_context_parallel_size > 1
1953	and vllm_config.parallel_config.dcp_comm_backend == "a2a"
1954	)
1955	if dcp_a2a:
1956	self.dcp_combine = partial(dcp_a2a_lse_reduce, is_lse_base_on_e=False)
1957	else:
1958	self.dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
1959
1960	def fused_output_quant_supported(self, quant_key: QuantKey):
1961	if quant_key == kNvfp4Dynamic and self.is_kvcache_nvfp4:
1962	logger.warning_once(
1963	"fuse_attn_quant: not registering the NVFP4-output attention "
1964	"fusion with an NVFP4 KV cache (no such trtllm-gen kernel). "
1965	"Layers whose o_proj takes NVFP4 input keep the unfused path "
1966	"(FP8 -> BF16 -> o_proj quant); FP8-output fusion is unaffected."
1967	)
1968	return False
1969	# XQA does not support FP8/NVFP4 output, so require trtllm-gen
1970	# (SM100+) here. Without that we cannot fuse the output quant.
1971	return (
1972	self.supports_xqa_or_trtllm_gen_decode
1973	and is_quantized_kv_cache(self.kv_cache_dtype)
1974	and current_platform.is_device_capability_family(100)
1975	and quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic)
1976	)
1977
1978	# FlashInfer requires attention sinks to be float32
1979	def process_weights_after_loading(self, act_dtype: torch.dtype):
1980	source_sinks = self._sinks_source
1981	if source_sinks is None:
1982	return
1983	if source_sinks.dtype == torch.float32:
1984	self.sinks = source_sinks
1985	elif self.sinks is None or self.sinks.dtype != torch.float32:
1986	self.sinks = source_sinks.to(torch.float32)
1987	else:
1988	self.sinks.copy_(source_sinks)
1989
1990	def get_xqa_bmm1_scale(self, layer: torch.nn.Module, q_data_type: torch.dtype):
1991	bmm1_scale = self.scale
1992	if is_quantized_kv_cache(self.kv_cache_dtype):
1993	if q_data_type in (torch.float8_e4m3fn, torch.float8_e5m2):
1994	bmm1_scale *= layer._q_scale_float
1995	bmm1_scale *= layer._k_scale_float
1996	return bmm1_scale
1997
1998	# SM90 may need FP8-Q for native prefill and BF16/FP16-Q for XQA decode,
1999	# so quantize only the slice whose target dtype differs.
2000	def maybe_quant_query(
2001	self,
2002	query: torch.Tensor,
2003	q_data_type: torch.dtype,
2004	scale: torch.Tensor,
2005	) -> torch.Tensor:
2006	if query.dtype != q_data_type:
2007	assert query.dtype in [torch.float16, torch.bfloat16]
2008	assert q_data_type in [torch.float8_e4m3fn, torch.float8_e5m2]
2009	assert query.dim() == 3
2010	num_tokens = query.shape[0]
2011	num_heads = query.shape[1]
2012	head_size = query.shape[2]
2013	assert query.stride(2) == 1 and query.stride(1) == head_size
2014	query_quantized, _ = custom_ops.scaled_fp8_quant(
2015	query.view(num_tokens, num_heads * head_size), scale=scale
2016	)
2017	return query_quantized.view(num_tokens, num_heads, head_size)
2018
2019	return query
2020
2021	def forward(
2022	self,
2023	layer: torch.nn.Module,
2024	query: torch.Tensor,
2025	key: torch.Tensor \| None,
2026	value: torch.Tensor \| None,
2027	kv_cache: torch.Tensor,
2028	attn_metadata: FlashInferMetadata,
2029	output: torch.Tensor,
2030	output_scale: torch.Tensor \| None = None,
2031	output_block_scale: torch.Tensor \| None = None,
2032	) -> torch.Tensor:
2033	"""Forward pass with FlashInfer.
2034
2035	Args:
2036	layer: The attention layer, providing the q/k/v quantization scales.
2037	query: shape = [num_tokens, num_heads, head_size]
2038	key: shape = [num_tokens, num_kv_heads, head_size], or None for a
2039	KV-sharing decoder layer.
2040	value: shape = [num_tokens, num_kv_heads, head_size], or None for a
2041	KV-sharing decoder layer.
2042	kv_cache: [num_blocks, num_kv_heads, block_size, 2*head_size]
2043	attn_metadata: Metadata for attention.
2044	output: Tensor that the attention result is written into.
2045	output_scale: Scale for fused output quantization. Enables the
2046	attention+quantization fusion path when provided.
2047	output_block_scale: Block scale for fused output quantization,
2048	required for nvfp4 output and rejected for fp8 output.
2049
2050	Returns:
2051	shape = [num_tokens, num_heads * head_size]
2052
2053	"""
2054	if attn_metadata is None:
2055	# Profiling run.
2056	return output.fill_(0)
2057
2058	if self.bmm1_scale is None:
2059	self.bmm1_scale = self.scale
2060	if is_quantized_kv_cache(self.kv_cache_dtype):
2061	self.bmm1_scale = layer._q_scale_float layer._k_scale_float
2062
2063	if self.bmm2_scale is None:
2064	self.bmm2_scale = 1.0
2065	if is_quantized_kv_cache(self.kv_cache_dtype):
2066	self.bmm2_scale *= layer._v_scale_float
2067
2068	prefill_use_trtllm = isinstance(attn_metadata.prefill, TRTLLMPrefill)
2069	decode_kernel = (
2070	attn_metadata.decode.kernel
2071	if isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode)
2072	else None
2073	)
2074	decode_with_xqa = decode_kernel == FlashInferDecodeKernel.XQA
2075	decode_with_trtllm_gen = decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN
2076	decode_with_flashinfer_trtllm_api = decode_with_xqa or decode_with_trtllm_gen
2077
2078	# The attn+quant fusion happens when output_scale is provided.
2079	if output_scale is None:
2080	assert output_block_scale is None, (
2081	"output_block_scale is not supported when fusion has not happened"
2082	)
2083	else:
2084	assert attn_metadata.q_data_type_prefill == FP8_DTYPE, (
2085	"Query must be FP8 when attn+quant fusion happened for prefill."
2086	)
2087	assert attn_metadata.q_data_type_decode == FP8_DTYPE, (
2088	"Query must be FP8 when attn+quant fusion happened for decode."
2089	)
2090	assert (attn_metadata.num_prefills == 0 or prefill_use_trtllm) and (
2091	attn_metadata.num_decodes == 0 or decode_with_trtllm_gen
2092	), "Output quant fusion requires TRTLLM prefill/trtllm-gen decode"
2093
2094	if output.dtype == FP8_DTYPE:
2095	assert output_block_scale is None, (
2096	"output_block_scale should not be provided for fp8 output"
2097	)
2098	elif output.dtype == FP4_DTYPE:
2099	assert output_block_scale is not None, (
2100	"output_block_scale is required for nvfp4 output"
2101	)
2102	else:
2103	raise ValueError(f"Unsupported output dtype: {output.dtype}")
2104
2105	# TRTLLM attn kernel requires to scale to pass as a host scalar,
2106	# store the o scale as a host scalar in warmup run with cuda graph
2107	# not enabled
2108	if layer._o_scale_float is None:
2109	layer._o_scale_float = output_scale.cpu().item()
2110	if output.dtype == FP8_DTYPE:
2111	self.bmm2_scale = self.bmm2_scale / layer._o_scale_float
2112	elif output.dtype == FP4_DTYPE:
2113	self.o_sf_scale = layer._o_scale_float
2114
2115	# IMPORTANT!
2116	# NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
2117	# eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
2118	# in this method. For example, `view` and `slice` (or `[:n]`) operations
2119	# are surprisingly slow even in the case they do not invoke any GPU ops.
2120	# Minimize the PyTorch ops in this method as much as possible.
2121	# Whenever making a change in this method, please benchmark the
2122	# performance to make sure it does not introduce any overhead.
2123
2124	num_actual_tokens = attn_metadata.num_actual_tokens
2125	kv_cache_layout = attn_metadata.kv_cache_layout
2126
2127	# FlashInfer treats uint8 KV cache as NVFP4. vLLM stores FP8 KV cache
2128	# as uint8 bytes, so pass FP8 caches with their logical dtype.
2129	if not self.is_kvcache_nvfp4 and kv_cache.dtype == torch.uint8:
2130	fp8_view_dtype = None
2131	if self.kv_cache_dtype in ("fp8", "fp8_e4m3", torch.float8_e4m3fn):
2132	fp8_view_dtype = torch.float8_e4m3fn
2133	elif self.kv_cache_dtype in ("fp8_e5m2", torch.float8_e5m2):
2134	fp8_view_dtype = torch.float8_e5m2
2135	if fp8_view_dtype is not None:
2136	kv_cache = kv_cache.view(fp8_view_dtype)
2137
2138	# Inputs and outputs may be padded for CUDA graphs. Keep the original
2139	# query for single-token XQA, whose non-ragged API requires one row for
2140	# every (including padded) decode request.
2141	query_padded = query
2142	query = query[:num_actual_tokens]
2143	output_padded = output
2144	output = output[:num_actual_tokens]
2145
2146	if attn_metadata.use_cascade:
2147	# Cascade attention (rare case).
2148	assert attn_metadata.cascade_wrapper is not None
2149	stride_order = kv_cache_layout.layer_view_order
2150	if self.is_kvcache_nvfp4:
2151	kv_cache_views = tuple(
2152	cache.permute(*stride_order)
2153	for cache in kv_cache.split(self.num_kv_heads, dim=1)
2154	)
2155	else:
2156	kv_perm = kv_cache.permute(*stride_order)
2157	kv_cache_views = kv_perm.split(self.head_size, dim=-1)
2158	kv_tuple = tuple(
2159	canonicalize_singleton_dim_strides(cache) for cache in kv_cache_views
2160	)
2161	output.copy_(attn_metadata.cascade_wrapper.run(query, kv_tuple))
2162	return output
2163
2164	# When using spec decoding, num_decodes can be < num_decode_tokens
2165	# because some decode requests may have more than one query token.
2166	num_decode_tokens = attn_metadata.num_decode_tokens
2167	num_prefill_tokens = attn_metadata.num_prefill_tokens
2168
2169	stride_order = kv_cache_layout.layer_view_order
2170	kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous
2171	# Fix degenerate strides on any size-1 dimension (e.g. num_kv_heads=1
2172	# with TP=8). PyTorch permits non-canonical strides on size-1 dims;
2173	# CUDA TMA requires ≥16-byte alignment on all non-outermost strides.
2174	# canonicalize_singleton_dim_strides patches metadata via as_strided —
2175	# zero-copy. See vllm.utils.torch_utils.
2176	fixed = canonicalize_singleton_dim_strides(kv_cache_permute)
2177	if fixed is not kv_cache_permute:
2178	logger.debug(
2179	"Canonicalized degenerate KV cache strides (FlashInfer): "
2180	"shape=%s, strides before=%s, strides after=%s",
2181	kv_cache_permute.shape,
2182	kv_cache_permute.stride(),
2183	fixed.stride(),
2184	)
2185	kv_cache_permute = fixed
2186
2187	# Split K/V — zero-copy views. NVFP4 stores K/V as separate head
2188	# groups; other dtypes pack K/V in the content dim.
2189	hs = self.head_size
2190	nvfp4_kv_data = None
2191	nvfp4_kv_block_scales = None
2192	if self.is_kvcache_nvfp4:
2193	k_cache, v_cache = kv_cache.split(self.num_kv_heads, dim=1)
2194	kv_cache_tuple = (
2195	canonicalize_singleton_dim_strides(k_cache.permute(*stride_order)),
2196	canonicalize_singleton_dim_strides(v_cache.permute(*stride_order)),
2197	)
2198	k_data, k_sf = nvfp4_split_data_scale(kv_cache_tuple[0])
2199	v_data, v_sf = nvfp4_split_data_scale(kv_cache_tuple[1])
2200	nvfp4_kv_data = (k_data, v_data)
2201	nvfp4_kv_block_scales = (k_sf, v_sf)
2202	else:
2203	kv_cache_tuple = kv_cache_permute.split(hs, dim=-1)
2204
2205	use_dcp = self.dcp_world_size > 1
2206	if decode_with_xqa:
2207	assert not use_dcp
2208	assert not self.is_kvcache_nvfp4
2209	assert self.o_sf_scale is None
2210	assert output.dtype != FP4_DTYPE
2211
2212	# Regular attention (common case).
2213	# Decodes are at the front and prefills are at the back.
2214	if num_prefill_tokens > 0:
2215	prefill_query = query[num_decode_tokens:]
2216	assert prefill_query.shape[0] == num_prefill_tokens
2217
2218	# Convert query to the expected dtype for prefill if needed.
2219	prefill_query = self.maybe_quant_query(
2220	prefill_query,
2221	attn_metadata.q_data_type_prefill,
2222	layer._q_scale,
2223	)
2224
2225	if not prefill_use_trtllm:
2226	assert isinstance(attn_metadata.prefill, FIPrefill)
2227	prefill_wrapper = attn_metadata.prefill.wrapper
2228	assert prefill_wrapper is not None
2229	if use_dcp:
2230	if key is None or value is None:
2231	raise NotImplementedError(
2232	"FlashInfer DCP prefill does not support KV-sharing layers"
2233	)
2234	assert isinstance(prefill_wrapper, BatchDCPPrefillWrapper)
2235	assert prefill_wrapper._context._window_left == self.window_left
2236	assert prefill_wrapper._context._logits_soft_cap == (
2237	self.logits_soft_cap or 0.0
2238	)
2239	assert prefill_wrapper._context._sm_scale == self.scale
2240	assert not prefill_wrapper._context._causal
2241	assert prefill_wrapper._new_tokens._window_left == self.window_left
2242	assert prefill_wrapper._new_tokens._logits_soft_cap == (
2243	self.logits_soft_cap or 0.0
2244	)
2245	assert prefill_wrapper._new_tokens._sm_scale == self.scale
2246	assert prefill_wrapper._new_tokens._causal
2247
2248	prefill_wrapper.run(
2249	layer,
2250	prefill_query,
2251	kv_cache_tuple,
2252	key[num_decode_tokens:num_actual_tokens],
2253	value[num_decode_tokens:num_actual_tokens],
2254	out=output[num_decode_tokens:],
2255	)
2256	else:
2257	assert isinstance(
2258	prefill_wrapper, BatchPrefillWithPagedKVCacheWrapper
2259	)
2260	assert prefill_wrapper._window_left == self.window_left
2261	assert prefill_wrapper._logits_soft_cap == (
2262	self.logits_soft_cap or 0.0
2263	)
2264	assert prefill_wrapper._sm_scale == self.scale
2265	assert prefill_wrapper._causal == attn_metadata.causal
2266
2267	if self.is_kvcache_nvfp4:
2268	kv_cache_for_fi = nvfp4_kv_data
2269	else:
2270	kv_cache_for_fi = kv_cache_tuple
2271	kv_cache_sf = (
2272	nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None
2273	)
2274
2275	# NVFP4 trtllm kernel only supports FP8 output.
2276	# Use a pre-allocated FP8 buffer and dequantize
2277	# afterwards.
2278	needs_fp8_out_prefill = (
2279	self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE
2280	)
2281	if needs_fp8_out_prefill:
2282	out_prefill = self._nvfp4_fp8_out[:num_prefill_tokens]
2283	else:
2284	out_prefill = output[num_decode_tokens:]
2285
2286	if isinstance(
2287	prefill_wrapper, BatchAttentionWithAttentionSinkWrapper
2288	):
2289	assert self.sinks is not None
2290	prefill_wrapper.run(
2291	prefill_query,
2292	kv_cache_for_fi,
2293	self.sinks,
2294	self.scale * layer._q_scale_float * layer._k_scale_float,
2295	v_scale=layer._v_scale_float,
2296	out=out_prefill,
2297	)
2298	else:
2299	prefill_wrapper.run(
2300	prefill_query,
2301	kv_cache_for_fi,
2302	q_scale=layer._q_scale_float,
2303	k_scale=layer._k_scale_float,
2304	v_scale=layer._v_scale_float,
2305	out=out_prefill,
2306	kv_cache_sf=kv_cache_sf,
2307	)
2308
2309	if needs_fp8_out_prefill:
2310	output[
2311	num_decode_tokens : num_decode_tokens + num_prefill_tokens
2312	].copy_(out_prefill)
2313	else:
2314	assert isinstance(attn_metadata.prefill, TRTLLMPrefill)
2315	# prefill_query may be non-contiguous or have degenerate strides
2316	# on size=1 dims. contiguous() ensures memory layout; then
2317	# canonicalize_singleton_dim_strides fixes any remaining
2318	# degenerate strides on size=1 dims for TMA alignment.
2319	prefill_query = prefill_query.contiguous()
2320	prefill_query = canonicalize_singleton_dim_strides(prefill_query)
2321	workspace_buffer = _get_trtllm_workspace_buffer()
2322	block_tables_prefill = attn_metadata.prefill.block_tables
2323	seq_lens_prefill = attn_metadata.prefill.seq_lens
2324
2325	# This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
2326	assert get_flashinfer_layout_string(kv_cache_layout) == "HND"
2327	assert is_strictly_contiguous(prefill_query)
2328	assert is_strictly_contiguous(workspace_buffer)
2329	assert is_strictly_contiguous(block_tables_prefill)
2330	assert is_strictly_contiguous(seq_lens_prefill)
2331
2332	if output.dtype == FP4_DTYPE:
2333	assert self.o_sf_scale is not None
2334	out = FP4Tensor(
2335	data=output[num_decode_tokens:],
2336	scale=output_block_scale,
2337	scale_start_index=num_decode_tokens,
2338	original_shape=prefill_query.shape,
2339	)
2340	else:
2341	assert self.o_sf_scale is None
2342	out = output[num_decode_tokens:]
2343
2344	# NVFP4 trtllm kernel only supports FP8 output.
2345	# Use a pre-allocated FP8 buffer and dequantize afterwards.
2346	needs_fp8_out = self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE
2347	if needs_fp8_out:
2348	out = self._nvfp4_fp8_out[:num_prefill_tokens]
2349
2350	prefill_kv_block_scales = None
2351	if self.is_kvcache_nvfp4:
2352	# NVFP4 trtllm-gen kernel requires FP8 query.
2353	assert attn_metadata.q_data_type_prefill == FP8_DTYPE, (
2354	"NVFP4 KV cache requires FP8 quantized queries for "
2355	"trtllm-gen prefill. Set "
2356	"disable_flashinfer_q_quantization=False."
2357	)
2358	mock_kv_cache = nvfp4_kv_data
2359	mock_block_table = block_tables_prefill
2360	prefill_kv_block_scales = nvfp4_kv_block_scales
2361	elif (
2362	attn_metadata.q_data_type_prefill != FP8_DTYPE
2363	and self.kv_cache_dtype.startswith("fp8")
2364	):
2365	# TRTLLM prefill attention does not support BF16 Q
2366	# and fp8 kv cache. So to enable prefill attention
2367	# with fp8 kv cache, we can construct a mock block
2368	# and mock kv cache with BF16 KV involved in the prefill.
2369	kv_cache_permute = canonicalize_singleton_dim_strides(
2370	kv_cache_permute
2371	)
2372	kv_strides = kv_cache_permute.stride()
2373	assert (
2374	kv_strides[-1] == 1
2375	and kv_strides[-2] == kv_cache_permute.shape[-1]
2376	), (
2377	"KV cache inner dims (block_size, head_size) must be "
2378	f"contiguous, got strides {kv_strides}"
2379	)
2380	# fp8 uses (B, H, N, 2*hs); reshape to (B, 2, H, N, hs)
2381	# for the dequant kernel — zero-copy view. The dequant
2382	# kernel handles the interleaved K/V block stride.
2383	B_kv, H_kv, N_kv = kv_cache_permute.shape[:3]
2384	kv_cache_5d = kv_cache_permute.view(B_kv, H_kv, N_kv, 2, hs)
2385	kv_cache_5d = kv_cache_5d.permute(0, 3, 1, 2, 4)
2386	mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
2387	kv_cache_5d,
2388	block_tables_prefill,
2389	layer._k_scale,
2390	layer._v_scale,
2391	attn_metadata.q_data_type_prefill,
2392	)
2393	else:
2394	mock_kv_cache = kv_cache_tuple
2395	mock_block_table = block_tables_prefill
2396
2397	trtllm_batch_context_with_kv_cache(
2398	query=prefill_query,
2399	kv_cache=mock_kv_cache,
2400	workspace_buffer=workspace_buffer,
2401	block_tables=mock_block_table,
2402	seq_lens=seq_lens_prefill,
2403	max_q_len=attn_metadata.prefill.max_q_len,
2404	max_kv_len=attn_metadata.prefill.max_seq_len,
2405	bmm1_scale=self.bmm1_scale,
2406	bmm2_scale=self.bmm2_scale,
2407	batch_size=attn_metadata.num_prefills,
2408	cum_seq_lens_q=attn_metadata.prefill.cum_seq_lens_q,
2409	cum_seq_lens_kv=attn_metadata.prefill.cum_seq_lens_kv,
2410	window_left=self.window_left,
2411	sinks=self.sinks,
2412	o_sf_scale=self.o_sf_scale,
2413	out=out,
2414	kv_cache_sf=prefill_kv_block_scales,
2415	)
2416
2417	if needs_fp8_out:
2418	output[
2419	num_decode_tokens : num_decode_tokens + num_prefill_tokens
2420	].copy_(out[:num_prefill_tokens])
2421
2422	if num_decode_tokens > 0:
2423	decode_query_tokens = num_decode_tokens
2424	if decode_with_xqa:
2425	assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode)
2426	if attn_metadata.decode.q_len_per_req == 1:
2427	decode_query_tokens = attn_metadata.num_decodes
2428	decode_query = query_padded[:decode_query_tokens]
2429	assert decode_query.shape[0] == decode_query_tokens
2430
2431	# Convert query to the expected dtype for decode if needed.
2432	decode_query = self.maybe_quant_query(
2433	decode_query,
2434	attn_metadata.q_data_type_decode,
2435	layer._q_scale,
2436	)
2437
2438	if not decode_with_flashinfer_trtllm_api:
2439	assert isinstance(attn_metadata.decode, FIDecode)
2440	decode_wrapper = attn_metadata.decode.wrapper
2441	assert decode_wrapper is not None
2442	assert decode_wrapper._window_left == self.window_left
2443	assert decode_wrapper._logits_soft_cap == (self.logits_soft_cap or 0.0)
2444	assert decode_wrapper._sm_scale == self.scale
2445
2446	if self.is_kvcache_nvfp4:
2447	kv_cache_for_fi = nvfp4_kv_data
2448	else:
2449	kv_cache_for_fi = kv_cache_tuple
2450	kv_cache_sf = nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None
2451
2452	# NVFP4 kernel only supports FP8 output.
2453	# Use a pre-allocated FP8 buffer and dequantize afterwards.
2454	needs_fp8_out = self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE
2455	if needs_fp8_out:
2456	out_decode = self._nvfp4_fp8_out[:num_decode_tokens]
2457	else:
2458	out_decode = output[:num_decode_tokens]
2459
2460	if use_dcp:
2461	decode_query = get_dcp_group().all_gather(
2462	decode_query.contiguous(), dim=-2
2463	)
2464	output_tmp = torch.empty_like(decode_query)
2465	lse = torch.empty(
2466	(decode_query.size(0), decode_query.size(1)),
2467	dtype=torch.float32,
2468	device=decode_query.device,
2469	)
2470	decode_wrapper.run(
2471	decode_query,
2472	kv_cache_for_fi,
2473	q_scale=layer._q_scale_float,
2474	k_scale=layer._k_scale_float,
2475	v_scale=layer._v_scale_float,
2476	out=output_tmp,
2477	lse=lse,
2478	return_lse=True,
2479	kv_cache_sf=kv_cache_sf,
2480	sinks=self.sinks,
2481	)
2482	output[:num_decode_tokens] = self.dcp_combine(
2483	output_tmp,
2484	lse,
2485	get_dcp_group(),
2486	)
2487	else:
2488	decode_wrapper.run(
2489	decode_query,
2490	kv_cache_for_fi,
2491	q_scale=layer._q_scale_float,
2492	k_scale=layer._k_scale_float,
2493	v_scale=layer._v_scale_float,
2494	out=out_decode,
2495	kv_cache_sf=kv_cache_sf,
2496	sinks=self.sinks,
2497	)
2498
2499	if needs_fp8_out:
2500	output[:num_decode_tokens].copy_(out_decode)
2501	else:
2502	assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode)
2503	# decode_query may be non-contiguous or have degenerate strides
2504	# on size=1 dims. contiguous() ensures memory layout; then
2505	# canonicalize_singleton_dim_strides fixes any remaining
2506	# degenerate strides on size=1 dims for TMA alignment.
2507	decode_query = decode_query.contiguous()
2508	decode_query = canonicalize_singleton_dim_strides(decode_query)
2509	workspace_buffer = _get_trtllm_workspace_buffer()
2510	block_tables_decode = attn_metadata.decode.block_tables
2511	seq_lens_decode = attn_metadata.decode.seq_lens
2512
2513	# trtllm-gen needs HND layout on SM100. XQA is selected
2514	# separately on SM90 and does not use this SM100 layout gate.
2515	if decode_with_trtllm_gen:
2516	assert get_flashinfer_layout_string(kv_cache_layout) == "HND"
2517	else:
2518	assert decode_with_xqa
2519	assert is_strictly_contiguous(decode_query)
2520	assert is_strictly_contiguous(workspace_buffer)
2521	assert is_strictly_contiguous(block_tables_decode)
2522	assert is_strictly_contiguous(seq_lens_decode)
2523	kv_cache_permute = canonicalize_singleton_dim_strides(kv_cache_permute)
2524	kv_strides = kv_cache_permute.stride()
2525	assert (
2526	kv_strides[-1] == 1 and kv_strides[-2] == kv_cache_permute.shape[-1]
2527	), (
2528	"KV cache inner dims (block_size, head_size) must be "
2529	f"contiguous, got strides {kv_strides}"
2530	)
2531
2532	if use_dcp:
2533	assert decode_with_trtllm_gen
2534	if output.dtype == FP4_DTYPE:
2535	raise NotImplementedError(
2536	"DCP decode with FlashInfer trtllm-gen does not support "
2537	"FP4 attention output yet."
2538	)
2539	decode_query = get_dcp_group().all_gather(
2540	decode_query.contiguous(), dim=-2
2541	)
2542	decode_query = canonicalize_singleton_dim_strides(decode_query)
2543
2544	if decode_with_xqa:
2545	bmm1_scale = self.get_xqa_bmm1_scale(
2546	layer, attn_metadata.q_data_type_decode
2547	)
2548	q_len_per_req: int \| None = attn_metadata.decode.q_len_per_req
2549
2550	flashinfer_xqa_batch_decode_with_kv_cache(
2551	query=decode_query,
2552	kv_cache=kv_cache_tuple,
2553	workspace_buffer=workspace_buffer,
2554	block_tables=block_tables_decode,
2555	seq_lens=seq_lens_decode,
2556	max_seq_len=attn_metadata.decode.max_seq_len,
2557	bmm1_scale=bmm1_scale,
2558	bmm2_scale=self.bmm2_scale,
2559	window_left=self.window_left,
2560	out=output_padded[:decode_query_tokens],
2561	sinks=self.sinks,
2562	kv_layout=get_flashinfer_layout_string(kv_cache_layout),
2563	q_len_per_req=q_len_per_req,
2564	mask=attn_metadata.decode.mask,
2565	q_cu_seq_lens=attn_metadata.decode.q_cu_seq_lens,
2566	)
2567	return output_padded
2568
2569	assert decode_with_trtllm_gen
2570	if output.dtype == FP4_DTYPE:
2571	assert self.o_sf_scale is not None
2572	out = FP4Tensor(
2573	data=output[:num_decode_tokens],
2574	scale=output_block_scale,
2575	scale_start_index=0,
2576	original_shape=decode_query.shape,
2577	)
2578	else:
2579	assert self.o_sf_scale is None
2580	out = output[:num_decode_tokens]
2581
2582	# NVFP4 trtllm kernel only supports FP8 output.
2583	# Use a pre-allocated FP8 buffer and dequantize afterwards.
2584	needs_fp8_out = self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE
2585	if needs_fp8_out:
2586	out = self._nvfp4_fp8_out[:num_decode_tokens]
2587
2588	q_cu_seq_lens = attn_metadata.decode.q_cu_seq_lens
2589	max_q_len = None
2590	q_len_per_req = None
2591	if q_cu_seq_lens is not None:
2592	# q_len_per_req must stay None for the kernel's varlen path.
2593	max_q_len = attn_metadata.decode.q_len_per_req
2594	elif num_decode_tokens % attn_metadata.num_decodes != 0:
2595	# This gets triggered when the dummy_run forces
2596	# attention to be initialized with q_len = 0
2597	q_len_per_req = 1
2598	else:
2599	q_len_per_req = num_decode_tokens // attn_metadata.num_decodes
2600
2601	lse = None
2602	if use_dcp:
2603	out = torch.empty(
2604	decode_query.shape,
2605	dtype=output.dtype,
2606	device=decode_query.device,
2607	)
2608	lse = torch.empty(
2609	(decode_query.size(0), decode_query.size(1)),
2610	dtype=torch.float32,
2611	device=decode_query.device,
2612	)
2613
2614	trtllm_batch_decode_with_kv_cache(
2615	query=decode_query,
2616	kv_cache=(
2617	nvfp4_kv_data if self.is_kvcache_nvfp4 else kv_cache_tuple
2618	),
2619	workspace_buffer=workspace_buffer,
2620	block_tables=block_tables_decode,
2621	seq_lens=seq_lens_decode,
2622	max_seq_len=attn_metadata.decode.max_seq_len,
2623	bmm1_scale=self.bmm1_scale,
2624	bmm2_scale=self.bmm2_scale,
2625	window_left=self.window_left,
2626	sinks=self.sinks,
2627	o_sf_scale=self.o_sf_scale,
2628	out=out,
2629	kv_layout=get_flashinfer_layout_string(kv_cache_layout),
2630	backend=attn_metadata.decode.kernel.value,
2631	q_len_per_req=q_len_per_req,
2632	max_q_len=max_q_len,
2633	cum_seq_lens_q=q_cu_seq_lens,
2634	kv_cache_sf=(
2635	nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None
2636	),
2637	lse=lse,
2638	return_lse=self.need_to_return_lse_for_decode,
2639	)
2640
2641	if use_dcp:
2642	assert isinstance(out, torch.Tensor)
2643	assert lse is not None
2644	output[:num_decode_tokens] = self.dcp_combine(
2645	out,
2646	lse,
2647	get_dcp_group(),
2648	seq_lens=seq_lens_decode,
2649	query_start_loc=attn_metadata.decode.dcp_query_start_loc,
2650	)
2651	elif needs_fp8_out:
2652	output[:num_decode_tokens].copy_(out)
2653	return output_padded
2654
2655	def do_kv_cache_update(
2656	self,
2657	layer: torch.nn.Module,
2658	key: torch.Tensor,
2659	value: torch.Tensor,
2660	kv_cache: torch.Tensor,
2661	slot_mapping: torch.Tensor,
2662	) -> None:
2663	if self.kv_sharing_target_layer_name is None:
2664	# Reshape the input keys and values and store them in the cache.
2665	# Skip this if sharing KV cache with an earlier attention layer.
2666	# NOTE(woosuk): Here, key and value are padded while slot_mapping is
2667	# not padded. However, we don't need to do key[:num_actual_tokens]
2668	# and value[:num_actual_tokens] because the reshape_and_cache_flash
2669	# op uses the slot_mapping's shape to determine the number of
2670	# actual tokens.
2671	if self.is_kvcache_nvfp4:
2672	# (B, 2*H, N, full_dim) -> ((B, N, H, full_dim),
2673	# (B, N, H, full_dim));
2674	# K heads first, then V heads.
2675	k_cache, v_cache = kv_cache.transpose(1, 2).split(
2676	self.num_kv_heads, dim=-2
2677	)
2678	else:
2679	# (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
2680	k_cache, v_cache = kv_cache.transpose(1, 2).split(
2681	self.head_size, dim=-1
2682	)
2683	torch.ops._C_cache_ops.reshape_and_cache_flash(
2684	key,
2685	value,
2686	k_cache,
2687	v_cache,
2688	slot_mapping,
2689	self.cache_dtype,
2690	layer._k_scale,
2691	layer._v_scale,
2692	)
2693
2694
2695	def fast_plan_decode(
2696	self, # decode wrapper
2697	indptr_cpu: torch.Tensor,
2698	indices: torch.Tensor,
2699	last_page_len_cpu: torch.Tensor,
2700	seq_lens_cpu: torch.Tensor,
2701	num_qo_heads: int,
2702	num_kv_heads: int,
2703	head_dim: int,
2704	page_size: int,
2705	pos_encoding_mode: str = "NONE",
2706	window_left: int = -1,
2707	logits_soft_cap: float \| None = None,
2708	q_data_type: str \| torch.dtype \| None = "float16",
2709	kv_data_type: str \| torch.dtype \| None = None,
2710	o_data_type: str \| torch.dtype \| None = None,
2711	data_type: str \| torch.dtype \| None = None,
2712	sm_scale: float \| None = None,
2713	rope_scale: float \| None = None,
2714	rope_theta: float \| None = None,
2715	non_blocking: bool = True,
2716	fixed_split_size: int = -1,
2717	disable_split_kv: bool = False,
2718	) -> None:
2719	"""A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for
2720	cudagraph capture/replay, while the no cudagraph version turns back
2721	to the original plan.
2722	using original plan after passing host-side buffers:
2723	- only host-to-device copy of indptr and last_page_len buffers
2724	Modifications for cudagraph:
2725	- only host-to-device copy of indptr and last_page_len buffers.
2726	- avoid device-to-device copy of indices buffer.
2727
2728	Part of the code get inspiration from the original plan from FlashInfer repo
2729	and the implementation of fast_decode_plan for FlashInfer in SGlang repo.
2730	"""
2731	# Warm up with the original plan if it is first call, and always run the
2732	# original plan if we run for dynamic shape. For fixed shape (cudagraph),
2733	# this warm up is to generate the _cached_module for the decode wrapper.
2734	if not self.is_cuda_graph_enabled or getattr(self, "vllm_first_call", True):
2735	# plan() copies qo_indptr from flashinfer's internal cached host
2736	# arange buffer (_get_range_buf) with non_blocking=True; pin the
2737	# cache entry so the copy stays async.
2738	pin_host_range_buf(len(last_page_len_cpu) + 1)
2739	self.plan(
2740	indptr=indptr_cpu,
2741	indices=indices,
2742	last_page_len=last_page_len_cpu,
2743	num_qo_heads=num_qo_heads,
2744	num_kv_heads=num_kv_heads,
2745	head_dim=head_dim,
2746	page_size=page_size,
2747	pos_encoding_mode=pos_encoding_mode,
2748	window_left=window_left,
2749	logits_soft_cap=logits_soft_cap,
2750	q_data_type=q_data_type,
2751	kv_data_type=kv_data_type,
2752	o_data_type=o_data_type,
2753	data_type=data_type,
2754	sm_scale=sm_scale,
2755	rope_scale=rope_scale,
2756	rope_theta=rope_theta,
2757	non_blocking=non_blocking,
2758	block_tables=None,
2759	seq_lens=seq_lens_cpu,
2760	fixed_split_size=fixed_split_size,
2761	disable_split_kv=disable_split_kv,
2762	)
2763	self.vllm_first_call = False
2764	return
2765
2766	assert self.is_cuda_graph_enabled, "Should be cudagraph only here"
2767
2768	fast_decode_plan(
2769	self,
2770	indptr=indptr_cpu,
2771	indices=indices,
2772	last_page_len=last_page_len_cpu,
2773	num_qo_heads=num_qo_heads,
2774	num_kv_heads=num_kv_heads,
2775	head_dim=head_dim,
2776	page_size=page_size,
2777	pos_encoding_mode=pos_encoding_mode,
2778	window_left=window_left,
2779	logits_soft_cap=logits_soft_cap,
2780	q_data_type=q_data_type,
2781	kv_data_type=kv_data_type,
2782	data_type=data_type,
2783	sm_scale=sm_scale,
2784	rope_scale=rope_scale,
2785	rope_theta=rope_theta,
2786	non_blocking=non_blocking,
2787	fixed_split_size=fixed_split_size,
2788	disable_split_kv=disable_split_kv,
2789	)
2790
2791
2792	@triton.jit
2793	def _copy_page_indices_kernel(
2794	page_indices,
2795	block_table,
2796	block_table_stride,
2797	cu_num_blocks,
2798	BLOCK_SIZE: tl.constexpr,
2799	):
2800	req_idx = tl.program_id(0)
2801	row_ptr = block_table + req_idx * block_table_stride
2802	start_idx = tl.load(cu_num_blocks + req_idx)
2803	end_idx = tl.load(cu_num_blocks + req_idx + 1)
2804	num_blocks = end_idx - start_idx
2805
2806	offset = tl.arange(0, BLOCK_SIZE)
2807	for i in tl.range(0, num_blocks, BLOCK_SIZE):
2808	block_ids = tl.load(row_ptr + i + offset, mask=i + offset < num_blocks)
2809	tl.store(
2810	page_indices + start_idx + i + offset,
2811	block_ids,
2812	mask=i + offset < num_blocks,
2813	)