A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 08:45 UTC
[Fix]Load kv-cache dtype from hf_quant_config.json automatically (fix for reverted PR) (#30785)
Signed-off-by: <> Co-authored-by: root <root@gpu-937.slurm-workers-slurm.slurm.svc.cluster.local>
danielafrimi committed
Dec 17, 2025 at 11:56 UTC
7b966ae2ba73b5391937907bfd8aaf63af033ff1
2 files changed
+83
-1
vllm/engine/arg_utils.py
+8
-1
@@ -93,6 +93,7 @@ from vllm.transformers_utils.utils import is_cloud_storage
93
from vllm.utils.argparse_utils import FlexibleArgumentParser
94
from vllm.utils.mem_constants import GiB_bytes
95
from vllm.utils.network_utils import get_ip
96
+from vllm.utils.torch_utils import resolve_kv_cache_dtype_string
97
from vllm.v1.sample.logits_processor import LogitsProcessor
98
99
if TYPE_CHECKING:
@@ -106,6 +107,7 @@ else:
107
LoadFormats = Any
108
UsageContext = Any
109
110
+
111
logger = init_logger(__name__)
112
113
# object is used to allow for special typing forms
@@ -1361,12 +1363,17 @@ class EngineArgs:
1363
f"dcp_size={self.decode_context_parallel_size}."
1364
)
1365
1366
+ # Resolve "auto" kv_cache_dtype to actual value from model config
1367
+ resolved_cache_dtype = resolve_kv_cache_dtype_string(
1368
+ self.kv_cache_dtype, model_config
1369
+ )
1370
+
1371
cache_config = CacheConfig(
1372
block_size=self.block_size,
1373
gpu_memory_utilization=self.gpu_memory_utilization,
1374
kv_cache_memory_bytes=self.kv_cache_memory_bytes,
1375
swap_space=self.swap_space,
1369
- cache_dtype=self.kv_cache_dtype,
1376
+ cache_dtype=resolved_cache_dtype,
1377
is_attention_free=model_config.is_attention_free,
1378
num_gpu_blocks_override=self.num_gpu_blocks_override,
1379
sliding_window=sliding_window,
vllm/utils/torch_utils.py
+75
@@ -24,6 +24,10 @@ else:
24
ModelConfig = object
25
IntermediateTensors = object
26
27
+import logging
28
+
29
+logger = logging.getLogger(__name__)
30
+
31
32
STR_DTYPE_TO_TORCH_DTYPE = {
33
"float32": torch.float32,
@@ -49,6 +53,13 @@ TORCH_DTYPE_TO_NUMPY_DTYPE = {
53
}
54
55
56
+MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP = {
57
+ # TODO: Add more modelopt kv cache dtype
58
+ # mappings here when it supported by some attention backend
59
+ # (for example supports nvfp4).
60
+ "fp8": "fp8_e4m3",
61
+}
62
+
63
T = TypeVar("T")
64
65
@@ -194,6 +205,70 @@ def get_kv_cache_torch_dtype(
205
return torch_dtype
206
207
208
+def get_kv_cache_quant_algo_string(quant_cfg: dict[str, Any]) -> str | None:
209
+ """Get the KV cache quantization algorithm string from the quantization config.
210
+
211
+ Maps various FP8 format names to vLLM's standard cache dtype strings.
212
+ Returns None if no kv_cache_quant_algo is specified.
213
+ Returns "auto" if the value is not recognized/supported.
214
+ """
215
+ # Mapping from model config values to vLLM cache_dtype strings
216
+
217
+ quant_method = quant_cfg.get("quant_method", "")
218
+ if quant_method.startswith("modelopt"):
219
+ quantization_inner = quant_cfg.get("quantization", quant_cfg)
220
+ # Check if quant config is specified and use kv cache quant algo
221
+ kv_algo = quantization_inner.get("kv_cache_quant_algo") or quant_cfg.get(
222
+ "kv_cache_quant_algo"
223
+ )
224
+ if isinstance(kv_algo, str):
225
+ kv_algo_lower = kv_algo.lower()
226
+
227
+ # Try to map to vLLM's standard format
228
+ if kv_algo_lower in MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP:
229
+ return MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP[kv_algo_lower]
230
+ else:
231
+ # Unknown/unsupported format - return "auto" as safe fallback
232
+ logger.warning(
233
+ "WARNING: Unknown kv_cache_quant_algo '%s' in model "
234
+ "config. Supported values: %s. Falling back to 'auto'.",
235
+ kv_algo,
236
+ list(MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP.keys()),
237
+ )
238
+ return "auto"
239
+ return None
240
+
241
+
242
+def get_kv_cache_quant_algo_dtype(quant_cfg: dict[str, Any]) -> torch.dtype | None:
243
+ """Get the KV cache quantization algorithm dtype from the quantization config."""
244
+ kv_algo_str = get_kv_cache_quant_algo_string(quant_cfg)
245
+ if kv_algo_str is not None and kv_algo_str != "auto":
246
+ # Only convert if we have a valid dtype string (not "auto" fallback)
247
+ return STR_DTYPE_TO_TORCH_DTYPE[kv_algo_str]
248
+ return None
249
+
250
+
251
+def resolve_kv_cache_dtype_string(
252
+ kv_cache_dtype: str, model_config: ModelConfig
253
+) -> str:
254
+ """Resolve 'auto' kv_cache_dtype to the actual string value from model config.
255
+ Returns the resolved cache_dtype string.
256
+ """
257
+ if kv_cache_dtype != "auto":
258
+ return kv_cache_dtype
259
+
260
+ hf_cfg = getattr(model_config, "hf_config", None)
261
+ if hf_cfg is not None:
262
+ quant_cfg = getattr(hf_cfg, "quantization_config", None)
263
+ if quant_cfg is not None:
264
+ kv_algo_str = get_kv_cache_quant_algo_string(quant_cfg)
265
+ if kv_algo_str is not None:
266
+ return kv_algo_str
267
+
268
+ # Default to auto (will be handled by downstream code)
269
+ return "auto"
270
+
271
+
272
def kv_cache_dtype_str_to_dtype(
273
kv_cache_dtype: str, model_config: ModelConfig
274
) -> torch.dtype: