[Attention] Extend XQA decode support on SM90 (#50439)
Signed-off-by: Andrii Skliar <askliar@nvidia.com> Co-authored-by: Andrii Skliar <askliar@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Andrii Skliar committed
Sep 11, 2026 at 08:19 UTC
4b839c3378c6a40fa613328a9b97069ed83b181b
3 files changed
+239
-105
tests/v1/attention/test_attention_backends.py
+147
-13
@@ -889,7 +889,7 @@ def test_flashinfer_xqa_query_lens_preserve_cudagraph_padding():
889
890
device = torch.device("cpu")
891
builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder)
892
- builder.use_dedicated_xqa = True
892
+ builder.use_xqa = True
893
qo_indptr = torch.tensor([0, 3, 9, 15, 15], dtype=torch.int32, device=device)
894
895
q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens(
@@ -905,6 +905,95 @@ def test_flashinfer_xqa_query_lens_preserve_cudagraph_padding():
905
assert q_cu_seq_lens.tolist() == [0, 3, 9, 15, 15]
906
907
908
+@pytest.mark.skipif(
909
+ AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
910
+ reason="FlashInfer is not available.",
911
+)
912
+def test_flashinfer_xqa_single_token_decode_preserves_cudagraph_padding(monkeypatch):
913
+ """Non-ragged XQA keeps one query/output row per padded request."""
914
+ import unittest.mock
915
+
916
+ from vllm.v1.attention.backends import flashinfer as flashinfer_backend
917
+
918
+ impl = object.__new__(flashinfer_backend.FlashInferImpl)
919
+ impl.scale = 1.0
920
+ impl.bmm1_scale = None
921
+ impl.bmm2_scale = None
922
+ impl.kv_cache_dtype = "auto"
923
+ impl.is_kvcache_nvfp4 = False
924
+ impl.head_size = 16
925
+ impl.dcp_world_size = 1
926
+ impl.o_sf_scale = None
927
+ impl.window_left = -1
928
+ impl.sinks = None
929
+ impl.cache_config = unittest.mock.Mock()
930
+ impl.cache_config.get_resolved_kv_cache_layout.return_value = KVCacheLayout.LBHNC
931
+
932
+ layer = unittest.mock.Mock(
933
+ _q_scale=torch.tensor(1.0),
934
+ _q_scale_float=1.0,
935
+ _k_scale_float=1.0,
936
+ _v_scale_float=1.0,
937
+ )
938
+ decode = flashinfer_backend.FlashInferTrtllmAPIDecode(
939
+ kernel=flashinfer_backend.FlashInferDecodeKernel.XQA,
940
+ block_tables=torch.zeros((4, 1), dtype=torch.int32),
941
+ seq_lens=torch.tensor([8, 8, 0, 0], dtype=torch.int32),
942
+ max_seq_len=8,
943
+ q_len_per_req=1,
944
+ )
945
+ attn_metadata = flashinfer_backend.FlashInferMetadata(
946
+ num_actual_tokens=2,
947
+ slot_mapping=torch.empty(0, dtype=torch.int64),
948
+ q_data_type_prefill=torch.bfloat16,
949
+ q_data_type_decode=torch.bfloat16,
950
+ num_decodes=4,
951
+ num_decode_tokens=2,
952
+ num_prefills=0,
953
+ num_prefill_tokens=0,
954
+ causal=True,
955
+ prefill=None,
956
+ decode=decode,
957
+ use_cascade=False,
958
+ cascade_wrapper=None,
959
+ )
960
+
961
+ seen_shapes = {}
962
+
963
+ def mock_xqa(**kwargs):
964
+ seen_shapes["query"] = kwargs["query"].shape
965
+ seen_shapes["out"] = kwargs["out"].shape
966
+ seen_shapes["block_tables"] = kwargs["block_tables"].shape
967
+
968
+ monkeypatch.setattr(
969
+ flashinfer_backend,
970
+ "_get_trtllm_workspace_buffer",
971
+ lambda: torch.empty(1, dtype=torch.uint8),
972
+ )
973
+ monkeypatch.setattr(
974
+ flashinfer_backend, "flashinfer_xqa_batch_decode_with_kv_cache", mock_xqa
975
+ )
976
+
977
+ query = torch.zeros((4, 1, 16), dtype=torch.bfloat16)
978
+ output = torch.empty_like(query)
979
+ result = impl.forward(
980
+ layer,
981
+ query,
982
+ torch.empty_like(query),
983
+ torch.empty_like(query),
984
+ torch.zeros((1, 1, 1, 32), dtype=torch.bfloat16),
985
+ attn_metadata,
986
+ output,
987
+ )
988
+
989
+ assert result is output
990
+ assert seen_shapes == {
991
+ "query": torch.Size([4, 1, 16]),
992
+ "out": torch.Size([4, 1, 16]),
993
+ "block_tables": torch.Size([4, 1]),
994
+ }
995
+
996
+
997
@pytest.mark.skipif(
998
AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
999
reason="FlashInfer is not available.",
@@ -913,7 +1002,7 @@ def test_flashinfer_xqa_query_lens_require_exact_uniform_product():
1002
from vllm.v1.attention.backends import flashinfer as flashinfer_backend
1003
1004
builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder)
916
- builder.use_dedicated_xqa = True
1005
+ builder.use_xqa = True
1006
qo_indptr = torch.tensor([0, 3, 3, 3], dtype=torch.int32)
1007
1008
q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens(
@@ -929,6 +1018,42 @@ def test_flashinfer_xqa_query_lens_require_exact_uniform_product():
1018
assert q_cu_seq_lens.tolist() == [0, 3, 3, 3]
1019
1020
1021
+@pytest.mark.skipif(
1022
+ AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
1023
+ reason="FlashInfer is not available.",
1024
+)
1025
+@pytest.mark.parametrize(
1026
+ ("qo_indptr_values", "num_decode_tokens", "expected_q_len"),
1027
+ [
1028
+ ([0, 3, 3], 3, 3),
1029
+ ([0, 2, 4, 4, 4], 4, 2),
1030
+ ],
1031
+)
1032
+def test_flashinfer_trtllm_gen_padded_decode_uses_varlen_offsets(
1033
+ qo_indptr_values: list[int],
1034
+ num_decode_tokens: int,
1035
+ expected_q_len: int,
1036
+):
1037
+ """Padded speculative decode keeps its actual packed query width."""
1038
+ from vllm.v1.attention.backends import flashinfer as flashinfer_backend
1039
+
1040
+ builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder)
1041
+ builder.use_xqa = False
1042
+ qo_indptr = torch.tensor(qo_indptr_values, dtype=torch.int32)
1043
+
1044
+ q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens(
1045
+ qo_indptr,
1046
+ qo_indptr,
1047
+ num_decodes=len(qo_indptr_values) - 1,
1048
+ num_decode_tokens=num_decode_tokens,
1049
+ )
1050
+
1051
+ assert q_len == expected_q_len
1052
+ assert q_lens is None
1053
+ assert q_cu_seq_lens is not None
1054
+ assert q_cu_seq_lens.tolist() == qo_indptr_values
1055
+
1056
+
1057
@pytest.mark.skipif(
1058
AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
1059
reason="FlashInfer is not available.",
@@ -961,10 +1086,12 @@ def test_flashinfer_attention_sinks_refreshed_after_reload(dtype):
1086
reason="FlashInfer is not available.",
1087
)
1088
def test_flashinfer_native_prefill_with_sinks(default_vllm_config):
964
- if not (
965
- current_platform.is_cuda() and current_platform.is_device_capability_family(120)
966
- ):
967
- pytest.skip("Native FlashInfer prefill with sinks requires SM12x.")
1089
+ supported = current_platform.is_cuda() and (
1090
+ current_platform.is_device_capability(90)
1091
+ or current_platform.is_device_capability_family(120)
1092
+ )
1093
+ if not supported:
1094
+ pytest.skip("Native FlashInfer prefill with sinks requires SM90 or SM12x.")
1095
1096
from vllm.v1.attention.backends.flashinfer import FlashInferBackend
1097
@@ -995,7 +1122,7 @@ def test_flashinfer_native_prefill_with_sinks(default_vllm_config):
1122
reason="FlashInfer is not available.",
1123
)
1124
def test_flashinfer_xqa_decode_correctness(default_vllm_config):
998
- """FlashInfer should route supported decode through XQA and match SDPA."""
1125
+ """FlashInfer should route SM90/SM12x decode through XQA and match SDPA."""
1126
supported = current_platform.is_cuda() and (
1127
current_platform.is_device_capability(90)
1128
or current_platform.is_device_capability_family(120)
@@ -1066,16 +1193,23 @@ def test_flashinfer_xqa_decode_correctness(default_vllm_config):
1193
)
1194
attn_metadata = builder.build(0, common_attn_metadata)
1195
1069
- expected_cg_support = (
1070
- AttentionCGSupport.UNIFORM_BATCH
1071
- if current_platform.is_device_capability_family(120)
1072
- else AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
1073
- )
1196
assert (
1197
flashinfer_backend.FlashInferMetadataBuilder.get_cudagraph_support(
1198
vllm_config, kv_cache_spec
1199
)
1078
- == expected_cg_support
1200
+ == AttentionCGSupport.UNIFORM_BATCH
1201
+ )
1202
+ wide_head_kv_cache_spec = FullAttentionSpec(
1203
+ block_size=vllm_config.cache_config.block_size,
1204
+ num_kv_heads=kv_cache_spec.num_kv_heads,
1205
+ head_size=512,
1206
+ dtype=vllm_config.model_config.dtype,
1207
+ )
1208
+ assert (
1209
+ flashinfer_backend.FlashInferMetadataBuilder.get_cudagraph_support(
1210
+ vllm_config, wide_head_kv_cache_spec
1211
+ )
1212
+ == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
1213
)
1214
assert isinstance(
1215
attn_metadata.decode,
tests/v1/attention/test_flashinfer_dcp_spec_reorder.py
+5
@@ -13,6 +13,7 @@ import torch
13
14
from tests.v1.attention.utils import create_vllm_config
15
from vllm.config import SpeculativeConfig, set_current_vllm_config
16
+from vllm.v1.attention.backend import AttentionCGSupport
17
from vllm.v1.attention.backends import flashinfer as flashinfer_backend
18
from vllm.v1.attention.backends.flashinfer import (
19
FlashInferDecodeKernel,
@@ -73,3 +74,7 @@ def test_flashinfer_gqa_dcp_spec_decode_clamps_reorder_threshold(monkeypatch):
74
builder.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN
75
)
76
assert builder.reorder_batch_threshold == 1
77
+ assert (
78
+ FlashInferMetadataBuilder.get_cudagraph_support(vllm_config, kv_cache_spec)
79
+ == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
80
+ )
vllm/v1/attention/backends/flashinfer.py
+87
-92
@@ -519,7 +519,7 @@ class FlashInferBackend(AttentionBackend):
519
520
@classmethod
521
def supports_sink(cls) -> bool:
522
- """FlashInfer supports sinks on SM12x XQA and SM100 trtllm-gen."""
522
+ """Whether FlashInfer can serve attention sinks on this platform."""
523
from vllm.utils.flashinfer import (
524
force_use_trtllm_attention,
525
)
@@ -527,7 +527,9 @@ class FlashInferBackend(AttentionBackend):
527
if force_use_trtllm_attention() is False:
528
return False
529
530
- if current_platform.is_device_capability_family(120):
530
+ if current_platform.is_device_capability(
531
+ 90
532
+ ) or current_platform.is_device_capability_family(120):
533
return supports_trtllm_attention(is_prefill=False)
534
535
if not current_platform.is_device_capability_family(100):
@@ -570,6 +572,10 @@ class FlashInferDecodeKernel(Enum):
572
TRTLLM_GEN = "trtllm-gen"
573
574
575
+def _is_xqa_head_dim_supported(head_dim: int) -> bool:
576
+ return 16 <= head_dim <= 256 and head_dim % 16 == 0
577
+
578
+
579
@dataclass
580
class TRTLLMPrefill:
581
"""Metadata for the TRTLLM prefill pathway."""
@@ -827,7 +833,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
833
# the wider groups fall back to native FlashInfer decode.
834
if (
835
self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
830
- and not (16 <= self.head_dim <= 256 and self.head_dim % 16 == 0)
836
+ and not _is_xqa_head_dim_supported(self.head_dim)
837
):
838
logger.warning_once(
839
"FlashInfer XQA decode does not support head_dim=%d; "
@@ -847,9 +853,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
853
)
854
self.use_trtllm_decode_attention = False
855
self.flashinfer_trtllm_api_decode_kernel = None
850
- self.use_dedicated_xqa = (
851
- current_platform.is_device_capability_family(120)
852
- and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
856
+ self.use_xqa = (
857
+ self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
858
)
859
# Adaptive verification trims drafts on device, so decode query lengths
860
# must come from the device qo_indptr; only trtllm-gen supports that
@@ -861,14 +866,11 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
866
== FlashInferDecodeKernel.TRTLLM_GEN
867
and not self.use_dcp
868
)
864
- supports_spec_as_decode = (
865
- self.flashinfer_trtllm_api_decode_kernel
866
- == FlashInferDecodeKernel.TRTLLM_GEN
867
- or self.use_dedicated_xqa
868
- )
869
self._init_reorder_batch_threshold(
870
1,
871
- supports_spec_as_decode=supports_spec_as_decode,
871
+ supports_spec_as_decode=(
872
+ self.flashinfer_trtllm_api_decode_kernel is not None
873
+ ),
874
# trtllm-gen decode receives no cp_rank/global-seq-len information,
875
# so its end-aligned causal mask is wrong for q_len > 1 over the
876
# DCP-interleaved local KV shard (spec token i misses up to
@@ -886,16 +888,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
888
per_layer_parameters = get_per_layer_parameters(
889
vllm_config, layer_names, FlashInferImpl
890
)
889
- if current_platform.is_device_capability(90) and any(
890
- params.window_left != -1 for params in per_layer_parameters.values()
891
- ):
892
- # FlashInfer SM90 sliding-window prefill is not reliable with FP8-Q:
893
- # https://github.com/flashinfer-ai/flashinfer/issues/3578
894
- raise NotImplementedError(
895
- "FlashInfer backend on SM90 currently crashes with "
896
- "sliding-window attention layers. Use the default attention "
897
- "backend."
898
- )
891
self.global_hyperparameters = infer_global_hyperparameters(per_layer_parameters)
892
self.sm_scale = self.global_hyperparameters.sm_scale
893
self.window_left = self.global_hyperparameters.window_left
@@ -950,7 +942,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
942
# if cache_config requests a quantized dtype globally.
943
cache_dtype = self.cache_dtype
944
953
- # XQA decode requires BF16/FP16-Q even with FP8 KV cache.
945
+ # On SM90/SM12x, XQA decode requires BF16/FP16-Q even with FP8 KV cache.
946
+ # FI native prefill on SM90 still uses FP8-Q in that case; SM12x prefill
947
+ # is fa2-only and keeps the model dtype (handled below).
948
if (
949
(
950
current_platform.is_device_capability(90)
@@ -964,7 +958,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
958
959
# Otherwise, match Q dtype to the KV cache dtype.
960
if cache_dtype.startswith("fp8"):
967
- # FP8-Q requires an fp8 tensor-core attention path.
961
+ # FP8-Q requires an fp8 tensor-core attention path. SM90 uses
962
+ # native FA3 for prefill; XQA decode on SM90/SM12x is handled
963
+ # above with model-dtype Q, while SM100 uses trtllm-gen.
964
# Architectures with only fa2 (e.g. SM89, SM120) cannot
965
# consume FP8 queries, so keep the model dtype for Q there.
966
if current_platform.is_device_capability(
@@ -983,24 +979,19 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
979
vllm_config: VllmConfig,
980
kv_cache_spec: KVCacheSpec,
981
) -> AttentionCGSupport:
986
- """Get the cudagraph support level for FlashInfer attention.
987
-
988
- SM90 XQA supports only single-token decode. SM12x uses the dedicated
989
- XQA API, which supports speculative and non-causal decode.
990
- """
991
- if current_platform.is_device_capability(90):
992
- return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
993
-
994
- is_sm12x = current_platform.is_device_capability_family(120)
995
- # XQA does not return LSE and therefore does not support DCP.
996
- if is_sm12x and vllm_config.parallel_config.decode_context_parallel_size > 1:
982
+ """Get the cudagraph support level for FlashInfer attention."""
983
+ # XQA lacks LSE for DCP; DCP also cannot graph variable-length trtllm-gen.
984
+ if vllm_config.parallel_config.decode_context_parallel_size > 1:
985
return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
986
987
kv_specs = iter_layer_specs(kv_cache_spec)
988
num_qo_heads = vllm_config.model_config.get_num_attention_heads(
989
vllm_config.parallel_config
990
)
1003
- has_trtllm_support: bool = len(kv_specs) > 0
991
+ is_xqa_arch = current_platform.is_device_capability(
992
+ 90
993
+ ) or current_platform.is_device_capability_family(120)
994
+ has_uniform_batch_support: bool = len(kv_specs) > 0
995
for spec in kv_specs:
996
if not isinstance(spec, AttentionSpec):
997
# FlashInfer only applies to attention, so we don't consider other types
@@ -1010,13 +1001,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1001
num_qo_heads=num_qo_heads,
1002
num_kv_heads=spec.num_kv_heads,
1003
is_prefill=False,
1013
- ):
1014
- has_trtllm_support = False
1004
+ ) or (is_xqa_arch and not _is_xqa_head_dim_supported(spec.head_size)):
1005
+ has_uniform_batch_support = False
1006
break
1007
1017
- if has_trtllm_support and (
1018
- is_sm12x or not vllm_config.attention_config.use_non_causal
1019
- ):
1008
+ use_non_causal = vllm_config.attention_config.use_non_causal
1009
+ if has_uniform_batch_support and (not use_non_causal or is_xqa_arch):
1010
return AttentionCGSupport.UNIFORM_BATCH
1011
else:
1012
return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
@@ -1066,10 +1056,27 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1056
num_decode_tokens: int,
1057
) -> tuple[int, torch.Tensor | None, list[int] | None]:
1058
"""Return the query width, ragged offsets, and effective query lengths."""
1069
- assert self.use_dedicated_xqa
1059
if num_decodes == 0 or num_decode_tokens == 0:
1060
return 1, None, None
1061
1062
+ if not self.use_xqa:
1063
+ decode_q_lens = (
1064
+ qo_indptr_cpu[1 : num_decodes + 1] - qo_indptr_cpu[:num_decodes]
1065
+ )
1066
+ q_len_per_req = num_decode_tokens // num_decodes
1067
+ if num_decode_tokens % num_decodes == 0 and bool(
1068
+ (decode_q_lens == q_len_per_req).all().item()
1069
+ ):
1070
+ return q_len_per_req, None, None
1071
+
1072
+ # CUDA-graph padding can leave zero-length requests after a
1073
+ # uniform speculative decode batch. Keep the packed query and
1074
+ # describe it through the trtllm-gen varlen API.
1075
+ max_q_len = int(decode_q_lens.max().item())
1076
+ if max_q_len > 1:
1077
+ return max_q_len, qo_indptr[: num_decodes + 1], None
1078
+ return 1, None, None
1079
+
1080
decode_q_lens = qo_indptr_cpu[1 : num_decodes + 1] - qo_indptr_cpu[:num_decodes]
1081
nonzero = decode_q_lens[decode_q_lens > 0]
1082
q_len_per_req = int(nonzero.max().item()) if nonzero.numel() > 0 else 1
@@ -1128,7 +1135,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1135
"NVFP4 KV cache."
1136
)
1137
if self._noncausal_prefill_wrapper is None:
1131
- if self.has_sinks and current_platform.is_device_capability_family(120):
1138
+ if self.has_sinks and (
1139
+ current_platform.is_device_capability(90)
1140
+ or current_platform.is_device_capability_family(120)
1141
+ ):
1142
self._noncausal_prefill_wrapper = (
1143
BatchAttentionWithAttentionSinkWrapper(
1144
self._get_workspace_buffer(),
@@ -1159,7 +1169,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1169
dcp_a2a=self.dcp_a2a,
1170
)
1171
else:
1162
- if self.has_sinks and current_platform.is_device_capability_family(120):
1172
+ if self.has_sinks and (
1173
+ current_platform.is_device_capability(90)
1174
+ or current_platform.is_device_capability_family(120)
1175
+ ):
1176
assert not self.is_kvcache_nvfp4
1177
self._prefill_wrapper = BatchAttentionWithAttentionSinkWrapper(
1178
self._get_workspace_buffer(),
@@ -1296,13 +1309,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1309
num_reqs = common_attn_metadata.num_reqs
1310
num_actual_tokens = common_attn_metadata.num_actual_tokens
1311
causal = common_attn_metadata.causal
1299
- route_decode = causal or self.use_dedicated_xqa
1312
+ route_decode = causal or self.use_xqa
1313
if route_decode:
1314
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
1315
split_decodes_and_prefills(
1316
common_attn_metadata,
1317
decode_threshold=self.reorder_batch_threshold,
1305
- require_uniform=not self.use_dedicated_xqa,
1318
+ require_uniform=not self.use_xqa,
1319
)
1320
)
1321
else:
@@ -1342,7 +1355,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1355
has_spec=uses_spec_reorder,
1356
)
1357
decode_with_flashinfer_trtllm_api = self.use_trtllm_decode_attention and (
1345
- causal or self.use_dedicated_xqa
1358
+ causal or self.use_xqa
1359
)
1360
1361
if not causal and self.use_dcp:
@@ -1360,13 +1373,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1373
)
1374
1375
if not all_uses_trtllm:
1363
- if self.has_sinks and (
1364
- not self.use_dedicated_xqa or self.use_dcp or use_cascade
1365
- ):
1376
+ sinks_would_be_dropped = self.use_dcp or use_cascade
1377
+ if self.has_sinks and sinks_would_be_dropped:
1378
raise NotImplementedError(
1367
- "FlashInfer backend currently does not support attention "
1368
- "sinks, please use trtllm on blackwell or flash attention "
1369
- "on earlier GPUs."
1379
+ "FlashInfer BatchDCPPrefillWrapper and "
1380
+ "MultiLevelCascadeAttentionWrapper do not support attention sinks. "
1381
+ "Please use TRTLLM on Blackwell or FlashAttention on earlier GPUs."
1382
)
1383
1384
if not self.global_hyperparameters.has_same_window_lefts:
@@ -1663,29 +1675,19 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1675
if num_decodes > 0:
1676
if decode_with_flashinfer_trtllm_api:
1677
assert self.flashinfer_trtllm_api_decode_kernel is not None
1666
- if not self.use_dedicated_xqa and not self.use_trtllm_gen_varlen_decode:
1667
- assert num_decode_tokens % num_decodes == 0, (
1668
- "XQA/trtllm-gen decode requires uniform query lengths "
1669
- f"per request. Got {num_decode_tokens=} and {num_decodes=}."
1670
- )
1678
seq_lens_decode = seq_lens[:num_decodes]
1679
if self.use_dcp:
1680
assert common_attn_metadata.dcp_local_seq_lens is not None
1681
seq_lens_decode = common_attn_metadata.dcp_local_seq_lens[
1682
:num_decodes
1683
]
1677
- q_len_per_req = 1
1678
- q_cu_seq_lens = None
1679
- decode_mask = None
1680
- if self.use_dedicated_xqa:
1681
- q_len_per_req, q_cu_seq_lens, ragged_q_lens = (
1682
- self._compute_decode_query_lens(
1683
- qo_indptr,
1684
- qo_indptr_cpu,
1685
- num_decodes,
1686
- num_decode_tokens,
1687
- )
1684
+ q_len_per_req, q_cu_seq_lens, ragged_q_lens = (
1685
+ self._compute_decode_query_lens(
1686
+ qo_indptr, qo_indptr_cpu, num_decodes, num_decode_tokens
1687
)
1688
+ )
1689
+ decode_mask = None
1690
+ if self.use_xqa:
1691
decode_mask = self._get_decode_mask(
1692
q_len_per_req,
1693
ragged_q_lens,
@@ -2066,7 +2068,10 @@ class FlashInferImpl(AttentionImpl):
2068
if fp8_view_dtype is not None:
2069
kv_cache = kv_cache.view(fp8_view_dtype)
2070
2069
- # Inputs and outputs may be padded for CUDA graphs
2071
+ # Inputs and outputs may be padded for CUDA graphs. Keep the original
2072
+ # query for single-token XQA, whose non-ragged API requires one row for
2073
+ # every (including padded) decode request.
2074
+ query_padded = query
2075
query = query[:num_actual_tokens]
2076
output_padded = output
2077
output = output[:num_actual_tokens]
@@ -2131,10 +2136,7 @@ class FlashInferImpl(AttentionImpl):
2136
kv_cache_tuple = kv_cache_permute.split(hs, dim=-1)
2137
2138
use_dcp = self.dcp_world_size > 1
2134
- decode_with_dedicated_xqa = (
2135
- decode_with_xqa and current_platform.is_device_capability_family(120)
2136
- )
2137
- if decode_with_dedicated_xqa:
2139
+ if decode_with_xqa:
2140
assert not use_dcp
2141
assert not self.is_kvcache_nvfp4
2142
assert self.o_sf_scale is None
@@ -2351,8 +2353,13 @@ class FlashInferImpl(AttentionImpl):
2353
].copy_(out[:num_prefill_tokens])
2354
2355
if num_decode_tokens > 0:
2354
- decode_query = query[:num_decode_tokens]
2355
- assert decode_query.shape[0] == num_decode_tokens
2356
+ decode_query_tokens = num_decode_tokens
2357
+ if decode_with_xqa:
2358
+ assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode)
2359
+ if attn_metadata.decode.q_len_per_req == 1:
2360
+ decode_query_tokens = attn_metadata.num_decodes
2361
+ decode_query = query_padded[:decode_query_tokens]
2362
+ assert decode_query.shape[0] == decode_query_tokens
2363
2364
# Convert query to the expected dtype for decode if needed.
2365
decode_query = self.maybe_quant_query(
@@ -2467,7 +2474,7 @@ class FlashInferImpl(AttentionImpl):
2474
)
2475
decode_query = canonicalize_singleton_dim_strides(decode_query)
2476
2470
- if decode_with_dedicated_xqa:
2477
+ if decode_with_xqa:
2478
bmm1_scale = self.get_xqa_bmm1_scale(
2479
layer, attn_metadata.q_data_type_decode
2480
)
@@ -2483,7 +2490,7 @@ class FlashInferImpl(AttentionImpl):
2490
bmm1_scale=bmm1_scale,
2491
bmm2_scale=self.bmm2_scale,
2492
window_left=self.window_left,
2486
- out=output[:num_decode_tokens],
2493
+ out=output_padded[:decode_query_tokens],
2494
sinks=self.sinks,
2495
kv_layout=get_flashinfer_layout_string(self.kv_cache_layout),
2496
q_len_per_req=q_len_per_req,
@@ -2492,6 +2499,7 @@ class FlashInferImpl(AttentionImpl):
2499
)
2500
return output_padded
2501
2502
+ assert decode_with_trtllm_gen
2503
if output.dtype == FP4_DTYPE:
2504
assert self.o_sf_scale is not None
2505
out = FP4Tensor(
@@ -2523,19 +2531,6 @@ class FlashInferImpl(AttentionImpl):
2531
else:
2532
q_len_per_req = num_decode_tokens // attn_metadata.num_decodes
2533
2526
- if decode_with_xqa and q_len_per_req is not None and q_len_per_req > 1:
2527
- raise NotImplementedError(
2528
- "FlashInfer XQA speculative decode is not wired in vLLM yet."
2529
- )
2530
-
2531
- # XQA decode can use model-dtype Q with FP8 KV, so only include
2532
- # q_scale when the decode query is actually FP8.
2533
- bmm1_scale = (
2534
- self.get_xqa_bmm1_scale(layer, attn_metadata.q_data_type_decode)
2535
- if decode_with_xqa
2536
- else self.bmm1_scale
2537
- )
2538
-
2534
lse = None
2535
if use_dcp:
2536
out = torch.empty(
@@ -2558,7 +2553,7 @@ class FlashInferImpl(AttentionImpl):
2553
block_tables=block_tables_decode,
2554
seq_lens=seq_lens_decode,
2555
max_seq_len=attn_metadata.decode.max_seq_len,
2561
- bmm1_scale=bmm1_scale,
2556
+ bmm1_scale=self.bmm1_scale,
2557
bmm2_scale=self.bmm2_scale,
2558
window_left=self.window_left,
2559
sinks=self.sinks,