[CI] Align extraction test with canonical auxiliary layer order (#55457)
Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: Codex <noreply@openai.com>
Kevin H. Luu committed
Sep 6, 2026 at 11:22 UTC
722d169339b10117cc65fec85c85279c00c82d8c
1 file changed
+6
-4
tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
+6
-4
@@ -75,6 +75,7 @@ PREDICTABLE_LLAMA_MODEL_CLS = (
75
)
76
77
78
+@create_new_process_for_each_test()
79
def test_extract_hidden_states_with_predictable_dummy_model(
80
predictable_llama_config_path, tmp_path, monkeypatch
81
):
@@ -82,7 +83,7 @@ def test_extract_hidden_states_with_predictable_dummy_model(
83
84
Tests 3 scenarios:
85
85
- 1. **Basic extraction**: non-sequential layer ordering, multiple prompts
86
+ 1. **Basic extraction**: unsorted layer selection, multiple prompts
87
of varying length — verifies correct layer association and
88
deterministic values.
89
2. **Chunked prefill**: max_num_batched_tokens=128 with ~500-token
@@ -95,6 +96,7 @@ def test_extract_hidden_states_with_predictable_dummy_model(
96
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork")
97
98
layer_ids = [5, 2, 10]
99
+ expected_layer_ids = sorted(layer_ids)
100
num_layers = len(layer_ids)
101
max_num_batched_tokens = 128
102
@@ -146,7 +148,7 @@ def test_extract_hidden_states_with_predictable_dummy_model(
148
)
149
_token_ids, hidden_states = get_and_check_output(output, expected_shape)
150
149
- for idx, layer_id in enumerate(layer_ids):
151
+ for idx, layer_id in enumerate(expected_layer_ids):
152
layer_hidden = hidden_states[:, idx, :]
153
assert torch.allclose(
154
layer_hidden,
@@ -174,7 +176,7 @@ def test_extract_hidden_states_with_predictable_dummy_model(
176
expected_shape = (prompt_len, num_layers, hidden_size)
177
_token_ids, hidden_states = get_and_check_output(output, expected_shape)
178
177
- for idx, layer_id in enumerate(layer_ids):
179
+ for idx, layer_id in enumerate(expected_layer_ids):
180
layer_hidden = hidden_states[:, idx, :]
181
assert torch.allclose(
182
layer_hidden,
@@ -236,7 +238,7 @@ def test_extract_hidden_states_with_predictable_dummy_model(
238
assert hidden_states.shape == (total_tokens, num_layers, hidden_size)
239
240
# Verify predictable layer values hold for all tokens (prompt + output)
239
- for idx, layer_id in enumerate(layer_ids):
241
+ for idx, layer_id in enumerate(expected_layer_ids):
242
layer_hidden = hidden_states[:, idx, :]
243
assert torch.allclose(
244
layer_hidden,