[BugFix] Derive FlashInfer Q dtype from resolved per-group builder state (#47485)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Michael Goin committed
Jul 2, 2026 at 22:33 UTC
41de1380c23d9993197f8103b4f6c5b63a5b0835
1 file changed
+21
-19
vllm/v1/attention/backends/flashinfer.py
+21
-19
@@ -713,12 +713,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
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
- )
716
+ self.q_data_type_prefill = self.get_q_data_type(is_prefill=True)
717
+ self.q_data_type_decode = self.get_q_data_type(is_prefill=False)
718
719
# Prefer TRTLLM/XQA for decoding whenever supported. The decode kernel
720
# must be selected statically for FULL cudagraph capture.
@@ -803,35 +799,41 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
799
self.paged_kv_last_page_len = self._make_buffer(max_num_reqs)
800
801
# 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:
802
+ def get_q_data_type(self, is_prefill: bool) -> torch.dtype:
803
# The user sets --attention-config.disable_flashinfer_q_quantization
804
# 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
805
+ if self.vllm_config.attention_config.disable_flashinfer_q_quantization:
806
+ return self.model_config.dtype
807
+
808
+ # self.cache_dtype is resolved per KV-cache group: it is "auto" when
809
+ # this group is unquantized (e.g. --kv-cache-dtype-skip-layers), even
810
+ # if cache_config requests a quantized dtype globally.
811
+ cache_dtype = self.cache_dtype
812
813
# On SM90, XQA decode requires BF16/FP16-Q even with FP8 KV cache.
814
# FI native prefill on SM90 still uses FP8-Q in that case.
820
- cache_dtype = vllm_config.cache_config.cache_dtype
815
if (
816
current_platform.is_device_capability(90)
817
and not is_prefill
818
and force_use_trtllm_attention() is not False
819
and cache_dtype.startswith("fp8")
820
):
827
- return vllm_config.model_config.dtype
821
+ return self.model_config.dtype
822
823
# Otherwise, match Q dtype to the KV cache dtype.
824
if cache_dtype.startswith("fp8"):
831
- return FlashInferBackend.get_dtype_for_flashinfer(cache_dtype)
825
+ # FP8-Q requires an fp8 tensor-core attention path
826
+ # (FI native fa3 on SM90, trtllm-gen/XQA on SM100).
827
+ # Architectures with only fa2 (e.g. SM89, SM120) cannot
828
+ # consume FP8 queries, so keep the model dtype for Q there.
829
+ if current_platform.is_device_capability(
830
+ 90
831
+ ) or current_platform.is_device_capability_family(100):
832
+ return FlashInferBackend.get_dtype_for_flashinfer(cache_dtype)
833
+ return self.model_config.dtype
834
if cache_dtype == "nvfp4":
835
return FlashInferBackend.get_dtype_for_flashinfer("fp8_e4m3")
834
- return kv_cache_spec.dtype
836
+ return self.kv_cache_spec.dtype
837
838
def _make_buffer(
839
self, *size: int | torch.SymInt, dtype: torch.dtype = torch.int32