@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1d2d83a07f

[Attention] Add FlashInfer XQA decode support on SM12x (#49718)

Signed-off-by: Andrii Skliar <askliar@nvidia.com> Signed-off-by: Andrii Skliar <andreyws96@gmail.com> Co-authored-by: Andrii Skliar <askliar@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Andrii Skliar committed Aug 12, 2026 at 01:32 UTC 1d2d83a07fd3180068917a031c28dcc83141d0be

4 files changed +413 -57

tests/kernels/attention/test_use_trtllm_attention.py

+14

@@ -101,6 +101,20 @@ def test_supports_sm90_decode_only(_art, _family, _cap):

101 assert supports_trtllm_attention(is_prefill=True) is False

102

103

104 +@patch("vllm.envs.VLLM_BATCH_INVARIANT", False)

105 +@patch(

106 + "vllm.utils.flashinfer.current_platform.is_device_capability", return_value=False

107 +)

108 +@patch(

109 + "vllm.utils.flashinfer.current_platform.is_device_capability_family",

110 + side_effect=lambda capability: capability == 120,

111 +)

112 +@patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)

113 +def test_supports_sm12x_decode_only(_art, _family, _cap):

114 + assert supports_trtllm_attention(is_prefill=False) is True

115 + assert supports_trtllm_attention(is_prefill=True) is False

116 +

117 +

118 @patch("vllm.envs.VLLM_BATCH_INVARIANT", False)

119 @patch(

120 "vllm.utils.flashinfer.current_platform.is_device_capability_family",

tests/v1/attention/test_attention_backends.py

+171 -17

@@ -250,6 +250,7 @@ def run_attention_backend(

250 attn_type: AttentionType = AttentionType.DECODER,

251 sliding_window: int | None = None,

252 kv_cache_dtype: str = "auto",

253 + sinks: torch.Tensor | None = None,

254 ) -> torch.Tensor:

255 """Run attention computation using the specified backend's AttentionImpl."""

256

@@ -277,6 +278,7 @@ def run_attention_backend(

278 window_left=-1, # No sliding window

279 logits_soft_cap=0.0, # No soft cap

280 sm_scale=1.0 / (head_size**0.5), # Standard scale

281 + has_sinks=sinks is not None,

282 )

283 for layer_name in layer_names

284 }

@@ -318,6 +320,7 @@ def run_attention_backend(

320 sliding_window=sliding_window,

321 attn_type=attn_type,

322 kv_cache_dtype=kv_cache_dtype,

323 + **({"sinks": sinks} if sinks is not None else {}),

324 )

325

326 # Create mock layer and output buffer

@@ -354,10 +357,11 @@ def _test_backend_correctness(

357 rtol: float = 1e-2,

358 tensor_parallel_size: int = 1,

359 kv_cache_dtype: str = "auto",

360 + use_sinks: bool = False,

361 ):

362 """

363 Test that all backends produce similar outputs to a reference implementation

360 - using torch.nn.functional.scaled_dot_product_attention.

364 + using FlexAttention or an explicit attention-sink reference.

365

366 This test works by:

367 1. Generating a batch of sequences with specified context and query lengths.

@@ -417,6 +421,11 @@ def _test_backend_correctness(

421 num_kv_heads = vllm_config.model_config.get_num_kv_heads(

422 vllm_config.parallel_config

423 )

424 + sinks = (

425 + torch.linspace(-1.0, 1.0, num_q_heads, dtype=torch.float32, device=device)

426 + if use_sinks

427 + else None

428 + )

429 head_size = vllm_config.model_config.get_head_size()

430 sliding_window = vllm_config.model_config.get_sliding_window()

431 dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)

@@ -471,17 +480,38 @@ def _test_backend_correctness(

480 kv_len = s_len

481

482 final_mask_mod = partial(mask_mod, context_len=context_len)

474 - block_mask = create_block_mask(

475 - final_mask_mod, B=None, H=None, Q_LEN=q_len, KV_LEN=kv_len, device=device

476 - )

477 - sdpa_out_i = flex_attention(

478 - q_sdpa_in,

479 - k_sdpa_in,

480 - v_sdpa_in,

481 - block_mask=block_mask,

482 - scale=scale,

483 - enable_gqa=True,

484 - )

483 + if sinks is None:

484 + block_mask = create_block_mask(

485 + final_mask_mod,

486 + B=None,

487 + H=None,

488 + Q_LEN=q_len,

489 + KV_LEN=kv_len,

490 + device=device,

491 + )

492 + sdpa_out_i = flex_attention(

493 + q_sdpa_in,

494 + k_sdpa_in,

495 + v_sdpa_in,

496 + block_mask=block_mask,

497 + scale=scale,

498 + enable_gqa=True,

499 + )

500 + else:

501 + scores = (

502 + torch.matmul(q_sdpa_in.float(), k_sdpa_in.float().transpose(-2, -1))

503 + * scale

504 + )

505 + q_idx = torch.arange(q_len, device=device).unsqueeze(1)

506 + kv_idx = torch.arange(kv_len, device=device).unsqueeze(0)

507 + zero = torch.zeros((), dtype=torch.int32, device=device)

508 + valid = final_mask_mod(zero, zero, q_idx, kv_idx)

509 + scores.masked_fill_(~valid.unsqueeze(0).unsqueeze(0), -torch.inf)

510 + sink_logits = sinks.view(1, num_q_heads, 1, 1).expand(1, -1, q_len, -1)

511 + weights = torch.softmax(torch.cat((scores, sink_logits), dim=-1), dim=-1)[

512 + ..., :kv_len

513 + ]

514 + sdpa_out_i = torch.matmul(weights, v_sdpa_in.float()).to(dtype)

515

516 all_sdpa_outputs.append(sdpa_out_i.transpose(1, 2).squeeze(0))

517

@@ -574,6 +604,7 @@ def _test_backend_correctness(

604 sliding_window=sliding_window,

605 attn_type=attn_type,

606 kv_cache_dtype=kv_cache_dtype,

607 + sinks=sinks,

608 )

609 finally:

610 if reset_kv_cache_layout:

@@ -706,6 +737,86 @@ def test_flashinfer_xqa_bmm1_scale_matches_decode_q_dtype():

737 assert impl.get_xqa_bmm1_scale(MockLayer, torch.float8_e4m3fn) == 3.0

738

739

740 +@pytest.mark.skipif(

741 + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,

742 + reason="FlashInfer is not available.",

743 +)

744 +def test_flashinfer_xqa_draft_masks():

745 + from vllm.v1.attention.backends import flashinfer as flashinfer_backend

746 +

747 + device = torch.device("cpu")

748 + causal = flashinfer_backend._make_xqa_draft_block_mask(3, True, device)

749 + full = flashinfer_backend._make_xqa_draft_block_mask(3, False, device)

750 + ragged = flashinfer_backend._make_xqa_ragged_draft_block_mask(

751 + [2, 3], 3, True, device

752 + )

753 +

754 + assert torch.equal(

755 + causal.view(torch.int16),

756 + torch.tensor([[1, 0], [3, 0], [7, 0]], dtype=torch.int16),

757 + )

758 + assert torch.equal(

759 + full.view(torch.int16), torch.tensor([[7, 0]] * 3, dtype=torch.int16)

760 + )

761 + assert torch.equal(

762 + ragged.view(torch.int16),

763 + torch.tensor(

764 + [[1, 0], [3, 0], [1, 0], [3, 0], [7, 0]],

765 + dtype=torch.int16,

766 + ),

767 + )

768 +

769 +

770 +@pytest.mark.skipif(

771 + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,

772 + reason="FlashInfer is not available.",

773 +)

774 +def test_flashinfer_xqa_query_lens_preserve_cudagraph_padding():

775 + """CUDA-graph padding stays as zero-length requests in ragged offsets."""

776 + from vllm.v1.attention.backends import flashinfer as flashinfer_backend

777 +

778 + device = torch.device("cpu")

779 + builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder)

780 + builder.use_dedicated_xqa = True

781 + qo_indptr = torch.tensor([0, 3, 9, 15, 15], dtype=torch.int32, device=device)

782 +

783 + q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens(

784 + qo_indptr,

785 + qo_indptr,

786 + num_decodes=4,

787 + num_decode_tokens=21,

788 + )

789 +

790 + assert q_len == 6

791 + assert q_lens == [3, 6, 6, 0]

792 + assert q_cu_seq_lens is not None

793 + assert q_cu_seq_lens.tolist() == [0, 3, 9, 15, 15]

794 +

795 +

796 +@pytest.mark.skipif(

797 + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,

798 + reason="FlashInfer is not available.",

799 +)

800 +def test_flashinfer_xqa_query_lens_require_exact_uniform_product():

801 + from vllm.v1.attention.backends import flashinfer as flashinfer_backend

802 +

803 + builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder)

804 + builder.use_dedicated_xqa = True

805 + qo_indptr = torch.tensor([0, 3, 3, 3], dtype=torch.int32)

806 +

807 + q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens(

808 + qo_indptr,

809 + qo_indptr,

810 + num_decodes=3,

811 + num_decode_tokens=6,

812 + )

813 +

814 + assert q_len == 3

815 + assert q_lens == [3, 0, 0]

816 + assert q_cu_seq_lens is not None

817 + assert q_cu_seq_lens.tolist() == [0, 3, 3, 3]

818 +

819 +

820 @pytest.mark.skipif(

821 AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,

822 reason="FlashInfer is not available.",

@@ -737,10 +848,48 @@ def test_flashinfer_attention_sinks_refreshed_after_reload(dtype):

848 AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,

849 reason="FlashInfer is not available.",

850 )

740 -def test_flashinfer_sm90_xqa_decode_correctness(default_vllm_config):

741 - """FlashInfer should route Hopper decode through XQA and match SDPA."""

742 - if not current_platform.is_cuda() or not current_platform.is_device_capability(90):

743 - pytest.skip("FlashInfer XQA decode requires SM90.")

851 +def test_flashinfer_native_prefill_with_sinks():

852 + if not (

853 + current_platform.is_cuda() and current_platform.is_device_capability_family(120)

854 + ):

855 + pytest.skip("Native FlashInfer prefill with sinks requires SM12x.")

856 +

857 + from vllm.v1.attention.backends.flashinfer import FlashInferBackend

858 +

859 + if not FlashInferBackend.supports_sink():

860 + pytest.skip("FlashInfer attention sinks are not available in this setup.")

861 +

862 + def causal_mask_mod(

863 + b: torch.Tensor,

864 + h: torch.Tensor,

865 + q_idx: torch.Tensor,

866 + kv_idx: torch.Tensor,

867 + *,

868 + context_len: int,

869 + ):

870 + return (q_idx + context_len) >= kv_idx

871 +

872 + _test_backend_correctness(

873 + BATCH_SPECS["small_prefill"],

874 + "meta-llama/Meta-Llama-3-8B",

875 + [AttentionBackendEnum.FLASHINFER],

876 + causal_mask_mod,

877 + use_sinks=True,

878 + )

879 +

880 +

881 +@pytest.mark.skipif(

882 + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,

883 + reason="FlashInfer is not available.",

884 +)

885 +def test_flashinfer_xqa_decode_correctness(default_vllm_config):

886 + """FlashInfer should route supported decode through XQA and match SDPA."""

887 + supported = current_platform.is_cuda() and (

888 + current_platform.is_device_capability(90)

889 + or current_platform.is_device_capability_family(120)

890 + )

891 + if not supported:

892 + pytest.skip("FlashInfer XQA decode requires SM90 or SM12x.")

893

894 import unittest.mock

895

@@ -805,11 +954,16 @@ def test_flashinfer_sm90_xqa_decode_correctness(default_vllm_config):

954 )

955 attn_metadata = builder.build(0, common_attn_metadata)

956

957 + expected_cg_support = (

958 + AttentionCGSupport.UNIFORM_BATCH

959 + if current_platform.is_device_capability_family(120)

960 + else AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE

961 + )

962 assert (

963 flashinfer_backend.FlashInferMetadataBuilder.get_cudagraph_support(

964 vllm_config, kv_cache_spec

965 )

812 - == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE

966 + == expected_cg_support

967 )

968 assert isinstance(

969 attn_metadata.decode,

vllm/utils/flashinfer.py

+16 -9

@@ -158,6 +158,10 @@ flashinfer_trtllm_batch_decode_sparse_mla_dsv4 = _lazy_import_wrapper(

158 "trtllm_batch_decode_sparse_mla_dsv4",

159 fallback_fn=_missing_sparse_mla,

160 )

161 +flashinfer_xqa_batch_decode_with_kv_cache = _lazy_import_wrapper(

162 + "flashinfer.decode",

163 + "xqa_batch_decode_with_kv_cache",

164 +)

165

166

167 # Special case for autotune since it returns a context manager

@@ -375,8 +379,8 @@ def supports_trtllm_attention(is_prefill: bool = False) -> bool:

379 """Return whether TRTLLM attention is available on the current platform

380 for the given attention phase.

381

378 - SM90 (Hopper) supports the XQA decode kernel but not TRTLLM prefill.

379 - SM100+ supports TRTLLM for both phases. All others are unsupported.

382 + SM90 (Hopper) and SM12x support the XQA decode kernel but not TRTLLM

383 + prefill. SM100+ supports TRTLLM for both phases. All others are unsupported.

384 """

385 # Batch-invariant mode disables TRTLLM attention

386 if envs.VLLM_BATCH_INVARIANT:

@@ -386,8 +390,10 @@ def supports_trtllm_attention(is_prefill: bool = False) -> bool:

390 if not has_nvidia_artifactory():

391 return False

392

389 - # SM90 has XQA decode; prefill is not supported.

390 - if current_platform.is_device_capability(90):

393 + # SM90 and SM12x have XQA decode only.

394 + if current_platform.is_device_capability(

395 + 90

396 + ) or current_platform.is_device_capability_family(120):

397 return not is_prefill

398

399 # SM100/SM103 has both prefill and decode TRTLLM kernels.

@@ -490,11 +496,11 @@ def use_trtllm_attention(

496 if is_prefill:

497 # Prefill auto-detection

498 use_trtllm = kv_cache_dtype == "auto"

493 - elif current_platform.is_device_capability(90) and kv_cache_dtype.startswith(

494 - "fp8"

495 - ):

496 - # SM90 + FP8 KV cache: prefer the XQA decode kernel. XQA does not

497 - # support NVFP4 KV (that is an SM100 trtllm-gen path only).

499 + elif (

500 + current_platform.is_device_capability(90)

501 + or current_platform.is_device_capability_family(120)

502 + ) and kv_cache_dtype.startswith("fp8"):

503 + # SM90/SM12x + FP8 KV cache: prefer the XQA decode kernel.

504 use_trtllm = True

505 else:

506 # Decode auto-detection

@@ -1036,6 +1042,7 @@ __all__ = [

1042 "trtllm_fp4_block_scale_moe",

1043 "flashinfer_trtllm_batch_decode_with_kv_cache_mla",

1044 "flashinfer_trtllm_batch_decode_sparse_mla_dsv4",

1045 + "flashinfer_xqa_batch_decode_with_kv_cache",

1046 "autotune",

1047 "has_flashinfer_moe",

1048 "has_flashinfer_comm",

vllm/v1/attention/backends/flashinfer.py

+212 -31

@@ -40,6 +40,7 @@ from vllm.platforms.interface import DeviceCapability

40 from vllm.triton_utils import tl, triton

41 from vllm.utils.flashinfer import (

42 can_use_trtllm_attention,

43 + flashinfer_xqa_batch_decode_with_kv_cache,

44 force_use_trtllm_attention,

45 supports_trtllm_attention,

46 use_trtllm_attention,

@@ -102,6 +103,50 @@ def _get_trtllm_workspace_buffer():

103 return trtllm_workspace_buffer

104

105

106 +def _pack_draft_block_bool_mask(

107 + bool_mask: torch.Tensor, num_packed: int

108 +) -> torch.Tensor:

109 + """Pack a boolean draft mask into XQA's uint16 layout."""

110 + num_rows = bool_mask.shape[0]

111 + bits = 1 << torch.arange(32, device=bool_mask.device, dtype=torch.int64)

112 + mask_u32 = (

113 + (bool_mask.view(num_rows, num_packed, 32).to(torch.int64) * bits)

114 + .sum(dim=-1)

115 + .to(torch.uint32)

116 + )

117 + return mask_u32.view(torch.uint16).reshape(num_rows, num_packed * 2)

118 +

119 +

120 +def _make_xqa_draft_block_mask(

121 + q_len: int, causal: bool, device: torch.device

122 +) -> torch.Tensor:

123 + """Build a packed XQA draft mask for a uniform request."""

124 + num_packed = (q_len + 31) // 32

125 + padded = num_packed * 32

126 + q_idx = torch.arange(q_len, device=device).unsqueeze(1)

127 + kv_idx = torch.arange(padded, device=device).unsqueeze(0)

128 + bool_mask = kv_idx <= q_idx if causal else (kv_idx < q_len).expand(q_len, padded)

129 + return _pack_draft_block_bool_mask(bool_mask, num_packed)

130 +

131 +

132 +def _make_xqa_ragged_draft_block_mask(

133 + q_lens: list[int], max_q_len: int, causal: bool, device: torch.device

134 +) -> torch.Tensor:

135 + """Build a packed XQA draft mask for a ragged batch."""

136 + num_packed = (max_q_len + 31) // 32

137 + padded = num_packed * 32

138 + q_lens_t = torch.tensor(q_lens, device=device)

139 + row_lens = torch.repeat_interleave(q_lens_t, q_lens_t)

140 + request_starts = torch.cumsum(q_lens_t, dim=0) - q_lens_t

141 + row_starts = torch.repeat_interleave(request_starts, q_lens_t)

142 + q_idx = torch.arange(sum(q_lens), device=device) - row_starts

143 + kv_idx = torch.arange(padded, device=device).unsqueeze(0)

144 + bool_mask = kv_idx < row_lens.unsqueeze(1)

145 + if causal:

146 + bool_mask &= kv_idx <= q_idx.unsqueeze(1)

147 + return _pack_draft_block_bool_mask(bool_mask, num_packed)

148 +

149 +

150 @triton.jit

151 def _trtllm_prefill_attn_kvfp8_dequant(

152 kv_cache_ptr,

@@ -472,20 +517,20 @@ class FlashInferBackend(AttentionBackend):

517

518 @classmethod

519 def supports_sink(cls) -> bool:

475 - """FlashInfer supports sinks only on the SM100 trtllm-gen path."""

520 + """FlashInfer supports sinks on SM12x XQA and SM100 trtllm-gen."""

521 from vllm.utils.flashinfer import (

522 force_use_trtllm_attention,

523 )

524

480 - # Respect explicit disable flag (e.g.,

481 - # --attention-config.use_trtllm_attention=0)

525 if force_use_trtllm_attention() is False:

526 return False

527

528 + if current_platform.is_device_capability_family(120):

529 + return supports_trtllm_attention(is_prefill=False)

530 +

531 if not current_platform.is_device_capability_family(100):

532 return False

533

488 - # Check if TRTLLM is supported on this platform

534 return supports_trtllm_attention(

535 is_prefill=False

536 ) and supports_trtllm_attention(is_prefill=True)

@@ -551,12 +596,7 @@ class TRTLLMPrefill:

596

597 @dataclass

598 class FlashInferTrtllmAPIDecode:

554 - """Metadata for decode paths using FlashInfer's TRTLLM decode API.

555 -

556 - FlashInfer exposes both XQA (SM90) and trtllm-gen (SM100) through

557 - ``trtllm_batch_decode_with_kv_cache``. Keep them as distinct vLLM

558 - decode kernels because their dtype/layout/output constraints differ.

559 - """

599 + """Metadata for XQA and trtllm-gen decode."""

600

601 kernel: FlashInferDecodeKernel

602

@@ -575,6 +615,15 @@ class FlashInferTrtllmAPIDecode:

615 max_seq_len: int

616 """The maximum sequence length for KV Cache."""

617

618 + q_len_per_req: int = 1

619 + """Query tokens per request for uniform speculative decode."""

620 +

621 + q_cu_seq_lens: torch.Tensor | None = None

622 + """Cumulative query lengths for ragged XQA decode."""

623 +

624 + mask: torch.Tensor | None = None

625 + """Packed XQA draft mask."""

626 +

627

628 @dataclass

629 class FlashInferMetadata:

@@ -659,7 +708,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

708 self.model_config.max_model_len, self.kv_cache_spec.block_size

709 )

710 max_num_reqs = vllm_config.scheduler_config.max_num_seqs

711 + self.max_num_reqs = max_num_reqs

712 max_num_pages = max_num_reqs * max_num_pages_per_req

713 + # Persistent uniform masks keep stable addresses for CUDA graphs.

714 + self._decode_mask_cache: dict[tuple[int, bool], torch.Tensor] = {}

715 speculative_config = vllm_config.speculative_config

716 num_spec_tokens = (

717 speculative_config.num_speculative_tokens

@@ -771,9 +823,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

823 )

824 self.use_trtllm_decode_attention = False

825 self.flashinfer_trtllm_api_decode_kernel = None

826 + self.use_dedicated_xqa = (

827 + current_platform.is_device_capability_family(120)

828 + and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA

829 + )

830 supports_spec_as_decode = (

831 self.flashinfer_trtllm_api_decode_kernel

832 == FlashInferDecodeKernel.TRTLLM_GEN

833 + or self.use_dedicated_xqa

834 )

835 self._init_reorder_batch_threshold(

836 1,

@@ -856,10 +913,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

913 # if cache_config requests a quantized dtype globally.

914 cache_dtype = self.cache_dtype

915

859 - # On SM90, XQA decode requires BF16/FP16-Q even with FP8 KV cache.

860 - # FI native prefill on SM90 still uses FP8-Q in that case.

916 + # XQA decode requires BF16/FP16-Q even with FP8 KV cache.

917 if (

862 - current_platform.is_device_capability(90)

918 + (

919 + current_platform.is_device_capability(90)

920 + or current_platform.is_device_capability_family(120)

921 + )

922 and not is_prefill

923 and force_use_trtllm_attention() is not False

924 and cache_dtype.startswith("fp8")

@@ -868,8 +927,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

927

928 # Otherwise, match Q dtype to the KV cache dtype.

929 if cache_dtype.startswith("fp8"):

871 - # FP8-Q requires an fp8 tensor-core attention path

872 - # (FI native fa3 on SM90, trtllm-gen/XQA on SM100).

930 + # FP8-Q requires an fp8 tensor-core attention path.

931 # Architectures with only fa2 (e.g. SM89, SM120) cannot

932 # consume FP8 queries, so keep the model dtype for Q there.

933 if current_platform.is_device_capability(

@@ -901,13 +959,17 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

959 ) -> AttentionCGSupport:

960 """Get the cudagraph support level for FlashInfer attention.

961

904 - The SM90 XQA integration only enables single-token decode today. Keep

905 - specdec CUDA graphs limited to trtllm-gen until vLLM wires the XQA

906 - specdec mask.

962 + SM90 XQA supports only single-token decode. SM12x uses the dedicated

963 + XQA API, which supports speculative and non-causal decode.

964 """

965 if current_platform.is_device_capability(90):

966 return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE

967

968 + is_sm12x = current_platform.is_device_capability_family(120)

969 + # XQA does not return LSE and therefore does not support DCP.

970 + if is_sm12x and vllm_config.parallel_config.decode_context_parallel_size > 1:

971 + return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE

972 +

973 # For UniformTypeKVCacheSpecs, check all contained specs

974 kv_specs = (

975 kv_cache_spec.kv_cache_specs.values()

@@ -931,8 +993,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

993 has_trtllm_support = False

994 break

995

934 - # trtllm-gen only supports causal attention.

935 - if has_trtllm_support and not vllm_config.attention_config.use_non_causal:

996 + if has_trtllm_support and (

997 + is_sm12x or not vllm_config.attention_config.use_non_causal

998 + ):

999 return AttentionCGSupport.UNIFORM_BATCH

1000 else:

1001 return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE

@@ -967,11 +1030,68 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1030

1031 @staticmethod

1032 def _get_flashinfer_trtllm_api_decode_kernel() -> FlashInferDecodeKernel:

970 - if current_platform.is_device_capability(90):

1033 + if current_platform.is_device_capability(

1034 + 90

1035 + ) or current_platform.is_device_capability_family(120):

1036 return FlashInferDecodeKernel.XQA

1037 assert current_platform.is_device_capability_family(100)

1038 return FlashInferDecodeKernel.TRTLLM_GEN

1039

1040 + def _compute_decode_query_lens(

1041 + self,

1042 + qo_indptr: torch.Tensor,

1043 + qo_indptr_cpu: torch.Tensor,

1044 + num_decodes: int,

1045 + num_decode_tokens: int,

1046 + ) -> tuple[int, torch.Tensor | None, list[int] | None]:

1047 + """Return the query width, ragged offsets, and effective query lengths."""

1048 + assert self.use_dedicated_xqa

1049 + if num_decodes == 0 or num_decode_tokens == 0:

1050 + return 1, None, None

1051 +

1052 + decode_q_lens = qo_indptr_cpu[1 : num_decodes + 1] - qo_indptr_cpu[:num_decodes]

1053 + nonzero = decode_q_lens[decode_q_lens > 0]

1054 + q_len_per_req = int(nonzero.max().item()) if nonzero.numel() > 0 else 1

1055 + uniform = nonzero.numel() <= 1 or bool((nonzero == nonzero[0]).all().item())

1056 + if uniform and num_decode_tokens == num_decodes * q_len_per_req:

1057 + return q_len_per_req, None, None

1058 +

1059 + if q_len_per_req <= 1:

1060 + return 1, None, None

1061 +

1062 + q_lens = decode_q_lens.tolist()

1063 + return (

1064 + q_len_per_req,

1065 + qo_indptr[: num_decodes + 1],

1066 + q_lens,

1067 + )

1068 +

1069 + def _get_decode_mask(

1070 + self,

1071 + q_len_per_req: int,

1072 + ragged_q_lens: list[int] | None,

1073 + num_decodes: int,

1074 + causal: bool,

1075 + ) -> torch.Tensor | None:

1076 + """Return the packed XQA draft mask for speculative decode."""

1077 + if q_len_per_req <= 1:

1078 + return None

1079 +

1080 + if ragged_q_lens is None:

1081 + key = (q_len_per_req, causal)

1082 + buf = self._decode_mask_cache.get(key)

1083 + if buf is None:

1084 + per_req = _make_xqa_draft_block_mask(q_len_per_req, causal, self.device)

1085 + buf = (

1086 + per_req.unsqueeze(0).expand(self.max_num_reqs, -1, -1).contiguous()

1087 + )

1088 + self._decode_mask_cache[key] = buf

1089 + return buf[:num_decodes]

1090 +

1091 + return _make_xqa_ragged_draft_block_mask(

1092 + ragged_q_lens, q_len_per_req, causal, self.device

1093 + )

1094 +

1095 def _get_prefill_wrapper(

1096 self,

1097 causal: bool = True,

@@ -1125,17 +1245,16 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1245 num_reqs = common_attn_metadata.num_reqs

1246 num_actual_tokens = common_attn_metadata.num_actual_tokens

1247 causal = common_attn_metadata.causal

1128 - if causal:

1248 + route_decode = causal or self.use_dedicated_xqa

1249 + if route_decode:

1250 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

1251 split_decodes_and_prefills(

1252 common_attn_metadata,

1253 decode_threshold=self.reorder_batch_threshold,

1133 - require_uniform=True,

1254 + require_uniform=not self.use_dedicated_xqa,

1255 )

1256 )

1257 else:

1137 - # FlashInfer decode/TRTLLM paths cannot express non-causal

1138 - # query-query attention, so DFlash runs as native prefill.

1258 num_decodes = 0

1259 num_prefills = num_reqs

1260 num_decode_tokens = 0

@@ -1171,7 +1290,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1290 has_sinks=self.has_sinks,

1291 has_spec=uses_spec_reorder,

1292 )

1174 - decode_with_flashinfer_trtllm_api = causal and self.use_trtllm_decode_attention

1293 + decode_with_flashinfer_trtllm_api = self.use_trtllm_decode_attention and (

1294 + causal or self.use_dedicated_xqa

1295 + )

1296

1297 if not causal and self.use_dcp:

1298 raise NotImplementedError(

@@ -1188,7 +1309,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1309 )

1310

1311 if not all_uses_trtllm:

1191 - if self.has_sinks:

1312 + if self.has_sinks and (

1313 + not self.use_dedicated_xqa or self.use_dcp or use_cascade

1314 + ):

1315 raise NotImplementedError(

1316 "FlashInfer backend currently does not support attention "

1317 "sinks, please use trtllm on blackwell or flash attention "

@@ -1455,22 +1578,44 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1578 ## DECODE PATHWAY

1579 if num_decodes > 0:

1580 if decode_with_flashinfer_trtllm_api:

1458 - assert num_decode_tokens % num_decodes == 0, (

1459 - "XQA/trtllm-gen decode requires uniform query lengths per request. "

1460 - f"Got {num_decode_tokens=} and {num_decodes=}."

1461 - )

1581 assert self.flashinfer_trtllm_api_decode_kernel is not None

1582 + if not self.use_dedicated_xqa:

1583 + assert num_decode_tokens % num_decodes == 0, (

1584 + "XQA/trtllm-gen decode requires uniform query lengths "

1585 + f"per request. Got {num_decode_tokens=} and {num_decodes=}."

1586 + )

1587 seq_lens_decode = seq_lens[:num_decodes]

1588 if self.use_dcp:

1589 assert common_attn_metadata.dcp_local_seq_lens is not None

1590 seq_lens_decode = common_attn_metadata.dcp_local_seq_lens[

1591 :num_decodes

1592 ]

1593 + q_len_per_req = 1

1594 + q_cu_seq_lens = None

1595 + decode_mask = None

1596 + if self.use_dedicated_xqa:

1597 + q_len_per_req, q_cu_seq_lens, ragged_q_lens = (

1598 + self._compute_decode_query_lens(

1599 + qo_indptr,

1600 + qo_indptr_cpu,

1601 + num_decodes,

1602 + num_decode_tokens,

1603 + )

1604 + )

1605 + decode_mask = self._get_decode_mask(

1606 + q_len_per_req,

1607 + ragged_q_lens,

1608 + num_decodes,

1609 + bool(causal),

1610 + )

1611 attn_metadata.decode = FlashInferTrtllmAPIDecode(

1612 kernel=self.flashinfer_trtllm_api_decode_kernel,

1613 block_tables=block_table_tensor[:num_decodes],

1614 seq_lens=seq_lens_decode,

1615 max_seq_len=max_seq_len,

1616 + q_len_per_req=q_len_per_req,

1617 + q_cu_seq_lens=q_cu_seq_lens,

1618 + mask=decode_mask,

1619 )

1620 else:

1621 assert seq_lens_cpu is not None

@@ -1856,6 +2001,14 @@ class FlashInferImpl(AttentionImpl):

2001 kv_cache_tuple = kv_cache_permute.split(hs, dim=-1)

2002

2003 use_dcp = self.dcp_world_size > 1

2004 + decode_with_dedicated_xqa = (

2005 + decode_with_xqa and current_platform.is_device_capability_family(120)

2006 + )

2007 + if decode_with_dedicated_xqa:

2008 + assert not use_dcp

2009 + assert not self.is_kvcache_nvfp4

2010 + assert self.o_sf_scale is None

2011 + assert output.dtype != FP4_DTYPE

2012

2013 # Regular attention (common case).

2014 # Decodes are at the front and prefills are at the back.

@@ -1935,6 +2088,7 @@ class FlashInferImpl(AttentionImpl):

2088 v_scale=layer._v_scale_float,

2089 out=out_prefill,

2090 kv_cache_sf=kv_cache_sf,

2091 + sinks=self.sinks,

2092 )

2093

2094 if needs_fp8_out_prefill:

@@ -2103,6 +2257,7 @@ class FlashInferImpl(AttentionImpl):

2257 lse=lse,

2258 return_lse=True,

2259 kv_cache_sf=kv_cache_sf,

2260 + sinks=self.sinks,

2261 )

2262 output[:num_decode_tokens] = self.dcp_combine(

2263 output_tmp,

@@ -2118,6 +2273,7 @@ class FlashInferImpl(AttentionImpl):

2273 v_scale=layer._v_scale_float,

2274 out=out_decode,

2275 kv_cache_sf=kv_cache_sf,

2276 + sinks=self.sinks,

2277 )

2278

2279 if needs_fp8_out:

@@ -2165,6 +2321,31 @@ class FlashInferImpl(AttentionImpl):

2321 )

2322 decode_query = canonicalize_singleton_dim_strides(decode_query)

2323

2324 + if decode_with_dedicated_xqa:

2325 + bmm1_scale = self.get_xqa_bmm1_scale(

2326 + layer, attn_metadata.q_data_type_decode

2327 + )

2328 + q_len_per_req = attn_metadata.decode.q_len_per_req

2329 +

2330 + flashinfer_xqa_batch_decode_with_kv_cache(

2331 + query=decode_query,

2332 + kv_cache=kv_cache_tuple,

2333 + workspace_buffer=workspace_buffer,

2334 + block_tables=block_tables_decode,

2335 + seq_lens=seq_lens_decode,

2336 + max_seq_len=attn_metadata.decode.max_seq_len,

2337 + bmm1_scale=bmm1_scale,

2338 + bmm2_scale=self.bmm2_scale,

2339 + window_left=self.window_left,

2340 + out=output[:num_decode_tokens],

2341 + sinks=self.sinks,

2342 + kv_layout=get_kv_cache_layout(),

2343 + q_len_per_req=q_len_per_req,

2344 + mask=attn_metadata.decode.mask,

2345 + q_cu_seq_lens=attn_metadata.decode.q_cu_seq_lens,

2346 + )

2347 + return output_padded

2348 +

2349 if output.dtype == FP4_DTYPE:

2350 assert self.o_sf_scale is not None

2351 out = FP4Tensor(