@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9571e99945
Fork 0

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 7, 17:45 UTC

[ROCm][CI] Extending attention backend coverage for Eagle spec decode tests (#35265)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Andreas Karatzas committed Feb 25, 2026 at 16:16 UTC 9571e999451a468423e97cd7e3f36e9d27a098cb
4 files changed +314 -150
.buildkite/test_areas/engine.yaml
+1 -1
@@ -30,7 +30,7 @@ steps:
30 - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
31 mirror:
32 amd:
33 - device: mi325_8
33 + device: mi325_1
34 depends_on:
35 - image-build-amd
36 commands:
tests/utils.py
+51
@@ -1327,6 +1327,57 @@ def multi_gpu_test(*, num_gpus: int):
1327 return wrapper
1328
1329
1330 +def gpu_tier_mark(*, min_gpus: int = 1, max_gpus: int | None = None):
1331 + """
1332 + Mark a test to only run when the GPU count falls within [min_gpus, max_gpus].
1333 +
1334 + Examples:
1335 + @gpu_tier_mark(min_gpus=2) # only on multi-GPU
1336 + @gpu_tier_mark(max_gpus=1) # only on single-GPU
1337 + @gpu_tier_mark(min_gpus=2, max_gpus=4) # 2-4 GPUs only
1338 + """
1339 + gpu_count = cuda_device_count_stateless()
1340 + marks = []
1341 +
1342 + if min_gpus > 1:
1343 + marks.append(pytest.mark.distributed(num_gpus=min_gpus))
1344 +
1345 + reasons = []
1346 + if gpu_count < min_gpus:
1347 + reasons.append(f"Need at least {min_gpus} GPUs (have {gpu_count})")
1348 + if max_gpus is not None and gpu_count > max_gpus:
1349 + reasons.append(f"Need at most {max_gpus} GPUs (have {gpu_count})")
1350 +
1351 + if reasons:
1352 + marks.append(pytest.mark.skipif(True, reason="; ".join(reasons)))
1353 +
1354 + return marks
1355 +
1356 +
1357 +def single_gpu_only(f=None):
1358 + """Skip this test when running in a multi-GPU environment."""
1359 + marks = gpu_tier_mark(max_gpus=1)
1360 +
1361 + def wrapper(func):
1362 + for mark in reversed(marks):
1363 + func = mark(func)
1364 + return func
1365 +
1366 + return wrapper(f) if f is not None else wrapper
1367 +
1368 +
1369 +def multi_gpu_only(*, num_gpus: int = 2):
1370 + """Skip this test when running on fewer than num_gpus GPUs."""
1371 + marks = gpu_tier_mark(min_gpus=num_gpus)
1372 +
1373 + def wrapper(f):
1374 + for mark in reversed(marks):
1375 + f = mark(f)
1376 + return f
1377 +
1378 + return wrapper
1379 +
1380 +
1381 async def completions_with_server_args(
1382 prompts: list[str],
1383 model_name: str,
tests/v1/e2e/test_async_scheduling.py
+4
@@ -6,6 +6,7 @@ from typing import Any
6 import pytest
7 import torch._dynamo.config as dynamo_config
8
9 +from tests.utils import large_gpu_mark, single_gpu_only
10 from vllm import SamplingParams
11 from vllm.logprobs import Logprob
12 from vllm.platforms import current_platform
@@ -36,6 +37,7 @@ default_params = dict(
37 )
38
39
40 +@single_gpu_only
41 def test_without_spec_decoding(
42 sample_json_schema,
43 monkeypatch: pytest.MonkeyPatch,
@@ -95,6 +97,8 @@ def test_without_spec_decoding(
97 run_tests(monkeypatch, MODEL, test_configs, test_sampling_params)
98
99
100 +@single_gpu_only
101 +@large_gpu_mark(min_gb=16)
102 def test_with_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch):
103 """Test consistency and acceptance rates with some different combos of
104 preemption, executor, async scheduling, prefill chunking,
tests/v1/e2e/test_spec_decode.py
+258 -149
@@ -9,7 +9,13 @@ import pytest
9 import torch
10
11 from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_offline
12 -from tests.utils import get_attn_backend_list_based_on_platform, large_gpu_mark
12 +from tests.utils import (
13 + get_attn_backend_list_based_on_platform,
14 + large_gpu_mark,
15 + multi_gpu_marks,
16 + multi_gpu_only,
17 + single_gpu_only,
18 +)
19 from vllm import LLM, SamplingParams
20 from vllm.assets.base import VLLM_S3_BUCKET_URL
21 from vllm.assets.image import VLM_IMAGES_DIR
@@ -160,6 +166,8 @@ def reset_torch_dynamo():
166 },
167 ],
168 )
169 +@single_gpu_only
170 +@large_gpu_mark(min_gb=20)
171 def test_ngram_and_suffix_correctness(
172 speculative_config: dict,
173 model_name: str,
@@ -175,6 +183,8 @@ def test_ngram_and_suffix_correctness(
183 cleanup_dist_env_and_memory()
184
185
186 +@single_gpu_only
187 +@large_gpu_mark(min_gb=20)
188 def test_suffix_decoding_acceptance(
189 monkeypatch: pytest.MonkeyPatch,
190 sampling_config: SamplingParams,
@@ -242,6 +252,8 @@ def test_suffix_decoding_acceptance(
252 ],
253 ids=["llama3_eagle3_speculator", "qwen3_eagle3_speculator"],
254 )
255 +@single_gpu_only
256 +@large_gpu_mark(min_gb=24)
257 def test_speculators_model_integration(
258 monkeypatch: pytest.MonkeyPatch,
259 sampling_config: SamplingParams,
@@ -319,137 +331,7 @@ def test_speculators_model_integration(
331 )
332
333
322 -@pytest.mark.parametrize(
323 - [
324 - "model_setup",
325 - "mm_enabled",
326 - "enable_chunked_prefill",
327 - "model_impl",
328 - "expected_accuracy_threshold",
329 - ],
330 - [
331 - (
332 - ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
333 - False,
334 - False,
335 - "auto",
336 - 0.8, # ref: 90%
337 - ),
338 - (
339 - ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
340 - False,
341 - False,
342 - "transformers",
343 - 0.8, # ref: 90%
344 - ),
345 - pytest.param(
346 - (
347 - "eagle3",
348 - "Qwen/Qwen3-VL-8B-Instruct",
349 - "taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
350 - 1,
351 - ),
352 - False,
353 - False,
354 - "auto",
355 - 0.8, # ref: 90%
356 - marks=pytest.mark.skip(
357 - reason="architecture of its eagle3 is LlamaForCausalLMEagle3"
358 - ),
359 - ),
360 - pytest.param(
361 - (
362 - "eagle3",
363 - "Qwen/Qwen2.5-VL-7B-Instruct",
364 - "Rayzl/qwen2.5-vl-7b-eagle3-sgl",
365 - 1,
366 - ),
367 - False,
368 - False,
369 - "auto",
370 - 0.7, # TODO, update this with a reference value when re-enabling this case
371 - marks=pytest.mark.skip(
372 - reason="Skipping due to its head_dim not being a a multiple of 32"
373 - ),
374 - ),
375 - pytest.param(
376 - (
377 - "eagle",
378 - "meta-llama/Llama-3.1-8B-Instruct",
379 - "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
380 - 1,
381 - ),
382 - False,
383 - True,
384 - "auto",
385 - 0.7, # ref: 75%-80%
386 - marks=large_gpu_mark(min_gb=40),
387 - ), # works on 4x H100
388 - (
389 - (
390 - "eagle3",
391 - "meta-llama/Llama-3.1-8B-Instruct",
392 - "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
393 - 1,
394 - ),
395 - False,
396 - False,
397 - "auto",
398 - 0.7, # ref: 75%-80%
399 - ),
400 - pytest.param(
401 - (
402 - "eagle",
403 - "meta-llama/Llama-4-Scout-17B-16E-Instruct",
404 - "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
405 - 4,
406 - ),
407 - False,
408 - False,
409 - "auto",
410 - 0.8, # ref: 90%
411 - # marks=large_gpu_mark(min_gb=80),
412 - ), # works on 4x H100
413 - pytest.param(
414 - (
415 - "eagle",
416 - "meta-llama/Llama-4-Scout-17B-16E-Instruct",
417 - "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
418 - 4,
419 - ),
420 - True,
421 - True,
422 - "auto",
423 - 0.8, # ref: 90%
424 - marks=large_gpu_mark(min_gb=80),
425 - ), # works on 4x H100
426 - (
427 - (
428 - "eagle",
429 - "eagle618/deepseek-v3-random",
430 - "eagle618/eagle-deepseek-v3-random",
431 - 1,
432 - ),
433 - False,
434 - False,
435 - "auto",
436 - 0.0, # dummy model, skip gsm8k check
437 - ),
438 - ],
439 - ids=[
440 - "qwen3_eagle3",
441 - "qwen3_eagle3-transformers",
442 - "qwen3_vl_eagle3",
443 - "qwen2_5_vl_eagle3",
444 - "llama3_eagle",
445 - "llama3_eagle3",
446 - "llama4_eagle",
447 - "llama4_eagle_mm",
448 - "deepseek_eagle",
449 - ],
450 -)
451 -@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
452 -def test_eagle_correctness(
334 +def _run_eagle_correctness(
335 monkeypatch: pytest.MonkeyPatch,
336 sampling_config: SamplingParams,
337 model_setup: tuple[str, str, str, int],
@@ -460,14 +342,10 @@ def test_eagle_correctness(
342 attn_backend: str,
343 ):
344 """
463 - Compare the outputs of a original LLM and a speculative LLM
464 - which should be the same when using eagle speculative decoding. Due to some variance
465 - in the engine, it is possible for some outputs to differ, so we expect that at least
466 - 6/10 output tokens match exactly, and that the GSM8k accuracy is above
467 - a precomputed reference threshold for each model.
345 + Compare the outputs of an original LLM and a speculative LLM
346 + which should be the same when using eagle speculative decoding.
347 """
348 if attn_backend == "TREE_ATTN":
470 - # TODO: Fix this flaky test
349 pytest.skip(
350 "TREE_ATTN is flaky in the test disable for now until it can be "
351 "resolved (see https://github.com/vllm-project/vllm/issues/22922)"
@@ -484,17 +362,17 @@ def test_eagle_correctness(
362 f"transformers>={required}, but got {installed}"
363 )
364
487 - # Generate test prompts inside the function instead of using fixture
365 test_prompts = get_test_prompts(mm_enabled)
489 - # Determine attention config
490 - # Scout requires default backend selection because vision encoder has
491 - # head_dim 88 being incompatible with FLASH_ATTN and needs to fall back
492 - # to Flex Attn
366 +
367 if "Llama-4-Scout" in model_setup[1] and attn_backend == "FLASH_ATTN":
368 if current_platform.is_rocm():
495 - # TODO: Enable Flex Attn for spec_decode on ROCm
496 - pytest.skip("Flex Attn for spec_decode not supported on ROCm currently")
497 - attention_config = None # Let it fall back to default
369 + print(
370 + "FLASH_ATTN for spec_decode not supported on "
371 + "ROCm currently. Changing to FLEX_ATTENTION backend."
372 + )
373 + attention_config = {"backend": "FLEX_ATTENTION"}
374 + else:
375 + attention_config = None
376 else:
377 attention_config = {"backend": attn_backend}
378
@@ -509,7 +387,9 @@ def test_eagle_correctness(
387
388 if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
389 if "deepseek" in model_setup[1].lower():
512 - pytest.skip("ROCM_AITER_FA for deepseek not supported on ROCm platform")
390 + m.setenv("VLLM_ROCM_USE_AITER", "1")
391 + m.delenv("VLLM_MLA_DISABLE", raising=False)
392 + attention_config = {"backend": "TRITON_MLA"}
393 else:
394 m.setenv("VLLM_ROCM_USE_AITER", "1")
395
@@ -563,14 +443,235 @@ def test_eagle_correctness(
443 print(f"ref_output: {ref_output.outputs[0].text}")
444 print(f"spec_output: {spec_output.outputs[0].text}")
445
566 - # Heuristic: expect at least 60% of the prompts to match exactly
567 - # Upon failure, inspect the outputs to check for inaccuracy.
446 assert matches > int(0.6 * len(ref_outputs))
447 del spec_llm
448 torch.cuda.empty_cache()
449 cleanup_dist_env_and_memory()
450
451
452 +@single_gpu_only
453 +@pytest.mark.parametrize(
454 + [
455 + "model_setup",
456 + "mm_enabled",
457 + "enable_chunked_prefill",
458 + "model_impl",
459 + "expected_accuracy_threshold",
460 + ],
461 + [
462 + (
463 + (
464 + "eagle",
465 + "eagle618/deepseek-v3-random",
466 + "eagle618/eagle-deepseek-v3-random",
467 + 1,
468 + ),
469 + False,
470 + False,
471 + "auto",
472 + 0.0,
473 + ),
474 + ],
475 + ids=["deepseek_eagle"],
476 +)
477 +@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
478 +def test_eagle_correctness_light(
479 + monkeypatch: pytest.MonkeyPatch,
480 + sampling_config: SamplingParams,
481 + model_setup: tuple[str, str, str, int],
482 + mm_enabled: bool,
483 + expected_accuracy_threshold: float,
484 + enable_chunked_prefill: bool,
485 + model_impl: str,
486 + attn_backend: str,
487 +):
488 + _run_eagle_correctness(
489 + monkeypatch,
490 + sampling_config,
491 + model_setup,
492 + mm_enabled,
493 + expected_accuracy_threshold,
494 + enable_chunked_prefill,
495 + model_impl,
496 + attn_backend,
497 + )
498 +
499 +
500 +@single_gpu_only
501 +@large_gpu_mark(min_gb=24)
502 +@pytest.mark.parametrize(
503 + [
504 + "model_setup",
505 + "mm_enabled",
506 + "enable_chunked_prefill",
507 + "model_impl",
508 + "expected_accuracy_threshold",
509 + ],
510 + [
511 + (
512 + ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
513 + False,
514 + False,
515 + "auto",
516 + 0.8,
517 + ),
518 + (
519 + ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
520 + False,
521 + False,
522 + "transformers",
523 + 0.8,
524 + ),
525 + pytest.param(
526 + (
527 + "eagle3",
528 + "Qwen/Qwen3-VL-8B-Instruct",
529 + "taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
530 + 1,
531 + ),
532 + False,
533 + False,
534 + "auto",
535 + 0.8,
536 + marks=pytest.mark.skip(
537 + reason="architecture of its eagle3 is LlamaForCausalLMEagle3"
538 + ),
539 + ),
540 + pytest.param(
541 + (
542 + "eagle3",
543 + "Qwen/Qwen2.5-VL-7B-Instruct",
544 + "Rayzl/qwen2.5-vl-7b-eagle3-sgl",
545 + 1,
546 + ),
547 + False,
548 + False,
549 + "auto",
550 + 0.7,
551 + marks=pytest.mark.skip(
552 + reason="Skipping due to its head_dim not being a multiple of 32"
553 + ),
554 + ),
555 + (
556 + (
557 + "eagle3",
558 + "meta-llama/Llama-3.1-8B-Instruct",
559 + "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
560 + 1,
561 + ),
562 + False,
563 + False,
564 + "auto",
565 + 0.7,
566 + ),
567 + ],
568 + ids=[
569 + "qwen3_eagle3",
570 + "qwen3_eagle3-transformers",
571 + "qwen3_vl_eagle3",
572 + "qwen2_5_vl_eagle3",
573 + "llama3_eagle3",
574 + ],
575 +)
576 +@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
577 +def test_eagle_correctness_medium(
578 + monkeypatch: pytest.MonkeyPatch,
579 + sampling_config: SamplingParams,
580 + model_setup: tuple[str, str, str, int],
581 + mm_enabled: bool,
582 + expected_accuracy_threshold: float,
583 + enable_chunked_prefill: bool,
584 + model_impl: str,
585 + attn_backend: str,
586 +):
587 + _run_eagle_correctness(
588 + monkeypatch,
589 + sampling_config,
590 + model_setup,
591 + mm_enabled,
592 + expected_accuracy_threshold,
593 + enable_chunked_prefill,
594 + model_impl,
595 + attn_backend,
596 + )
597 +
598 +
599 +@pytest.mark.parametrize(
600 + [
601 + "model_setup",
602 + "mm_enabled",
603 + "enable_chunked_prefill",
604 + "model_impl",
605 + "expected_accuracy_threshold",
606 + ],
607 + [
608 + pytest.param(
609 + (
610 + "eagle",
611 + "meta-llama/Llama-3.1-8B-Instruct",
612 + "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
613 + 1,
614 + ),
615 + False,
616 + True,
617 + "auto",
618 + 0.7,
619 + marks=large_gpu_mark(min_gb=40),
620 + id="llama3_eagle",
621 + ),
622 + pytest.param(
623 + (
624 + "eagle",
625 + "meta-llama/Llama-4-Scout-17B-16E-Instruct",
626 + "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
627 + 4,
628 + ),
629 + False,
630 + False,
631 + "auto",
632 + 0.8,
633 + marks=multi_gpu_marks(num_gpus=4),
634 + id="llama4_eagle",
635 + ),
636 + pytest.param(
637 + (
638 + "eagle",
639 + "meta-llama/Llama-4-Scout-17B-16E-Instruct",
640 + "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
641 + 4,
642 + ),
643 + True,
644 + True,
645 + "auto",
646 + 0.8,
647 + marks=[*multi_gpu_marks(num_gpus=4), large_gpu_mark(min_gb=80)],
648 + id="llama4_eagle_mm",
649 + ),
650 + ],
651 +)
652 +@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
653 +def test_eagle_correctness_heavy(
654 + monkeypatch: pytest.MonkeyPatch,
655 + sampling_config: SamplingParams,
656 + model_setup: tuple[str, str, str, int],
657 + mm_enabled: bool,
658 + expected_accuracy_threshold: float,
659 + enable_chunked_prefill: bool,
660 + model_impl: str,
661 + attn_backend: str,
662 +):
663 + _run_eagle_correctness(
664 + monkeypatch,
665 + sampling_config,
666 + model_setup,
667 + mm_enabled,
668 + expected_accuracy_threshold,
669 + enable_chunked_prefill,
670 + model_impl,
671 + attn_backend,
672 + )
673 +
674 +
675 @pytest.mark.parametrize(
676 ["model_setup", "mm_enabled", "expected_accuracy_threshold"],
677 [
@@ -579,6 +680,8 @@ def test_eagle_correctness(
680 ],
681 ids=["mimo", "deepseek"],
682 )
683 +@single_gpu_only
684 +@large_gpu_mark(min_gb=20)
685 def test_mtp_correctness(
686 monkeypatch: pytest.MonkeyPatch,
687 sampling_config: SamplingParams,
@@ -694,11 +797,13 @@ cases = [
797
798 @pytest.mark.parametrize("args", cases)
799 @pytest.mark.parametrize("enforce_eager", [True, False])
800 +@single_gpu_only
801 def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
802 args.enforce_eager = enforce_eager
803 assert_draft_model_correctness(args)
804
805
806 +@single_gpu_only
807 def test_draft_model_realistic_example():
808 args = ArgsTest(
809 target_model="Qwen/Qwen3-1.7B",
@@ -713,6 +818,7 @@ def test_draft_model_realistic_example():
818 assert_draft_model_correctness(args)
819
820
821 +@single_gpu_only
822 def test_draft_model_parallel_drafting():
823 args = ArgsTest(
824 target_model="Qwen/Qwen3-1.7B",
@@ -738,6 +844,7 @@ def test_draft_model_parallel_drafting():
844 ids=["target_quantized", "draft_quantized"],
845 )
846 @pytest.mark.parametrize("enforce_eager", [True, False])
847 +@single_gpu_only
848 def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
849 tgt_model, draft_model = models
850 sd_case = ArgsTest(
@@ -749,6 +856,7 @@ def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
856 assert_draft_model_correctness(sd_case)
857
858
859 +@multi_gpu_only(num_gpus=2)
860 def test_draft_model_tensor_parallelism():
861 """Ensure spec decode works when running with TP > 1."""
862 _skip_if_insufficient_gpus_for_tp(2)
@@ -764,6 +872,7 @@ def test_draft_model_tensor_parallelism():
872 assert_draft_model_correctness(sd_case)
873
874
875 +@multi_gpu_only(num_gpus=2)
876 def test_draft_model_engine_args_tensor_parallelism():
877 """Ensure the vllm_config for the draft model is created correctly,
878 and independently of the target model (quantization, TP, etc.)"""