[CI] Fix stale/fragile untethered kernels-root tests (#49423)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

tests/kernels/conftest.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch
+
+
+@pytest.fixture(autouse=True)
+def reset_default_torch_device():
+    """Several kernel tests call torch.set_default_device without restoring
+    it, which poisons subsequent tests in the same pytest run (e.g. CPU
+    tensors silently created on CUDA). Restore the factory default after
+    every test.
+    """
+    yield
+    torch.set_default_device(None)

tests/kernels/test_flex_attention.py

@@ -264,7 +264,7 @@ def test_block_mask_direct_vs_slow_path():
     device = torch.device("cuda")
 
     vllm_config = create_vllm_config(
-        model_name="meta-llama/Meta-Llama-3-8B", block_size=16, max_model_len=1024
+        model_name="Qwen/Qwen2.5-1.5B-Instruct", block_size=16, max_model_len=1024
     )
     kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
 

tests/kernels/test_fused_inv_rope_fp8_quant.py

@@ -726,14 +726,19 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
     This catches stride/layout bugs that only manifest when the einsum
     kernel actually consumes the quantized activations.
     """
-    from deep_gemm.utils.math import ceil_div
-
     from vllm.utils.deep_gemm import (
         fp8_einsum,
+        is_deep_gemm_supported,
         per_block_cast_to_fp8,
         transform_sf_into_required_layout,
     )
 
+    if not is_deep_gemm_supported():
+        pytest.skip("DeepGEMM not supported on this platform")
+
+    def ceil_div(a: int, b: int) -> int:
+        return (a + b - 1) // b
+
     heads_per_group = num_heads // n_groups
     d = heads_per_group * HEAD_DIM
     o_lora_rank = 1024
@@ -809,8 +814,12 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
     # -- Checks --
     # Einsum output: Triton and CUDA both rotate in fp32 now, so diffs
     # come from fp32 ordering and UE8M0 boundary shifts only.
-    # Use relative diff (same metric as test_fp8_einsum.py).
-    from deep_gemm.testing import calc_diff
+    # Use relative diff (same metric as deep_gemm.testing.calc_diff).
+    def calc_diff(x, y):
+        x, y = x.double(), y.double()
+        denominator = (x * x + y * y).sum()
+        sim = 2 * (x * y).sum() / denominator
+        return 1 - sim
 
     z_diff = calc_diff(z_fused, z_ref)
     assert z_diff < 0.01, (

tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py

@@ -84,6 +84,22 @@ def norm_rope_ref(x, weight, positions, cos_sin_cache, eps):
     return roped
 
 
+def assert_fp8_cache_close(kv_cache, expected_kv_cache):
+    """Compare two e4m3 caches allowing 1 ulp.
+
+    On CUDA the fused kernel quantizes K from its fp32 intermediate, while the
+    reshape_and_cache_flash reference quantizes the bf16-materialized value, so
+    rounding-boundary values may differ by one e4m3 code.
+    """
+    byte_diff = (kv_cache.int() - expected_kv_cache.int()).abs()
+    got = kv_cache.view(torch.float8_e4m3fn).float()
+    exp = expected_kv_cache.view(torch.float8_e4m3fn).float()
+    ok = (byte_diff <= 1) | ((got == 0) & (exp == 0))
+    assert bool(ok.all()), (
+        f"fp8 cache differs by more than 1 ulp in {int((~ok).sum())} elements"
+    )
+
+
 # ── Test 1: dense mode (norm+rope only, no index, no insert) ─────────────────
 
 
@@ -265,7 +281,7 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
             scale,
             scale,
         )
-        torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
+        assert_fp8_cache_close(kv_cache, expected_kv_cache)
     else:
         for t in range(num_tokens):
             s = slot_mapping[t].item()
@@ -383,7 +399,7 @@ def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype):
             scale,
             scale,
         )
-        torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
+        assert_fp8_cache_close(kv_cache, expected_kv_cache)
     else:
         k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
         v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM)

tests/kernels/test_fused_recurrent_packed_decode.py

@@ -41,11 +41,12 @@ def test_fused_recurrent_packed_decode_matches_reference(
     A_log = torch.randn((HV,), device=device, dtype=dtype)
     dt_bias = torch.randn((HV,), device=device, dtype=dtype)
 
-    # Continuous batching indices (include PAD_SLOT_ID=-1 cases).
-    ssm_state_indices = torch.arange(B, device=device, dtype=torch.int32)
+    # Continuous batching indices (include PAD_SLOT_ID=-1 cases). Index 0 is
+    # reserved as NULL_BLOCK_ID (CUDA graph padding), so valid slots start at 1.
+    ssm_state_indices = torch.arange(1, B + 1, device=device, dtype=torch.int32)
     ssm_state_indices[-3:] = -1
 
-    state0 = torch.randn((B, HV, V, K), device=device, dtype=dtype)
+    state0 = torch.randn((B + 1, HV, V, K), device=device, dtype=dtype)
     state_ref = state0.clone()
     state_packed = state0.clone()
 
@@ -94,5 +95,8 @@ def test_fused_recurrent_packed_decode_matches_reference(
 
     atol = 2e-2 if dtype != torch.float32 else 1e-4
     rtol = 1e-2 if dtype != torch.float32 else 1e-4
-    torch.testing.assert_close(out_packed, out_ref, rtol=rtol, atol=atol)
+    # Output rows for PAD_SLOT_ID entries are never written (uninitialized in
+    # both paths), so compare only the valid rows.
+    valid = ssm_state_indices > 0
+    torch.testing.assert_close(out_packed[valid], out_ref[valid], rtol=rtol, atol=atol)
     torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol)

tests/kernels/test_fused_sigmoid_gating_delta_rule.py

@@ -58,10 +58,12 @@ def test_fused_sigmoid_gating_delta_rule_update_non_spec(
     dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
     a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
     b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
+    # Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
+    # state indices start at 1.
     ssm_state = torch.rand(
-        total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
+        total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
     )
-    state_indices = torch.randperm(total_entries, dtype=torch.int32)[:num_tokens]
+    state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[:num_tokens]
     cu_seqlens = torch.arange(0, num_tokens + 1, dtype=torch.int32)
 
     beta = b.sigmoid()
@@ -144,13 +146,14 @@ def test_fused_sigmoid_gating_delta_rule_update_spec(
     dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
     a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
     b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
+    # Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
+    # state indices start at 1.
     ssm_state = torch.rand(
-        total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
+        total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
     )
-    state_indices = torch.randperm(
-        total_entries,
-        dtype=torch.int32,
-    )[:num_tokens].view(num_reqs, num_speculative_tokens + 1)
+    state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[
+        :num_tokens
+    ].view(num_reqs, num_speculative_tokens + 1)
     num_accepted_tokens = torch.randint(
         1, num_speculative_tokens + 1, (num_reqs,), dtype=torch.int32
     )