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 10, 22:45 UTC
[BugFix] Fix cache issue in compilation_config (#31376)
Signed-off-by: Boyuan Feng <boyuan@meta.com>
Boyuan Feng committed
Dec 27, 2025 at 06:30 UTC
2f12cd32c06647f5f3f8d3e020eaed4cd290424a
2 files changed
+47
tests/compile/test_config.py
+42
@@ -428,3 +428,45 @@ def test_cudagraph_sizes_post_init(
428
vllm_config.compilation_config.max_cudagraph_capture_size
429
== expected_max_size
430
)
431
+
432
+
433
+def test_cached_compilation_config():
434
+ import torch
435
+ from torch._inductor.utils import run_and_get_code
436
+
437
+ from vllm.config import get_cached_compilation_config, set_current_vllm_config
438
+ from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
439
+ from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
440
+
441
+ dtype = torch.bfloat16
442
+ device = torch.device("cuda:0")
443
+ batch_size, num_qo_heads, head_size = 8, 16, 128
444
+
445
+ # access and cache default compilation config
446
+ # default compilation config does not contain +quant_fp8 custom op. If this is
447
+ # used, the generated code would use inductor-generated triton kernel instead
448
+ # of the custom op `torch.ops._C.static_scaled_fp8_quant`.
449
+ get_cached_compilation_config()
450
+
451
+ vllm_config = VllmConfig(
452
+ compilation_config=CompilationConfig(
453
+ mode=CompilationMode.VLLM_COMPILE,
454
+ custom_ops=["+quant_fp8"],
455
+ )
456
+ )
457
+
458
+ # set_current_vllm_config should clear cached compilation config and
459
+ # use the new compilation_config in vllm_config
460
+ with set_current_vllm_config(vllm_config):
461
+ query_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
462
+ query_quant = torch.compile(query_quant)
463
+
464
+ _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
465
+ query = torch.randn(
466
+ batch_size, num_qo_heads * head_size, dtype=dtype, device=device
467
+ )
468
+
469
+ _, code = run_and_get_code(query_quant, query, _q_scale)
470
+
471
+ code = " ".join(code)
472
+ assert "torch.ops._C.static_scaled_fp8_quant.default(" in code
vllm/config/vllm.py
+5
@@ -1360,6 +1360,11 @@ def set_current_vllm_config(
1360
1361
num_models_seen = compilation_counter.num_models_seen
1362
try:
1363
+ # Clear the compilation config cache when context changes.
1364
+ # This is needed since the old config may have been accessed
1365
+ # and cached before the new config is set.
1366
+ get_cached_compilation_config.cache_clear()
1367
+
1368
_current_vllm_config = vllm_config
1369
_current_prefix = prefix
1370
yield