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, 15:15 UTC
[CI] Add warmup run in test_fusion_attn (#31183)
Signed-off-by: angelayi <yiangela7@gmail.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Angela Yi committed
Jan 6, 2026 at 16:31 UTC
9a1d20a89c3b1f2c2687dee585b22c93f05b2310
1 file changed
+16
-8
tests/compile/test_fusion_attn.py
+16
-8
@@ -305,8 +305,12 @@ def test_attention_quant_pattern(
305
model_class: type[AttentionQuantPatternModel],
306
backend: AttentionBackendEnum,
307
dist_init,
308
+ monkeypatch,
309
+ use_fresh_inductor_cache,
310
):
311
"""Test AttentionStaticQuantPattern fusion pass"""
312
+ monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
313
+
314
if backend == AttentionBackendEnum.FLASHINFER and (
315
not current_platform.is_device_capability((10, 0)) or not has_flashinfer()
316
):
@@ -363,13 +367,15 @@ def test_attention_quant_pattern(
367
vllm_config=vllm_config_unfused,
368
)
369
model_unfused = model_unfused.to(device)
370
+ result_unfused_0 = model_unfused(q, k, v) # noqa: F841 HACK: See #131044
371
372
forward_ctx = get_forward_context()
373
forward_ctx.attn_metadata = model_unfused.build_attn_metadata(batch_size)
374
375
# Run model directly without fusion
376
# Still compile so query QuantFP8 has closer numerics
372
- result_unfused = torch.compile(model_unfused, fullgraph=True)(q, k, v)
377
+ compiled_unfused = torch.compile(model_unfused, fullgraph=True)
378
+ result_unfused = compiled_unfused(q, k, v)
379
380
# Run model with attn fusion enabled
381
vllm_config.compilation_config.pass_config = PassConfig(
@@ -399,24 +405,26 @@ def test_attention_quant_pattern(
405
cleanup_pass = PostCleanupPass(vllm_config)
406
407
test_backend = TestBackend(noop_pass, attn_pass, cleanup_pass)
408
+ # HACK: See https://github.com/vllm-project/vllm/issues/31044
409
+ result_fused_0 = model_fused(q, k, v) # noqa: F841
410
411
# Compile model with fusion enabled
404
- model_compiled = torch.compile(
412
+ compiled_fused = torch.compile(
413
model_fused, backend=test_backend, fullgraph=True
414
)
407
- assert model_compiled.attn._o_scale_float is None
415
+ assert compiled_fused.attn._o_scale_float is None
416
409
- result_fused_1 = model_compiled(q, k, v)
417
+ result_fused = compiled_fused(q, k, v)
418
419
if backend == AttentionBackendEnum.FLASHINFER:
420
# With the Flashinfer backend after the 1st round of the forward
421
# pass, output quant scale should be loaded into the attn layer's
422
# _o_scale_float, the 2nd round should reuse the loaded
423
# _o_scale_float
416
- assert model_compiled.attn._o_scale_float is not None
417
- result_fused_2 = model_compiled(q, k, v)
424
+ assert compiled_fused.attn._o_scale_float is not None
425
+ result_fused_2 = compiled_fused(q, k, v)
426
419
- assert model_compiled.attn._o_scale_float is not None
427
+ assert compiled_fused.attn._o_scale_float is not None
428
429
torch.testing.assert_close(
430
result_unfused, result_fused_2, atol=1e-2, rtol=1e-2
@@ -474,4 +482,4 @@ def test_attention_quant_pattern(
482
)
483
484
# Check that results are close
477
- torch.testing.assert_close(result_unfused, result_fused_1, atol=1e-2, rtol=1e-2)
485
+ torch.testing.assert_close(result_unfused, result_fused, atol=1e-2, rtol=1e-2)