From 1d2d83a07fd3180068917a031c28dcc83141d0be Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Wed, 12 Aug 2026 01:32:49 +0200 Subject: [PATCH] [Attention] Add FlashInfer XQA decode support on SM12x (#49718) Signed-off-by: Andrii Skliar Signed-off-by: Andrii Skliar Co-authored-by: Andrii Skliar Co-authored-by: OpenAI Codex --- .../attention/test_use_trtllm_attention.py | 14 + tests/v1/attention/test_attention_backends.py | 188 ++++++++++++-- vllm/utils/flashinfer.py | 25 +- vllm/v1/attention/backends/flashinfer.py | 243 +++++++++++++++--- 4 files changed, 413 insertions(+), 57 deletions(-) diff --git a/tests/kernels/attention/test_use_trtllm_attention.py b/tests/kernels/attention/test_use_trtllm_attention.py index 8ae4ac7a32..c311824618 100644 --- a/tests/kernels/attention/test_use_trtllm_attention.py +++ b/tests/kernels/attention/test_use_trtllm_attention.py @@ -101,6 +101,20 @@ def test_supports_sm90_decode_only(_art, _family, _cap): assert supports_trtllm_attention(is_prefill=True) is False +@patch("vllm.envs.VLLM_BATCH_INVARIANT", False) +@patch( + "vllm.utils.flashinfer.current_platform.is_device_capability", return_value=False +) +@patch( + "vllm.utils.flashinfer.current_platform.is_device_capability_family", + side_effect=lambda capability: capability == 120, +) +@patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True) +def test_supports_sm12x_decode_only(_art, _family, _cap): + assert supports_trtllm_attention(is_prefill=False) is True + assert supports_trtllm_attention(is_prefill=True) is False + + @patch("vllm.envs.VLLM_BATCH_INVARIANT", False) @patch( "vllm.utils.flashinfer.current_platform.is_device_capability_family", diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 87a9c80942..3d9ab6df8d 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -250,6 +250,7 @@ def run_attention_backend( attn_type: AttentionType = AttentionType.DECODER, sliding_window: int | None = None, kv_cache_dtype: str = "auto", + sinks: torch.Tensor | None = None, ) -> torch.Tensor: """Run attention computation using the specified backend's AttentionImpl.""" @@ -277,6 +278,7 @@ def run_attention_backend( window_left=-1, # No sliding window logits_soft_cap=0.0, # No soft cap sm_scale=1.0 / (head_size**0.5), # Standard scale + has_sinks=sinks is not None, ) for layer_name in layer_names } @@ -318,6 +320,7 @@ def run_attention_backend( sliding_window=sliding_window, attn_type=attn_type, kv_cache_dtype=kv_cache_dtype, + **({"sinks": sinks} if sinks is not None else {}), ) # Create mock layer and output buffer @@ -354,10 +357,11 @@ def _test_backend_correctness( rtol: float = 1e-2, tensor_parallel_size: int = 1, kv_cache_dtype: str = "auto", + use_sinks: bool = False, ): """ Test that all backends produce similar outputs to a reference implementation - using torch.nn.functional.scaled_dot_product_attention. + using FlexAttention or an explicit attention-sink reference. This test works by: 1. Generating a batch of sequences with specified context and query lengths. @@ -417,6 +421,11 @@ def _test_backend_correctness( num_kv_heads = vllm_config.model_config.get_num_kv_heads( vllm_config.parallel_config ) + sinks = ( + torch.linspace(-1.0, 1.0, num_q_heads, dtype=torch.float32, device=device) + if use_sinks + else None + ) head_size = vllm_config.model_config.get_head_size() sliding_window = vllm_config.model_config.get_sliding_window() dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype) @@ -471,17 +480,38 @@ def _test_backend_correctness( kv_len = s_len final_mask_mod = partial(mask_mod, context_len=context_len) - block_mask = create_block_mask( - final_mask_mod, B=None, H=None, Q_LEN=q_len, KV_LEN=kv_len, device=device - ) - sdpa_out_i = flex_attention( - q_sdpa_in, - k_sdpa_in, - v_sdpa_in, - block_mask=block_mask, - scale=scale, - enable_gqa=True, - ) + if sinks is None: + block_mask = create_block_mask( + final_mask_mod, + B=None, + H=None, + Q_LEN=q_len, + KV_LEN=kv_len, + device=device, + ) + sdpa_out_i = flex_attention( + q_sdpa_in, + k_sdpa_in, + v_sdpa_in, + block_mask=block_mask, + scale=scale, + enable_gqa=True, + ) + else: + scores = ( + torch.matmul(q_sdpa_in.float(), k_sdpa_in.float().transpose(-2, -1)) + * scale + ) + q_idx = torch.arange(q_len, device=device).unsqueeze(1) + kv_idx = torch.arange(kv_len, device=device).unsqueeze(0) + zero = torch.zeros((), dtype=torch.int32, device=device) + valid = final_mask_mod(zero, zero, q_idx, kv_idx) + scores.masked_fill_(~valid.unsqueeze(0).unsqueeze(0), -torch.inf) + sink_logits = sinks.view(1, num_q_heads, 1, 1).expand(1, -1, q_len, -1) + weights = torch.softmax(torch.cat((scores, sink_logits), dim=-1), dim=-1)[ + ..., :kv_len + ] + sdpa_out_i = torch.matmul(weights, v_sdpa_in.float()).to(dtype) all_sdpa_outputs.append(sdpa_out_i.transpose(1, 2).squeeze(0)) @@ -574,6 +604,7 @@ def _test_backend_correctness( sliding_window=sliding_window, attn_type=attn_type, kv_cache_dtype=kv_cache_dtype, + sinks=sinks, ) finally: if reset_kv_cache_layout: @@ -706,6 +737,86 @@ def test_flashinfer_xqa_bmm1_scale_matches_decode_q_dtype(): assert impl.get_xqa_bmm1_scale(MockLayer, torch.float8_e4m3fn) == 3.0 +@pytest.mark.skipif( + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, + reason="FlashInfer is not available.", +) +def test_flashinfer_xqa_draft_masks(): + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + device = torch.device("cpu") + causal = flashinfer_backend._make_xqa_draft_block_mask(3, True, device) + full = flashinfer_backend._make_xqa_draft_block_mask(3, False, device) + ragged = flashinfer_backend._make_xqa_ragged_draft_block_mask( + [2, 3], 3, True, device + ) + + assert torch.equal( + causal.view(torch.int16), + torch.tensor([[1, 0], [3, 0], [7, 0]], dtype=torch.int16), + ) + assert torch.equal( + full.view(torch.int16), torch.tensor([[7, 0]] * 3, dtype=torch.int16) + ) + assert torch.equal( + ragged.view(torch.int16), + torch.tensor( + [[1, 0], [3, 0], [1, 0], [3, 0], [7, 0]], + dtype=torch.int16, + ), + ) + + +@pytest.mark.skipif( + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, + reason="FlashInfer is not available.", +) +def test_flashinfer_xqa_query_lens_preserve_cudagraph_padding(): + """CUDA-graph padding stays as zero-length requests in ragged offsets.""" + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + device = torch.device("cpu") + builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder) + builder.use_dedicated_xqa = True + qo_indptr = torch.tensor([0, 3, 9, 15, 15], dtype=torch.int32, device=device) + + q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens( + qo_indptr, + qo_indptr, + num_decodes=4, + num_decode_tokens=21, + ) + + assert q_len == 6 + assert q_lens == [3, 6, 6, 0] + assert q_cu_seq_lens is not None + assert q_cu_seq_lens.tolist() == [0, 3, 9, 15, 15] + + +@pytest.mark.skipif( + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, + reason="FlashInfer is not available.", +) +def test_flashinfer_xqa_query_lens_require_exact_uniform_product(): + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder) + builder.use_dedicated_xqa = True + qo_indptr = torch.tensor([0, 3, 3, 3], dtype=torch.int32) + + q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens( + qo_indptr, + qo_indptr, + num_decodes=3, + num_decode_tokens=6, + ) + + assert q_len == 3 + assert q_lens == [3, 0, 0] + assert q_cu_seq_lens is not None + assert q_cu_seq_lens.tolist() == [0, 3, 3, 3] + + @pytest.mark.skipif( AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, reason="FlashInfer is not available.", @@ -737,10 +848,48 @@ def test_flashinfer_attention_sinks_refreshed_after_reload(dtype): AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, reason="FlashInfer is not available.", ) -def test_flashinfer_sm90_xqa_decode_correctness(default_vllm_config): - """FlashInfer should route Hopper decode through XQA and match SDPA.""" - if not current_platform.is_cuda() or not current_platform.is_device_capability(90): - pytest.skip("FlashInfer XQA decode requires SM90.") +def test_flashinfer_native_prefill_with_sinks(): + if not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(120) + ): + pytest.skip("Native FlashInfer prefill with sinks requires SM12x.") + + from vllm.v1.attention.backends.flashinfer import FlashInferBackend + + if not FlashInferBackend.supports_sink(): + pytest.skip("FlashInfer attention sinks are not available in this setup.") + + def causal_mask_mod( + b: torch.Tensor, + h: torch.Tensor, + q_idx: torch.Tensor, + kv_idx: torch.Tensor, + *, + context_len: int, + ): + return (q_idx + context_len) >= kv_idx + + _test_backend_correctness( + BATCH_SPECS["small_prefill"], + "meta-llama/Meta-Llama-3-8B", + [AttentionBackendEnum.FLASHINFER], + causal_mask_mod, + use_sinks=True, + ) + + +@pytest.mark.skipif( + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, + reason="FlashInfer is not available.", +) +def test_flashinfer_xqa_decode_correctness(default_vllm_config): + """FlashInfer should route supported decode through XQA and match SDPA.""" + supported = current_platform.is_cuda() and ( + current_platform.is_device_capability(90) + or current_platform.is_device_capability_family(120) + ) + if not supported: + pytest.skip("FlashInfer XQA decode requires SM90 or SM12x.") import unittest.mock @@ -805,11 +954,16 @@ def test_flashinfer_sm90_xqa_decode_correctness(default_vllm_config): ) attn_metadata = builder.build(0, common_attn_metadata) + expected_cg_support = ( + AttentionCGSupport.UNIFORM_BATCH + if current_platform.is_device_capability_family(120) + else AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + ) assert ( flashinfer_backend.FlashInferMetadataBuilder.get_cudagraph_support( vllm_config, kv_cache_spec ) - == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + == expected_cg_support ) assert isinstance( attn_metadata.decode, diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 8b999c8478..121621ba20 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -158,6 +158,10 @@ flashinfer_trtllm_batch_decode_sparse_mla_dsv4 = _lazy_import_wrapper( "trtllm_batch_decode_sparse_mla_dsv4", fallback_fn=_missing_sparse_mla, ) +flashinfer_xqa_batch_decode_with_kv_cache = _lazy_import_wrapper( + "flashinfer.decode", + "xqa_batch_decode_with_kv_cache", +) # Special case for autotune since it returns a context manager @@ -375,8 +379,8 @@ def supports_trtllm_attention(is_prefill: bool = False) -> bool: """Return whether TRTLLM attention is available on the current platform for the given attention phase. - SM90 (Hopper) supports the XQA decode kernel but not TRTLLM prefill. - SM100+ supports TRTLLM for both phases. All others are unsupported. + SM90 (Hopper) and SM12x support the XQA decode kernel but not TRTLLM + prefill. SM100+ supports TRTLLM for both phases. All others are unsupported. """ # Batch-invariant mode disables TRTLLM attention if envs.VLLM_BATCH_INVARIANT: @@ -386,8 +390,10 @@ def supports_trtllm_attention(is_prefill: bool = False) -> bool: if not has_nvidia_artifactory(): return False - # SM90 has XQA decode; prefill is not supported. - if current_platform.is_device_capability(90): + # SM90 and SM12x have XQA decode only. + if current_platform.is_device_capability( + 90 + ) or current_platform.is_device_capability_family(120): return not is_prefill # SM100/SM103 has both prefill and decode TRTLLM kernels. @@ -490,11 +496,11 @@ def use_trtllm_attention( if is_prefill: # Prefill auto-detection use_trtllm = kv_cache_dtype == "auto" - elif current_platform.is_device_capability(90) and kv_cache_dtype.startswith( - "fp8" - ): - # SM90 + FP8 KV cache: prefer the XQA decode kernel. XQA does not - # support NVFP4 KV (that is an SM100 trtllm-gen path only). + elif ( + current_platform.is_device_capability(90) + or current_platform.is_device_capability_family(120) + ) and kv_cache_dtype.startswith("fp8"): + # SM90/SM12x + FP8 KV cache: prefer the XQA decode kernel. use_trtllm = True else: # Decode auto-detection @@ -1036,6 +1042,7 @@ __all__ = [ "trtllm_fp4_block_scale_moe", "flashinfer_trtllm_batch_decode_with_kv_cache_mla", "flashinfer_trtllm_batch_decode_sparse_mla_dsv4", + "flashinfer_xqa_batch_decode_with_kv_cache", "autotune", "has_flashinfer_moe", "has_flashinfer_comm", diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 8e70440e77..e074afb3ea 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -40,6 +40,7 @@ from vllm.platforms.interface import DeviceCapability from vllm.triton_utils import tl, triton from vllm.utils.flashinfer import ( can_use_trtllm_attention, + flashinfer_xqa_batch_decode_with_kv_cache, force_use_trtllm_attention, supports_trtllm_attention, use_trtllm_attention, @@ -102,6 +103,50 @@ def _get_trtllm_workspace_buffer(): return trtllm_workspace_buffer +def _pack_draft_block_bool_mask( + bool_mask: torch.Tensor, num_packed: int +) -> torch.Tensor: + """Pack a boolean draft mask into XQA's uint16 layout.""" + num_rows = bool_mask.shape[0] + bits = 1 << torch.arange(32, device=bool_mask.device, dtype=torch.int64) + mask_u32 = ( + (bool_mask.view(num_rows, num_packed, 32).to(torch.int64) * bits) + .sum(dim=-1) + .to(torch.uint32) + ) + return mask_u32.view(torch.uint16).reshape(num_rows, num_packed * 2) + + +def _make_xqa_draft_block_mask( + q_len: int, causal: bool, device: torch.device +) -> torch.Tensor: + """Build a packed XQA draft mask for a uniform request.""" + num_packed = (q_len + 31) // 32 + padded = num_packed * 32 + q_idx = torch.arange(q_len, device=device).unsqueeze(1) + kv_idx = torch.arange(padded, device=device).unsqueeze(0) + bool_mask = kv_idx <= q_idx if causal else (kv_idx < q_len).expand(q_len, padded) + return _pack_draft_block_bool_mask(bool_mask, num_packed) + + +def _make_xqa_ragged_draft_block_mask( + q_lens: list[int], max_q_len: int, causal: bool, device: torch.device +) -> torch.Tensor: + """Build a packed XQA draft mask for a ragged batch.""" + num_packed = (max_q_len + 31) // 32 + padded = num_packed * 32 + q_lens_t = torch.tensor(q_lens, device=device) + row_lens = torch.repeat_interleave(q_lens_t, q_lens_t) + request_starts = torch.cumsum(q_lens_t, dim=0) - q_lens_t + row_starts = torch.repeat_interleave(request_starts, q_lens_t) + q_idx = torch.arange(sum(q_lens), device=device) - row_starts + kv_idx = torch.arange(padded, device=device).unsqueeze(0) + bool_mask = kv_idx < row_lens.unsqueeze(1) + if causal: + bool_mask &= kv_idx <= q_idx.unsqueeze(1) + return _pack_draft_block_bool_mask(bool_mask, num_packed) + + @triton.jit def _trtllm_prefill_attn_kvfp8_dequant( kv_cache_ptr, @@ -472,20 +517,20 @@ class FlashInferBackend(AttentionBackend): @classmethod def supports_sink(cls) -> bool: - """FlashInfer supports sinks only on the SM100 trtllm-gen path.""" + """FlashInfer supports sinks on SM12x XQA and SM100 trtllm-gen.""" from vllm.utils.flashinfer import ( force_use_trtllm_attention, ) - # Respect explicit disable flag (e.g., - # --attention-config.use_trtllm_attention=0) if force_use_trtllm_attention() is False: return False + if current_platform.is_device_capability_family(120): + return supports_trtllm_attention(is_prefill=False) + if not current_platform.is_device_capability_family(100): return False - # Check if TRTLLM is supported on this platform return supports_trtllm_attention( is_prefill=False ) and supports_trtllm_attention(is_prefill=True) @@ -551,12 +596,7 @@ class TRTLLMPrefill: @dataclass class FlashInferTrtllmAPIDecode: - """Metadata for decode paths using FlashInfer's TRTLLM decode API. - - FlashInfer exposes both XQA (SM90) and trtllm-gen (SM100) through - ``trtllm_batch_decode_with_kv_cache``. Keep them as distinct vLLM - decode kernels because their dtype/layout/output constraints differ. - """ + """Metadata for XQA and trtllm-gen decode.""" kernel: FlashInferDecodeKernel @@ -575,6 +615,15 @@ class FlashInferTrtllmAPIDecode: max_seq_len: int """The maximum sequence length for KV Cache.""" + q_len_per_req: int = 1 + """Query tokens per request for uniform speculative decode.""" + + q_cu_seq_lens: torch.Tensor | None = None + """Cumulative query lengths for ragged XQA decode.""" + + mask: torch.Tensor | None = None + """Packed XQA draft mask.""" + @dataclass class FlashInferMetadata: @@ -659,7 +708,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): self.model_config.max_model_len, self.kv_cache_spec.block_size ) max_num_reqs = vllm_config.scheduler_config.max_num_seqs + self.max_num_reqs = max_num_reqs max_num_pages = max_num_reqs * max_num_pages_per_req + # Persistent uniform masks keep stable addresses for CUDA graphs. + self._decode_mask_cache: dict[tuple[int, bool], torch.Tensor] = {} speculative_config = vllm_config.speculative_config num_spec_tokens = ( speculative_config.num_speculative_tokens @@ -771,9 +823,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): ) self.use_trtllm_decode_attention = False self.flashinfer_trtllm_api_decode_kernel = None + self.use_dedicated_xqa = ( + current_platform.is_device_capability_family(120) + and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA + ) supports_spec_as_decode = ( self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN + or self.use_dedicated_xqa ) self._init_reorder_batch_threshold( 1, @@ -856,10 +913,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): # if cache_config requests a quantized dtype globally. cache_dtype = self.cache_dtype - # On SM90, XQA decode requires BF16/FP16-Q even with FP8 KV cache. - # FI native prefill on SM90 still uses FP8-Q in that case. + # XQA decode requires BF16/FP16-Q even with FP8 KV cache. if ( - current_platform.is_device_capability(90) + ( + current_platform.is_device_capability(90) + or current_platform.is_device_capability_family(120) + ) and not is_prefill and force_use_trtllm_attention() is not False and cache_dtype.startswith("fp8") @@ -868,8 +927,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): # Otherwise, match Q dtype to the KV cache dtype. if cache_dtype.startswith("fp8"): - # FP8-Q requires an fp8 tensor-core attention path - # (FI native fa3 on SM90, trtllm-gen/XQA on SM100). + # FP8-Q requires an fp8 tensor-core attention path. # Architectures with only fa2 (e.g. SM89, SM120) cannot # consume FP8 queries, so keep the model dtype for Q there. if current_platform.is_device_capability( @@ -901,13 +959,17 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): ) -> AttentionCGSupport: """Get the cudagraph support level for FlashInfer attention. - The SM90 XQA integration only enables single-token decode today. Keep - specdec CUDA graphs limited to trtllm-gen until vLLM wires the XQA - specdec mask. + SM90 XQA supports only single-token decode. SM12x uses the dedicated + XQA API, which supports speculative and non-causal decode. """ if current_platform.is_device_capability(90): return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + is_sm12x = current_platform.is_device_capability_family(120) + # XQA does not return LSE and therefore does not support DCP. + if is_sm12x and vllm_config.parallel_config.decode_context_parallel_size > 1: + return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + # For UniformTypeKVCacheSpecs, check all contained specs kv_specs = ( kv_cache_spec.kv_cache_specs.values() @@ -931,8 +993,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): has_trtllm_support = False break - # trtllm-gen only supports causal attention. - if has_trtllm_support and not vllm_config.attention_config.use_non_causal: + if has_trtllm_support and ( + is_sm12x or not vllm_config.attention_config.use_non_causal + ): return AttentionCGSupport.UNIFORM_BATCH else: return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE @@ -967,11 +1030,68 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): @staticmethod def _get_flashinfer_trtllm_api_decode_kernel() -> FlashInferDecodeKernel: - if current_platform.is_device_capability(90): + if current_platform.is_device_capability( + 90 + ) or current_platform.is_device_capability_family(120): return FlashInferDecodeKernel.XQA assert current_platform.is_device_capability_family(100) return FlashInferDecodeKernel.TRTLLM_GEN + def _compute_decode_query_lens( + self, + qo_indptr: torch.Tensor, + qo_indptr_cpu: torch.Tensor, + num_decodes: int, + num_decode_tokens: int, + ) -> tuple[int, torch.Tensor | None, list[int] | None]: + """Return the query width, ragged offsets, and effective query lengths.""" + assert self.use_dedicated_xqa + if num_decodes == 0 or num_decode_tokens == 0: + return 1, None, None + + decode_q_lens = qo_indptr_cpu[1 : num_decodes + 1] - qo_indptr_cpu[:num_decodes] + nonzero = decode_q_lens[decode_q_lens > 0] + q_len_per_req = int(nonzero.max().item()) if nonzero.numel() > 0 else 1 + uniform = nonzero.numel() <= 1 or bool((nonzero == nonzero[0]).all().item()) + if uniform and num_decode_tokens == num_decodes * q_len_per_req: + return q_len_per_req, None, None + + if q_len_per_req <= 1: + return 1, None, None + + q_lens = decode_q_lens.tolist() + return ( + q_len_per_req, + qo_indptr[: num_decodes + 1], + q_lens, + ) + + def _get_decode_mask( + self, + q_len_per_req: int, + ragged_q_lens: list[int] | None, + num_decodes: int, + causal: bool, + ) -> torch.Tensor | None: + """Return the packed XQA draft mask for speculative decode.""" + if q_len_per_req <= 1: + return None + + if ragged_q_lens is None: + key = (q_len_per_req, causal) + buf = self._decode_mask_cache.get(key) + if buf is None: + per_req = _make_xqa_draft_block_mask(q_len_per_req, causal, self.device) + buf = ( + per_req.unsqueeze(0).expand(self.max_num_reqs, -1, -1).contiguous() + ) + self._decode_mask_cache[key] = buf + return buf[:num_decodes] + + return _make_xqa_ragged_draft_block_mask( + ragged_q_lens, q_len_per_req, causal, self.device + ) + def _get_prefill_wrapper( self, causal: bool = True, @@ -1125,17 +1245,16 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): num_reqs = common_attn_metadata.num_reqs num_actual_tokens = common_attn_metadata.num_actual_tokens causal = common_attn_metadata.causal - if causal: + route_decode = causal or self.use_dedicated_xqa + if route_decode: num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( split_decodes_and_prefills( common_attn_metadata, decode_threshold=self.reorder_batch_threshold, - require_uniform=True, + require_uniform=not self.use_dedicated_xqa, ) ) else: - # FlashInfer decode/TRTLLM paths cannot express non-causal - # query-query attention, so DFlash runs as native prefill. num_decodes = 0 num_prefills = num_reqs num_decode_tokens = 0 @@ -1171,7 +1290,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): has_sinks=self.has_sinks, has_spec=uses_spec_reorder, ) - decode_with_flashinfer_trtllm_api = causal and self.use_trtllm_decode_attention + decode_with_flashinfer_trtllm_api = self.use_trtllm_decode_attention and ( + causal or self.use_dedicated_xqa + ) if not causal and self.use_dcp: raise NotImplementedError( @@ -1188,7 +1309,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): ) if not all_uses_trtllm: - if self.has_sinks: + if self.has_sinks and ( + not self.use_dedicated_xqa or self.use_dcp or use_cascade + ): raise NotImplementedError( "FlashInfer backend currently does not support attention " "sinks, please use trtllm on blackwell or flash attention " @@ -1455,22 +1578,44 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): ## DECODE PATHWAY if num_decodes > 0: if decode_with_flashinfer_trtllm_api: - assert num_decode_tokens % num_decodes == 0, ( - "XQA/trtllm-gen decode requires uniform query lengths per request. " - f"Got {num_decode_tokens=} and {num_decodes=}." - ) assert self.flashinfer_trtllm_api_decode_kernel is not None + if not self.use_dedicated_xqa: + assert num_decode_tokens % num_decodes == 0, ( + "XQA/trtllm-gen decode requires uniform query lengths " + f"per request. Got {num_decode_tokens=} and {num_decodes=}." + ) seq_lens_decode = seq_lens[:num_decodes] if self.use_dcp: assert common_attn_metadata.dcp_local_seq_lens is not None seq_lens_decode = common_attn_metadata.dcp_local_seq_lens[ :num_decodes ] + q_len_per_req = 1 + q_cu_seq_lens = None + decode_mask = None + if self.use_dedicated_xqa: + q_len_per_req, q_cu_seq_lens, ragged_q_lens = ( + self._compute_decode_query_lens( + qo_indptr, + qo_indptr_cpu, + num_decodes, + num_decode_tokens, + ) + ) + decode_mask = self._get_decode_mask( + q_len_per_req, + ragged_q_lens, + num_decodes, + bool(causal), + ) attn_metadata.decode = FlashInferTrtllmAPIDecode( kernel=self.flashinfer_trtllm_api_decode_kernel, block_tables=block_table_tensor[:num_decodes], seq_lens=seq_lens_decode, max_seq_len=max_seq_len, + q_len_per_req=q_len_per_req, + q_cu_seq_lens=q_cu_seq_lens, + mask=decode_mask, ) else: assert seq_lens_cpu is not None @@ -1856,6 +2001,14 @@ class FlashInferImpl(AttentionImpl): kv_cache_tuple = kv_cache_permute.split(hs, dim=-1) use_dcp = self.dcp_world_size > 1 + decode_with_dedicated_xqa = ( + decode_with_xqa and current_platform.is_device_capability_family(120) + ) + if decode_with_dedicated_xqa: + assert not use_dcp + assert not self.is_kvcache_nvfp4 + assert self.o_sf_scale is None + assert output.dtype != FP4_DTYPE # Regular attention (common case). # Decodes are at the front and prefills are at the back. @@ -1935,6 +2088,7 @@ class FlashInferImpl(AttentionImpl): v_scale=layer._v_scale_float, out=out_prefill, kv_cache_sf=kv_cache_sf, + sinks=self.sinks, ) if needs_fp8_out_prefill: @@ -2103,6 +2257,7 @@ class FlashInferImpl(AttentionImpl): lse=lse, return_lse=True, kv_cache_sf=kv_cache_sf, + sinks=self.sinks, ) output[:num_decode_tokens] = self.dcp_combine( output_tmp, @@ -2118,6 +2273,7 @@ class FlashInferImpl(AttentionImpl): v_scale=layer._v_scale_float, out=out_decode, kv_cache_sf=kv_cache_sf, + sinks=self.sinks, ) if needs_fp8_out: @@ -2165,6 +2321,31 @@ class FlashInferImpl(AttentionImpl): ) decode_query = canonicalize_singleton_dim_strides(decode_query) + if decode_with_dedicated_xqa: + bmm1_scale = self.get_xqa_bmm1_scale( + layer, attn_metadata.q_data_type_decode + ) + q_len_per_req = attn_metadata.decode.q_len_per_req + + flashinfer_xqa_batch_decode_with_kv_cache( + query=decode_query, + kv_cache=kv_cache_tuple, + workspace_buffer=workspace_buffer, + block_tables=block_tables_decode, + seq_lens=seq_lens_decode, + max_seq_len=attn_metadata.decode.max_seq_len, + bmm1_scale=bmm1_scale, + bmm2_scale=self.bmm2_scale, + window_left=self.window_left, + out=output[:num_decode_tokens], + sinks=self.sinks, + kv_layout=get_kv_cache_layout(), + q_len_per_req=q_len_per_req, + mask=attn_metadata.decode.mask, + q_cu_seq_lens=attn_metadata.decode.q_cu_seq_lens, + ) + return output_padded + if output.dtype == FP4_DTYPE: assert self.o_sf_scale is not None out = FP4Tensor(