[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>

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",

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,

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",

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(