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 15, 07:15 UTC
[TPU] Update dynamo dump file name in compilation test (#19108)
Signed-off-by: Siyuan Liu <lsiyuan@google.com>
Siyuan Liu committed
Jun 4, 2025 at 13:13 UTC
7ee2590478e088e8df56652ba2607cecb5123b6b
3 files changed
+19
-13
.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh
+1
-1
@@ -150,7 +150,7 @@ run_and_track_test 9 "test_multimodal.py" \
150
run_and_track_test 10 "test_pallas.py" \
151
"python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_pallas.py"
152
run_and_track_test 11 "test_struct_output_generate.py" \
153
- "python3 -m pytest -s -v /workspace/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py -k 'not test_structured_output_with_reasoning_matrices'"
153
+ "python3 -m pytest -s -v /workspace/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py -k \"not test_structured_output_with_reasoning_matrices\""
154
run_and_track_test 12 "test_moe_pallas.py" \
155
"python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_moe_pallas.py"
156
run_and_track_test 13 "test_lora.py" \
tests/tpu/test_compilation.py
+3
-2
@@ -64,9 +64,10 @@ def test_tpu_compilation():
64
numbers = [int(part) for part in parts if part.isdigit()]
65
return numbers[0]
66
67
- # Check all the compilations are as expected
67
+ # Check all the compilations are as expected. The dump files include the
68
+ # captured graph for the forward function of the nn.Module.
69
compiled_fns = sorted(glob.glob(
69
- os.path.join(temp_dir, "__compiled_fn*Captured*.py")),
70
+ os.path.join(temp_dir, "__compiled_fn*Forward_graph*.py")),
71
key=lambda s: extract_compiled_index(s))
72
73
for i, compiled_fn in enumerate(compiled_fns):
tests/v1/tpu/worker/test_tpu_model_runner.py
+15
-10
@@ -370,6 +370,7 @@ def test_get_req_paddings():
370
assert _get_req_paddings(8, 36) == [8, 16, 32, 36]
371
372
373
+@pytest.mark.skip(reason="Test is broken on TPU when it's added.")
374
def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():
375
layer_0 = "model.layers.0.self_attn.attn"
376
layer_1 = "model.layers.1.self_attn.attn"
@@ -381,7 +382,7 @@ def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():
382
layer_0:
383
Attention(
384
num_heads=8,
384
- head_size=64,
385
+ head_size=128,
386
scale=1.0,
387
prefix=layer_0,
388
kv_sharing_target_layer_name=layer_1,
@@ -389,7 +390,7 @@ def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():
390
layer_1:
391
Attention(
392
num_heads=8,
392
- head_size=64,
393
+ head_size=128,
394
scale=1.0,
395
prefix=layer_1,
396
)
@@ -398,6 +399,7 @@ def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():
399
assert fwd_context is not None
400
401
402
+@pytest.mark.skip(reason="Test is broken on TPU when it's added.")
403
def test_init_kv_cache_with_kv_sharing_target_layer_not_exist():
404
layer_0 = "model.layers.0.self_attn.attn"
405
layer_1 = "model.layers.1.self_attn.attn"
@@ -408,14 +410,14 @@ def test_init_kv_cache_with_kv_sharing_target_layer_not_exist():
410
layer_0:
411
Attention(
412
num_heads=8,
411
- head_size=64,
413
+ head_size=128,
414
scale=1.0,
415
prefix=layer_0,
416
),
417
layer_1:
418
Attention(
419
num_heads=8,
418
- head_size=64,
420
+ head_size=128,
421
scale=1.0,
422
prefix=layer_1,
423
# invalid layer: cross_attn.atn doesn't exist!
@@ -426,6 +428,7 @@ def test_init_kv_cache_with_kv_sharing_target_layer_not_exist():
428
assert fwd_context is not None
429
430
431
+@pytest.mark.skip(reason="Test is broken on TPU when it's added.")
432
def test_init_kv_cache_with_kv_sharing_target_same_as_current():
433
layer_0 = "model.layers.0.self_attn.attn"
434
layer_1 = "model.layers.1.self_attn.attn"
@@ -437,14 +440,14 @@ def test_init_kv_cache_with_kv_sharing_target_same_as_current():
440
layer_0:
441
Attention(
442
num_heads=8,
440
- head_size=64,
443
+ head_size=128,
444
scale=1.0,
445
prefix=layer_0,
446
),
447
layer_1:
448
Attention(
449
num_heads=8,
447
- head_size=64,
450
+ head_size=128,
451
scale=1.0,
452
prefix=layer_1,
453
kv_sharing_target_layer_name=layer_1,
@@ -454,6 +457,7 @@ def test_init_kv_cache_with_kv_sharing_target_same_as_current():
457
assert fwd_context is not None
458
459
460
+@pytest.mark.skip(reason="Test is broken on TPU when it's added.")
461
def test_init_kv_cache_without_kv_sharing(model_runner):
462
layer_0 = "model.layers.0.self_attn.attn"
463
layer_1 = "model.layers.1.self_attn.attn"
@@ -463,14 +467,14 @@ def test_init_kv_cache_without_kv_sharing(model_runner):
467
layer_0:
468
Attention(
469
num_heads=8,
466
- head_size=64,
470
+ head_size=128,
471
scale=1.0,
472
prefix=layer_0,
473
),
474
layer_1:
475
Attention(
476
num_heads=8,
473
- head_size=64,
477
+ head_size=128,
478
scale=1.0,
479
prefix=layer_1,
480
)
@@ -520,6 +524,7 @@ def test_init_kv_cache_without_kv_sharing(model_runner):
524
assert kv_cache_config.kv_cache_groups[0].layer_names[1] == layer_1
525
526
527
+@pytest.mark.skip(reason="Test is broken on TPU when it's added.")
528
def test_init_kv_cache_with_kv_sharing_valid(model_runner):
529
layer_0 = "model.layers.0.self_attn.attn"
530
layer_1 = "model.layers.1.self_attn.attn"
@@ -529,14 +534,14 @@ def test_init_kv_cache_with_kv_sharing_valid(model_runner):
534
layer_0:
535
Attention(
536
num_heads=8,
532
- head_size=64,
537
+ head_size=128,
538
scale=1.0,
539
prefix=layer_0,
540
),
541
layer_1:
542
Attention(
543
num_heads=8,
539
- head_size=64,
544
+ head_size=128,
545
scale=1.0,
546
prefix=layer_1,
547
kv_sharing_target_layer_name="model.layers.0.self_attn.attn",