From 431934522b65f126c43f74b19cdd8f3bf63f9747 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Wed, 22 Jul 2026 21:43:07 +0100 Subject: [PATCH] [CI] Fix stale/fragile untethered kernels-root tests (#49423) Signed-off-by: Nick Hill --- tests/kernels/conftest.py | 15 ++++++++++++++ tests/kernels/test_flex_attention.py | 2 +- .../kernels/test_fused_inv_rope_fp8_quant.py | 17 ++++++++++++---- ..._fused_minimax_m3_qknorm_rope_kv_insert.py | 20 +++++++++++++++++-- .../test_fused_recurrent_packed_decode.py | 12 +++++++---- .../test_fused_sigmoid_gating_delta_rule.py | 17 +++++++++------- 6 files changed, 65 insertions(+), 18 deletions(-) create mode 100644 tests/kernels/conftest.py diff --git a/tests/kernels/conftest.py b/tests/kernels/conftest.py new file mode 100644 index 0000000000..290d155342 --- /dev/null +++ b/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) diff --git a/tests/kernels/test_flex_attention.py b/tests/kernels/test_flex_attention.py index 86f26cfe8c..f01dca373a 100644 --- a/tests/kernels/test_flex_attention.py +++ b/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) diff --git a/tests/kernels/test_fused_inv_rope_fp8_quant.py b/tests/kernels/test_fused_inv_rope_fp8_quant.py index b5b81efa77..5b71e59baa 100644 --- a/tests/kernels/test_fused_inv_rope_fp8_quant.py +++ b/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, ( diff --git a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py index 5c762b63df..626b06290e 100644 --- a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py +++ b/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) diff --git a/tests/kernels/test_fused_recurrent_packed_decode.py b/tests/kernels/test_fused_recurrent_packed_decode.py index 128a020600..928c3e5bd7 100644 --- a/tests/kernels/test_fused_recurrent_packed_decode.py +++ b/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) diff --git a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py index 82a5a6f4ca..fc923cb945 100644 --- a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py +++ b/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 )