@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d29125c085

Xqa decode kernels (#43232)

Signed-off-by: Dan Blanaru <48605845+DanBlanaru@users.noreply.github.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>

Dan Blanaru committed Jul 2, 2026 at 21:32 UTC d29125c0852eb61efef060d55675010d1abf51fe
7 files changed +538 -143
docs/design/attention_backends.md
+3 -2
@@ -161,7 +161,8 @@ Priority is **1 = highest** (tried first).
161 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
162 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ✅ | ❌ | ❌ | All | N/A |
163 | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x |
164 -| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
164 +| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
165 +| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
166 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
167 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
168 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
@@ -175,7 +176,7 @@ Priority is **1 = highest** (tried first).
176 | `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
177 | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
178
178 -> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
179 +> **†** FlashInfer Native is the regular FlashInfer path. XQA is the SM90 decode path exposed through FlashInfer's TRTLLM decode API. trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen via `--attention-config.use_trtllm_attention=0`.
180 >
181 > **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.
182
tests/kernels/attention/test_use_trtllm_attention.py
+16 -1
@@ -78,14 +78,29 @@ def test_supports_sm100_with_artifactory(_art, _cap):
78
79
80 @patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
81 +@patch(
82 + "vllm.utils.flashinfer.current_platform.is_device_capability", return_value=False
83 +)
84 @patch(
85 "vllm.utils.flashinfer.current_platform.is_device_capability_family",
86 return_value=False,
87 )
85 -def test_supports_non_sm100_platform(_cap):
88 +def test_supports_unsupported_platform(_family, _cap):
89 assert supports_trtllm_attention() is False
90
91
92 +@patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
93 +@patch("vllm.utils.flashinfer.current_platform.is_device_capability", return_value=True)
94 +@patch(
95 + "vllm.utils.flashinfer.current_platform.is_device_capability_family",
96 + return_value=False,
97 +)
98 +@patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)
99 +def test_supports_sm90_decode_only(_art, _family, _cap):
100 + assert supports_trtllm_attention(is_prefill=False) is True
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_family",
tests/v1/attention/test_attention_backends.py
+120 -4
@@ -16,7 +16,7 @@ from tests.v1.attention.utils import (
16 try_backend_includes_kv_cache_update,
17 try_get_attention_backend,
18 )
19 -from vllm.config import ModelConfig
19 +from vllm.config import ModelConfig, set_current_vllm_config
20 from vllm.platforms import current_platform
21 from vllm.utils.math_utils import cdiv
22 from vllm.utils.torch_utils import (
@@ -24,7 +24,11 @@ from vllm.utils.torch_utils import (
24 is_torch_equal_or_newer,
25 set_random_seed,
26 )
27 -from vllm.v1.attention.backend import AttentionType, CommonAttentionMetadata
27 +from vllm.v1.attention.backend import (
28 + AttentionCGSupport,
29 + AttentionType,
30 + CommonAttentionMetadata,
31 +)
32 from vllm.v1.attention.backends.registry import AttentionBackendEnum
33 from vllm.v1.attention.backends.utils import (
34 set_kv_cache_layout,
@@ -626,6 +630,118 @@ def test_causal_backend_correctness(
630 )
631
632
633 +@pytest.mark.skipif(
634 + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
635 + reason="FlashInfer is not available.",
636 +)
637 +def test_flashinfer_xqa_bmm1_scale_matches_decode_q_dtype():
638 + """XQA decode should only apply q_scale when decode Q is FP8."""
639 + from vllm.v1.attention.backends import flashinfer as flashinfer_backend
640 +
641 + class MockLayer:
642 + _q_scale_float = 2.0
643 + _k_scale_float = 3.0
644 +
645 + impl = object.__new__(flashinfer_backend.FlashInferImpl)
646 + impl.scale = 0.5
647 + impl.kv_cache_dtype = "fp8"
648 +
649 + assert impl.get_xqa_bmm1_scale(MockLayer, torch.bfloat16) == 1.5
650 + assert impl.get_xqa_bmm1_scale(MockLayer, torch.float8_e4m3fn) == 3.0
651 +
652 +
653 +@pytest.mark.skipif(
654 + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
655 + reason="FlashInfer is not available.",
656 +)
657 +def test_flashinfer_sm90_xqa_decode_correctness(default_vllm_config):
658 + """FlashInfer should route Hopper decode through XQA and match SDPA."""
659 + if not current_platform.is_cuda() or not current_platform.is_device_capability(90):
660 + pytest.skip("FlashInfer XQA decode requires SM90.")
661 +
662 + import unittest.mock
663 +
664 + from vllm.utils.flashinfer import can_use_trtllm_attention
665 + from vllm.v1.attention.backends import flashinfer as flashinfer_backend
666 + from vllm.v1.attention.backends.utils import PerLayerParameters
667 +
668 + def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
669 + return {
670 + "placeholder": PerLayerParameters(
671 + window_left=-1,
672 + logits_soft_cap=0.0,
673 + sm_scale=1.0,
674 + )
675 + }
676 +
677 + def causal_mask_mod(
678 + b: torch.Tensor,
679 + h: torch.Tensor,
680 + q_idx: torch.Tensor,
681 + kv_idx: torch.Tensor,
682 + *,
683 + context_len: int,
684 + ):
685 + return (q_idx + context_len) >= kv_idx
686 +
687 + batch_spec = BATCH_SPECS["small_decode"]
688 + vllm_config = create_vllm_config(
689 + model_name="meta-llama/Meta-Llama-3-8B",
690 + max_model_len=max(batch_spec.seq_lens),
691 + block_size=16,
692 + )
693 + device = torch.device(f"{DEVICE_TYPE}:0")
694 + kv_cache_spec = FullAttentionSpec(
695 + block_size=vllm_config.cache_config.block_size,
696 + num_kv_heads=vllm_config.model_config.get_num_kv_heads(
697 + vllm_config.parallel_config
698 + ),
699 + head_size=vllm_config.model_config.get_head_size(),
700 + dtype=vllm_config.model_config.dtype,
701 + )
702 +
703 + with set_current_vllm_config(vllm_config):
704 + if not can_use_trtllm_attention(
705 + vllm_config.model_config.get_num_attention_heads(
706 + vllm_config.parallel_config
707 + ),
708 + kv_cache_spec.num_kv_heads,
709 + is_prefill=False,
710 + ):
711 + pytest.skip("FlashInfer XQA decode is not available in this setup.")
712 +
713 + with unittest.mock.patch(
714 + "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
715 + mock_get_per_layer_parameters,
716 + ):
717 + builder = flashinfer_backend.FlashInferMetadataBuilder(
718 + kv_cache_spec, ["placeholder"], vllm_config, device
719 + )
720 + common_attn_metadata = create_common_attn_metadata(
721 + batch_spec, vllm_config.cache_config.block_size, device
722 + )
723 + attn_metadata = builder.build(0, common_attn_metadata)
724 +
725 + assert (
726 + flashinfer_backend.FlashInferMetadataBuilder.get_cudagraph_support(
727 + vllm_config, kv_cache_spec
728 + )
729 + == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
730 + )
731 + assert isinstance(
732 + attn_metadata.decode,
733 + flashinfer_backend.FlashInferTrtllmAPIDecode,
734 + )
735 + assert attn_metadata.decode.kernel == flashinfer_backend.FlashInferDecodeKernel.XQA
736 +
737 + _test_backend_correctness(
738 + batch_spec,
739 + "meta-llama/Meta-Llama-3-8B",
740 + [AttentionBackendEnum.FLASHINFER],
741 + causal_mask_mod,
742 + )
743 +
744 +
745 if current_platform.is_rocm():
746 # FLASH_ATTN is not supported on ROCm
747 SLIDING_WINDOW_BACKENDS_TO_TEST = [
@@ -656,7 +772,7 @@ else:
772 @pytest.mark.parametrize("model", ["microsoft/Phi-tiny-MoE-instruct"])
773 @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
774 def test_sliding_window_backend_correctness(
659 - batch_spec_name: str, model: str, tensor_parallel_size: int
775 + default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
776 ):
777 """Test backend's correctness with sliding window attention."""
778
@@ -718,7 +834,7 @@ def test_sliding_window_backend_correctness(
834 @pytest.mark.parametrize("model", ["google/embeddinggemma-300m"])
835 @pytest.mark.parametrize("tensor_parallel_size", [1, 2])
836 def test_sliding_window_encoder_backend_correctness(
721 - batch_spec_name: str, model: str, tensor_parallel_size: int
837 + default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
838 ):
839 """Test backend's correctness with sliding window attention."""
840
tests/v1/attention/test_trtllm_attention_integration.py
+6 -3
@@ -32,9 +32,10 @@ if not current_platform.is_device_capability_family(100):
32 )
33
34 from vllm.v1.attention.backends.flashinfer import ( # noqa: E402
35 + FlashInferDecodeKernel,
36 FlashInferImpl,
37 FlashInferMetadataBuilder,
37 - TRTLLMDecode,
38 + FlashInferTrtllmAPIDecode,
39 TRTLLMPrefill,
40 )
41
@@ -435,9 +436,11 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL)
436 f"Expected TRTLLMPrefill, got {type(attn_metadata.prefill)}"
437 )
438 if has_decodes:
438 - assert isinstance(attn_metadata.decode, TRTLLMDecode), (
439 - f"Expected TRTLLMDecode, got {type(attn_metadata.decode)}"
439 + assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode), (
440 + "Expected FlashInferTrtllmAPIDecode, got "
441 + f"{type(attn_metadata.decode)}"
442 )
443 + assert attn_metadata.decode.kernel == FlashInferDecodeKernel.TRTLLM_GEN
444
445 impl = FlashInferImpl(
446 num_heads=num_q_heads,
tools/pre_commit/generate_attention_backend_docs.py
+70 -23
@@ -261,6 +261,25 @@ def _find_cc_in_function(tree: ast.AST, func_name: str) -> str | None:
261 return None
262
263
264 +def _find_exact_cc_in_function(tree: ast.AST, func_name: str) -> str | None:
265 + """Find a compute capability from is_device_capability() calls in a function."""
266 + for node in ast.walk(tree):
267 + if not isinstance(node, ast.FunctionDef) or node.name != func_name:
268 + continue
269 + for n in ast.walk(node):
270 + if (
271 + isinstance(n, ast.Call)
272 + and isinstance(n.func, ast.Attribute)
273 + and n.func.attr == "is_device_capability"
274 + and n.args
275 + and isinstance(n.args[0], ast.Constant)
276 + and isinstance(n.args[0].value, int)
277 + ):
278 + capability = n.args[0].value
279 + return f"{capability // 10}.{capability % 10}"
280 + return None
281 +
282 +
283 # ---------------------------------------------------------------------------
284 # Registry and file resolution
285 # ---------------------------------------------------------------------------
@@ -1048,10 +1067,11 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
1067
1068
1069 def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:
1051 - """Parse flashinfer.py to detect TRTLLM-specific features.
1070 + """Parse flashinfer.py to detect FlashInfer TRTLLM API variants.
1071
1053 - FLASHINFER uses TRTLLM attention on SM100 (Blackwell), which has different
1054 - capabilities (e.g., sink support) than native FlashInfer on earlier GPUs.
1072 + FLASHINFER uses XQA on SM90 and trtllm-gen on SM100 through FlashInfer's
1073 + TRTLLM decode API. These variants have different capabilities than native
1074 + FlashInfer.
1075 """
1076 if not FLASHINFER_UTILS_FILE.exists():
1077 return {}
@@ -1061,9 +1081,10 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:
1081 except Exception:
1082 return {}
1083
1064 - trtllm_compute_cap = _find_cc_in_function(tree, "supports_trtllm_attention")
1084 + xqa_compute_cap = _find_exact_cc_in_function(tree, "supports_trtllm_attention")
1085 + trtllm_gen_compute_cap = _find_cc_in_function(tree, "supports_trtllm_attention")
1086
1066 - if not trtllm_compute_cap:
1087 + if not xqa_compute_cap and not trtllm_gen_compute_cap:
1088 return {}
1089
1090 # KV cache dtypes that only work with a dedicated kernel (e.g. nvfp4
@@ -1073,12 +1094,17 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:
1094
1095 return {
1096 "native": {
1076 - # Native FlashInfer: everything except SM100
1097 + # Native FlashInfer path.
1098 "supports_sink": False,
1099 },
1079 - "trtllm": {
1080 - # TRTLLM pathway on Blackwell
1081 - "compute_capability": trtllm_compute_cap,
1100 + "xqa": {
1101 + # XQA decode path on Hopper.
1102 + "compute_capability": xqa_compute_cap,
1103 + "supports_sink": False,
1104 + },
1105 + "trtllm_gen": {
1106 + # trtllm-gen pathway on Blackwell.
1107 + "compute_capability": trtllm_gen_compute_cap,
1108 "supports_sink": True,
1109 },
1110 "exclude_kv_dtypes": kernel_only_kv_dtypes,
@@ -1086,7 +1112,7 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:
1112
1113
1114 # ---------------------------------------------------------------------------
1089 -# Backend variant expansion (FA2/FA3/FA4, FlashInfer native/TRTLLM)
1115 +# Backend variant expansion (FA2/FA3/FA4, FlashInfer native/XQA/trtllm-gen)
1116 # ---------------------------------------------------------------------------
1117
1118
@@ -1146,7 +1172,7 @@ def _expand_flashinfer_variants(
1172 all_backends: list[dict[str, Any]],
1173 fi_features: dict[str, dict[str, Any]],
1174 ) -> list[dict[str, Any]]:
1149 - """Expand FLASHINFER into native and TRTLLM variants."""
1175 + """Expand FLASHINFER into native, XQA, and trtllm-gen variants."""
1176 expanded = []
1177 for backend in all_backends:
1178 if backend["name"] != "FLASHINFER":
@@ -1157,9 +1183,8 @@ def _expand_flashinfer_variants(
1183 orig_cap = backend["compute_capability"]
1184 parts = orig_cap.replace(".x", "").split("-")
1185 min_cc = parts[0] if parts else "7"
1160 - trtllm_cc = fi_features["trtllm"]["compute_capability"]
1186
1162 - # Create native entry (pre-Blackwell GPUs)
1187 + # Create native entry.
1188 native = backend.copy()
1189 native["version"] = "Native†"
1190 native["_sort_key"] = "FLASHINFER"
@@ -1176,16 +1201,36 @@ def _expand_flashinfer_variants(
1201 if d not in exclude
1202 )
1203
1179 - # Create TRTLLM entry
1180 - trtllm = backend.copy()
1181 - trtllm["version"] = "TRTLLM†"
1182 - trtllm["_sort_key"] = "FLASHINFER"
1183 - trtllm["_sort_order"] = 1
1184 - trtllm["compute_capability"] = trtllm_cc
1185 - trtllm["supports_sink"] = fi_features["trtllm"]["supports_sink"]
1204 + # Create XQA entry.
1205 + xqa = backend.copy()
1206 + xqa["version"] = "XQA†"
1207 + xqa["_sort_key"] = "FLASHINFER"
1208 + xqa["_sort_order"] = 1
1209 + xqa["compute_capability"] = fi_features["xqa"]["compute_capability"]
1210 + xqa["supports_sink"] = fi_features["xqa"]["supports_sink"]
1211 + xqa["supports_non_causal"] = False
1212 + if exclude:
1213 + xqa["kv_cache_dtypes"] = ", ".join(
1214 + d
1215 + for d in (d.strip() for d in xqa["kv_cache_dtypes"].split(","))
1216 + if d not in exclude
1217 + )
1218 +
1219 + # Create trtllm-gen entry.
1220 + trtllm_gen = backend.copy()
1221 + trtllm_gen["version"] = "trtllm-gen†"
1222 + trtllm_gen["_sort_key"] = "FLASHINFER"
1223 + trtllm_gen["_sort_order"] = 2
1224 + trtllm_gen["compute_capability"] = fi_features["trtllm_gen"][
1225 + "compute_capability"
1226 + ]
1227 + trtllm_gen["supports_sink"] = fi_features["trtllm_gen"]["supports_sink"]
1228
1229 expanded.append(native)
1188 - expanded.append(trtllm)
1230 + if fi_features["xqa"]["compute_capability"]:
1231 + expanded.append(xqa)
1232 + if fi_features["trtllm_gen"]["compute_capability"]:
1233 + expanded.append(trtllm_gen)
1234 return expanded
1235
1236
@@ -1788,8 +1833,10 @@ def generate_docs() -> str:
1833 footnotes = []
1834 if fi_features:
1835 footnotes.append(
1791 - "> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which "
1792 - "supports sinks. Disable via `--attention-config.use_trtllm_attention=0`."
1836 + "> **†** FlashInfer Native is the regular FlashInfer path. XQA is the "
1837 + "SM90 decode path exposed through FlashInfer's TRTLLM decode API. "
1838 + "trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen "
1839 + "via `--attention-config.use_trtllm_attention=0`."
1840 )
1841 if fa_features:
1842 footnotes.append(
vllm/utils/flashinfer.py
+37 -18
@@ -371,19 +371,27 @@ def has_nvidia_artifactory() -> bool:
371
372
373 @functools.cache
374 -def supports_trtllm_attention() -> bool:
375 - """
376 - TRTLLM attention is supported if the platform is SM100,
377 - NVIDIA artifactory is accessible, and batch-invariant mode is not enabled.
374 +def supports_trtllm_attention(is_prefill: bool = False) -> bool:
375 + """Return whether TRTLLM attention is available on the current platform
376 + for the given attention phase.
377 +
378 + SM90 (Hopper) supports the XQA decode kernel but not TRTLLM prefill.
379 + SM100+ supports TRTLLM for both phases. All others are unsupported.
380 """
381 # Batch-invariant mode disables TRTLLM attention
382 if envs.VLLM_BATCH_INVARIANT:
383 return False
384
383 - # Requires SM100 and NVIDIA artifactory to be accessible to download cubins
384 - return (
385 - current_platform.is_device_capability_family(100) and has_nvidia_artifactory()
386 - )
385 + # Requires NVIDIA artifactory to be accessible to download cubins
386 + if not has_nvidia_artifactory():
387 + return False
388 +
389 + # SM90 has XQA decode; prefill is not supported.
390 + if current_platform.is_device_capability(90):
391 + return not is_prefill
392 +
393 + # SM100/SM103 has both prefill and decode TRTLLM kernels.
394 + return current_platform.is_device_capability_family(100)
395
396
397 def force_use_trtllm_attention() -> bool | None:
@@ -400,12 +408,15 @@ def force_use_trtllm_attention() -> bool | None:
408 return vllm_config.attention_config.use_trtllm_attention
409
410
403 -def can_use_trtllm_attention(num_qo_heads: int, num_kv_heads: int) -> bool:
411 +def can_use_trtllm_attention(
412 + num_qo_heads: int, num_kv_heads: int, is_prefill: bool = False
413 +) -> bool:
414 """Check if the current configuration supports TRTLLM attention."""
415 if force_use_trtllm_attention() is False:
416 return False
407 - has_trtllm = supports_trtllm_attention()
408 - return has_trtllm and (num_qo_heads % num_kv_heads == 0)
417 + return supports_trtllm_attention(is_prefill=is_prefill) and (
418 + num_qo_heads % num_kv_heads == 0
419 + )
420
421
422 def use_trtllm_attention(
@@ -437,11 +448,12 @@ def use_trtllm_attention(
448 return False
449
450 # The platform is not supported
440 - if not supports_trtllm_attention():
451 + if not supports_trtllm_attention(is_prefill=is_prefill):
452 if force_use_trtllm:
453 logger.warning_once(
443 - "TRTLLM attention is not supported on this platform, "
444 - "but --attention-config.use_trtllm_attention is set to 1"
454 + "TRTLLM attention is not supported on this platform for %s, "
455 + "but --attention-config.use_trtllm_attention is set to 1",
456 + "prefill" if is_prefill else "decode",
457 )
458 return False
459
@@ -476,13 +488,20 @@ def use_trtllm_attention(
488 if is_prefill:
489 # Prefill auto-detection
490 use_trtllm = kv_cache_dtype == "auto"
479 - if use_trtllm:
480 - logger.warning_once("Using TRTLLM prefill attention (auto-detected).")
491 + elif current_platform.is_device_capability(90) and kv_cache_dtype.startswith(
492 + "fp8"
493 + ):
494 + # SM90 + FP8 KV cache: prefer the XQA decode kernel. XQA does not
495 + # support NVFP4 KV (that is an SM100 trtllm-gen path only).
496 + use_trtllm = True
497 else:
498 # Decode auto-detection
499 use_trtllm = num_tokens <= 256 and kv_cache_dtype == "auto"
484 - if use_trtllm:
485 - logger.warning_once("Using TRTLLM decode attention (auto-detected).")
500 + if use_trtllm:
501 + logger.warning_once(
502 + "Using TRTLLM %s attention (auto-detected).",
503 + "prefill" if is_prefill else "decode",
504 + )
505 return use_trtllm
506
507 # CLI argument is set to 1 - respect it
vllm/v1/attention/backends/flashinfer.py
+286 -92
@@ -3,6 +3,7 @@
3 """Attention layer with FlashInfer."""
4
5 from dataclasses import dataclass
6 +from enum import Enum
7 from functools import partial
8 from typing import ClassVar
9
@@ -19,6 +20,7 @@ from flashinfer.prefill import trtllm_batch_context_with_kv_cache
20 from flashinfer.utils import FP4Tensor
21 from typing_extensions import override
22
23 +from vllm import _custom_ops as custom_ops
24 from vllm import envs
25 from vllm.config import (
26 CUDAGraphMode,
@@ -38,6 +40,8 @@ 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 + force_use_trtllm_attention,
44 + supports_trtllm_attention,
45 use_trtllm_attention,
46 )
47 from vllm.utils.math_utils import cdiv
@@ -84,16 +88,16 @@ FP4_DTYPE = torch.uint8
88
89 logger = init_logger(__name__)
90
87 -trtllm_gen_workspace_buffer = None
91 +trtllm_workspace_buffer = None
92
93
90 -def _get_trtllm_gen_workspace_buffer():
91 - global trtllm_gen_workspace_buffer
92 - if trtllm_gen_workspace_buffer is None:
93 - trtllm_gen_workspace_buffer = torch.zeros(
94 +def _get_trtllm_workspace_buffer():
95 + global trtllm_workspace_buffer
96 + if trtllm_workspace_buffer is None:
97 + trtllm_workspace_buffer = torch.zeros(
98 envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8, device="cuda"
99 )
96 - return trtllm_gen_workspace_buffer
100 + return trtllm_workspace_buffer
101
102
103 @triton.jit
@@ -350,6 +354,7 @@ class FlashInferBackend(AttentionBackend):
354 use_large_pages = (
355 num_kv_heads > 0
356 and num_qo_heads // num_kv_heads > 1
357 + and current_platform.is_device_capability_family(100)
358 and can_use_trtllm_attention(num_qo_heads, num_kv_heads)
359 )
360 if not use_large_pages:
@@ -418,6 +423,16 @@ class FlashInferBackend(AttentionBackend):
423 else:
424 raise ValueError(f"Unrecognized dtype: {kv_cache_dtype}")
425
426 + @classmethod
427 + def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
428 + if kv_cache_dtype == "nvfp4":
429 + return (
430 + current_platform.is_device_capability_family(100)
431 + and supports_trtllm_attention(is_prefill=True)
432 + and supports_trtllm_attention(is_prefill=False)
433 + )
434 + return super().supports_kv_cache_dtype(kv_cache_dtype)
435 +
436 @classmethod
437 def get_supported_head_sizes(cls) -> list[int]:
438 # https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157
@@ -436,10 +451,9 @@ class FlashInferBackend(AttentionBackend):
451
452 @classmethod
453 def supports_sink(cls) -> bool:
439 - """FlashInfer supports sinks when TRTLLM attention is available (SM100)."""
454 + """FlashInfer supports sinks only on the SM100 trtllm-gen path."""
455 from vllm.utils.flashinfer import (
456 force_use_trtllm_attention,
442 - supports_trtllm_attention,
457 )
458
459 # Respect explicit disable flag (e.g.,
@@ -447,8 +461,13 @@ class FlashInferBackend(AttentionBackend):
461 if force_use_trtllm_attention() is False:
462 return False
463
464 + if not current_platform.is_device_capability_family(100):
465 + return False
466 +
467 # Check if TRTLLM is supported on this platform
451 - return supports_trtllm_attention()
468 + return supports_trtllm_attention(
469 + is_prefill=False
470 + ) and supports_trtllm_attention(is_prefill=True)
471
472 @classmethod
473 def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None:
@@ -474,6 +493,13 @@ class FIDecode:
493 wrapper: BatchDecodeWithPagedKVCacheWrapper
494
495
496 +class FlashInferDecodeKernel(Enum):
497 + """Decode kernels selected inside the FlashInfer backend."""
498 +
499 + XQA = "xqa"
500 + TRTLLM_GEN = "trtllm-gen"
501 +
502 +
503 @dataclass
504 class TRTLLMPrefill:
505 """Metadata for the TRTLLM prefill pathway."""
@@ -503,8 +529,15 @@ class TRTLLMPrefill:
529
530
531 @dataclass
506 -class TRTLLMDecode:
507 - """Metadata for the TRTLLM decode pathway."""
532 +class FlashInferTrtllmAPIDecode:
533 + """Metadata for decode paths using FlashInfer's TRTLLM decode API.
534 +
535 + FlashInfer exposes both XQA (SM90) and trtllm-gen (SM100) through
536 + ``trtllm_batch_decode_with_kv_cache``. Keep them as distinct vLLM
537 + decode kernels because their dtype/layout/output constraints differ.
538 + """
539 +
540 + kernel: FlashInferDecodeKernel
541
542 block_tables: torch.Tensor
543 """
@@ -530,7 +563,10 @@ class FlashInferMetadata:
563 slot_mapping: torch.Tensor
564 """Tensor for writing K/V to the cache. Shape: [num_actual_tokens]"""
565
533 - q_data_type: torch.dtype
566 + # The data types of the query for prefill and decode.
567 + # On SM90, these two data types may be different.
568 + q_data_type_prefill: torch.dtype
569 + q_data_type_decode: torch.dtype
570
571 num_decodes: int
572 num_decode_tokens: int
@@ -544,7 +580,7 @@ class FlashInferMetadata:
580 Will be `None` if `num_prefill_tokens == 0`.
581 """
582
547 - decode: FIDecode | TRTLLMDecode | None
583 + decode: FIDecode | FlashInferTrtllmAPIDecode | None
584 """
585 Holds the metadata for the decode portion of the batch.
586 Will be `None` if `num_decode_tokens == 0`.
@@ -651,12 +687,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
687 # storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3)
688 self.is_kvcache_nvfp4 = self.cache_dtype == "nvfp4"
689 if self.is_kvcache_nvfp4:
654 - # trtllm-gen FP4 FMHA kernels only exist for sm100f (sm_100/sm_103).
655 - # Fail fast at init rather than crashing on the first request.
656 - if not current_platform.is_device_capability_family(100):
690 + if (
691 + force_use_trtllm_attention() is False
692 + or not supports_trtllm_attention(is_prefill=True)
693 + or not supports_trtllm_attention(is_prefill=False)
694 + ):
695 raise ValueError(
658 - "--kv-cache-dtype nvfp4 requires sm100f, "
659 - "please try a different dtype or remove"
696 + "--kv-cache-dtype nvfp4 requires the SM100 trtllm-gen "
697 + "FlashInfer path."
698 )
699 # For NVFP4, kv_cache_dtype stays as the string "nvfp4"
700 # which is passed to FlashInferImpl
@@ -671,54 +709,87 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
709 assert self.kv_cache_spec.dtype == self.model_config.dtype
710 self.kv_cache_dtype = self.kv_cache_spec.dtype
711
674 - # Use model dtype as q dtype when TRTLLM attn is not supported, or
675 - # --attention-config.disable_flashinfer_q_quantization is set to 1. Otherwise,
676 - # try to use fp8 q if kv cache is fp8, and will fall back to model dtype
677 - # if TRTLLM attention kernel is not used when building attn metadata
678 - can_use_trtllm = can_use_trtllm_attention(self.num_qo_heads, self.num_kv_heads)
712 + # Compute per-phase Q dtype. On SM90 (XQA decode), the prefill and
713 + # decode phases require different Q dtypes when the KV cache is FP8
714 + # (FP8-Q for the FI native prefill, BF16/FP16-Q for XQA decode),
715 + # so both values must be tracked independently.
716 + self.q_data_type_prefill = self.get_q_data_type(
717 + vllm_config, self.kv_cache_spec, is_prefill=True
718 + )
719 + self.q_data_type_decode = self.get_q_data_type(
720 + vllm_config, self.kv_cache_spec, is_prefill=False
721 + )
722
723 + # Prefer TRTLLM/XQA for decoding whenever supported. The decode kernel
724 + # must be selected statically for FULL cudagraph capture.
725 + can_use_xqa_or_trtllm_gen_decode = can_use_trtllm_attention(
726 + self.num_qo_heads, self.num_kv_heads, is_prefill=False
727 + )
728 # Page sizes >= 128 require the trtllm-gen GQA/MQA path (guaranteed by
729 # get_supported_kernel_block_sizes).
730 assert self.page_size <= 64 or (
683 - can_use_trtllm and self.num_qo_heads // self.num_kv_heads > 1
731 + current_platform.is_device_capability_family(100)
732 + and can_use_xqa_or_trtllm_gen_decode
733 + and self.num_qo_heads // self.num_kv_heads > 1
734 ), f"Unexpected FlashInfer page size {self.page_size} without trtllm-gen GQA"
685 -
686 - if (
687 - can_use_trtllm
688 - and not vllm_config.attention_config.disable_flashinfer_q_quantization
689 - ):
690 - if self.is_kvcache_nvfp4:
691 - # NVFP4 KV cache uses FP8 quantized queries
692 - self.q_data_type = FlashInferBackend.get_dtype_for_flashinfer(
693 - "fp8_e4m3"
694 - )
695 - else:
696 - self.q_data_type = self.kv_cache_dtype
697 - else:
698 - self.q_data_type = self.model_config.dtype
699 -
700 - # Prefer TRTLLM attention for decoding in all cases.
701 - # This allows us to use AttentionCGSupport.UNIFORM_BATCH mode.
702 - self.use_trtllm_decode_attention = can_use_trtllm
703 - self._init_reorder_batch_threshold(1, supports_spec_as_decode=can_use_trtllm)
735 + self.use_trtllm_decode_attention = can_use_xqa_or_trtllm_gen_decode
736 + self.flashinfer_trtllm_api_decode_kernel: FlashInferDecodeKernel | None = (
737 + self._get_flashinfer_trtllm_api_decode_kernel()
738 + if can_use_xqa_or_trtllm_gen_decode
739 + else None
740 + )
741 + supports_spec_as_decode = (
742 + self.flashinfer_trtllm_api_decode_kernel
743 + == FlashInferDecodeKernel.TRTLLM_GEN
744 + )
745 + self._init_reorder_batch_threshold(
746 + 1, supports_spec_as_decode=supports_spec_as_decode
747 + )
748
749 self._cascade_wrapper = None # Wrapper for cascade attention
750
751 # Global hyperparameters shared by all attention layers
752 # TODO: discard this for trtllm-gen backend
709 - self.global_hyperparameters = infer_global_hyperparameters(
710 - get_per_layer_parameters(vllm_config, layer_names, FlashInferImpl)
753 + per_layer_parameters = get_per_layer_parameters(
754 + vllm_config, layer_names, FlashInferImpl
755 )
756 + if current_platform.is_device_capability(90) and any(
757 + params.window_left != -1 for params in per_layer_parameters.values()
758 + ):
759 + # FlashInfer SM90 sliding-window prefill is not reliable with FP8-Q:
760 + # https://github.com/flashinfer-ai/flashinfer/issues/3578
761 + raise NotImplementedError(
762 + "FlashInfer backend on SM90 currently crashes with "
763 + "sliding-window attention layers. Use the default attention "
764 + "backend."
765 + )
766 + self.global_hyperparameters = infer_global_hyperparameters(per_layer_parameters)
767 self.sm_scale = self.global_hyperparameters.sm_scale
768 self.window_left = self.global_hyperparameters.window_left
769 self.logits_soft_cap = self.global_hyperparameters.logits_soft_cap
770 self.has_sinks = self.global_hyperparameters.has_sinks
716 - if self.has_sinks and not can_use_trtllm:
771 + if self.has_sinks and not FlashInferBackend.supports_sink():
772 raise NotImplementedError(
773 "FlashInfer backend currently does not support attention "
774 "sinks, please use trtllm on blackwell or flash attention on "
775 "earlier GPUs."
776 )
777 + capability = current_platform.get_device_capability()
778 + arch = f"sm{capability.major}{capability.minor}" if capability else "unknown"
779 + decode_backend = (
780 + self.flashinfer_trtllm_api_decode_kernel.value
781 + if self.flashinfer_trtllm_api_decode_kernel is not None
782 + else "flashinfer-native"
783 + )
784 + logger.info_once(
785 + "FlashInfer resolved query dtypes: prefill=%s, decode=%s, "
786 + "decode_backend=%s, kv_cache_dtype=%s, arch=%s",
787 + self.q_data_type_prefill,
788 + self.q_data_type_decode,
789 + decode_backend,
790 + self.kv_cache_dtype,
791 + arch,
792 + )
793 # Preparing persistent buffers
794 # Since we do not have explicit synchronization in ModelRunnerV2, we do not pin
795 # reused CPU buffers to avoid a race condition between step N async copies to
@@ -731,6 +802,37 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
802 self.paged_kv_indices = self._make_buffer(max_num_pages)
803 self.paged_kv_last_page_len = self._make_buffer(max_num_reqs)
804
805 + # Keep SM90 prefill/decode Q dtype selection in one place.
806 + @classmethod
807 + def get_q_data_type(
808 + cls,
809 + vllm_config: VllmConfig,
810 + kv_cache_spec: AttentionSpec,
811 + is_prefill: bool,
812 + ) -> torch.dtype:
813 + # The user sets --attention-config.disable_flashinfer_q_quantization
814 + # to 1 explicitly, use model dtype for query.
815 + if vllm_config.attention_config.disable_flashinfer_q_quantization:
816 + return vllm_config.model_config.dtype
817 +
818 + # On SM90, XQA decode requires BF16/FP16-Q even with FP8 KV cache.
819 + # FI native prefill on SM90 still uses FP8-Q in that case.
820 + cache_dtype = vllm_config.cache_config.cache_dtype
821 + if (
822 + current_platform.is_device_capability(90)
823 + and not is_prefill
824 + and force_use_trtllm_attention() is not False
825 + and cache_dtype.startswith("fp8")
826 + ):
827 + return vllm_config.model_config.dtype
828 +
829 + # Otherwise, match Q dtype to the KV cache dtype.
830 + if cache_dtype.startswith("fp8"):
831 + return FlashInferBackend.get_dtype_for_flashinfer(cache_dtype)
832 + if cache_dtype == "nvfp4":
833 + return FlashInferBackend.get_dtype_for_flashinfer("fp8_e4m3")
834 + return kv_cache_spec.dtype
835 +
836 def _make_buffer(
837 self, *size: int | torch.SymInt, dtype: torch.dtype = torch.int32
838 ) -> CpuGpuBuffer:
@@ -751,12 +853,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
853 ) -> AttentionCGSupport:
854 """Get the cudagraph support level for FlashInfer attention.
855
754 - This depends on whether we can use TRTLLM attention for decodes, since we can
755 - only do UNIFORM_SINGLE_TOKEN_DECODE if it is unavailable.
756 - To check this, we must call can_use_trtllm_attention with the number of KV
757 - heads from the kv_cache_spec. We check all available KV cache specs and
758 - only return UNIFORM_BATCH if all of them support TRTLLM attention.
856 + The SM90 XQA integration only enables single-token decode today. Keep
857 + specdec CUDA graphs limited to trtllm-gen until vLLM wires the XQA
858 + specdec mask.
859 """
860 + if current_platform.is_device_capability(90):
861 + return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
862 +
863 # For UniformTypeKVCacheSpecs, check all contained specs
864 kv_specs = (
865 kv_cache_spec.kv_cache_specs.values()
@@ -775,6 +878,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
878 if not can_use_trtllm_attention(
879 num_qo_heads=num_qo_heads,
880 num_kv_heads=spec.num_kv_heads,
881 + is_prefill=False,
882 ):
883 has_trtllm_support = False
884 break
@@ -797,6 +901,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
901 def set_workspace_buffer(self, workspace_buffer: torch.Tensor):
902 self._workspace_buffer = workspace_buffer
903
904 + @staticmethod
905 + def _get_flashinfer_trtllm_api_decode_kernel() -> FlashInferDecodeKernel:
906 + if current_platform.is_device_capability(90):
907 + return FlashInferDecodeKernel.XQA
908 + assert current_platform.is_device_capability_family(100)
909 + return FlashInferDecodeKernel.TRTLLM_GEN
910 +
911 def _get_prefill_wrapper(
912 self,
913 causal: bool = True,
@@ -976,7 +1087,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1087 # Step 1: Decide which dispatch modes to use:
1088 # - Cascade attention (distinct mode)
1089 # - Prefill (FI native or TRTLLM)
979 - # - Decode (FI native or TRTLLM)
1090 + # - Decode (FI native, XQA, or trtllm-gen)
1091 use_cascade = common_prefix_len > 0
1092 uses_spec_reorder = self.reorder_batch_threshold > 1
1093 # Page sizes >= 128 must use trtllm-gen; force it for prefill too.
@@ -990,13 +1101,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1101 max_seq_len,
1102 self.dcp_world_size,
1103 self.cache_dtype,
993 - self.q_data_type,
1104 + self.q_data_type_prefill,
1105 is_prefill=True,
1106 force_use_trtllm=prefill_force_trtllm,
1107 has_sinks=self.has_sinks,
1108 has_spec=uses_spec_reorder,
1109 )
999 - decode_use_trtllm = (
1110 + decode_with_flashinfer_trtllm_api = (
1111 causal and self.use_trtllm_decode_attention and self.dcp_world_size <= 1
1112 )
1113
@@ -1011,7 +1122,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1122 )
1123 all_uses_trtllm = causal and (
1124 (num_prefills == 0 or prefill_use_trtllm)
1014 - and (num_decodes == 0 or decode_use_trtllm)
1125 + and (num_decodes == 0 or decode_with_flashinfer_trtllm_api)
1126 )
1127
1128 if not all_uses_trtllm:
@@ -1035,17 +1146,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1146 "`sm_scale`."
1147 )
1148
1038 - # The q quantization is not supported for non-trtllm attention,
1039 - # fall back to model dtype.
1040 - self.q_data_type = self.model_config.dtype
1041 -
1149 # Step 2: Initialize the output metadata
1150 # Leave prefill/decode/cascade_wrapper empty, to be populated
1151 # case by case depending on the batch contents and backend selection.
1152 attn_metadata = FlashInferMetadata(
1153 num_actual_tokens=num_actual_tokens,
1154 slot_mapping=common_attn_metadata.slot_mapping,
1048 - q_data_type=self.q_data_type,
1155 + q_data_type_prefill=self.q_data_type_prefill,
1156 + q_data_type_decode=self.q_data_type_decode,
1157 num_decodes=num_decodes,
1158 num_decode_tokens=num_decode_tokens,
1159 num_prefills=num_prefills,
@@ -1101,7 +1209,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1209
1210 # Compute paged_kv_indices if necessary
1211 # paged_kv_indices is only needed for FlashInfer native paths;
1104 - # TRTLLM paths use block_tables directly on GPU.
1212 + # XQA/trtllm-gen paths use block_tables directly on GPU.
1213 needs_paged_kv_indices = use_cascade or not all_uses_trtllm
1214 if needs_paged_kv_indices:
1215 assert num_blocks_np is not None
@@ -1143,6 +1251,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1251 paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_reqs]
1252
1253 attn_metadata.cascade_wrapper = self._get_cascade_wrapper()
1254 + # Cascade attention must use the same q dtype for prefill and decode
1255 + # because it does not support FP8 kv-cache or FP8 query yet.
1256 + assert self.q_data_type_prefill == self.q_data_type_decode
1257 attn_metadata.cascade_wrapper.plan(
1258 qo_indptr_arr=[shared_qo_indptr_cpu, qo_indptr_cpu],
1259 paged_kv_indptr_arr=[shared_kv_page_indptr_cpu, paged_kv_indptr_cpu],
@@ -1159,7 +1270,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1270 sm_scale=self.sm_scale,
1271 window_left=self.window_left,
1272 logits_soft_cap=self.logits_soft_cap,
1162 - q_data_type=self.q_data_type,
1273 + q_data_type=self.q_data_type_prefill,
1274 kv_data_type=self.kv_cache_dtype,
1275 )
1276 return attn_metadata
@@ -1233,7 +1344,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1344 sm_scale=self.sm_scale,
1345 window_left=self.window_left,
1346 logits_soft_cap=self.logits_soft_cap,
1236 - q_data_type=self.q_data_type,
1347 + q_data_type=self.q_data_type_prefill,
1348 kv_cache_dtype=self.kv_cache_dtype,
1349 prefill_fixed_split_size=self.prefill_fixed_split_size,
1350 disable_split_kv=self.disable_split_kv,
@@ -1262,7 +1373,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1373 sm_scale=self.sm_scale,
1374 window_left=self.window_left,
1375 logits_soft_cap=self.logits_soft_cap,
1265 - q_data_type=self.q_data_type,
1376 + q_data_type=self.q_data_type_prefill,
1377 kv_data_type=self.kv_cache_dtype,
1378 o_data_type=o_dtype,
1379 fixed_split_size=self.prefill_fixed_split_size,
@@ -1272,12 +1383,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1383
1384 ## DECODE PATHWAY
1385 if num_decodes > 0:
1275 - if decode_use_trtllm:
1386 + if decode_with_flashinfer_trtllm_api:
1387 assert num_decode_tokens % num_decodes == 0, (
1277 - "TRTLLM decode requires uniform query lengths per request. "
1388 + "XQA/trtllm-gen decode requires uniform query lengths per request. "
1389 f"Got {num_decode_tokens=} and {num_decodes=}."
1390 )
1280 - attn_metadata.decode = TRTLLMDecode(
1391 + assert self.flashinfer_trtllm_api_decode_kernel is not None
1392 + attn_metadata.decode = FlashInferTrtllmAPIDecode(
1393 + kernel=self.flashinfer_trtllm_api_decode_kernel,
1394 block_tables=block_table_tensor[:num_decodes],
1395 seq_lens=seq_lens[:num_decodes],
1396 max_seq_len=max_seq_len,
@@ -1320,7 +1433,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1433 sm_scale=self.sm_scale,
1434 window_left=self.window_left,
1435 logits_soft_cap=self.logits_soft_cap,
1323 - q_data_type=self.q_data_type,
1436 + q_data_type=self.q_data_type_decode,
1437 kv_data_type=self.kv_cache_dtype,
1438 o_data_type=o_dtype,
1439 fixed_split_size=self.decode_fixed_split_size,
@@ -1396,10 +1509,17 @@ class FlashInferImpl(AttentionImpl):
1509 )
1510 self.sinks = sinks
1511
1399 - self.support_trtllm_attn = can_use_trtllm_attention(num_heads, num_kv_heads)
1512 + self.supports_xqa_or_trtllm_gen_decode = can_use_trtllm_attention(
1513 + num_heads, num_kv_heads, is_prefill=False
1514 + )
1515 vllm_config = get_current_vllm_config_or_none()
1516 + # Query pre-quantization needs a single dtype for the whole query tensor.
1517 + # SM90 XQA needs BF16/FP16-Q for decode and FP8 for prefill,
1518 + # so only enable this for SM100 trtllm-gen where both use FP8-Q.
1519 self.supports_quant_query_input = (
1402 - self.support_trtllm_attn
1520 + self.supports_xqa_or_trtllm_gen_decode
1521 + and is_quantized_kv_cache(self.kv_cache_dtype)
1522 + and current_platform.is_device_capability_family(100)
1523 and vllm_config is not None
1524 and not vllm_config.attention_config.disable_flashinfer_q_quantization
1525 )
@@ -1429,9 +1549,12 @@ class FlashInferImpl(AttentionImpl):
1549 self.dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
1550
1551 def fused_output_quant_supported(self, quant_key: QuantKey):
1552 + # XQA does not support FP8/NVFP4 output, so require trtllm-gen
1553 + # (SM100+) here. Without that we cannot fuse the output quant.
1554 return (
1433 - self.support_trtllm_attn
1555 + self.supports_xqa_or_trtllm_gen_decode
1556 and is_quantized_kv_cache(self.kv_cache_dtype)
1557 + and current_platform.is_device_capability_family(100)
1558 and quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic)
1559 )
1560
@@ -1440,6 +1563,37 @@ class FlashInferImpl(AttentionImpl):
1563 if self.sinks is not None and self.sinks.dtype != torch.float32:
1564 self.sinks = self.sinks.to(torch.float32)
1565
1566 + def get_xqa_bmm1_scale(self, layer: torch.nn.Module, q_data_type: torch.dtype):
1567 + bmm1_scale = self.scale
1568 + if is_quantized_kv_cache(self.kv_cache_dtype):
1569 + if q_data_type in (torch.float8_e4m3fn, torch.float8_e5m2):
1570 + bmm1_scale *= layer._q_scale_float
1571 + bmm1_scale *= layer._k_scale_float
1572 + return bmm1_scale
1573 +
1574 + # SM90 may need FP8-Q for native prefill and BF16/FP16-Q for XQA decode,
1575 + # so quantize only the slice whose target dtype differs.
1576 + def maybe_quant_query(
1577 + self,
1578 + query: torch.Tensor,
1579 + q_data_type: torch.dtype,
1580 + scale: torch.Tensor,
1581 + ) -> torch.Tensor:
1582 + if query.dtype != q_data_type:
1583 + assert query.dtype in [torch.float16, torch.bfloat16]
1584 + assert q_data_type in [torch.float8_e4m3fn, torch.float8_e5m2]
1585 + assert query.is_contiguous()
1586 + assert query.dim() == 3
1587 + num_tokens = query.shape[0]
1588 + num_heads = query.shape[1]
1589 + head_size = query.shape[2]
1590 + query_quantized, _ = custom_ops.scaled_fp8_quant(
1591 + query.view(num_tokens, num_heads * head_size), scale=scale
1592 + )
1593 + return query_quantized.view(num_tokens, num_heads, head_size)
1594 +
1595 + return query
1596 +
1597 def forward(
1598 self,
1599 layer: torch.nn.Module,
@@ -1469,12 +1623,6 @@ class FlashInferImpl(AttentionImpl):
1623 # Profiling run.
1624 return output.fill_(0)
1625
1472 - # Ensure query dtype matches the expected dtype from attention metadata
1473 - assert attn_metadata.q_data_type == query.dtype, (
1474 - f"Query dtype mismatch: expected {attn_metadata.q_data_type}, "
1475 - f"got {query.dtype}"
1476 - )
1477 -
1626 if self.bmm1_scale is None:
1627 self.bmm1_scale = self.scale
1628 if is_quantized_kv_cache(self.kv_cache_dtype):
@@ -1486,7 +1634,14 @@ class FlashInferImpl(AttentionImpl):
1634 self.bmm2_scale *= layer._v_scale_float
1635
1636 prefill_use_trtllm = isinstance(attn_metadata.prefill, TRTLLMPrefill)
1489 - decode_use_trtllm = isinstance(attn_metadata.decode, TRTLLMDecode)
1637 + decode_kernel = (
1638 + attn_metadata.decode.kernel
1639 + if isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode)
1640 + else None
1641 + )
1642 + decode_with_xqa = decode_kernel == FlashInferDecodeKernel.XQA
1643 + decode_with_trtllm_gen = decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN
1644 + decode_with_flashinfer_trtllm_api = decode_with_xqa or decode_with_trtllm_gen
1645
1646 # The attn+quant fusion happens when output_scale is provided.
1647 if output_scale is None:
@@ -1494,12 +1649,15 @@ class FlashInferImpl(AttentionImpl):
1649 "output_block_scale is not supported when fusion has not happened"
1650 )
1651 else:
1497 - assert attn_metadata.q_data_type == FP8_DTYPE, (
1498 - "Query must be FP8 when attn+quant fusion happened."
1652 + assert attn_metadata.q_data_type_prefill == FP8_DTYPE, (
1653 + "Query must be FP8 when attn+quant fusion happened for prefill."
1654 + )
1655 + assert attn_metadata.q_data_type_decode == FP8_DTYPE, (
1656 + "Query must be FP8 when attn+quant fusion happened for decode."
1657 )
1658 assert (attn_metadata.num_prefills == 0 or prefill_use_trtllm) and (
1501 - attn_metadata.num_decodes == 0 or decode_use_trtllm
1502 - ), "Must use TRT-LLM attn"
1659 + attn_metadata.num_decodes == 0 or decode_with_trtllm_gen
1660 + ), "Output quant fusion requires TRTLLM prefill/trtllm-gen decode"
1661
1662 if output.dtype == FP8_DTYPE:
1663 assert output_block_scale is None, (
@@ -1597,6 +1755,13 @@ class FlashInferImpl(AttentionImpl):
1755 prefill_query = query[num_decode_tokens:]
1756 assert prefill_query.shape[0] == num_prefill_tokens
1757
1758 + # Convert query to the expected dtype for prefill if needed.
1759 + prefill_query = self.maybe_quant_query(
1760 + prefill_query,
1761 + attn_metadata.q_data_type_prefill,
1762 + layer._q_scale,
1763 + )
1764 +
1765 if not prefill_use_trtllm:
1766 assert isinstance(attn_metadata.prefill, FIPrefill)
1767 prefill_wrapper = attn_metadata.prefill.wrapper
@@ -1655,6 +1820,7 @@ class FlashInferImpl(AttentionImpl):
1820 prefill_wrapper.run(
1821 prefill_query,
1822 kv_cache_permute,
1823 + q_scale=layer._q_scale_float,
1824 k_scale=layer._k_scale_float,
1825 v_scale=layer._v_scale_float,
1826 out=out_prefill,
@@ -1673,7 +1839,7 @@ class FlashInferImpl(AttentionImpl):
1839 # degenerate strides on size=1 dims for TMA alignment.
1840 prefill_query = prefill_query.contiguous()
1841 prefill_query = canonicalize_singleton_dim_strides(prefill_query)
1676 - workspace_buffer = _get_trtllm_gen_workspace_buffer()
1842 + workspace_buffer = _get_trtllm_workspace_buffer()
1843 block_tables_prefill = attn_metadata.prefill.block_tables
1844 seq_lens_prefill = attn_metadata.prefill.seq_lens
1845
@@ -1705,7 +1871,7 @@ class FlashInferImpl(AttentionImpl):
1871 prefill_kv_block_scales = None
1872 if self.is_kvcache_nvfp4:
1873 # NVFP4 trtllm-gen kernel requires FP8 query.
1708 - assert attn_metadata.q_data_type == FP8_DTYPE, (
1874 + assert attn_metadata.q_data_type_prefill == FP8_DTYPE, (
1875 "NVFP4 KV cache requires FP8 quantized queries for "
1876 "trtllm-gen prefill. Set "
1877 "disable_flashinfer_q_quantization=False."
@@ -1714,7 +1880,7 @@ class FlashInferImpl(AttentionImpl):
1880 mock_block_table = block_tables_prefill
1881 prefill_kv_block_scales = nvfp4_kv_block_scales
1882 elif (
1717 - attn_metadata.q_data_type != FP8_DTYPE
1883 + attn_metadata.q_data_type_prefill != FP8_DTYPE
1884 and self.kv_cache_dtype.startswith("fp8")
1885 ):
1886 # TRTLLM prefill attention does not support BF16 Q
@@ -1738,7 +1904,7 @@ class FlashInferImpl(AttentionImpl):
1904 block_tables_prefill,
1905 layer._k_scale,
1906 layer._v_scale,
1741 - attn_metadata.q_data_type,
1907 + attn_metadata.q_data_type_prefill,
1908 )
1909 else:
1910 mock_kv_cache = kv_cache_permute
@@ -1773,7 +1939,14 @@ class FlashInferImpl(AttentionImpl):
1939 decode_query = query[:num_decode_tokens]
1940 assert decode_query.shape[0] == num_decode_tokens
1941
1776 - if not decode_use_trtllm:
1942 + # Convert query to the expected dtype for decode if needed.
1943 + decode_query = self.maybe_quant_query(
1944 + decode_query,
1945 + attn_metadata.q_data_type_decode,
1946 + layer._q_scale,
1947 + )
1948 +
1949 + if not decode_with_flashinfer_trtllm_api:
1950 assert isinstance(attn_metadata.decode, FIDecode)
1951 decode_wrapper = attn_metadata.decode.wrapper
1952 assert decode_wrapper is not None
@@ -1806,6 +1979,7 @@ class FlashInferImpl(AttentionImpl):
1979 decode_wrapper.run(
1980 decode_query,
1981 kv_cache_permute,
1982 + q_scale=layer._q_scale_float,
1983 k_scale=layer._k_scale_float,
1984 v_scale=layer._v_scale_float,
1985 out=output_tmp,
@@ -1822,6 +1996,7 @@ class FlashInferImpl(AttentionImpl):
1996 decode_wrapper.run(
1997 decode_query,
1998 kv_cache_permute,
1999 + q_scale=layer._q_scale_float,
2000 k_scale=layer._k_scale_float,
2001 v_scale=layer._v_scale_float,
2002 out=out_decode,
@@ -1831,19 +2006,23 @@ class FlashInferImpl(AttentionImpl):
2006 if needs_fp8_out:
2007 output[:num_decode_tokens].copy_(out_decode.to(output.dtype))
2008 else:
1834 - assert isinstance(attn_metadata.decode, TRTLLMDecode)
2009 + assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode)
2010 # decode_query may be non-contiguous or have degenerate strides
2011 # on size=1 dims. contiguous() ensures memory layout; then
2012 # canonicalize_singleton_dim_strides fixes any remaining
2013 # degenerate strides on size=1 dims for TMA alignment.
2014 decode_query = decode_query.contiguous()
2015 decode_query = canonicalize_singleton_dim_strides(decode_query)
1841 - workspace_buffer = _get_trtllm_gen_workspace_buffer()
2016 + workspace_buffer = _get_trtllm_workspace_buffer()
2017 block_tables_decode = attn_metadata.decode.block_tables
2018 seq_lens_decode = attn_metadata.decode.seq_lens
2019
1845 - # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
1846 - assert get_kv_cache_layout() == "HND"
2020 + # trtllm-gen needs HND layout on SM100. XQA is selected
2021 + # separately on SM90 and does not use this SM100 layout gate.
2022 + if decode_with_trtllm_gen:
2023 + assert get_kv_cache_layout() == "HND"
2024 + else:
2025 + assert decode_with_xqa
2026 assert is_strictly_contiguous(decode_query)
2027 assert is_strictly_contiguous(workspace_buffer)
2028 assert is_strictly_contiguous(block_tables_decode)
@@ -1882,6 +2061,19 @@ class FlashInferImpl(AttentionImpl):
2061 else:
2062 q_len_per_req = num_decode_tokens // attn_metadata.num_decodes
2063
2064 + if decode_with_xqa and q_len_per_req > 1:
2065 + raise NotImplementedError(
2066 + "FlashInfer XQA speculative decode is not wired in vLLM yet."
2067 + )
2068 +
2069 + # XQA decode can use model-dtype Q with FP8 KV, so only include
2070 + # q_scale when the decode query is actually FP8.
2071 + bmm1_scale = (
2072 + self.get_xqa_bmm1_scale(layer, attn_metadata.q_data_type_decode)
2073 + if decode_with_xqa
2074 + else self.bmm1_scale
2075 + )
2076 +
2077 trtllm_batch_decode_with_kv_cache(
2078 query=decode_query,
2079 kv_cache=(
@@ -1891,12 +2083,14 @@ class FlashInferImpl(AttentionImpl):
2083 block_tables=block_tables_decode,
2084 seq_lens=seq_lens_decode,
2085 max_seq_len=attn_metadata.decode.max_seq_len,
1894 - bmm1_scale=self.bmm1_scale,
2086 + bmm1_scale=bmm1_scale,
2087 bmm2_scale=self.bmm2_scale,
2088 window_left=self.window_left,
2089 sinks=self.sinks,
2090 o_sf_scale=self.o_sf_scale,
2091 out=out,
2092 + kv_layout=get_kv_cache_layout(),
2093 + backend=attn_metadata.decode.kernel.value,
2094 q_len_per_req=q_len_per_req,
2095 kv_cache_sf=(
2096 nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None