[CI] Fix stale/fragile untethered kernels-root tests (#49423)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jul 22, 2026 at 21:43 UTC
431934522b65f126c43f74b19cdd8f3bf63f9747
6 files changed
+65
-18
tests/kernels/conftest.py
new
+15
@@ -0,0 +1,15 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+
6
+
7
+@pytest.fixture(autouse=True)
8
+def reset_default_torch_device():
9
+ """Several kernel tests call torch.set_default_device without restoring
10
+ it, which poisons subsequent tests in the same pytest run (e.g. CPU
11
+ tensors silently created on CUDA). Restore the factory default after
12
+ every test.
13
+ """
14
+ yield
15
+ torch.set_default_device(None)
tests/kernels/test_flex_attention.py
+1
-1
@@ -264,7 +264,7 @@ def test_block_mask_direct_vs_slow_path():
264
device = torch.device("cuda")
265
266
vllm_config = create_vllm_config(
267
- model_name="meta-llama/Meta-Llama-3-8B", block_size=16, max_model_len=1024
267
+ model_name="Qwen/Qwen2.5-1.5B-Instruct", block_size=16, max_model_len=1024
268
)
269
kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
270
tests/kernels/test_fused_inv_rope_fp8_quant.py
+13
-4
@@ -726,14 +726,19 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
726
This catches stride/layout bugs that only manifest when the einsum
727
kernel actually consumes the quantized activations.
728
"""
729
- from deep_gemm.utils.math import ceil_div
730
-
729
from vllm.utils.deep_gemm import (
730
fp8_einsum,
731
+ is_deep_gemm_supported,
732
per_block_cast_to_fp8,
733
transform_sf_into_required_layout,
734
)
735
736
+ if not is_deep_gemm_supported():
737
+ pytest.skip("DeepGEMM not supported on this platform")
738
+
739
+ def ceil_div(a: int, b: int) -> int:
740
+ return (a + b - 1) // b
741
+
742
heads_per_group = num_heads // n_groups
743
d = heads_per_group * HEAD_DIM
744
o_lora_rank = 1024
@@ -809,8 +814,12 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
814
# -- Checks --
815
# Einsum output: Triton and CUDA both rotate in fp32 now, so diffs
816
# come from fp32 ordering and UE8M0 boundary shifts only.
812
- # Use relative diff (same metric as test_fp8_einsum.py).
813
- from deep_gemm.testing import calc_diff
817
+ # Use relative diff (same metric as deep_gemm.testing.calc_diff).
818
+ def calc_diff(x, y):
819
+ x, y = x.double(), y.double()
820
+ denominator = (x * x + y * y).sum()
821
+ sim = 2 * (x * y).sum() / denominator
822
+ return 1 - sim
823
824
z_diff = calc_diff(z_fused, z_ref)
825
assert z_diff < 0.01, (
tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
+18
-2
@@ -84,6 +84,22 @@ def norm_rope_ref(x, weight, positions, cos_sin_cache, eps):
84
return roped
85
86
87
+def assert_fp8_cache_close(kv_cache, expected_kv_cache):
88
+ """Compare two e4m3 caches allowing 1 ulp.
89
+
90
+ On CUDA the fused kernel quantizes K from its fp32 intermediate, while the
91
+ reshape_and_cache_flash reference quantizes the bf16-materialized value, so
92
+ rounding-boundary values may differ by one e4m3 code.
93
+ """
94
+ byte_diff = (kv_cache.int() - expected_kv_cache.int()).abs()
95
+ got = kv_cache.view(torch.float8_e4m3fn).float()
96
+ exp = expected_kv_cache.view(torch.float8_e4m3fn).float()
97
+ ok = (byte_diff <= 1) | ((got == 0) & (exp == 0))
98
+ assert bool(ok.all()), (
99
+ f"fp8 cache differs by more than 1 ulp in {int((~ok).sum())} elements"
100
+ )
101
+
102
+
103
# ── Test 1: dense mode (norm+rope only, no index, no insert) ─────────────────
104
105
@@ -265,7 +281,7 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
281
scale,
282
scale,
283
)
268
- torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
284
+ assert_fp8_cache_close(kv_cache, expected_kv_cache)
285
else:
286
for t in range(num_tokens):
287
s = slot_mapping[t].item()
@@ -383,7 +399,7 @@ def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype):
399
scale,
400
scale,
401
)
386
- torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
402
+ assert_fp8_cache_close(kv_cache, expected_kv_cache)
403
else:
404
k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
405
v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM)
tests/kernels/test_fused_recurrent_packed_decode.py
+8
-4
@@ -41,11 +41,12 @@ def test_fused_recurrent_packed_decode_matches_reference(
41
A_log = torch.randn((HV,), device=device, dtype=dtype)
42
dt_bias = torch.randn((HV,), device=device, dtype=dtype)
43
44
- # Continuous batching indices (include PAD_SLOT_ID=-1 cases).
45
- ssm_state_indices = torch.arange(B, device=device, dtype=torch.int32)
44
+ # Continuous batching indices (include PAD_SLOT_ID=-1 cases). Index 0 is
45
+ # reserved as NULL_BLOCK_ID (CUDA graph padding), so valid slots start at 1.
46
+ ssm_state_indices = torch.arange(1, B + 1, device=device, dtype=torch.int32)
47
ssm_state_indices[-3:] = -1
48
48
- state0 = torch.randn((B, HV, V, K), device=device, dtype=dtype)
49
+ state0 = torch.randn((B + 1, HV, V, K), device=device, dtype=dtype)
50
state_ref = state0.clone()
51
state_packed = state0.clone()
52
@@ -94,5 +95,8 @@ def test_fused_recurrent_packed_decode_matches_reference(
95
96
atol = 2e-2 if dtype != torch.float32 else 1e-4
97
rtol = 1e-2 if dtype != torch.float32 else 1e-4
97
- torch.testing.assert_close(out_packed, out_ref, rtol=rtol, atol=atol)
98
+ # Output rows for PAD_SLOT_ID entries are never written (uninitialized in
99
+ # both paths), so compare only the valid rows.
100
+ valid = ssm_state_indices > 0
101
+ torch.testing.assert_close(out_packed[valid], out_ref[valid], rtol=rtol, atol=atol)
102
torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol)
tests/kernels/test_fused_sigmoid_gating_delta_rule.py
+10
-7
@@ -58,10 +58,12 @@ def test_fused_sigmoid_gating_delta_rule_update_non_spec(
58
dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
59
a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
60
b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
61
+ # Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
62
+ # state indices start at 1.
63
ssm_state = torch.rand(
62
- total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
64
+ total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
65
)
64
- state_indices = torch.randperm(total_entries, dtype=torch.int32)[:num_tokens]
66
+ state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[:num_tokens]
67
cu_seqlens = torch.arange(0, num_tokens + 1, dtype=torch.int32)
68
69
beta = b.sigmoid()
@@ -144,13 +146,14 @@ def test_fused_sigmoid_gating_delta_rule_update_spec(
146
dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
147
a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
148
b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
149
+ # Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
150
+ # state indices start at 1.
151
ssm_state = torch.rand(
148
- total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
152
+ total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
153
)
150
- state_indices = torch.randperm(
151
- total_entries,
152
- dtype=torch.int32,
153
- )[:num_tokens].view(num_reqs, num_speculative_tokens + 1)
154
+ state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[
155
+ :num_tokens
156
+ ].view(num_reqs, num_speculative_tokens + 1)
157
num_accepted_tokens = torch.randint(
158
1, num_speculative_tokens + 1, (num_reqs,), dtype=torch.int32
159
)