From 9c3753d69bd8a9926422baad9a08231364987fa5 Mon Sep 17 00:00:00 2001 From: Kai Ma Date: Fri, 17 Jul 2026 13:45:19 -0400 Subject: [PATCH] fix(collector): align model dtype for vLLM FlashInfer fp8 KV attention FlashInferMetadataBuilder asserts kv_cache_spec.dtype == model_config.dtype. During B200 vLLM 0.22.0 power collection, fp8 KV + head_dim=128 routes to FLASHINFER while model_config stayed bfloat16, aborting every case. Set model_config.dtype from the KV cache spec before builder init when fp8 KV is enabled. --- collector/vllm/collect_attn.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/collector/vllm/collect_attn.py b/collector/vllm/collect_attn.py index 925d439c2..b0d1a24d4 100644 --- a/collector/vllm/collect_attn.py +++ b/collector/vllm/collect_attn.py @@ -207,6 +207,10 @@ def run_attention_torch( # Mock flashinfer's get_per_layer_parameters if needed if backend_name_str == "FLASHINFER": + if use_fp8_kv_cache: + # FlashInferMetadataBuilder asserts kv_cache_spec.dtype == + # model_config.dtype; KV spec uses fp8 while model_config stays bf16. + vllm_config.model_config.dtype = kv_cache_spec.dtype import unittest.mock from vllm.v1.attention.backends.utils import PerLayerParameters