From eaf074d50eef0bdb99bf6282fa1ffc4793413b7b Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 6 May 2026 18:46:52 -0700 Subject: [PATCH] propagate pytest exit code from test __main__ entries (#24487) --- .../test_norm_tanh_mul_add_norm_scale.py | 88 -- python/sglang/jit_kernel/tests/test_cast.py | 314 ---- .../tests/test_flash_attention_3.py | 1359 ----------------- .../tests/test_fused_qknorm_rope.py | 448 ------ .../test/server/test_tracing.py | 155 -- .../test/unit/manual/test_patch_embed.py | 292 ---- test/registered/lora/test_lora_moe_runner.py | 788 ---------- .../lora/test_marlin_lora_correctness.py | 289 ---- .../lora/test_sgemm_sorted_by_adapter.py | 236 --- .../unit/test_no_bare_pytest_main.py | 90 ++ 10 files changed, 90 insertions(+), 3969 deletions(-) delete mode 100644 python/sglang/jit_kernel/tests/diffusion/test_norm_tanh_mul_add_norm_scale.py delete mode 100644 python/sglang/jit_kernel/tests/test_cast.py delete mode 100644 python/sglang/jit_kernel/tests/test_flash_attention_3.py delete mode 100644 python/sglang/jit_kernel/tests/test_fused_qknorm_rope.py delete mode 100644 python/sglang/multimodal_gen/test/server/test_tracing.py delete mode 100644 python/sglang/multimodal_gen/test/unit/manual/test_patch_embed.py delete mode 100644 test/registered/lora/test_lora_moe_runner.py delete mode 100644 test/registered/lora/test_marlin_lora_correctness.py delete mode 100644 test/registered/lora/test_sgemm_sorted_by_adapter.py create mode 100644 test/registered/unit/test_no_bare_pytest_main.py diff --git a/python/sglang/jit_kernel/tests/diffusion/test_norm_tanh_mul_add_norm_scale.py b/python/sglang/jit_kernel/tests/diffusion/test_norm_tanh_mul_add_norm_scale.py deleted file mode 100644 index e36736917..000000000 --- a/python/sglang/jit_kernel/tests/diffusion/test_norm_tanh_mul_add_norm_scale.py +++ /dev/null @@ -1,88 +0,0 @@ -import pytest -import torch - -from sglang.jit_kernel.diffusion.cutedsl.norm_tanh_mul_add_norm_scale import ( - fused_norm_tanh_mul_add, - fused_norm_tanh_mul_add_norm_scale, -) -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=45, suite="stage-b-kernel-unit-1-gpu-large") -register_cuda_ci(est_time=180, suite="nightly-kernel-1-gpu", nightly=True) - -BSD_CONFIG = [ - (1, 3648, 3840), # Z-image - (1, 4128, 3840), # Z-image - (3, 7, 256), # bound - (7, 1, 8192), # bound -] - - -@pytest.mark.parametrize("B,S,D", BSD_CONFIG) -@pytest.mark.parametrize("norm_type", ["rms", "layer"]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -def test_norm_tanh_mul_add(B: int, S: int, D: int, norm_type: str, dtype: str) -> None: - device = "cuda" - eps = 1e-5 - x = torch.randn(B, S, D, device=device, dtype=dtype) - weight = torch.randn(D, device=device, dtype=dtype) - bias = torch.randn(D, device=device, dtype=dtype) if norm_type == "layer" else None - scale = torch.randn(B, 1, D, device=device, dtype=dtype) - shift = torch.randn(B, 1, D, device=device, dtype=dtype) - - y = fused_norm_tanh_mul_add(x, weight, bias, scale, shift, norm_type, eps) - if norm_type == "rms": - normed = torch.rms_norm(x, x.shape[-1:], weight=weight, eps=eps) - else: - normed = torch.layer_norm(x, x.shape[-1:], weight=weight, bias=bias, eps=eps) - ref_y = normed * torch.tanh(scale) + shift - # Accuracy check - if dtype == "float32": - torch.testing.assert_close(y, ref_y, atol=1e-5, rtol=1e-5) - else: - torch.testing.assert_close(y, ref_y, atol=5e-2, rtol=5e-2) - - -@pytest.mark.parametrize("B,S,D", BSD_CONFIG) -@pytest.mark.parametrize("norm_type", ["rms", "layer"]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -def test_norm_tanh_mul_add_norm_scale( - B: int, S: int, D: int, norm_type: str, dtype: str -) -> None: - device = "cuda" - eps = 1e-5 - x = torch.randn(B, S, D, device=device, dtype=dtype) - weight = torch.randn(D, device=device, dtype=dtype) - bias = torch.randn(D, device=device, dtype=dtype) if norm_type == "layer" else None - scale = torch.randn(B, 1, D, device=device, dtype=dtype) - shift = torch.randn(B, 1, D, device=device, dtype=dtype) - weight2 = torch.randn(D, device=device, dtype=dtype) - bias2 = torch.randn(D, device=device, dtype=dtype) if norm_type == "layer" else None - scale2 = torch.randn(B, 1, D, device=device, dtype=dtype) - - y, y2 = fused_norm_tanh_mul_add_norm_scale( - x, weight, bias, scale, shift, weight2, bias2, scale2, norm_type, eps - ) - if norm_type == "rms": - normed = torch.rms_norm(x, x.shape[-1:], weight=weight, eps=eps) - else: - normed = torch.layer_norm(x, x.shape[-1:], weight=weight, bias=bias, eps=eps) - ref_y = normed * torch.tanh(scale) + shift - if norm_type == "rms": - normed2 = torch.rms_norm(ref_y, ref_y.shape[-1:], weight=weight2, eps=eps) - else: - normed2 = torch.layer_norm( - ref_y, ref_y.shape[-1:], weight=weight2, bias=bias2, eps=eps - ) - ref_y2 = normed2 * (1 + scale2) - # Accuracy check - if dtype == "float32": - torch.testing.assert_close(y, ref_y, atol=1e-5, rtol=1e-5) - torch.testing.assert_close(y2, ref_y2, atol=1e-5, rtol=1e-5) - else: - torch.testing.assert_close(y, ref_y, atol=5e-2, rtol=5e-2) - torch.testing.assert_close(y2, ref_y2, atol=5e-2, rtol=5e-2) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/python/sglang/jit_kernel/tests/test_cast.py b/python/sglang/jit_kernel/tests/test_cast.py deleted file mode 100644 index a63b4023c..000000000 --- a/python/sglang/jit_kernel/tests/test_cast.py +++ /dev/null @@ -1,314 +0,0 @@ -import pytest -import torch - -from sglang.jit_kernel.cast import downcast_fp8 -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=15, suite="stage-b-kernel-unit-1-gpu-large") -register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) - -DTYPES = [torch.bfloat16, torch.float16] - -# FP8 E4M3 representable range (matches kFP8E4M3Max in type.cuh) -_FP8_E4M3_MAX = 448.0 - - -def _run(input_sl, head, dim, out_sl, dtype): - k = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda") - v = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda") - k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - k_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda") - v_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda") - loc = torch.arange(input_sl, dtype=torch.int64, device="cuda") - downcast_fp8(k, v, k_out, v_out, k_scale, v_scale, loc) - return k_out, v_out - - -def _ref_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: - """Reference: replicate kernel precision — scale_inv in dtype T, then to fp8. - - Mirrors the kernel logic: - scale_inv = cast(1.0f) / cast(scale[0]) - out[j] = cast(clamp(x[j] * scale_inv)) - """ - dtype = x.dtype - scale_inv = x.new_ones(1) / scale[0].to(dtype) - x_scaled = (x * scale_inv).clamp(-_FP8_E4M3_MAX, _FP8_E4M3_MAX) - return x_scaled.to(torch.float8_e4m3fn).view(torch.uint8) - - -def _ref_downcast( - x: torch.Tensor, - scale: torch.Tensor, - loc: torch.Tensor, - out_sl: int, - mult: int = 1, - offset: int = 0, -) -> torch.Tensor: - """Scatter _ref_fp8 output to the correct output slots via loc/mult/offset.""" - head, dim = x.shape[1], x.shape[2] - out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device=x.device) - fp8 = _ref_fp8(x, scale) - for i, dst in enumerate(loc.tolist()): - out[dst * mult + offset] = fp8[i] - return out - - -# --------------------------------------------------------------------------- -# Existing sanity test -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("input_sl,head,dim,out_sl", [(4, 8, 128, 16)]) -def test_downcast_fp8(input_sl, head, dim, out_sl, dtype): - k = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda") - v = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda") - k_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda") - v_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda") - loc = torch.arange(input_sl, dtype=torch.int64, device="cuda") - - k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - downcast_fp8(k, v, k_out, v_out, k_scale, v_scale, loc) - - # Verify written slots are non-zero (fp8 of random non-zero values) - assert k_out[:input_sl].any(), "k_out should have non-zero fp8 values" - assert v_out[:input_sl].any(), "v_out should have non-zero fp8 values" - # Verify unwritten slots remain zero - assert not k_out[input_sl:].any(), "k_out slots beyond input_sl should be zero" - assert not v_out[input_sl:].any(), "v_out slots beyond input_sl should be zero" - - -# --------------------------------------------------------------------------- -# Numerical correctness: kernel output must match PyTorch fp8 reference. -# This verifies that cast(float) and cast(T) produce the -# same bit patterns as the removed ConvertFromFloat / ConvertToFP8 structs. -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("input_sl,head,dim,out_sl", [(4, 8, 128, 16), (1, 4, 64, 8)]) -def test_downcast_fp8_matches_reference(input_sl, head, dim, out_sl, dtype): - torch.manual_seed(42) - k = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda") - v = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda") - k_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda") - v_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda") - loc = torch.arange(input_sl, dtype=torch.int64, device="cuda") - - k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - downcast_fp8(k, v, k_out, v_out, k_scale, v_scale, loc) - - k_ref = _ref_downcast(k, k_scale, loc, out_sl) - v_ref = _ref_downcast(v, v_scale, loc, out_sl) - - torch.testing.assert_close(k_out, k_ref, msg="k: kernel vs reference mismatch") - torch.testing.assert_close(v_out, v_ref, msg="v: kernel vs reference mismatch") - - -# --------------------------------------------------------------------------- -# Scale: a non-unit scale divides the values before fp8 conversion. -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("scale_val", [0.5, 2.0, 0.1]) -def test_downcast_fp8_scale(scale_val, dtype): - torch.manual_seed(0) - input_sl, head, dim, out_sl = 4, 4, 64, 8 - - k = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda") - v = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda") - k_scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda") - v_scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda") - loc = torch.arange(input_sl, dtype=torch.int64, device="cuda") - - k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - downcast_fp8(k, v, k_out, v_out, k_scale, v_scale, loc) - - k_ref = _ref_downcast(k, k_scale, loc, out_sl) - v_ref = _ref_downcast(v, v_scale, loc, out_sl) - - torch.testing.assert_close( - k_out, k_ref, msg=f"scale={scale_val}: kernel vs reference mismatch" - ) - torch.testing.assert_close( - v_out, v_ref, msg=f"scale={scale_val}: kernel vs reference mismatch" - ) - - -# --------------------------------------------------------------------------- -# Clamping: values exceeding ±448 must be saturated to fp8 max/min. -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("dtype", DTYPES) -def test_downcast_fp8_clamp(dtype): - input_sl, head, dim, out_sl = 2, 1, 8, 4 - - # All values well outside fp8 range so clamping is unavoidable. - k = torch.full((input_sl, head, dim), 1000.0, dtype=dtype, device="cuda") - v = torch.full((input_sl, head, dim), -1000.0, dtype=dtype, device="cuda") - scale = torch.tensor([1.0], dtype=torch.float32, device="cuda") - loc = torch.arange(input_sl, dtype=torch.int64, device="cuda") - - k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - downcast_fp8(k, v, k_out, v_out, scale, scale, loc) - - # Reference fp8 max/min byte values (E4M3: 0x7e = 448.0, 0xfe = -448.0) - fp8_pos_max = ( - torch.tensor([_FP8_E4M3_MAX], dtype=dtype, device="cuda") - .to(torch.float8_e4m3fn) - .view(torch.uint8) - .item() - ) - fp8_neg_max = ( - torch.tensor([-_FP8_E4M3_MAX], dtype=dtype, device="cuda") - .to(torch.float8_e4m3fn) - .view(torch.uint8) - .item() - ) - - assert ( - k_out[:input_sl] == fp8_pos_max - ).all(), "large positive values should clamp to fp8 max" - assert ( - v_out[:input_sl] == fp8_neg_max - ).all(), "large negative values should clamp to fp8 min" - - -# --------------------------------------------------------------------------- -# Scatter: loc controls which output rows receive the converted values. -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("dtype", DTYPES) -def test_downcast_fp8_loc(dtype): - torch.manual_seed(7) - input_sl, head, dim, out_sl = 3, 2, 32, 10 - - k = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda") - v = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda") - scale = torch.tensor([1.0], dtype=torch.float32, device="cuda") - - # Write to non-contiguous output positions: 0, 5, 9 - loc = torch.tensor([0, 5, 9], dtype=torch.int64, device="cuda") - - k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - downcast_fp8(k, v, k_out, v_out, scale, scale, loc) - - k_ref = _ref_downcast(k, scale, loc, out_sl) - v_ref = _ref_downcast(v, scale, loc, out_sl) - - torch.testing.assert_close( - k_out, k_ref, msg="loc scatter: kernel vs reference mismatch" - ) - torch.testing.assert_close( - v_out, v_ref, msg="loc scatter: kernel vs reference mismatch" - ) - - # Slots not in loc must remain zero - written = {0, 5, 9} - for i in range(out_sl): - if i not in written: - assert not k_out[i].any(), f"k_out[{i}] should be zero (not a loc target)" - assert not v_out[i].any(), f"v_out[{i}] should be zero (not a loc target)" - - -# --------------------------------------------------------------------------- -# mult/offset: output index = loc[i] * mult + offset -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("mult,offset", [(2, 0), (1, 3), (2, 1)]) -def test_downcast_fp8_mult_offset(mult, offset, dtype): - torch.manual_seed(3) - input_sl, head, dim = 2, 2, 32 - out_sl = input_sl * mult + offset + 4 # ensure output is large enough - - k = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda") - v = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda") - scale = torch.tensor([1.0], dtype=torch.float32, device="cuda") - loc = torch.arange(input_sl, dtype=torch.int64, device="cuda") - - k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - downcast_fp8(k, v, k_out, v_out, scale, scale, loc, mult=mult, offset=offset) - - k_ref = _ref_downcast(k, scale, loc, out_sl, mult=mult, offset=offset) - v_ref = _ref_downcast(v, scale, loc, out_sl, mult=mult, offset=offset) - - torch.testing.assert_close( - k_out, k_ref, msg=f"mult={mult},offset={offset}: kernel vs reference mismatch" - ) - torch.testing.assert_close( - v_out, v_ref, msg=f"mult={mult},offset={offset}: kernel vs reference mismatch" - ) - - -# --------------------------------------------------------------------------- -# static_cast conversion: verify static_cast matches PyTorch fp8 -# for a comprehensive sweep including values near and at the fp8 boundary. -# This specifically validates that the static_cast fallback (used after -# removing explicit __nv_cvt_*raw_to_fp8 from dtype_trait) produces the -# same bit patterns as the reference path. -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("dtype", DTYPES) -def test_downcast_fp8_static_cast_boundary(dtype): - """Test conversion accuracy near ±448 fp8 boundary using static_cast path.""" - torch.manual_seed(0) - # Values specifically chosen to stress the static_cast conversion path: - # - exactly at ±448 (representable fp8 max) - # - just inside the range - # - just outside (must saturate) - # - zero, small, and mid-range values - boundary_vals = [ - 0.0, - 1.0, - -1.0, - 100.0, - -100.0, - 447.0, - -447.0, - 448.0, - -448.0, - 449.0, - -449.0, - 1000.0, - -1000.0, - ] - input_sl = len(boundary_vals) - head, dim, out_sl = 1, 8, input_sl - - base = torch.tensor(boundary_vals, dtype=dtype, device="cuda") - k = base.unsqueeze(1).unsqueeze(2).expand(input_sl, head, dim).contiguous() - v = (-base).unsqueeze(1).unsqueeze(2).expand(input_sl, head, dim).contiguous() - scale = torch.tensor([1.0], dtype=torch.float32, device="cuda") - loc = torch.arange(input_sl, dtype=torch.int64, device="cuda") - - k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda") - downcast_fp8(k, v, k_out, v_out, scale, scale, loc) - - k_ref = _ref_downcast(k, scale, loc, out_sl) - v_ref = _ref_downcast(v, scale, loc, out_sl) - - torch.testing.assert_close( - k_out, k_ref, msg="boundary values: k static_cast vs reference mismatch" - ) - torch.testing.assert_close( - v_out, v_ref, msg="boundary values: v static_cast vs reference mismatch" - ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/python/sglang/jit_kernel/tests/test_flash_attention_3.py b/python/sglang/jit_kernel/tests/test_flash_attention_3.py deleted file mode 100644 index 3f8447008..000000000 --- a/python/sglang/jit_kernel/tests/test_flash_attention_3.py +++ /dev/null @@ -1,1359 +0,0 @@ -# Adapted from https://github.com/Dao-AILab/flash-attention/blob/main/hopper/test_flash_attn.py -import itertools -import math -from typing import Optional - -import pytest -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - -apply_rotary_emb = None - -from sglang.jit_kernel.flash_attention_v3 import _is_fa3_supported -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=120, suite="stage-b-kernel-unit-1-gpu-large") -register_cuda_ci(est_time=900, suite="nightly-kernel-1-gpu", nightly=True) - - -def is_hopper(): - # Only Hopper supports different V headdim - return torch.cuda.get_device_properties(0).major == 9 - - -DISABLE_BACKWARD = True -# For CI test, we close them to True. -# DISABLE_SPLIT = os.getenv("FLASH_ATTENTION_DISABLE_SPLIT", "FALSE") == "TRUE" -# DISABLE_PAGEDKV = os.getenv("FLASH_ATTENTION_DISABLE_PAGEDKV", "FALSE") == "TRUE" -# DISABLE_APPENDKV = os.getenv("FLASH_ATTENTION_DISABLE_APPENDKV", "FALSE") == "TRUE" -# DISABLE_LOCAL = os.getenv("FLASH_ATTENTION_DISABLE_LOCAL", "FALSE") == "TRUE" -# DISABLE_SOFTCAP = os.getenv("FLASH_ATTENTION_DISABLE_SOFTCAP", "FALSE") == "TRUE" -# DISABLE_PACKGQA = os.getenv("FLASH_ATTENTION_DISABLE_PACKGQA", "FALSE") == "TRUE" -# DISABLE_FP16 = os.getenv("FLASH_ATTENTION_DISABLE_FP16", "FALSE") == "TRUE" -# DISABLE_FP8 = ( -# os.getenv("FLASH_ATTENTION_DISABLE_FP8", "FALSE") == "TRUE" -# or torch.cuda.get_device_capability("cuda")[0] < 9 -# ) - -DISABLE_SPLIT = False -DISABLE_PAGEDKV = True -DISABLE_APPENDKV = False -DISABLE_LOCAL = False -DISABLE_SOFTCAP = True -DISABLE_PACKGQA = False -DISABLE_FP16 = True -DISABLE_FP8 = True - - -# Adapted from https://github.com/Dao-AILab/flash-attention/blob/main/hopper/padding.py -def unpad_input(hidden_states, attention_mask, unused_mask=None): - """ - Arguments: - hidden_states: (batch, seqlen, ...) - attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. - unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. - Return: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. - indices: (total_nnz), the indices of masked tokens from the flattened input sequence. - cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. - max_seqlen_in_batch: int - seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. - """ - all_masks = ( - (attention_mask + unused_mask) if unused_mask is not None else attention_mask - ) - seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) - used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) - # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the - # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim - # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to - # index with integer indices. - return ( - rearrange(hidden_states, "b s ... -> (b s) ...")[indices], - indices, - cu_seqlens, - max_seqlen_in_batch, - used_seqlens_in_batch, - ) - - -def generate_random_padding_mask( - max_seqlen, batch_size, device, mode="random", zero_lengths=False -): - assert mode in ["full", "random", "third"] - if mode == "full": - lengths = torch.full( - (batch_size, 1), max_seqlen, device=device, dtype=torch.int32 - ) - elif mode == "random": - lengths = torch.randint( - max(0 if zero_lengths else 1, max_seqlen - 20), - max_seqlen + 1, - (batch_size, 1), - device=device, - ) - elif mode == "third": - lengths = torch.randint( - max_seqlen // 3, max_seqlen + 1, (batch_size, 1), device=device - ) - - if zero_lengths: - # Generate zero-lengths every 5 batches and the last batch. - for i in range(batch_size): - if i % 5 == 0: - lengths[i] = 0 - lengths[-1] = 0 - padding_mask = ( - repeat(torch.arange(max_seqlen, device=device), "s -> b s", b=batch_size) - < lengths - ) - return padding_mask - - -def pad_input(hidden_states, indices, batch, seqlen): - """ - Arguments: - hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. - indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. - batch: int, batch size for the padded sequence. - seqlen: int, maximum sequence length for the padded sequence. - Return: - hidden_states: (batch, seqlen, ...) - """ - dim = hidden_states.shape[1:] - output = torch.zeros( - (batch * seqlen), *dim, device=hidden_states.device, dtype=hidden_states.dtype - ) - output[indices] = hidden_states - return rearrange(output, "(b s) ... -> b s ...", b=batch) - - -def construct_local_mask( - seqlen_q, - seqlen_k, - window_size=(-1, -1), # -1 means infinite window size - sink_token_length=0, - query_padding_mask=None, - key_padding_mask=None, - key_leftpad=None, - device=None, -): - row_idx = rearrange( - torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1" - ) - col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long) - if key_leftpad is not None: - key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1") - col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0]) - col_idx = torch.where(col_idx >= key_leftpad, col_idx - key_leftpad, 2**32) - sk = ( - seqlen_k - if key_padding_mask is None - else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1") - ) - sq = ( - seqlen_q - if query_padding_mask is None - else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1") - ) - if window_size[0] < 0: - return col_idx > row_idx + sk - sq + window_size[1] - else: - sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk - return torch.logical_or( - col_idx > torch.minimum(row_idx + sk - sq + window_size[1], sk), - torch.logical_and( - col_idx < row_idx + sk - sq - window_size[0], - col_idx >= sink_token_length, - ), - ) - - -def attention_ref( - q, - k, - v, - query_padding_mask=None, - key_padding_mask=None, - key_leftpad=None, - attn_bias=None, - dropout_p=0.0, - dropout_mask=None, - causal=False, - qv=None, - q_descale=None, - k_descale=None, - v_descale=None, - window_size=(-1, -1), # -1 means infinite window size - sink_token_length=0, - sinks: Optional[torch.Tensor] = None, - softcap=0.0, - upcast=True, - reorder_ops=False, - intermediate_dtype=None, -): - """ - Arguments: - q: (batch_size, seqlen_q, nheads, head_dim) - k: (batch_size, seqlen_k, nheads, head_dim) - v: (batch_size, seqlen_k, nheads, head_dim_v) - qv: (batch_size, seqlen_q, nheads, head_dim_v) - query_padding_mask: (batch_size, seqlen_q) - key_padding_mask: (batch_size, seqlen_k) - attn_bias: broadcastable to (batch_size, nheads, seqlen_q, seqlen_k) - dropout_p: float - dropout_mask: (batch_size, nheads, seqlen_q, seqlen_k) - causal: whether to apply causal masking - upcast: whether to cast all inputs to fp32, do all computation in fp32, then cast - output back to fp16/bf16. - reorder_ops: whether to change the order of operations (scaling k instead of scaling k, etc.) - without changing the math. This is to estimate the numerical error from operation - reordering. - Output: - output: (batch_size, seqlen_q, nheads, head_dim_v) - attention: (batch_size, nheads, seqlen_q, seqlen_k), softmax after dropout - """ - if causal: - window_size = (window_size[0], 0) - dtype_og = q.dtype - if upcast: - q, k, v = q.float(), k.float(), v.float() - qv = qv.float() if qv is not None else None - if q_descale is not None: - q_descale = repeat(q_descale, "b h -> b 1 (h g) 1", g=q.shape[2] // k.shape[2]) - q = (q.float() * q_descale).to(q.dtype) - qv = (qv.float() * q_descale).to(qv.dtype) if qv is not None else None - if k_descale is not None: - k = (k.float() * rearrange(k_descale, "b h -> b 1 h 1")).to(dtype=k.dtype) - if v_descale is not None: - v = (v.float() * rearrange(v_descale, "b h -> b 1 h 1")).to(dtype=v.dtype) - seqlen_q, seqlen_k = q.shape[1], k.shape[1] - k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2]) - v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2]) - d = q.shape[-1] - dv = v.shape[-1] - softmax_scale = 1.0 / math.sqrt(d if qv is None else d + dv) - if not reorder_ops: - scores = torch.einsum("bthd,bshd->bhts", q * softmax_scale, k) - else: - scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale) - if qv is not None: - scores = scores + torch.einsum("bthd,bshd->bhts", qv * softmax_scale, v) - if softcap > 0: - scores = torch.tanh(scores / softcap) * softcap - if key_padding_mask is not None: - scores.masked_fill_( - rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf") - ) - if window_size[0] >= 0 or window_size[1] >= 0: - local_mask = construct_local_mask( - seqlen_q, - seqlen_k, - window_size, - sink_token_length, - query_padding_mask, - key_padding_mask, - key_leftpad=key_leftpad, - device=q.device, - ) - scores.masked_fill_(local_mask, float("-inf")) - if attn_bias is not None: - scores = scores + attn_bias - if sinks is None: - attention = torch.softmax(scores, dim=-1).to(v.dtype) - else: - scores_fp32 = scores.to(torch.float32) - logits_max = torch.amax(scores_fp32, dim=-1, keepdim=True) - sinks = rearrange(sinks, "h -> h 1 1") - logits_or_sinks_max = torch.maximum(sinks, logits_max) - unnormalized_scores = torch.exp(scores_fp32 - logits_or_sinks_max) - normalizer = unnormalized_scores.sum(dim=-1, keepdim=True) + torch.exp( - sinks - logits_or_sinks_max - ) - attention = (unnormalized_scores / normalizer).to(v.dtype) - # We want to mask here so that the attention matrix doesn't have any NaNs - # Otherwise we'll get NaN in dV - if query_padding_mask is not None: - attention = attention.masked_fill( - rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0 - ) - # Without this we might get NaN in dv - if key_padding_mask is not None: - attention = attention.masked_fill( - rearrange(~key_padding_mask, "b s -> b 1 1 s"), 0.0 - ) - # Some rows might be completely masked out so we fill them with zero instead of NaN - if window_size[0] >= 0 or window_size[1] >= 0: - attention = attention.masked_fill( - torch.all(local_mask, dim=-1, keepdim=True), 0.0 - ) - dropout_scaling = 1.0 / (1 - dropout_p) - # attention_drop = attention.masked_fill(~dropout_mask, 0.0) * dropout_scaling - # output = torch.einsum('bhts,bshd->bthd', attention_drop , v) - if dropout_mask is not None: - attention_drop = attention.masked_fill(~dropout_mask, 0.0) - else: - attention_drop = attention - if intermediate_dtype is not None: - attention_drop = attention_drop.to(intermediate_dtype).to(attention_drop.dtype) - output = torch.einsum("bhts,bshd->bthd", attention_drop, v * dropout_scaling) - if query_padding_mask is not None: - output.masked_fill_(rearrange(~query_padding_mask, "b s -> b s 1 1"), 0.0) - return output.to(dtype=dtype_og), attention.to(dtype=dtype_og) - - -def generate_qkv( - q, - k, - v, - query_padding_mask=None, - key_padding_mask=None, - kvpacked=False, - qkvpacked=False, - add_unused_qkv=False, - query_unused_mask=None, - key_unused_mask=None, -): - """ - Arguments: - q: (batch_size, seqlen_q, nheads, d) - k: (batch_size, seqlen_k, nheads_k, d) - v: (batch_size, seqlen_k, nheads_k, d) - query_padding_mask: (batch_size, seqlen), bool - key_padding_mask: (batch_size, seqlen), bool - """ - assert not (kvpacked and qkvpacked) - batch_size, seqlen_q, nheads, d = q.shape - _, seqlen_k, nheads_k, _ = k.shape - assert k.shape == (batch_size, seqlen_k, nheads_k, d) - assert v.shape == (batch_size, seqlen_k, nheads_k, d) - if query_unused_mask is not None or key_unused_mask is not None: - assert not kvpacked - assert not qkvpacked - - if query_padding_mask is not None: - q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, seqused_q = unpad_input( - q, - query_padding_mask, - query_unused_mask, - ) - output_pad_fn = lambda output_unpad: pad_input( - output_unpad, indices_q, batch_size, seqlen_q - ) - else: - q_unpad = rearrange(q, "b s h d -> (b s) h d") - cu_seqlens_q = torch.arange( - 0, - (batch_size + 1) * seqlen_q, - step=seqlen_q, - dtype=torch.int32, - device=q_unpad.device, - ) - seqused_q = None - max_seqlen_q = seqlen_q - output_pad_fn = lambda output_unpad: rearrange( - output_unpad, "(b s) h d -> b s h d", b=batch_size - ) - - if key_padding_mask is not None: - k_unpad, indices_k, cu_seqlens_k, max_seqlen_k, seqused_k = unpad_input( - k, key_padding_mask, key_unused_mask - ) - v_unpad, _, _, _, _ = unpad_input(v, key_padding_mask, key_unused_mask) - else: - k_unpad = rearrange(k, "b s h d -> (b s) h d") - v_unpad = rearrange(v, "b s h d -> (b s) h d") - cu_seqlens_k = torch.arange( - 0, - (batch_size + 1) * seqlen_k, - step=seqlen_k, - dtype=torch.int32, - device=k_unpad.device, - ) - seqused_k = None - max_seqlen_k = seqlen_k - - if qkvpacked: - assert (query_padding_mask == key_padding_mask).all() - assert nheads == nheads_k - qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1) - qkv = torch.stack([q, k, v], dim=2) - if query_padding_mask is not None: - dqkv_pad_fn = lambda dqkv_unpad: pad_input( - dqkv_unpad, indices_q, batch_size, seqlen_q - ) - else: - dqkv_pad_fn = lambda dqkv_unpad: rearrange( - dqkv_unpad, "(b s) t h d -> b s t h d", b=batch_size - ) - return ( - qkv_unpad.detach().requires_grad_(), - cu_seqlens_q, - max_seqlen_q, - qkv.detach().requires_grad_(), - output_pad_fn, - dqkv_pad_fn, - ) - elif kvpacked: - kv_unpad = torch.stack([k_unpad, v_unpad], dim=1) - kv = torch.stack([k, v], dim=2) - dq_pad_fn = output_pad_fn - if key_padding_mask is not None: - dkv_pad_fn = lambda dkv_unpad: pad_input( - dkv_unpad, indices_k, batch_size, seqlen_k - ) - else: - dkv_pad_fn = lambda dkv_unpad: rearrange( - dkv_unpad, "(b s) t h d -> b s t h d", b=batch_size - ) - return ( - q_unpad.detach().requires_grad_(), - kv_unpad.detach().requires_grad_(), - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - q.detach().requires_grad_(), - kv.detach().requires_grad_(), - output_pad_fn, - dq_pad_fn, - dkv_pad_fn, - ) - else: - dq_pad_fn = output_pad_fn - if key_padding_mask is not None: - dk_pad_fn = lambda dk_unpad: pad_input( - dk_unpad, indices_k, batch_size, seqlen_k - ) - else: - dk_pad_fn = lambda dk_unpad: rearrange( - dk_unpad, "(b s) h d -> b s h d", b=batch_size - ) - return ( - q_unpad.detach().requires_grad_(), - k_unpad.detach().requires_grad_(), - v_unpad.detach().requires_grad_(), - cu_seqlens_q, - cu_seqlens_k, - seqused_q, - seqused_k, - max_seqlen_q, - max_seqlen_k, - q.detach().requires_grad_(), - k.detach().requires_grad_(), - v.detach().requires_grad_(), - output_pad_fn, - dq_pad_fn, - dk_pad_fn, - ) - - -@pytest.mark.skipif( - not _is_fa3_supported(), - reason="flash_attn at sgl-kernel is only supported on CUDA sm90, sm80 or MUSA >= mp31", -) -# @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn]) -@pytest.mark.parametrize( - "dtype", [torch.bfloat16] + ([torch.float8_e4m3fn] if not DISABLE_FP8 else []) -) -# @pytest.mark.parametrize("dtype", [torch.bfloat16]) -# @pytest.mark.parametrize("dtype", [torch.float8_e4m3fn]) -@pytest.mark.parametrize("mha_type", ["mha", "mqa", "gqa"]) -# @pytest.mark.parametrize("mha_type", ["mha"]) -@pytest.mark.parametrize("has_sink", [False, True]) -# @pytest.mark.parametrize("has_sink", [False]) -@pytest.mark.parametrize("new_kv", [False] + ([True] if not DISABLE_APPENDKV else [])) -# @pytest.mark.parametrize("new_kv", [True]) -# @pytest.mark.parametrize( -# "causal,local", -# [(False, False), (True, False)] + ([(False, True)] if not DISABLE_LOCAL else []), -# ) -# @pytest.mark.parametrize("causal,local", [(False, False), (True, False)]) -@pytest.mark.parametrize("causal,local", [(False, False)]) -@pytest.mark.parametrize( - "seqlen_new_eq_seqlen_q", [True, False] if not DISABLE_APPENDKV else [True] -) -# @pytest.mark.parametrize("seqlen_new_eq_seqlen_q", [True]) -# @pytest.mark.parametrize("has_rotary_seqlens", [False, True]) -@pytest.mark.parametrize("has_rotary_seqlens", [False]) -@pytest.mark.parametrize( - "rotary_interleaved", [False, True] if not DISABLE_APPENDKV else [False] -) -# @pytest.mark.parametrize("rotary_interleaved", [True]) -@pytest.mark.parametrize( - "rotary_fraction", - ( - [0.0, 0.5, 1.0] - if (not DISABLE_APPENDKV) and (apply_rotary_emb is not None) - else [0.0] - ), -) -# @pytest.mark.parametrize("rotary_fraction", [0.0]) -@pytest.mark.parametrize( - "page_size", [None] + ([1, 4, 128] if not DISABLE_PAGEDKV else []) -) -# @pytest.mark.parametrize("page_size", [None]) -# @pytest.mark.parametrize("has_leftpad", [False, True]) -@pytest.mark.parametrize("has_leftpad", [False]) -# @pytest.mark.parametrize("has_batch_idx", [False, True]) -@pytest.mark.parametrize("has_batch_idx", [False]) -# @pytest.mark.parametrize("varlen_q", [False, True]) -@pytest.mark.parametrize("varlen_q", [False]) -# @pytest.mark.parametrize("d", [32, 59, 64, 80, 128, 256]) -# @pytest.mark.parametrize("d", [32, 64, 96, 128, 160, 192, 224, 256]) -# @pytest.mark.parametrize('d', [32, 40, 64, 80, 96, 128, 160, 192]) -# @pytest.mark.parametrize('d', [56, 80]) -@pytest.mark.parametrize("d", [64]) -# @pytest.mark.parametrize("d", [192]) -@pytest.mark.parametrize( - "seqlen_q,seqlen_k", - [ - (1, 128), - (1, 339), - (3, 1024), - (64, 800), - (64, 256), - (3, 799), - (64, 2048), - (16, 20000), - # (1, 128 * 1024), - # (16, 128 * 1024), - (128, 128), - (256, 512), # To test appending KV with more than 1 block - (2048, 3577), # Enough tile to test persistent scheduler - ], -) -# @pytest.mark.parametrize('seqlen_q,seqlen_k', [(256, 128)]) -def test_flash_attn_kvcache( - seqlen_q, - seqlen_k, - d, - varlen_q, - has_batch_idx, - has_leftpad, - page_size, - rotary_fraction, - rotary_interleaved, - has_rotary_seqlens, - seqlen_new_eq_seqlen_q, - causal, - local, - new_kv, - mha_type, - dtype, - has_sink, -): - from sgl_kernel.flash_attn import flash_attn_with_kvcache - - if page_size is not None and seqlen_k % page_size != 0: - pytest.skip() - if seqlen_q > seqlen_k and new_kv: - pytest.skip() - if not new_kv and rotary_fraction > 0.0: - pytest.skip() - if rotary_fraction == 0.0 and has_rotary_seqlens: - pytest.skip() - device = "cuda" - # set seed - torch.random.manual_seed(0) - batch_size = 5 - # batch_size = 1 - batch_size_cache = batch_size if not has_batch_idx else batch_size * 2 - nheads = 6 - # nheads = 1 - # rotary_dim must be a multiple of 16, and must be <= d - rotary_dim = math.floor(int(rotary_fraction * d) / 16) * 16 - nheads_k = nheads if mha_type == "mha" else (1 if mha_type == "mqa" else 3) - assert nheads % nheads_k == 0 - dtype_ref = torch.bfloat16 if dtype == torch.float8_e4m3fn else dtype - dv_vals = [128, d] if d > 128 and d <= 192 else ([256, 512, d] if d <= 64 else [d]) - - if has_sink: - sinks = torch.randn(nheads, dtype=torch.bfloat16, device=device) - else: - sinks = None - - if dtype == torch.float8_e4m3fn or not is_hopper(): - # for fp8 and ampere arch, we not support v head dim != qk head dim - dv_vals = [d] - for dv in dv_vals: - has_qv = d == 64 and dv >= 256 - q = ( - torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype_ref) - .to(dtype) - .to(dtype_ref) - ) - if has_qv: - qv = ( - torch.randn( - batch_size, seqlen_q, nheads, dv, device=device, dtype=dtype_ref - ) - .to(dtype) - .to(dtype_ref) - ) - else: - qv = None - if varlen_q: - query_padding_mask = generate_random_padding_mask( - seqlen_q, batch_size, device, mode="random" - ) - q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, *rest = unpad_input( - q, query_padding_mask - ) - output_pad_fn = lambda output_unpad: pad_input( - output_unpad, indices_q, batch_size, seqlen_q - ) - qv_unpad = ( - rearrange(qv, "b s ... -> (b s) ...")[indices_q] if has_qv else None - ) - else: - query_padding_mask = None - q_unpad = q - qv_unpad = qv - cu_seqlens_q, max_seqlen_q = None, None - # Put window_size after QKV randn so that window_size changes from test to test - window_size = (-1, -1) if not local else torch.randint(0, seqlen_k, (2,)) - - seqlen_new = ( - seqlen_q - if seqlen_new_eq_seqlen_q - else torch.randint(1, seqlen_q + 1, (1,)).item() - ) - cu_seqlens_k_new = None - key_new_padding_mask = None - if new_kv: - k = ( - torch.randn( - batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype_ref - ) - .to(dtype) - .to(dtype_ref) - ) - v = ( - torch.randn( - batch_size, seqlen_new, nheads_k, dv, device=device, dtype=dtype_ref - ) - .to(dtype) - .to(dtype_ref) - ) - if varlen_q: # k & v are also varlen - key_new_padding_mask = generate_random_padding_mask( - seqlen_new, batch_size, device, mode="random" - ) - k_unpad, indices_k, cu_seqlens_k_new, *rest = unpad_input( - k, key_new_padding_mask - ) - v_unpad, *rest = unpad_input(v, key_new_padding_mask) - else: - k_unpad, v_unpad = k, v - else: - k, v, k_unpad, v_unpad = None, None, None, None - if page_size is None: - k_cache = ( - torch.randn( - batch_size_cache, - seqlen_k, - nheads_k, - d, - device=device, - dtype=dtype_ref, - ) - .to(dtype) - .to(dtype_ref) - ) - v_cache = ( - torch.randn( - batch_size_cache, - seqlen_k, - nheads_k, - dv, - device=device, - dtype=dtype_ref, - ) - .to(dtype) - .to(dtype_ref) - ) - page_table = None - else: - ( - k_cache, - v_cache, - page_table, - k_cache_paged, - v_cache_paged, - num_blocks, - ) = _generate_block_kvcache( - seqlen_k, - page_size, - batch_size_cache, - nheads_k, - d, - dv, - device, - dtype, - dtype_ref, - ) - cache_seqlens = torch.randint( - 0 if new_kv else 1, - # If we don't use seqlen_q in the case of causal and rotary, cos/sin won't be long enough - ( - ( - seqlen_k - - (seqlen_q if (causal or local) and rotary_dim > 1 else seqlen_new) - + 1 - ) - if new_kv - else (seqlen_k + 1) - ), - (batch_size,), - dtype=torch.int32, - device=device, - ) - if has_leftpad: - cache_leftpad = torch.cat( - [ - ( - torch.randint( - 0, - cache_seqlens[i].item(), - (1,), - dtype=torch.int32, - device=device, - ) - if cache_seqlens[i].item() > 0 - else torch.zeros(1, dtype=torch.int32, device=device) - ) - for i in range(batch_size) - ] - ) - else: - cache_leftpad = None - if has_batch_idx: - cache_batch_idx = torch.randperm( - batch_size_cache, dtype=torch.int32, device=device - )[:batch_size] - else: - cache_batch_idx = None - arange = rearrange(torch.arange(seqlen_k, device=device), "s -> 1 s") - cache_seqlens_expanded = rearrange(cache_seqlens, "b -> b 1") - if not new_kv: - key_padding_mask = arange < cache_seqlens_expanded - else: - k_new_seqlens = ( - key_new_padding_mask.sum(-1, keepdims=True) if varlen_q else seqlen_new - ) - key_padding_mask = arange < cache_seqlens_expanded + k_new_seqlens - if has_leftpad: - key_padding_mask = torch.logical_and( - key_padding_mask, - arange >= cache_leftpad.unsqueeze(-1).expand(-1, seqlen_k), - ) - # cache_seqlens = torch.tensor([64], dtype=torch.int32, device=device) - rotary_seqlens = cache_seqlens if not has_rotary_seqlens else cache_seqlens // 2 - if rotary_dim > 0: - angle = ( - torch.rand( - seqlen_k if page_size is None else num_blocks * page_size, - rotary_dim // 2, - device=device, - ) - * 2 - * math.pi - ) - cos = torch.cos(angle).to(dtype=dtype_ref).to(dtype).to(dtype_ref) - sin = torch.sin(angle).to(dtype=dtype_ref).to(dtype).to(dtype_ref) - if causal or local: - q_ro = apply_rotary_emb( - q, - cos, - sin, - seqlen_offsets=rotary_seqlens, - interleaved=rotary_interleaved, - ) - else: - q_ro = rearrange( - apply_rotary_emb( - rearrange(q, "b s h d -> b 1 (s h) d"), - cos, - sin, - seqlen_offsets=rotary_seqlens, - interleaved=rotary_interleaved, - ), - "b 1 (s h) d -> b s h d", - s=seqlen_q, - ) - # q_ro = q - k_ro = apply_rotary_emb( - k, - cos, - sin, - seqlen_offsets=rotary_seqlens, - interleaved=rotary_interleaved, - ) - else: - cos, sin = None, None - q_ro, k_ro = q, k - # k_cache[:, 64:] = -1 - k_cache_ref = ( - k_cache if not has_batch_idx else k_cache[cache_batch_idx] - ).clone() - v_cache_ref = ( - v_cache if not has_batch_idx else v_cache[cache_batch_idx] - ).clone() - if new_kv: - update_mask = torch.logical_and( - cache_seqlens_expanded <= arange, - arange < cache_seqlens_expanded + k_new_seqlens, - ) - k_to_update = rearrange(k_ro, "b s ... -> (b s) ...") - v_to_update = rearrange(v, "b s ... -> (b s) ...") - if varlen_q: - k_to_update = k_to_update[indices_k] - v_to_update = v_to_update[indices_k] - k_cache_ref[update_mask] = k_to_update - v_cache_ref[update_mask] = v_to_update - k_cache_rep = repeat( - k_cache_ref, "b s h d -> b s (h g) d", g=nheads // nheads_k - ) - v_cache_rep = repeat( - v_cache_ref, "b s h d -> b s (h g) d", g=nheads // nheads_k - ) - out_ref, _ = attention_ref( - q_ro, - k_cache_rep, - v_cache_rep, - query_padding_mask, - key_padding_mask, - causal=causal, - qv=qv, - window_size=window_size, - key_leftpad=cache_leftpad, - sinks=sinks, - ) - out_pt, _ = attention_ref( - q_ro, - k_cache_rep, - v_cache_rep, - query_padding_mask, - key_padding_mask, - causal=causal, - qv=qv, - window_size=window_size, - upcast=False, - reorder_ops=True, - key_leftpad=cache_leftpad, - intermediate_dtype=dtype if dtype == torch.float8_e4m3fn else None, - sinks=sinks, - ) - q = q.to(dtype) - q_unpad = q_unpad.to(dtype) if varlen_q else None - k_cache = k_cache.to(dtype) - v_cache = v_cache.to(dtype) - k_cache_paged = k_cache_paged.to(dtype) if page_size is not None else None - v_cache_paged = v_cache_paged.to(dtype) if page_size is not None else None - k = k.to(dtype) if k is not None else None - v = v.to(dtype) if v is not None else None - k_unpad = k_unpad.to(dtype) if k_unpad is not None else None - v_unpad = v_unpad.to(dtype) if v_unpad is not None else None - qv = qv.to(dtype) if qv is not None else None - qv_unpad = qv_unpad.to(dtype) if (varlen_q and qv is not None) else None - cos = cos.to(dtype) if cos is not None else None - sin = sin.to(dtype) if sin is not None else None - k_cache_saved = k_cache.clone() if page_size is None else k_cache_paged.clone() - v_cache_saved = v_cache.clone() if page_size is None else v_cache_paged.clone() - num_splits_vals = [1, 0] if not DISABLE_SPLIT else [1] - precompute_metadata_vals = [False] - for num_splits, precompute_metadata in itertools.product( - num_splits_vals, precompute_metadata_vals - ): - scheduler_metadata = None - # Repeat to test metadata reuse - for _ in range(1 if not precompute_metadata else 2): - if page_size is None: - k_cache.copy_(k_cache_saved) - v_cache.copy_(v_cache_saved) - else: - k_cache_paged.copy_(k_cache_saved) - v_cache_paged.copy_(v_cache_saved) - out, lse, *rest = flash_attn_with_kvcache( - q if not varlen_q else q_unpad, - k_cache if page_size is None else k_cache_paged, - v_cache if page_size is None else v_cache_paged, - k if not new_kv or not varlen_q else k_unpad, - v if not new_kv or not varlen_q else v_unpad, - qv=qv if not varlen_q else qv_unpad, - rotary_cos=cos, - rotary_sin=sin, - cache_seqlens=cache_seqlens, - cache_batch_idx=cache_batch_idx, - cache_leftpad=cache_leftpad, - page_table=page_table, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k_new=cu_seqlens_k_new, - max_seqlen_q=max_seqlen_q, - rotary_seqlens=rotary_seqlens, - causal=causal, - window_size=window_size, - rotary_interleaved=rotary_interleaved, - scheduler_metadata=scheduler_metadata, - num_splits=num_splits, - return_softmax_lse=True, - sinks=sinks, - ) - if varlen_q: - out = output_pad_fn(out) - # out = flash_attn_with_kvcache( - # q, k_cache, v_cache, cache_seqlens=cache_seqlens, causal=causal, window_size=window_size - # ) - # out = flash_attn_with_kvcache(q, k_cache, v_cache, causal=causal, window_size=window_size) - # qk = torch.einsum("bqhd,bkhd->bhqk", q, k_cache_ref) - # m = qk.amax(-1, keepdim=True) - # s_tmp = torch.exp((qk - m) / math.sqrt(d)) - # o1 = torch.einsum('bhst,bthd->bshd', s_tmp, v_cache_ref) - # lse_ref = torch.logsumexp(qk / math.sqrt(d), -1) - # probs = torch.softmax(qk, dim=-1) - print(f"Output max diff: {(out - out_ref).abs().max().item()}") - print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") - print(f"Pytorch max diff: {(out_pt - out_ref).abs().max().item()}") - print(f"Pytorch mean diff: {(out_pt - out_ref).abs().mean().item()}") - # breakpoint() - - # Check that FlashAttention's numerical error is at most twice the numerical error - # of a Pytorch implementation. - if new_kv: - if page_size is None: - k_cache_select = ( - k_cache.to(dtype_ref) - if not has_batch_idx - else k_cache.to(dtype_ref)[cache_batch_idx] - ) - v_cache_select = ( - v_cache.to(dtype_ref) - if not has_batch_idx - else v_cache.to(dtype_ref)[cache_batch_idx] - ) - else: - k_cache_select = rearrange( - k_cache_paged.to(dtype_ref)[ - ( - page_table - if not has_batch_idx - else page_table[cache_batch_idx] - ).flatten() - ], - "(b nblocks) block_size ... -> b (nblocks block_size) ...", - b=batch_size, - )[:, :seqlen_k].to(dtype_ref) - v_cache_select = rearrange( - v_cache_paged.to(dtype_ref)[ - ( - page_table - if not has_batch_idx - else page_table[cache_batch_idx] - ).flatten() - ], - "(b nblocks) block_size ... -> b (nblocks block_size) ...", - b=batch_size, - )[:, :seqlen_k].to(dtype_ref) - k_cache_ref = k_cache_ref.to(dtype).to(dtype_ref) - v_cache_ref = v_cache_ref.to(dtype).to(dtype_ref) - if dtype is not torch.float8_e4m3fn: - assert torch.equal(v_cache_select, v_cache_ref) - else: - assert torch.allclose( - v_cache_select, v_cache_ref, rtol=1e-3, atol=1e-3 - ) - # breakpoint() - # if rotary_dim == 0 and dtype is not torch.float8_e4m3fn: - if rotary_dim == 0: - assert torch.equal(k_cache_select, k_cache_ref) - else: - # if not torch.allclose(k_cache_select, k_cache_ref, rtol=1e-3, atol=1e-3): - # breakpoint() - if dtype is not torch.float8_e4m3fn: - assert torch.allclose( - k_cache_select, k_cache_ref, rtol=1e-3, atol=1e-3 - ) - else: - assert torch.allclose( - k_cache_select, k_cache_ref, rtol=1e-1, atol=1e-1 - ) - mult = 4 if dtype == torch.float8_e4m3fn else 2 - assert (out - out_ref).abs().max().item() <= mult * ( - out_pt - out_ref - ).abs().max().item() + 1e-5 - mult_mean = 3 if dtype == torch.float8_e4m3fn else 1.5 - assert (out - out_ref).abs().mean().item() <= mult_mean * ( - out_pt - out_ref - ).abs().mean().item() - - -def _generate_block_kvcache( - seqlen_k, page_size, batch_size, nheads_k, d, dv, device, dtype, dtype_ref -): - num_blocks = math.ceil(seqlen_k / page_size) * batch_size * 3 - k_cache_paged = ( - torch.randn(num_blocks, page_size, nheads_k, d, device=device, dtype=dtype_ref) - .to(dtype) - .to(dtype_ref) - ) - v_cache_paged = ( - torch.randn(num_blocks, page_size, nheads_k, dv, device=device, dtype=dtype_ref) - .to(dtype) - .to(dtype_ref) - ) - page_table = rearrange( - torch.randperm(num_blocks, dtype=torch.int32, device=device), - "(b nblocks) -> b nblocks", - b=batch_size, - ) - k_cache = rearrange( - k_cache_paged[page_table.flatten()], - "(b nblocks) block_size ... -> b (nblocks block_size) ...", - b=batch_size, - )[:, :seqlen_k] - v_cache = rearrange( - v_cache_paged[page_table.flatten()], - "(b nblocks) block_size ... -> b (nblocks block_size) ...", - b=batch_size, - )[:, :seqlen_k] - return k_cache, v_cache, page_table, k_cache_paged, v_cache_paged, num_blocks - - -@pytest.mark.skipif( - not _is_fa3_supported(), - reason="flash_attn at sgl-kernel is only supported on CUDA sm90, sm80 or MUSA >= mp31", -) -# @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn]) -@pytest.mark.parametrize( - "dtype", [torch.bfloat16] + ([torch.float8_e4m3fn] if not DISABLE_FP8 else []) -) -# @pytest.mark.parametrize("dtype", [torch.bfloat16]) -# @pytest.mark.parametrize("dtype", [torch.float8_e4m3fn]) -@pytest.mark.parametrize("mha_type", ["mha", "mqa", "gqa"]) -# @pytest.mark.parametrize("mha_type", ["mha"]) -@pytest.mark.parametrize("has_sink", [False, True]) -# @pytest.mark.parametrize("has_sink", [False]) -# @pytest.mark.parametrize("has_qv", [False, True]) -@pytest.mark.parametrize("has_qv", [False]) -# @pytest.mark.parametrize("deterministic", [False, True]) -@pytest.mark.parametrize("deterministic", [False]) -@pytest.mark.parametrize("softcap", [0.0] + ([15.0] if not DISABLE_SOFTCAP else [])) -# @pytest.mark.parametrize("softcap", [0.0]) -@pytest.mark.parametrize("local", [False]) -# @pytest.mark.parametrize("local", [False]) -@pytest.mark.parametrize("causal", [False, True]) -# @pytest.mark.parametrize("causal", [False]) -@pytest.mark.parametrize("add_unused_qkv", [False, True]) -# @pytest.mark.parametrize("add_unused_qkv", [True]) -# @pytest.mark.parametrize("d", [32, 64, 96, 128, 160, 192, 224, 256]) -# @pytest.mark.parametrize('d', [32, 40, 64, 80, 96, 128, 160, 192, 256]) -# @pytest.mark.parametrize('d', [32, 64, 96, 128, 160, 192]) -# @pytest.mark.parametrize('d', [56, 80]) -# @pytest.mark.parametrize('d', [32, 40, 64, 80, 96, 128]) -# @pytest.mark.parametrize("d", [64, 96, 128]) -# @pytest.mark.parametrize("d", COMPILED_HDIMS) -@pytest.mark.parametrize("d", [128]) -@pytest.mark.parametrize( - "seqlen_q,seqlen_k", - [ - (1, 1), - (1, 3), - (2, 1), - (511, 1), - (3, 513), - (64, 128), - (128, 128), - (256, 256), - (113, 203), - (128, 217), - (113, 211), - (108, 256), - (256, 512), - (307, 256), - (640, 128), - (512, 256), - (1024, 1024), - (1023, 1024), - (1024, 1023), - (2048, 2048), - ], -) -def test_flash_attn_varlen_output( - seqlen_q, - seqlen_k, - d, - add_unused_qkv, - causal, - local, - softcap, - deterministic, - has_qv, - mha_type, - dtype, - has_sink, -): - from sglang.jit_kernel.flash_attention import flash_attn_varlen_func - - device = "cuda" - # set seed - torch.random.manual_seed(seqlen_q + seqlen_k + d + int(causal) * 2 + int(local)) - # batch_size = 40 - # nheads = 16 - batch_size = 9 if seqlen_q <= 2048 else 2 - nheads = 6 - # batch_size = 2 - # nheads = 1 - nheads_kv = nheads if mha_type == "mha" else (2 if mha_type == "gqa" else 1) - dtype_ref = torch.bfloat16 if dtype == torch.float8_e4m3fn else dtype - dv_vals = [128, d] if d > 128 and d <= 192 else ([256, 512, d] if d <= 64 else [d]) - if dtype == torch.float8_e4m3fn: - dv_vals = [d] - for dv in dv_vals: - q_ref = torch.randn( - batch_size, seqlen_q, nheads, d, device=device, dtype=dtype_ref - ) - if softcap > 0.0: - # Ensure the values of qk are at least within softcap range. - q_ref = (q_ref * softcap / 4).detach().requires_grad_() - q_ref = q_ref.to(dtype).to(dtype_ref).requires_grad_() - k_ref = ( - torch.randn( - batch_size, seqlen_k, nheads_kv, d, device=device, dtype=dtype_ref - ) - .to(dtype) - .to(dtype_ref) - .requires_grad_() - ) - v_ref = ( - torch.randn( - batch_size, seqlen_k, nheads_kv, dv, device=device, dtype=dtype_ref - ) - .to(dtype) - .to(dtype_ref) - .requires_grad_() - ) - if has_qv: - qv_ref = ( - torch.randn( - batch_size, seqlen_q, nheads, dv, device=device, dtype=dtype_ref - ) - .to(dtype) - .to(dtype_ref) - ) - else: - qv_ref = None - # Put window_size after QKV randn so that window_size changes from test to test - window_size = (-1, -1) if not local else torch.randint(0, seqlen_k, (2,)) - - if has_sink: - sinks = torch.randn(nheads, dtype=torch.bfloat16, device=device) - else: - sinks = None - - if dtype == torch.float8_e4m3fn: - q_descale, k_descale, v_descale = [ - torch.rand(batch_size, nheads_kv, device=device, dtype=torch.float32) - * 2 - for _ in range(3) - ] - else: - q_descale, k_descale, v_descale = None, None, None - q, k, v = [x.detach().requires_grad_() for x in (q_ref, k_ref, v_ref)] - qv = qv_ref.detach() if has_qv else None - query_padding_mask = generate_random_padding_mask( - seqlen_q, batch_size, device, mode="random", zero_lengths=False - ) - key_padding_mask = generate_random_padding_mask( - seqlen_k, batch_size, device, mode="random", zero_lengths=True - ) - - def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): - if add_unused: - another_mask = generate_random_padding_mask(max_seq_len, bs, device) - attn_mask = torch.logical_and(padding_mask, another_mask) - unused_mask = torch.logical_xor( - torch.logical_or(padding_mask, another_mask), attn_mask - ) - else: - attn_mask = padding_mask - unused_mask = None - return attn_mask, unused_mask - - query_padding_mask, query_unused_mask = _gen_unused_masks( - query_padding_mask, add_unused_qkv, seqlen_q, batch_size, q.device - ) - key_padding_mask, key_unused_mask = _gen_unused_masks( - key_padding_mask, add_unused_qkv, seqlen_k, batch_size, k.device - ) - - ( - q_unpad, - k_unpad, - v_unpad, - cu_seqlens_q, - cu_seqlens_k, - seqused_q, - seqused_k, - max_seqlen_q, - max_seqlen_k, - q, - k, - v, - output_pad_fn, - dq_pad_fn, - dk_pad_fn, - ) = generate_qkv( - q, - k, - v, - query_padding_mask, - key_padding_mask, - kvpacked=False, - query_unused_mask=query_unused_mask, - key_unused_mask=key_unused_mask, - ) - q_unpad, k_unpad, v_unpad = [ - x.detach().to(dtype).requires_grad_() for x in (q_unpad, k_unpad, v_unpad) - ] - out_ref, attn_ref = attention_ref( - q_ref, - k_ref, - v_ref, - query_padding_mask, - key_padding_mask, - causal=causal, - qv=qv_ref, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - window_size=window_size, - softcap=softcap, - sinks=sinks, - ) - out_pt, attn_pt = attention_ref( - q_ref, - k_ref, - v_ref, - query_padding_mask, - key_padding_mask, - causal=causal, - qv=qv_ref, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - window_size=window_size, - softcap=softcap, - upcast=False, - reorder_ops=True, - intermediate_dtype=dtype if dtype == torch.float8_e4m3fn else None, - sinks=sinks, - ) - - print(f"Pytorch max diff: {(out_pt - out_ref).abs().max().item()}") - print(f"Pytorch mean diff: {(out_pt - out_ref).abs().mean().item()}") - - if query_unused_mask is not None: - q_zero_masking = rearrange(query_unused_mask, "b s -> b s 1 1") - - # Numerical error if we just do any arithmetic on out_ref - fwd_atol = 2 * (out_ref + 0.3 - 0.3 - out_ref).abs().max().item() - rtol = 2 if softcap == 0.0 else 3 - - pack_gqa_vals = [False, True] if not DISABLE_PACKGQA else [False] - num_splits_vals = [1, 3] if not DISABLE_SPLIT else [1] - for pack_gqa, num_splits in itertools.product(pack_gqa_vals, num_splits_vals): - out_unpad, lse, *rest = flash_attn_varlen_func( - q_unpad, - k_unpad, - v_unpad, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - seqused_q=seqused_q, - seqused_k=seqused_k, - causal=causal, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - window_size=window_size, - softcap=softcap, - return_softmax_lse=True, - sinks=sinks, - ) - out = output_pad_fn(out_unpad) - if query_unused_mask is not None: - out.masked_fill_(q_zero_masking, 0.0) - print(f"Output max diff: {(out - out_ref).abs().max().item()}") - print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") - - # Check that FlashAttention's numerical error is at most 3x the numerical error - # of a Pytorch implementation. - assert (out - out_ref).abs().max().item() <= rtol * ( - out_pt - out_ref - ).abs().max().item() + fwd_atol - - if not DISABLE_BACKWARD and dtype != torch.float8_e4m3fn and not has_qv: - g_unpad = torch.randn_like(out_unpad) - do_o = ((g_unpad.float() * out_unpad.float()).sum(-1)).transpose(-1, -2) - dq_unpad, dk_unpad, dv_unpad = torch.autograd.grad( - out_unpad, (q_unpad, k_unpad, v_unpad), g_unpad - ) - dq = dq_pad_fn(dq_unpad) - dk = dk_pad_fn(dk_unpad) - dv = dk_pad_fn(dv_unpad) - if key_unused_mask is not None: - k_zero_masking = rearrange(key_unused_mask, "b s -> b s 1 1") - dk.masked_fill_(k_zero_masking, 0.0) - dv.masked_fill_(k_zero_masking, 0.0) - if query_unused_mask is not None: - dq.masked_fill_(q_zero_masking, 0.0) - # print(f"dO_O max diff: {(softmax_d - do_o).abs().max().item()}") - # assert (softmax_d - do_o).abs().max().item() <= 1e-5 - # assert dq_accum.abs().max().item() == 0.0 - g = output_pad_fn(g_unpad) - - # dq, dk, dv = torch.autograd.grad(out, (q, k, v), g) - dq_ref, dk_ref, dv_ref = torch.autograd.grad(out_ref, (q_ref, k_ref, v_ref), g) - dq_pt, dk_pt, dv_pt = torch.autograd.grad(out_pt, (q_ref, k_ref, v_ref), g) - print(f"dQ max diff: {(dq - dq_ref).abs().max().item()}") - print(f"dK max diff: {(dk - dk_ref).abs().max().item()}") - print(f"dV max diff: {(dv - dv_ref).abs().max().item()}") - print(f"dQ mean diff: {(dq - dq_ref).abs().mean().item()}") - print(f"dK mean diff: {(dk - dk_ref).abs().mean().item()}") - print(f"dV mean diff: {(dv - dv_ref).abs().mean().item()}") - print(f"dQ Pytorch max diff: {(dq_pt - dq_ref).abs().max().item()}") - print(f"dK Pytorch max diff: {(dk_pt - dk_ref).abs().max().item()}") - print(f"dV Pytorch max diff: {(dv_pt - dv_ref).abs().max().item()}") - print(f"dQ Pytorch mean diff: {(dq_pt - dq_ref).abs().mean().item()}") - print(f"dK Pytorch mean diff: {(dk_pt - dk_ref).abs().mean().item()}") - print(f"dV Pytorch mean diff: {(dv_pt - dv_ref).abs().mean().item()}") - - if not DISABLE_BACKWARD and dtype != torch.float8_e4m3fn and not has_qv: - dq_atol = 2 * (dq_ref + 0.3 - 0.3 - dq_ref).abs().max().item() + ( - 0 if softcap == 0 else 3e-4 - ) - assert (dq - dq_ref).abs().max().item() <= rtol * ( - dq_pt - dq_ref - ).abs().max().item() + dq_atol - dk_atol = 2 * (dk_ref + 0.3 - 0.3 - dk_ref).abs().max().item() + ( - 0 if softcap == 0 else 3e-4 - ) - assert (dk - dk_ref).abs().max().item() <= rtol * ( - dk_pt - dk_ref - ).abs().max().item() + dk_atol - dv_atol = 2 * (dv_ref + 0.3 - 0.3 - dv_ref).abs().max().item() + ( - 0 if softcap == 0 else 3e-4 - ) - assert (dv - dv_ref).abs().max().item() <= rtol * ( - dv_pt - dv_ref - ).abs().max().item() + dv_atol - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/python/sglang/jit_kernel/tests/test_fused_qknorm_rope.py b/python/sglang/jit_kernel/tests/test_fused_qknorm_rope.py deleted file mode 100644 index 0843db13f..000000000 --- a/python/sglang/jit_kernel/tests/test_fused_qknorm_rope.py +++ /dev/null @@ -1,448 +0,0 @@ -""" -Correctness tests for the fused_qknorm_rope JIT kernel. - -Validates fused_qk_norm_rope against a pure-PyTorch reference and (when -available) the sgl_kernel AOT implementation. -""" - -import pytest -import torch - -from sglang.jit_kernel.fused_qknorm_rope import fused_qk_norm_rope -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=35, suite="stage-b-kernel-unit-1-gpu-large") -register_cuda_ci(est_time=256, suite="nightly-kernel-1-gpu", nightly=True) - -try: - from sgl_kernel import fused_qk_norm_rope as fused_qk_norm_rope_aot - - AOT_AVAILABLE = True -except ImportError: - AOT_AVAILABLE = False - -HEAD_DIMS = [64, 128, 256] -NUM_TOKENS = [1, 16, 128] - - -# --------------------------------------------------------------------------- -# Pure-PyTorch reference -# --------------------------------------------------------------------------- - - -def _compute_inv_freq_yarn(base, rotary_dim, factor, low, high, device): - """Compute YaRN-adjusted inverse frequencies for rotary_dim//2 positions.""" - half_dims = torch.arange(rotary_dim // 2, dtype=torch.float32, device=device) - inv_freq = base ** (-2.0 * half_dims / rotary_dim) - - if factor != 1.0: - inv_freq_interp = inv_freq / factor - inv_freq_extrap = inv_freq - high_adj = high if abs(high - low) > 1e-6 else high + 0.001 - linear = (half_dims - low) / (high_adj - low) - ramp = linear.clamp(0.0, 1.0) - extrap_factor = 1.0 - ramp - inv_freq = ( - inv_freq_interp * (1 - extrap_factor) + inv_freq_extrap * extrap_factor - ) - - return inv_freq - - -def fused_qk_norm_rope_ref( - qkv, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - eps, - q_weight, - k_weight, - base, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - rotary_dim, -): - """ - Pure-PyTorch reference: RMSNorm per head, then RoPE on Q and K. - - Returns a new tensor (same shape as qkv) with the transformation applied. - """ - num_tokens = qkv.shape[0] - total_heads = num_heads_q + num_heads_k + num_heads_v - - qkv_f = qkv.float() - qw = q_weight.float() - kw = k_weight.float() - - # Reshape to [num_tokens, total_heads, head_dim] - qkv_3d = qkv_f.view(num_tokens, total_heads, head_dim) - q = qkv_3d[:, :num_heads_q].clone() # [num_tokens, nq, head_dim] - k = qkv_3d[:, num_heads_q : num_heads_q + num_heads_k].clone() - - # RMSNorm per head - def rms_norm_heads(x, w): - # x: [num_tokens, n_heads, head_dim], w: [head_dim] - rms = (x**2).mean(-1, keepdim=True) - return x * torch.rsqrt(rms + eps) * w - - q = rms_norm_heads(q, qw) - k = rms_norm_heads(k, kw) - - # Compute frequencies - inv_freq = _compute_inv_freq_yarn(base, rotary_dim, factor, low, high, qkv.device) - # theta: [num_tokens, rotary_dim//2] - theta = position_ids.float().unsqueeze(1) * inv_freq.unsqueeze(0) - cos = torch.cos(theta) # [num_tokens, rotary_dim//2] - sin = torch.sin(theta) - # Broadcast across heads: [num_tokens, 1, rotary_dim//2] - c = cos.unsqueeze(1) - s = sin.unsqueeze(1) - - if not is_neox: - # Interleave (GPT-J) style: rotate pairs (x[2i], x[2i+1]) - def apply_interleave(x): - # x: [num_tokens, n_heads, head_dim] - x_rot = x[:, :, :rotary_dim] # [num_tokens, n_heads, rotary_dim] - x_pairs = x_rot.view(num_tokens, -1, rotary_dim // 2, 2) - x0, x1 = x_pairs[..., 0], x_pairs[..., 1] - x0_new = x0 * c - x1 * s - x1_new = x1 * c + x0 * s - x_rot_new = torch.stack([x0_new, x1_new], dim=-1).view( - num_tokens, -1, rotary_dim - ) - result = x.clone() - result[:, :, :rotary_dim] = x_rot_new * attention_factor - return result - - q = apply_interleave(q) - k = apply_interleave(k) - else: - # NeoX style: first half × cos − second half × sin (and vice versa) - def apply_neox(x): - # x: [num_tokens, n_heads, head_dim] - x1 = x[:, :, : rotary_dim // 2] - x2 = x[:, :, rotary_dim // 2 : rotary_dim] - x1_new = x1 * c - x2 * s - x2_new = x2 * c + x1 * s - result = x.clone() - result[:, :, : rotary_dim // 2] = x1_new * attention_factor - result[:, :, rotary_dim // 2 : rotary_dim] = x2_new * attention_factor - return result - - q = apply_neox(q) - k = apply_neox(k) - - # Write back into a copy of the full QKV - result_3d = qkv_f.view(num_tokens, total_heads, head_dim).clone() - result_3d[:, :num_heads_q] = q - result_3d[:, num_heads_q : num_heads_q + num_heads_k] = k - return result_3d.view(num_tokens, -1).bfloat16() - - -# --------------------------------------------------------------------------- -# Tests: correctness vs PyTorch reference -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("head_dim", HEAD_DIMS) -@pytest.mark.parametrize("num_tokens", NUM_TOKENS) -@pytest.mark.parametrize("is_neox", [False, True]) -def test_fused_qknorm_rope_vs_ref(head_dim, num_tokens, is_neox): - torch.manual_seed(head_dim * num_tokens + int(is_neox)) - device = "cuda" - num_heads_q, num_heads_k, num_heads_v = 4, 2, 2 - total_heads = num_heads_q + num_heads_k + num_heads_v - rotary_dim = head_dim # full rotary - - qkv = torch.randn( - (num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device - ) - q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device) - k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device) - position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device) - - eps = 1e-5 - base = 10000.0 - factor = 1.0 # no YaRN - low, high = 1.0, 32.0 - attention_factor = 1.0 - - ref = fused_qk_norm_rope_ref( - qkv, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - eps, - q_weight, - k_weight, - base, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - rotary_dim, - ) - - qkv_jit = qkv.clone() - fused_qk_norm_rope( - qkv_jit, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - eps, - q_weight, - k_weight, - base, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - rotary_dim, - ) - - assert torch.allclose(qkv_jit.float(), ref.float(), atol=5e-3, rtol=1e-2), ( - f"mismatch: head_dim={head_dim}, num_tokens={num_tokens}, " - f"is_neox={is_neox}, " - f"max_err={( qkv_jit.float() - ref.float()).abs().max().item():.4e}" - ) - - -@pytest.mark.parametrize("head_dim", HEAD_DIMS) -@pytest.mark.parametrize("is_neox", [False, True]) -def test_fused_qknorm_rope_partial_rotary(head_dim, is_neox): - """Test with rotary_dim < head_dim: non-rotary elements should be RMSNorm-only.""" - torch.manual_seed(42 + head_dim + int(is_neox)) - device = "cuda" - num_tokens = 16 - num_heads_q, num_heads_k, num_heads_v = 2, 2, 2 - total_heads = num_heads_q + num_heads_k + num_heads_v - rotary_dim = head_dim // 2 # half of head_dim - - # NeoX requires half_rotary_lanes to be power of 2. - # half_rotary_lanes = rotary_dim / (head_dim / 32) / 2 = (head_dim//2) / (head_dim/32) / 2 - # = 16 / 2 = 8 → power of 2, OK for all supported head_dims. - - qkv = torch.randn( - (num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device - ) - q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device) - k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device) - position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device) - - ref = fused_qk_norm_rope_ref( - qkv, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - 1e-5, - q_weight, - k_weight, - 10000.0, - is_neox, - position_ids, - 1.0, - 1.0, - 32.0, - 1.0, - rotary_dim, - ) - - qkv_jit = qkv.clone() - fused_qk_norm_rope( - qkv_jit, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - 1e-5, - q_weight, - k_weight, - 10000.0, - is_neox, - position_ids, - 1.0, - 1.0, - 32.0, - 1.0, - rotary_dim, - ) - - assert torch.allclose(qkv_jit.float(), ref.float(), atol=5e-3, rtol=1e-2), ( - f"partial rotary mismatch: head_dim={head_dim}, is_neox={is_neox}, " - f"max_err={(qkv_jit.float() - ref.float()).abs().max().item():.4e}" - ) - - -@pytest.mark.parametrize("head_dim", HEAD_DIMS) -def test_fused_qknorm_rope_yarn_scaling(head_dim): - """Test with YaRN scaling (factor != 1.0).""" - torch.manual_seed(99 + head_dim) - device = "cuda" - num_tokens = 32 - num_heads_q, num_heads_k, num_heads_v = 2, 2, 2 - total_heads = num_heads_q + num_heads_k + num_heads_v - rotary_dim = head_dim - - qkv = torch.randn( - (num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device - ) - q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device) - k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device) - position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device) - - factor = 2.5 - low, high = 4.0, 32.0 - attention_factor = 0.9 - is_neox = False # test with interleave; NeoX also tested in other tests - - ref = fused_qk_norm_rope_ref( - qkv, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - 1e-5, - q_weight, - k_weight, - 500000.0, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - rotary_dim, - ) - - qkv_jit = qkv.clone() - fused_qk_norm_rope( - qkv_jit, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - 1e-5, - q_weight, - k_weight, - 500000.0, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - rotary_dim, - ) - - assert torch.allclose(qkv_jit.float(), ref.float(), atol=5e-3, rtol=1e-2), ( - f"YaRN mismatch: head_dim={head_dim}, " - f"max_err={(qkv_jit.float() - ref.float()).abs().max().item():.4e}" - ) - - -def test_fused_qknorm_rope_default_rotary_dim(): - """rotary_dim=None should default to head_dim.""" - device = "cuda" - num_tokens = 8 - num_heads_q, num_heads_k, num_heads_v = 2, 2, 2 - head_dim = 128 - total_heads = num_heads_q + num_heads_k + num_heads_v - - torch.manual_seed(0) - qkv1 = torch.randn( - (num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device - ) - qkv2 = qkv1.clone() - q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device) - k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device) - position_ids = torch.zeros(num_tokens, dtype=torch.int32, device=device) - - common_kwargs = dict( - num_heads_q=num_heads_q, - num_heads_k=num_heads_k, - num_heads_v=num_heads_v, - head_dim=head_dim, - eps=1e-5, - q_weight=q_weight, - k_weight=k_weight, - base=10000.0, - is_neox=False, - position_ids=position_ids, - factor=1.0, - low=1.0, - high=32.0, - attention_factor=1.0, - ) - - fused_qk_norm_rope(qkv1, **common_kwargs, rotary_dim=None) - fused_qk_norm_rope(qkv2, **common_kwargs, rotary_dim=head_dim) - - assert torch.equal(qkv1, qkv2), "rotary_dim=None must equal rotary_dim=head_dim" - - -# --------------------------------------------------------------------------- -# Cross-validation against AOT sgl_kernel -# --------------------------------------------------------------------------- - - -@pytest.mark.skipif(not AOT_AVAILABLE, reason="sgl_kernel not available") -@pytest.mark.parametrize("head_dim", [64, 128, 256]) -@pytest.mark.parametrize("is_neox", [False, True]) -def test_fused_qknorm_rope_vs_aot(head_dim, is_neox): - torch.manual_seed(head_dim * 7 + int(is_neox)) - device = "cuda" - num_tokens = 32 - num_heads_q, num_heads_k, num_heads_v = 4, 2, 2 - total_heads = num_heads_q + num_heads_k + num_heads_v - - qkv = torch.randn( - (num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device - ) - q_weight = torch.randn(head_dim, dtype=torch.bfloat16, device=device).abs() + 0.5 - k_weight = torch.randn(head_dim, dtype=torch.bfloat16, device=device).abs() + 0.5 - position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device) - - common = dict( - num_heads_q=num_heads_q, - num_heads_k=num_heads_k, - num_heads_v=num_heads_v, - head_dim=head_dim, - eps=1e-5, - q_weight=q_weight, - k_weight=k_weight, - base=10000.0, - is_neox=is_neox, - position_ids=position_ids, - factor=1.0, - low=1.0, - high=32.0, - attention_factor=1.0, - rotary_dim=head_dim, - ) - - qkv_jit = qkv.clone() - fused_qk_norm_rope(qkv_jit, **common) - - qkv_aot = qkv.clone() - fused_qk_norm_rope_aot(qkv_aot, **common) - - assert torch.allclose(qkv_jit.float(), qkv_aot.float(), atol=1e-2, rtol=1e-2), ( - f"JIT vs AOT mismatch: head_dim={head_dim}, is_neox={is_neox}, " - f"max_err={(qkv_jit.float() - qkv_aot.float()).abs().max().item():.4e}" - ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/python/sglang/multimodal_gen/test/server/test_tracing.py b/python/sglang/multimodal_gen/test/server/test_tracing.py deleted file mode 100644 index 5da984e3b..000000000 --- a/python/sglang/multimodal_gen/test/server/test_tracing.py +++ /dev/null @@ -1,155 +0,0 @@ -"""Integration test for OpenTelemetry tracing in the diffusion pipeline. - -Spins up a lightweight in-process OTLP collector, launches a diffusion server -with ``--enable-trace``, sends an image-generation request with a -``traceparent`` header, and asserts that the expected spans -(``scheduler_dispatch``, ``gpu_forward``) are exported. -""" - -import os - -# Configure OTLP exporter for faster test execution. -# Must be set before importing any sglang trace module. -os.environ.setdefault("SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS", "50") -os.environ.setdefault("SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE", "4") - -import logging -import time - -import pytest -import requests - -from sglang.multimodal_gen.test.server.test_server_utils import ServerManager -from sglang.multimodal_gen.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST -from sglang.test.otel_collector import LightweightOtlpCollector - -logger = logging.getLogger(__name__) - -# Expected diffusion trace span names (from DiffStage in trace_wrapper.py) -EXPECTED_DIFF_SPANS = ["scheduler_dispatch", "gpu_forward"] - -COLLECTOR_PORT = 4317 -SERVER_PORT = 39812 - - -@pytest.fixture(scope="module") -def tracing_env(): - """Start the OTLP collector and diffusion server once for all tests.""" - collector = LightweightOtlpCollector(port=COLLECTOR_PORT) - collector.start() - time.sleep(0.3) - - mgr = ServerManager( - model=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, - port=SERVER_PORT, - extra_args=f"--enable-trace --otlp-traces-endpoint 127.0.0.1:{COLLECTOR_PORT}", - ) - ctx = mgr.start() - - # Clear any warmup spans - time.sleep(2) - collector.clear() - - yield collector, ctx - - ctx.cleanup() - collector.stop() - - -def _generate_image(headers=None): - """Send a single image-generation request.""" - resp = requests.post( - f"http://127.0.0.1:{SERVER_PORT}/v1/images/generations", - json={ - "model": DEFAULT_SMALL_MODEL_NAME_FOR_TEST, - "prompt": "A white cat", - "size": "256x256", - "n": 1, - }, - headers=headers or {}, - timeout=300, - ) - assert resp.status_code == 200, f"Generation failed: {resp.text}" - return resp - - -def _wait_for_spans(collector, required_names=None, min_count=1, timeout=30): - """Wait until collector has the required span names (or at least ``min_count`` spans).""" - deadline = time.time() + timeout - while time.time() < deadline: - if required_names: - if all(collector.has_span(n) for n in required_names): - return - elif collector.count_spans() >= min_count: - return - time.sleep(0.5) - - -def test_spans_exported(tracing_env): - """After a generation request the expected diffusion spans appear.""" - collector, _ = tracing_env - collector.clear() - - # W3C Trace Context traceparent header - _generate_image( - headers={ - "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - } - ) - _wait_for_spans(collector, required_names=EXPECTED_DIFF_SPANS) - - span_names = collector.get_span_names() - for expected in EXPECTED_DIFF_SPANS: - assert ( - expected in span_names - ), f"Missing span '{expected}'. Collected: {sorted(span_names)}" - - -def test_spans_without_traceparent(tracing_env): - """Requests without a traceparent header still produce spans as a new - root trace (not linked to any upstream).""" - collector, _ = tracing_env - collector.clear() - - _generate_image() - _wait_for_spans(collector, required_names=EXPECTED_DIFF_SPANS) - - span_names = collector.get_span_names() - for expected in EXPECTED_DIFF_SPANS: - assert ( - expected in span_names - ), f"Missing span '{expected}'. Collected: {sorted(span_names)}" - - -def test_batch_requests(tracing_env): - """Multiple requests each produce their own set of spans.""" - collector, _ = tracing_env - collector.clear() - - batch_size = 3 - for i in range(batch_size): - # Each request gets a unique trace-id - trace_id = f"0af7651916cd43dd8448eb211c8031{i:02x}" - _generate_image(headers={"traceparent": f"00-{trace_id}-b7ad6b7169203331-01"}) - - # Wait until all scheduler_dispatch spans have arrived (they come from a - # separate process so may lag behind gpu_forward). - deadline = time.time() + 60 - while time.time() < deadline: - if all( - len(collector.get_spans_by_name(n)) >= batch_size - for n in EXPECTED_DIFF_SPANS - ): - break - time.sleep(0.5) - - for span_name in EXPECTED_DIFF_SPANS: - matching = collector.get_spans_by_name(span_name) - assert len(matching) >= batch_size, ( - f"Expected at least {batch_size} '{span_name}' spans, " - f"got {len(matching)}" - ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/python/sglang/multimodal_gen/test/unit/manual/test_patch_embed.py b/python/sglang/multimodal_gen/test/unit/manual/test_patch_embed.py deleted file mode 100644 index 59c8c4d37..000000000 --- a/python/sglang/multimodal_gen/test/unit/manual/test_patch_embed.py +++ /dev/null @@ -1,292 +0,0 @@ -""" -Test that the optimized PatchEmbed (reshape + F.linear) is equivalent -to the original Conv3d-based PatchEmbed from upstream/main. - -The opt_krea branch replaces Conv3d forward with manual -reshape + permute + F.linear for 5D input. This is valid because -Conv3d with stride==kernel_size is a non-overlapping patch extraction -followed by linear projection, which is exactly what the manual path does. - -We disable TF32 so cuDNN (Conv3d) and cuBLAS (F.linear) both use -full FP32 precision, enabling strict numerical comparison. -""" - -import pytest -import torch -import torch.nn as nn -import torch.nn.functional as F - - -class PatchEmbed3D(nn.Module): - """PatchEmbed from upstream/main: uses Conv3d directly.""" - - def __init__( - self, patch_size, in_chans, embed_dim, flatten=True, bias=True, dtype=None - ): - super().__init__() - if isinstance(patch_size, list | tuple): - if len(patch_size) == 1: - patch_size = (patch_size[0], patch_size[0]) - else: - patch_size = (patch_size, patch_size) - self.patch_size = patch_size - self.flatten = flatten - self.proj = nn.Conv3d( - in_chans, - embed_dim, - kernel_size=patch_size, - stride=patch_size, - bias=bias, - dtype=dtype, - ) - self.norm = nn.Identity() - - def forward(self, x): - x = self.proj(x) - if self.flatten: - x = x.flatten(2).transpose(1, 2) - x = self.norm(x) - return x - - -class PatchEmbed(nn.Module): - """PatchEmbed from opt_krea: replaces Conv3d with reshape + F.linear for 5D input.""" - - def __init__( - self, patch_size, in_chans, embed_dim, flatten=True, bias=True, dtype=None - ): - super().__init__() - if isinstance(patch_size, list | tuple): - if len(patch_size) == 1: - patch_size = (1, patch_size[0], patch_size[0]) - elif len(patch_size) == 2: - patch_size = (1, patch_size[0], patch_size[1]) - else: - patch_size = (1, patch_size, patch_size) - self.patch_size = patch_size - self.flatten = flatten - self.proj = nn.Conv3d( - in_chans, - embed_dim, - kernel_size=patch_size, - stride=patch_size, - bias=bias, - dtype=dtype, - ) - self.norm = nn.Identity() - - def forward(self, x): - if x.dim() == 5: - B, C, T, H, W = x.shape - pt, ph, pw = self.patch_size - T_ = T // pt - H_ = H // ph - W_ = W // pw - x = x.reshape(B, C, T_, pt, H_, ph, W_, pw) - x = x.permute(0, 2, 4, 6, 1, 3, 5, 7).contiguous() - x = x.reshape(B, T_ * H_ * W_, C * pt * ph * pw) - w = self.proj.weight.reshape(self.proj.weight.shape[0], -1) - x = F.linear(x, w, self.proj.bias) - if not self.flatten: - x = x.reshape(B, T_, H_, W_, -1).permute(0, 4, 1, 2, 3).contiguous() - x = self.norm(x) - return x - x = self.proj(x) - if self.flatten: - x = x.flatten(2).transpose(1, 2) - x = self.norm(x) - return x - - -DEVICE = "cuda" if torch.cuda.is_available() else "cpu" - - -def _copy_weights(src, dst): - dst.proj.weight.data.copy_(src.proj.weight.data) - if src.proj.bias is not None: - dst.proj.bias.data.copy_(src.proj.bias.data) - - -def _run_equivalence( - patch_size, - in_chans, - embed_dim, - flatten, - bias, - weight_dtype, - input_dtype, - B, - T, - H, - W, - atol, - rtol, -): - """Helper: build both models with shared weights, run forward, compare. - - Args: - weight_dtype: dtype for Conv3d weights (None = FP32). - input_dtype: dtype for the input tensor (None = FP32). - """ - torch.manual_seed(42) - main = ( - PatchEmbed3D(patch_size, in_chans, embed_dim, flatten, bias, dtype=weight_dtype) - .to(DEVICE) - .eval() - ) - opt = ( - PatchEmbed(patch_size, in_chans, embed_dim, flatten, bias, dtype=weight_dtype) - .to(DEVICE) - .eval() - ) - _copy_weights(main, opt) - - x = torch.randn( - B, in_chans, T, H, W, device=DEVICE, dtype=input_dtype or torch.float32 - ) - with torch.no_grad(): - out_main = main(x) - out_opt = opt(x) - - assert ( - out_main.shape == out_opt.shape - ), f"Shape mismatch: {out_main.shape} vs {out_opt.shape}" - assert ( - out_main.dtype == out_opt.dtype - ), f"Dtype mismatch: {out_main.dtype} vs {out_opt.dtype}" - torch.testing.assert_close(out_main, out_opt, atol=atol, rtol=rtol) - - -@pytest.fixture(autouse=True) -def _disable_tf32(): - prev_cudnn = torch.backends.cudnn.allow_tf32 - prev_matmul = torch.backends.cuda.matmul.allow_tf32 - torch.backends.cudnn.allow_tf32 = False - torch.backends.cuda.matmul.allow_tf32 = False - yield - torch.backends.cudnn.allow_tf32 = prev_cudnn - torch.backends.cuda.matmul.allow_tf32 = prev_matmul - - -# ── Wan2.1 / Wan2.2 / CausalWan / Helios ──────────────────────────────────── -# patch_size=(1,2,2), in_channels=16, embed_dim=5120, flatten=False -# Real usage: weight=FP32 (no dtype passed), input=BF16 from VAE latent - - -@pytest.mark.parametrize( - "dtype,atol,rtol", - [ - (None, 1e-4, 1e-4), # weight=FP32, input=FP32 - (torch.bfloat16, 1e-2, 1e-2), # weight=BF16, input=BF16 - (torch.float16, 1e-2, 1e-2), # weight=FP16, input=FP16 - ], - ids=["fp32", "bf16", "fp16"], -) -@pytest.mark.parametrize( - "B,T,H,W", - [ - (1, 21, 60, 104), # 480p typical - (2, 9, 40, 64), # smaller resolution, batch=2 - (1, 33, 90, 160), # 720p longer video - ], - ids=["480p-B1", "small-B2", "720p-B1"], -) -def test_wan_helios(dtype, atol, rtol, B, T, H, W): - _run_equivalence( - patch_size=(1, 2, 2), - in_chans=16, - embed_dim=5120, - flatten=False, - bias=True, - weight_dtype=dtype, - input_dtype=dtype, - B=B, - T=T, - H=H, - W=W, - atol=atol, - rtol=rtol, - ) - - -# ── HunyuanVideo ───────────────────────────────────────────────────────────── -# patch_size=[1,2,2] (list!), in_channels=16, embed_dim=3072, flatten=True -# Real usage: dtype passed to PatchEmbed, so weight & input share same dtype - - -@pytest.mark.parametrize( - "dtype,atol,rtol", - [ - (None, 1e-4, 1e-4), # weight=FP32, input=FP32 - (torch.bfloat16, 1e-2, 1e-2), # weight=BF16, input=BF16 - (torch.float16, 1e-2, 1e-2), # weight=FP16, input=FP16 - ], - ids=["fp32", "bf16", "fp16"], -) -@pytest.mark.parametrize( - "B,T,H,W", - [ - (1, 21, 60, 104), - (2, 9, 40, 64), - ], - ids=["480p-B1", "small-B2"], -) -def test_hunyuanvideo(dtype, atol, rtol, B, T, H, W): - _run_equivalence( - patch_size=[1, 2, 2], - in_chans=16, - embed_dim=3072, - flatten=True, - bias=True, - weight_dtype=dtype, - input_dtype=dtype, - B=B, - T=T, - H=H, - W=W, - atol=atol, - rtol=rtol, - ) - - -# ── No-bias variants ───────────────────────────────────────────────────────── - - -def test_wan_no_bias(): - _run_equivalence( - patch_size=(1, 2, 2), - in_chans=16, - embed_dim=5120, - flatten=False, - bias=False, - weight_dtype=None, - input_dtype=None, - B=1, - T=21, - H=60, - W=104, - atol=1e-4, - rtol=1e-4, - ) - - -def test_hunyuanvideo_no_bias(): - _run_equivalence( - patch_size=[1, 2, 2], - in_chans=16, - embed_dim=3072, - flatten=True, - bias=False, - weight_dtype=None, - input_dtype=None, - B=1, - T=21, - H=60, - W=104, - atol=1e-4, - rtol=1e-4, - ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "--tb=short"]) diff --git a/test/registered/lora/test_lora_moe_runner.py b/test/registered/lora/test_lora_moe_runner.py deleted file mode 100644 index 253fdf1e9..000000000 --- a/test/registered/lora/test_lora_moe_runner.py +++ /dev/null @@ -1,788 +0,0 @@ -# Copyright 2023-2025 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -import random -from unittest.mock import patch - -import pytest -import torch - -from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig -from sglang.srt.layers.moe.moe_runner.runner import MoeRunner -from sglang.srt.layers.moe.moe_runner.triton import ( - TritonMoeQuantInfo, -) -from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput -from sglang.srt.layers.moe.topk import StandardTopKOutput -from sglang.srt.layers.moe.utils import MoeRunnerBackend -from sglang.srt.lora.lora_moe_runners import LoRAInfo -from sglang.srt.utils import set_random_seed -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=17, suite="stage-b-test-1-gpu-large") - - -def generate_request_data( - num_tokens: int, num_sequences: int, max_loras: int, device="cuda" -): - """ - Generates segment-based request data instead of token-based data. - """ - assert num_sequences > 0 and max_loras > 0 - assert num_tokens >= num_sequences, "num_tokens must be >= num_sequences" - - # 1. Generate random segment lengths - remaining = num_tokens - seg_lens = [] - for _ in range(num_sequences - 1): - # Ensure at least 1 token per sequence - max_len = remaining - (num_sequences - len(seg_lens)) + 1 - length = random.randint(1, min(max_len, num_tokens // num_sequences * 2)) - seg_lens.append(length) - remaining -= length - seg_lens.append(remaining) # Last segment gets the rest - - # 2. Build seg_indptr [0, len1, len1+len2, ...] - seg_indptr = torch.cumsum( - torch.tensor([0] + seg_lens, dtype=torch.int32, device=device), - dim=0, - dtype=torch.int32, - ) - - # 3. Assign one LoRA ID per Request - req_to_lora = torch.randint( - 0, max_loras, (num_sequences,), dtype=torch.int32, device=device - ) - - # 4. Create dense mapping for the Naive verification function - # (Expand req_to_lora based on seg_lens) - token_lora_mapping = torch.repeat_interleave( - req_to_lora, torch.tensor(seg_lens, device=device) - ) - - return seg_indptr, req_to_lora, token_lora_mapping - - -def assign_experts_to_tokens( - num_tokens: int, num_experts: int, top_k_num: int, dtype=torch.float32 -): - assert top_k_num <= num_experts, "top_k_num must be <= num_experts" - - expert_indices = torch.empty((num_tokens, top_k_num), dtype=torch.int32) - for i in range(num_tokens): - selected = torch.randperm(num_experts)[:top_k_num] - expert_indices[i] = selected - - expert_weights = torch.rand((num_tokens, top_k_num), dtype=dtype) - expert_weights = expert_weights / expert_weights.sum(dim=1, keepdim=True) - - return expert_indices, expert_weights - - -def sample_data( - num_tokens: int, - num_sequences: int, - max_loras: int, - num_experts: int, - top_k_num: int, - dtype=torch.float32, - device="cuda", -): - topk_ids, topk_weights = assign_experts_to_tokens( - num_tokens, num_experts, top_k_num, dtype - ) - seg_indptr, req_to_lora, token_lora_mapping = generate_request_data( - num_tokens, num_sequences, max_loras, device - ) - return topk_ids, topk_weights, seg_indptr, req_to_lora, token_lora_mapping - - -def create_lora_info( - seg_indptr, - weight_indices, - topk_ids, - max_loras, - num_experts, - max_lora_rank, - hidden_dim, - intermediate_dim, - gate_up_dim, - dtype, - device, - lora_use_virtual_experts=False, -): - # ------------------------------------------------------------------------- - # 1. Deterministic LoRA A Initialization - # ------------------------------------------------------------------------- - - val_gate_up_a = 0.1 - gate_up_lora_a_weights = torch.full( - (max_loras, num_experts, max_lora_rank * 2, hidden_dim), - val_gate_up_a, - dtype=dtype, - device=device, - ) - - val_down_a = 1.0 / intermediate_dim - down_lora_a_weights = torch.full( - (max_loras, num_experts, max_lora_rank, intermediate_dim), - val_down_a, - dtype=dtype, - device=device, - ) - - # ------------------------------------------------------------------------- - # 2. Deterministic LoRA B Initialization - # ------------------------------------------------------------------------- - base_target = 0.05 - - gate_up_lora_b_weights = torch.zeros( - (max_loras, num_experts, gate_up_dim, max_lora_rank), - dtype=dtype, - device=device, - ) - down_lora_b_weights = torch.zeros( - (max_loras, num_experts, hidden_dim, max_lora_rank), dtype=dtype, device=device - ) - - for i in range(num_experts): - expert_multiplier = i + 1 - divisor = max(1, max_lora_rank) - fill_val = (base_target * expert_multiplier) / divisor - - gate_up_lora_b_weights[:, i, :, :] = fill_val - down_lora_b_weights[:, i, :, :] = fill_val - - # ------------------------------------------------------------------------- - # 3. Setup Metadata - # ------------------------------------------------------------------------- - lora_ranks = torch.full( - (max_loras,), max_lora_rank, dtype=torch.int32, device=device - ) - - # Enable all adapters referenced in weight_indices - adapter_enabled = torch.zeros(max_loras + 1, dtype=torch.int32, device=device) - adapter_enabled.index_fill_(0, weight_indices.long(), 1) - - return LoRAInfo( - gate_up_lora_a_weights=gate_up_lora_a_weights, - gate_up_lora_b_weights=gate_up_lora_b_weights, - down_lora_a_weights=down_lora_a_weights, - down_lora_b_weights=down_lora_b_weights, - # UPDATED FIELDS - seg_indptr=seg_indptr, - req_to_lora=weight_indices, - lora_ranks=lora_ranks, - adapter_enabled=adapter_enabled, - max_lora_rank=max_lora_rank, - num_experts=num_experts, - lora_use_virtual_experts=lora_use_virtual_experts, - ) - - -def torch_naive_moe_with_lora( - hidden_states, - w13, - w2, - b13, - b2, - topk_weights, - topk_ids, - lora_info, - token_lora_mapping, -): - """ - Naive implementation. Note: We pass 'token_lora_mapping' explicitly because - lora_info no longer contains it, but the naive token-loop logic needs it. - """ - num_tokens, hidden_dim = hidden_states.shape - top_k = topk_ids.shape[1] - num_experts = w13.shape[0] - - # Expand hidden states for top-k routing - hidden_expanded = ( - hidden_states.unsqueeze(1).expand(-1, top_k, -1).reshape(-1, hidden_dim) - ) - - # 1. Gate/Up Projection (Base) - gate_up_out = torch.zeros( - num_tokens * top_k, - w13.shape[1], - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - - for expert_id in range(num_experts): - mask = (topk_ids == expert_id).flatten() - if mask.any(): - expert_result = hidden_expanded[mask] @ w13[expert_id].T - gate_up_out[mask] = expert_result - if b13 is not None: - gate_up_out[mask] += b13[expert_id] - - gate_up_out = gate_up_out.view(num_tokens, top_k, -1) - - # 1.5. LoRA Gate/Up Delta - # gate_up_lora_a is packed as [gate_a; up_a] along rank dim → [2*r, hidden_dim] - # gate_up_lora_b is packed as [gate_b; up_b] along output dim → [2*inter, r] - # Correct computation splits them: gate uses first r rows of A with first half of B, - # up uses last r rows of A with second half of B. - if lora_info.max_lora_rank > 0: - r = lora_info.max_lora_rank - for i in range(num_tokens): - for k in range(top_k): - expert_id = topk_ids[i, k] - lora_id = token_lora_mapping[i] - - if lora_id < len(lora_info.lora_ranks): - lora_a = lora_info.gate_up_lora_a_weights[lora_id, expert_id] - lora_b = lora_info.gate_up_lora_b_weights[lora_id, expert_id] - half = lora_b.shape[0] // 2 - lora_a_result = lora_a @ hidden_states[i] - gate_delta = lora_b[:half, :] @ lora_a_result[:r] - up_delta = lora_b[half:, :] @ lora_a_result[r:] - gate_up_out[i, k] += torch.cat([gate_delta, up_delta]) - - # 2. Activation - gate_up_dim = gate_up_out.shape[-1] - gate_dim = gate_up_dim // 2 - gate = gate_up_out[..., :gate_dim] - up = gate_up_out[..., gate_dim:] - - silu_gate = torch.nn.functional.silu(gate) - intermediate_out = silu_gate * up - - # 3. Down Projection (Base) - down_out = torch.zeros( - num_tokens, - top_k, - hidden_dim, - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - - for expert_id in range(num_experts): - mask = topk_ids == expert_id - if mask.any(): - masked_intermediate = intermediate_out[mask] - expert_down_result = masked_intermediate @ w2[expert_id].T - down_out[mask] = expert_down_result - if b2 is not None: - down_out[mask] += b2[expert_id] - - # 3.5. LoRA Down Delta - if lora_info.max_lora_rank > 0: - for i in range(num_tokens): - for k in range(top_k): - expert_id = topk_ids[i, k] - lora_id = token_lora_mapping[i] # Use explicit mapping - - if lora_id < len(lora_info.lora_ranks): - lora_a = lora_info.down_lora_a_weights[lora_id, expert_id] - lora_b = lora_info.down_lora_b_weights[lora_id, expert_id] - lora_a_result = lora_a @ intermediate_out[i, k] - lora_b_result = lora_b @ lora_a_result - down_out[i, k] += lora_b_result - - # 4. Final Reduction - weighted_out = down_out * topk_weights.unsqueeze(-1) - final_out = weighted_out.sum(dim=1) - - return final_out - - -@pytest.mark.parametrize("num_tokens", [32, 64]) -@pytest.mark.parametrize("top_k_num", [1, 2]) -@pytest.mark.parametrize("num_experts", [8, 20]) -@pytest.mark.parametrize("max_lora_rank", [8, 16]) -def test_lora_moe_runner_multi_expert( - num_tokens, top_k_num, num_experts, max_lora_rank -): - # Fixed parameters - max_loras = 2 - hidden_dim = 512 - intermediate_dim = 1024 - - dtype = torch.float32 - device = "cuda:0" - seed = 42 - - torch.set_default_device(device) - set_random_seed(seed) - - num_sequences = 4 - - # Generate Data using the new Request-Based generator - topk_ids, topk_weights, seg_indptr, req_to_lora, token_lora_mapping = sample_data( - num_tokens, num_sequences, max_loras, num_experts, top_k_num, dtype, device - ) - - gate_up_dim = intermediate_dim * 2 - - # Initialize experts - w13 = torch.randn(num_experts, gate_up_dim, hidden_dim, dtype=dtype) * 0.1 - w2 = torch.randn(num_experts, hidden_dim, intermediate_dim, dtype=dtype) * 0.1 - b13 = torch.randn(num_experts, gate_up_dim, dtype=dtype) * 0.1 - b2 = torch.randn(num_experts, hidden_dim, dtype=dtype) * 0.1 - - hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype) - - # Create LoRA Info using the new fields - lora_info_delta = create_lora_info( - seg_indptr=seg_indptr, - weight_indices=req_to_lora, - topk_ids=topk_ids, - max_loras=max_loras, - num_experts=num_experts, - max_lora_rank=max_lora_rank, - hidden_dim=hidden_dim, - intermediate_dim=intermediate_dim, - gate_up_dim=gate_up_dim, - dtype=dtype, - device=device, - ) - - lora_info_baseline = create_lora_info( - seg_indptr=seg_indptr, - weight_indices=req_to_lora, - topk_ids=topk_ids, - max_loras=max_loras, - num_experts=num_experts, - max_lora_rank=0, # Set rank to 0 for baseline - hidden_dim=hidden_dim, - intermediate_dim=intermediate_dim, - gate_up_dim=gate_up_dim, - dtype=dtype, - device=device, - ) - - # Sort tokens for the runner - topk_ids_flat = topk_ids.flatten() - sorted_indices = torch.argsort(topk_ids_flat) - sorted_token_ids = sorted_indices // top_k_num - expert_ids = topk_ids_flat[sorted_indices] - - num_dispatched = num_tokens * top_k_num - num_tokens_post_padded = torch.tensor( - [num_dispatched], dtype=torch.int32, device=device - ) - - quant_info = TritonMoeQuantInfo( - w13_weight=w13, - w2_weight=w2, - b13=b13, - b2=b2, - ) - - config = MoeRunnerConfig( - activation="silu", - is_gated=True, - inplace=False, - no_combine=False, - gemm1_alpha=None, - gemm1_clamp_limit=None, - routed_scaling_factor=1.0, - apply_router_weight_on_input=False, - num_local_experts=num_experts, - ) - - # Create StandardTopKOutput - router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device=device) - topk_output = StandardTopKOutput( - topk_weights=topk_weights, - topk_ids=topk_ids, - router_logits=router_logits, - ) - - # Create StandardDispatchOutput - dispatch_output = StandardDispatchOutput( - hidden_states=hidden_states, - hidden_states_scale=None, - topk_output=topk_output, - ) - - class MockServerArgs: - enable_deterministic_inference = False - - with patch( - "sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config.get_global_server_args", - return_value=MockServerArgs(), - ): - runner = MoeRunner(MoeRunnerBackend.TRITON, config, lora_enabled=True) - - # 3. Get outputs for both scenarios - output_with_lora = runner.run( - dispatch_output, quant_info, lora_info_delta - ).hidden_states - output_baseline = runner.run( - dispatch_output, quant_info, lora_info_baseline - ).hidden_states - - # Run Naive Torch Implementation (Uses dense mapping for verification) - torch_output_lora = torch_naive_moe_with_lora( - hidden_states, - w13, - w2, - b13, - b2, - topk_weights, - topk_ids, - lora_info_delta, - token_lora_mapping, - ) - - torch_output_base = torch_naive_moe_with_lora( - hidden_states, - w13, - w2, - b13, - b2, - topk_weights, - topk_ids, - lora_info_baseline, - token_lora_mapping, - ) - - # The actual "Delta" (LoRA effect) for both - sglang_delta = output_with_lora - output_baseline - torch_delta = torch_output_lora - torch_output_base - - # Larger expert counts accumulate more numerical drift in Triton kernels on GB300 - tol = 0.15 if num_experts >= 20 else 5e-2 - torch.testing.assert_close(sglang_delta, torch_delta, atol=tol, rtol=tol) - - -@pytest.mark.parametrize("num_tokens", [32, 64]) -@pytest.mark.parametrize("top_k_num", [1, 2]) -@pytest.mark.parametrize("num_experts", [8, 20]) -@pytest.mark.parametrize("max_lora_rank", [8, 16]) -def test_lora_moe_runner_virtual_experts( - num_tokens, top_k_num, num_experts, max_lora_rank -): - # Fixed parameters - max_loras = 2 - hidden_dim = 512 - intermediate_dim = 1024 - - dtype = torch.float32 - device = "cuda:0" - seed = 42 - - torch.set_default_device(device) - set_random_seed(seed) - - num_sequences = 4 - - # Generate Data using the new Request-Based generator - topk_ids, topk_weights, seg_indptr, req_to_lora, token_lora_mapping = sample_data( - num_tokens, num_sequences, max_loras, num_experts, top_k_num, dtype, device - ) - - gate_up_dim = intermediate_dim * 2 - - # Initialize experts - w13 = torch.randn(num_experts, gate_up_dim, hidden_dim, dtype=dtype) * 0.1 - w2 = torch.randn(num_experts, hidden_dim, intermediate_dim, dtype=dtype) * 0.1 - b13 = torch.randn(num_experts, gate_up_dim, dtype=dtype) * 0.1 - b2 = torch.randn(num_experts, hidden_dim, dtype=dtype) * 0.1 - - hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype) - - # Create LoRA Info with virtual experts enabled - lora_info_delta = create_lora_info( - seg_indptr=seg_indptr, - weight_indices=req_to_lora, - topk_ids=topk_ids, - max_loras=max_loras, - num_experts=num_experts, - max_lora_rank=max_lora_rank, - hidden_dim=hidden_dim, - intermediate_dim=intermediate_dim, - gate_up_dim=gate_up_dim, - dtype=dtype, - device=device, - lora_use_virtual_experts=True, - ) - - lora_info_baseline = create_lora_info( - seg_indptr=seg_indptr, - weight_indices=req_to_lora, - topk_ids=topk_ids, - max_loras=max_loras, - num_experts=num_experts, - max_lora_rank=0, - hidden_dim=hidden_dim, - intermediate_dim=intermediate_dim, - gate_up_dim=gate_up_dim, - dtype=dtype, - device=device, - lora_use_virtual_experts=True, - ) - - quant_info = TritonMoeQuantInfo( - w13_weight=w13, - w2_weight=w2, - b13=b13, - b2=b2, - ) - - config = MoeRunnerConfig( - activation="silu", - is_gated=True, - inplace=False, - no_combine=False, - gemm1_alpha=None, - gemm1_clamp_limit=None, - routed_scaling_factor=1.0, - apply_router_weight_on_input=False, - num_local_experts=num_experts, - ) - - router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device=device) - topk_output = StandardTopKOutput( - topk_weights=topk_weights, - topk_ids=topk_ids, - router_logits=router_logits, - ) - - dispatch_output = StandardDispatchOutput( - hidden_states=hidden_states, - hidden_states_scale=None, - topk_output=topk_output, - ) - - class MockServerArgs: - enable_deterministic_inference = False - - with patch( - "sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config.get_global_server_args", - return_value=MockServerArgs(), - ): - runner = MoeRunner(MoeRunnerBackend.TRITON, config, lora_enabled=True) - - output_with_lora = runner.run( - dispatch_output, quant_info, lora_info_delta - ).hidden_states - output_baseline = runner.run( - dispatch_output, quant_info, lora_info_baseline - ).hidden_states - - # Run Naive Torch Implementation (Uses dense mapping for verification) - torch_output_lora = torch_naive_moe_with_lora( - hidden_states, - w13, - w2, - b13, - b2, - topk_weights, - topk_ids, - lora_info_delta, - token_lora_mapping, - ) - - torch_output_base = torch_naive_moe_with_lora( - hidden_states, - w13, - w2, - b13, - b2, - topk_weights, - topk_ids, - lora_info_baseline, - token_lora_mapping, - ) - - # The actual "Delta" (LoRA effect) for both - sglang_delta = output_with_lora - output_baseline - torch_delta = torch_output_lora - torch_output_base - - # Larger expert counts accumulate more numerical drift in Triton kernels on GB300 - tol = 0.15 if num_experts >= 20 else 5e-2 - torch.testing.assert_close(sglang_delta, torch_delta, atol=tol, rtol=tol) - - -def _setup_marlin_moe_weights(num_experts, n, k, dtype): - """Quantize float weights into AWQ Marlin format for testing.""" - from sgl_kernel.scalar_type import scalar_types - - from sglang.test.test_marlin_utils import awq_marlin_quantize - - group_size = 128 - quant_type = scalar_types.uint4 - - w = torch.randn((num_experts, n, k), device="cuda", dtype=dtype) / 20 - - w_ref_l, qweight_l, scales_l, zeros_l = [], [], [], [] - for i in range(num_experts): - w_ref, qweight, scales, zeros = awq_marlin_quantize( - w[i].transpose(1, 0), quant_type, group_size - ) - w_ref_l.append(w_ref.T) - qweight_l.append(qweight) - scales_l.append(scales) - zeros_l.append(zeros) - - def _stack(tensors): - dev = tensors[0].device - return torch.stack(tensors, dim=0).to(dev) - - return ( - _stack(w_ref_l), - _stack(qweight_l).contiguous(), - _stack(scales_l), - _stack(zeros_l), - ) - - -@pytest.mark.parametrize("num_tokens", [32, 64]) -@pytest.mark.parametrize("top_k_num", [1, 2]) -@pytest.mark.parametrize("num_experts", [8]) -@pytest.mark.parametrize("max_lora_rank", [8, 16]) -def test_lora_moe_runner_marlin(num_tokens, top_k_num, num_experts, max_lora_rank): - from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo - - max_loras = 2 - hidden_dim = 512 - intermediate_dim = 1024 - gate_up_dim = intermediate_dim * 2 - - dtype = torch.float16 - device = "cuda:0" - seed = 42 - - torch.set_default_device(device) - set_random_seed(seed) - - num_sequences = 4 - - topk_ids, topk_weights, seg_indptr, req_to_lora, token_lora_mapping = sample_data( - num_tokens, - num_sequences, - max_loras, - num_experts, - top_k_num, - dtype, - device, - ) - - # Quantize base weights to Marlin format - _, w13_qweight, w13_scales, w13_qzeros = _setup_marlin_moe_weights( - num_experts, gate_up_dim, hidden_dim, dtype - ) - _, w2_qweight, w2_scales, w2_qzeros = _setup_marlin_moe_weights( - num_experts, hidden_dim, intermediate_dim, dtype - ) - - hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=device) - - lora_info_delta = create_lora_info( - seg_indptr=seg_indptr, - weight_indices=req_to_lora, - topk_ids=topk_ids, - max_loras=max_loras, - num_experts=num_experts, - max_lora_rank=max_lora_rank, - hidden_dim=hidden_dim, - intermediate_dim=intermediate_dim, - gate_up_dim=gate_up_dim, - dtype=dtype, - device=device, - lora_use_virtual_experts=True, - ) - - lora_info_baseline = create_lora_info( - seg_indptr=seg_indptr, - weight_indices=req_to_lora, - topk_ids=topk_ids, - max_loras=max_loras, - num_experts=num_experts, - max_lora_rank=0, - hidden_dim=hidden_dim, - intermediate_dim=intermediate_dim, - gate_up_dim=gate_up_dim, - dtype=dtype, - device=device, - lora_use_virtual_experts=True, - ) - - quant_info = MarlinMoeQuantInfo( - w13_qweight=w13_qweight, - w2_qweight=w2_qweight, - w13_scales=w13_scales, - w2_scales=w2_scales, - w13_qzeros=w13_qzeros, - w2_qzeros=w2_qzeros, - w13_g_idx=None, - w2_g_idx=None, - w13_g_idx_sort_indices=None, - w2_g_idx_sort_indices=None, - weight_bits=4, - ) - - config = MoeRunnerConfig( - activation="silu", - is_gated=True, - inplace=False, - no_combine=False, - gemm1_alpha=None, - gemm1_clamp_limit=None, - routed_scaling_factor=1.0, - apply_router_weight_on_input=False, - num_local_experts=num_experts, - ) - - router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device=device) - topk_output = StandardTopKOutput( - topk_weights=topk_weights, - topk_ids=topk_ids, - router_logits=router_logits, - ) - dispatch_output = StandardDispatchOutput( - hidden_states=hidden_states, - hidden_states_scale=None, - topk_output=topk_output, - ) - - class MockServerArgs: - enable_deterministic_inference = False - - with patch( - "sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config.get_global_server_args", - return_value=MockServerArgs(), - ): - runner = MoeRunner(MoeRunnerBackend.MARLIN, config, lora_enabled=True) - output_with_lora = runner.run( - dispatch_output, quant_info, lora_info_delta - ).hidden_states - output_baseline = runner.run( - dispatch_output, quant_info, lora_info_baseline - ).hidden_states - - marlin_delta = output_with_lora - output_baseline - - # Verify the LoRA hooks fired and produced a non-trivial delta - assert marlin_delta.abs().max().item() > 1e-4, ( - f"LoRA delta is too small ({marlin_delta.abs().max().item():.6f}), " - "hooks may not be firing" - ) - assert torch.isfinite( - output_with_lora - ).all(), "Marlin+LoRA output contains non-finite values" - assert torch.isfinite( - output_baseline - ).all(), "Marlin baseline output contains non-finite values" - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/test/registered/lora/test_marlin_lora_correctness.py b/test/registered/lora/test_marlin_lora_correctness.py deleted file mode 100644 index 44fc970db..000000000 --- a/test/registered/lora/test_marlin_lora_correctness.py +++ /dev/null @@ -1,289 +0,0 @@ -# Copyright 2023-2025 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -""" -Correctness test: Marlin (int4 base + LoRA) vs Triton (dequantized base + LoRA). - -Fake-quantizes random weights to int4/Marlin format and dequantizes them with the -same path, then runs both backends through MoeRunner and compares LoRA deltas. -""" - -from unittest.mock import patch - -import pytest -import torch - -from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig -from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo -from sglang.srt.layers.moe.moe_runner.runner import MoeRunner -from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo -from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput -from sglang.srt.layers.moe.topk import StandardTopKOutput -from sglang.srt.layers.moe.utils import MoeRunnerBackend -from sglang.srt.lora.lora_moe_runners import LoRAInfo -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=129, suite="stage-b-test-1-gpu-large") - - -# --------------------------------------------------------------------------- -# Fake quantization helpers (symmetric int4, matching Marlin's dequant path) -# --------------------------------------------------------------------------- - - -def _quantize_per_expert(w_float: torch.Tensor, K: int, group_size: int): - """Quantize [N, K] float weight to int4. Returns (q_int [N,K], scales_bf16 [N,groups]).""" - N = w_float.shape[0] - num_groups = K // group_size - - w_grouped = w_float.reshape(N, num_groups, group_size) - scales_fp32 = w_grouped.abs().amax(dim=-1) / 7.0 - scales_fp32 = scales_fp32.clamp(min=1e-6) - scales_bf16 = scales_fp32.to(torch.bfloat16) - - scales_for_quant = scales_bf16.float() - q_int = torch.zeros(N, K, dtype=torch.int32, device=w_float.device) - for g in range(num_groups): - s = scales_for_quant[:, g : g + 1] - sl = slice(g * group_size, (g + 1) * group_size) - q_int[:, sl] = torch.round(w_float[:, sl] / s).clamp(-8, 7).to(torch.int32) + 8 - - return q_int, scales_bf16 - - -def _fake_quantize_to_marlin_int4(weight_bf16: torch.Tensor): - """Fake-quantize [E, N, K] bf16 weight to Marlin int4 format. - - Returns: (qweight, scales, g_idx, g_idx_sort_indices) - """ - from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack - from sglang.srt.layers.quantization.marlin_utils import marlin_permute_scales - from sglang.srt.layers.quantization.utils import pack_rows - - E, N, K = weight_bf16.shape - num_bits = 4 - group_size = 128 - device = weight_bf16.device - - all_qweight, all_scales = [], [] - for e in range(E): - q_int, scales_bf16 = _quantize_per_expert(weight_bf16[e].float(), K, group_size) - w_quant_t = q_int.t().contiguous() - packed = pack_rows(w_quant_t, num_bits, K, N) - perm = torch.arange(K, device=device, dtype=torch.int32) - all_qweight.append(gptq_marlin_repack(packed.to(device), perm, K, N, num_bits)) - all_scales.append( - marlin_permute_scales( - scales_bf16.t().contiguous().to(device), K, N, group_size - ) - ) - - g_idx = ( - (torch.arange(K, device=device, dtype=torch.int32) // group_size) - .unsqueeze(0) - .expand(E, -1) - .contiguous() - ) - sort_indices = ( - torch.arange(K, device=device, dtype=torch.int32) - .unsqueeze(0) - .expand(E, -1) - .contiguous() - ) - - return torch.stack(all_qweight), torch.stack(all_scales), g_idx, sort_indices - - -def _dequantize_from_marlin_int4(weight_bf16_orig: torch.Tensor, group_size: int = 128): - """Dequantize using the same path as _fake_quantize, so Triton reference matches Marlin.""" - E, N, K = weight_bf16_orig.shape - result = torch.zeros_like(weight_bf16_orig) - for e in range(E): - q_int, scales_bf16 = _quantize_per_expert( - weight_bf16_orig[e].float(), K, group_size - ) - num_groups = K // group_size - for g in range(num_groups): - sl = slice(g * group_size, (g + 1) * group_size) - s = scales_bf16[:, g : g + 1] - result[e, :, sl] = (q_int[:, sl] - 8).to(torch.bfloat16) * s - return result - - -# --------------------------------------------------------------------------- -# Test -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("num_tokens", [1, 8, 32]) -@pytest.mark.parametrize("top_k", [2, 8]) -def test_marlin_vs_triton_lora_correctness(num_tokens, top_k): - torch.manual_seed(42) - - device = "cuda" - dtype = torch.bfloat16 - - hidden_dim = 7168 - intermediate_dim = 2048 - gate_up_dim = 2 * intermediate_dim - num_experts = 64 - lora_rank = 32 - num_loras = 1 - - hidden = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=device) - topk_weights = torch.randn( - num_tokens, top_k, dtype=torch.float32, device=device - ).softmax(dim=-1) - topk_ids = torch.randint( - 0, num_experts, (num_tokens, top_k), dtype=torch.int32, device=device - ) - - # Base weights (random bf16) - w13_bf16 = ( - torch.randn(num_experts, gate_up_dim, hidden_dim, dtype=dtype, device=device) - * 0.01 - ) - w2_bf16 = ( - torch.randn( - num_experts, hidden_dim, intermediate_dim, dtype=dtype, device=device - ) - * 0.01 - ) - - # LoRA weights (shared across both paths) - gu_lora_a = ( - torch.randn(num_loras, 1, lora_rank * 2, hidden_dim, dtype=dtype, device=device) - * 0.01 - ) - gu_lora_b = ( - torch.randn( - num_loras, num_experts, gate_up_dim, lora_rank, dtype=dtype, device=device - ) - * 0.01 - ) - dn_lora_a = ( - torch.randn( - num_loras, - num_experts, - lora_rank, - intermediate_dim, - dtype=dtype, - device=device, - ) - * 0.01 - ) - dn_lora_b = ( - torch.randn(num_loras, 1, hidden_dim, lora_rank, dtype=dtype, device=device) - * 0.01 - ) - - # Token-to-LoRA mapping: all tokens use adapter 0 - seg_indptr = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) - req_to_lora = torch.tensor([0], dtype=torch.int32, device=device) - - def _make_lora_info(rank): - return LoRAInfo( - gate_up_lora_a_weights=gu_lora_a if rank > 0 else gu_lora_a[:, :, :0, :], - gate_up_lora_b_weights=gu_lora_b if rank > 0 else gu_lora_b[:, :, :, :0], - down_lora_a_weights=dn_lora_a if rank > 0 else dn_lora_a[:, :, :0, :], - down_lora_b_weights=dn_lora_b if rank > 0 else dn_lora_b[:, :, :, :0], - seg_indptr=seg_indptr, - req_to_lora=req_to_lora, - lora_ranks=torch.full((num_loras,), rank, dtype=torch.int32, device=device), - adapter_enabled=torch.ones(num_loras + 1, dtype=torch.int32, device=device), - max_lora_rank=rank, - num_experts=num_experts, - lora_use_virtual_experts=True, - experts_shared_outer_loras=True, - ) - - lora_info = _make_lora_info(lora_rank) - lora_baseline = _make_lora_info(0) - - # Quantize for Marlin, dequantize for Triton reference - w13_qw, w13_sc, w13_gidx, w13_si = _fake_quantize_to_marlin_int4(w13_bf16) - w2_qw, w2_sc, w2_gidx, w2_si = _fake_quantize_to_marlin_int4(w2_bf16) - w13_deq = _dequantize_from_marlin_int4(w13_bf16) - w2_deq = _dequantize_from_marlin_int4(w2_bf16) - - marlin_qi = MarlinMoeQuantInfo( - w13_qweight=w13_qw, - w2_qweight=w2_qw, - w13_scales=w13_sc, - w2_scales=w2_sc, - w13_g_idx=w13_gidx, - w2_g_idx=w2_gidx, - w13_g_idx_sort_indices=w13_si, - w2_g_idx_sort_indices=w2_si, - weight_bits=4, - ) - triton_qi = TritonMoeQuantInfo( - w13_weight=w13_deq, w2_weight=w2_deq, b13=None, b2=None - ) - - config = MoeRunnerConfig( - activation="silu", - is_gated=True, - inplace=False, - no_combine=False, - gemm1_alpha=None, - gemm1_clamp_limit=None, - routed_scaling_factor=1.0, - apply_router_weight_on_input=False, - num_local_experts=num_experts, - ) - - router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device=device) - topk_output = StandardTopKOutput( - topk_weights=topk_weights, topk_ids=topk_ids, router_logits=router_logits - ) - dispatch_output = StandardDispatchOutput( - hidden_states=hidden, hidden_states_scale=None, topk_output=topk_output - ) - - class MockServerArgs: - enable_deterministic_inference = False - - with patch( - "sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config.get_global_server_args", - return_value=MockServerArgs(), - ): - marlin_runner = MoeRunner(MoeRunnerBackend.MARLIN, config, lora_enabled=True) - triton_runner = MoeRunner(MoeRunnerBackend.TRITON, config, lora_enabled=True) - - marlin_out = marlin_runner.run( - dispatch_output, marlin_qi, lora_info - ).hidden_states - marlin_base = marlin_runner.run( - dispatch_output, marlin_qi, lora_baseline - ).hidden_states - triton_out = triton_runner.run( - dispatch_output, triton_qi, lora_info - ).hidden_states - triton_base = triton_runner.run( - dispatch_output, triton_qi, lora_baseline - ).hidden_states - - marlin_delta = marlin_out - marlin_base - triton_delta = triton_out - triton_base - - # Remaining error is from kernel-level accumulation differences - # (Marlin fp32 reduce vs Triton bf16 dot), not from quantization mismatch. - torch.testing.assert_close( - marlin_delta.float(), triton_delta.float(), atol=0.01, rtol=0.05 - ) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/test/registered/lora/test_sgemm_sorted_by_adapter.py b/test/registered/lora/test_sgemm_sorted_by_adapter.py deleted file mode 100644 index d1f7a0baf..000000000 --- a/test/registered/lora/test_sgemm_sorted_by_adapter.py +++ /dev/null @@ -1,236 +0,0 @@ -"""Test that sgemm kernels produce identical results with and without SORTED_BY_ADAPTER.""" - -from typing import Any - -import pytest -import torch - -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=11, suite="stage-b-test-1-gpu-large") - - -def _make_batch_info( - bs: int, - weight_indices: list[int], - lora_ranks: list[int], - scalings: list[float], - device: str = "cuda", -) -> Any: - """Build a per-sequence LoRABatchInfo (no permutation).""" - from sglang.srt.lora.utils import LoRABatchInfo - - seg_lens = torch.ones(bs, dtype=torch.int32, device=device) - seg_indptr = torch.zeros(bs + 1, dtype=torch.int32, device=device) - seg_indptr[1:] = torch.cumsum(seg_lens, dim=0) - return LoRABatchInfo( - bs=bs, - use_cuda_graph=False, - num_segments=bs, - seg_lens=seg_lens, - seg_indptr=seg_indptr, - max_len=1, - weight_indices=torch.tensor(weight_indices, dtype=torch.int32, device=device), - lora_ranks=torch.tensor(lora_ranks, dtype=torch.int32, device=device), - scalings=torch.tensor(scalings, dtype=torch.float, device=device), - permutation=None, - ) - - -def _make_sorted_batch_info( - weight_indices: list[int], - lora_ranks: list[int], - scalings: list[float], - max_loras: int, - device: str = "cuda", -) -> Any: - from sglang.srt.lora.utils import LoRABatchInfo - - """Build a merged-by-adapter LoRABatchInfo (with permutation).""" - wi = torch.tensor(weight_indices, dtype=torch.int32, device=device) - bs = wi.shape[0] - - perm = torch.argsort(wi, stable=True).to(torch.int32) - sorted_wi = wi[perm] - adapter_ids = torch.arange(max_loras, device=device, dtype=torch.int32) - seg_starts = torch.searchsorted(sorted_wi, adapter_ids) - seg_ends = torch.searchsorted(sorted_wi, adapter_ids, right=True) - seg_lens = seg_ends - seg_starts - - seg_indptr = torch.zeros(max_loras + 1, dtype=torch.int32, device=device) - seg_indptr[1:] = torch.cumsum(seg_lens, dim=0) - - return LoRABatchInfo( - bs=max_loras, - use_cuda_graph=False, - num_segments=max_loras, - seg_lens=seg_lens, - seg_indptr=seg_indptr, - max_len=bs, - weight_indices=adapter_ids, - lora_ranks=torch.tensor(lora_ranks, dtype=torch.int32, device=device), - scalings=torch.tensor(scalings, dtype=torch.float, device=device), - permutation=perm, - ) - - -def _check_close( - a: torch.Tensor, b: torch.Tensor, name: str, atol: float = 1e-4, rtol: float = 1e-3 -) -> None: - diff = (a - b).abs().max().item() - assert torch.allclose(a, b, atol=atol, rtol=rtol), f"{name}: max diff = {diff}" - - -def test_sgemm_lora_a(): - from sglang.srt.lora.triton_ops import sgemm_lora_a_fwd - - torch.manual_seed(42) - bs, input_dim, rank, num_loras = 8, 256, 16, 3 - x = torch.randn(bs, input_dim, device="cuda", dtype=torch.bfloat16) - weights = torch.randn( - num_loras, rank, input_dim, device="cuda", dtype=torch.bfloat16 - ) - wi = [i % num_loras for i in range(bs)] - lora_ranks = [rank] * num_loras - scalings = [1.0] * num_loras - - bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings) - bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras) - - out_plain = sgemm_lora_a_fwd(x, weights, bi_plain) - out_sorted = sgemm_lora_a_fwd(x, weights, bi_sorted) - _check_close(out_plain, out_sorted, "sgemm_lora_a") - - -def test_sgemm_lora_b(): - from sglang.srt.lora.triton_ops import sgemm_lora_b_fwd - - torch.manual_seed(42) - bs, output_dim, rank, num_loras = 8, 256, 16, 3 - x = torch.randn(bs, rank, device="cuda", dtype=torch.bfloat16) - weights = torch.randn( - num_loras, output_dim, rank, device="cuda", dtype=torch.bfloat16 - ) - wi = [i % num_loras for i in range(bs)] - lora_ranks = [rank] * num_loras - scalings = [0.5] * num_loras - - bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings) - bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras) - - base_plain = torch.randn(bs, output_dim, device="cuda", dtype=torch.bfloat16) - base_sorted = base_plain.clone() - - out_plain = sgemm_lora_b_fwd(x, weights, bi_plain, base_plain) - out_sorted = sgemm_lora_b_fwd(x, weights, bi_sorted, base_sorted) - _check_close(out_plain, out_sorted, "sgemm_lora_b") - - -def test_qkv_lora_b(): - from sglang.srt.lora.triton_ops import qkv_lora_b_fwd - - torch.manual_seed(42) - bs, rank, num_loras = 8, 16, 3 - n_slices = 3 - q_dim, kv_dim = 128, 64 - total_out = q_dim + 2 * kv_dim - x = torch.randn(bs, n_slices * rank, device="cuda", dtype=torch.bfloat16) - weights = torch.randn( - num_loras, total_out, rank, device="cuda", dtype=torch.bfloat16 - ) - output_offset = torch.tensor( - [0, q_dim, q_dim + kv_dim, total_out], device="cuda", dtype=torch.int32 - ) - wi = [i % num_loras for i in range(bs)] - lora_ranks = [rank] * num_loras - scalings = [1.0] * num_loras - - bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings) - bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras) - - base_plain = torch.randn(bs, total_out, device="cuda", dtype=torch.bfloat16) - base_sorted = base_plain.clone() - - max_qkv_out_dim = max(q_dim, kv_dim) - out_plain = qkv_lora_b_fwd( - x, weights, bi_plain, output_offset, max_qkv_out_dim, base_plain - ) - out_sorted = qkv_lora_b_fwd( - x, weights, bi_sorted, output_offset, max_qkv_out_dim, base_sorted - ) - _check_close(out_plain, out_sorted, "qkv_lora_b") - - -def test_gate_up_lora_b(): - from sglang.srt.lora.triton_ops import gate_up_lora_b_fwd - - torch.manual_seed(42) - bs, rank, num_loras = 8, 16, 3 - output_dim = 128 - x = torch.randn(bs, 2 * rank, device="cuda", dtype=torch.bfloat16) - weights = torch.randn( - num_loras, 2 * output_dim, rank, device="cuda", dtype=torch.bfloat16 - ) - wi = [i % num_loras for i in range(bs)] - lora_ranks = [rank] * num_loras - scalings = [1.0] * num_loras - - bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings) - bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras) - - base_plain = torch.randn(bs, 2 * output_dim, device="cuda", dtype=torch.bfloat16) - base_sorted = base_plain.clone() - - out_plain = gate_up_lora_b_fwd(x, weights, bi_plain, output_dim, base_plain) - out_sorted = gate_up_lora_b_fwd(x, weights, bi_sorted, output_dim, base_sorted) - _check_close(out_plain, out_sorted, "gate_up_lora_b") - - -def test_mixed_ranks(): - """Test with different LoRA ranks per adapter.""" - from sglang.srt.lora.triton_ops import sgemm_lora_a_fwd - - torch.manual_seed(42) - bs, input_dim, num_loras = 12, 256, 4 - max_rank = 32 - lora_ranks = [8, 16, 32, 16] - scalings = [0.25, 0.5, 1.0, 2.0] - # Use max_rank for weight shape, kernel handles per-adapter rank - weights = torch.randn( - num_loras, max_rank, input_dim, device="cuda", dtype=torch.bfloat16 - ) - x = torch.randn(bs, input_dim, device="cuda", dtype=torch.bfloat16) - wi = [i % num_loras for i in range(bs)] - - bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings) - bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras) - - out_plain = sgemm_lora_a_fwd(x, weights, bi_plain) - out_sorted = sgemm_lora_a_fwd(x, weights, bi_sorted) - _check_close(out_plain, out_sorted, "sgemm_lora_a_mixed_ranks") - - -def test_single_adapter(): - """All sequences use the same adapter.""" - from sglang.srt.lora.triton_ops import sgemm_lora_a_fwd - - torch.manual_seed(42) - bs, input_dim, rank, num_loras = 16, 256, 16, 2 - x = torch.randn(bs, input_dim, device="cuda", dtype=torch.bfloat16) - weights = torch.randn( - num_loras, rank, input_dim, device="cuda", dtype=torch.bfloat16 - ) - wi = [0] * bs # all adapter 0 - lora_ranks = [rank, rank] - scalings = [1.0, 1.0] - - bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings) - bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras) - - out_plain = sgemm_lora_a_fwd(x, weights, bi_plain) - out_sorted = sgemm_lora_a_fwd(x, weights, bi_sorted) - _check_close(out_plain, out_sorted, "sgemm_lora_a_single_adapter") - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/test/registered/unit/test_no_bare_pytest_main.py b/test/registered/unit/test_no_bare_pytest_main.py new file mode 100644 index 000000000..5dda5c34f --- /dev/null +++ b/test/registered/unit/test_no_bare_pytest_main.py @@ -0,0 +1,90 @@ +import ast +import pathlib +import unittest + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="stage-a-test-cpu") + + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +_SCAN_ROOTS = [_REPO_ROOT / "python", _REPO_ROOT / "test"] + + +class TestNoBarePytestMain(CustomTestCase): + def test_no_bare_pytest_main_in_repo(self): + offenders = [] + for root in _SCAN_ROOTS: + if not root.exists(): + continue + for path in root.rglob("*.py"): + violation = _find_bare_pytest_main(path) + if violation is not None: + offenders.append(violation) + + self.assertFalse( + offenders, + msg=( + "Found bare `pytest.main(...)` in __main__ blocks (must be " + "wrapped in sys.exit(...) so failing tests propagate the exit " + "code to the CI runner):\n " + "\n ".join(offenders) + ), + ) + + +def _find_bare_pytest_main(path: pathlib.Path): + """Return `:` if `path` has a bare pytest.main(...) call + inside `if __name__ == "__main__":`, else None.""" + try: + source = path.read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError): + return None + try: + tree = ast.parse(source, filename=str(path)) + except SyntaxError: + return None + + for node in ast.walk(tree): + if not isinstance(node, ast.If): + continue + if not _is_main_guard(node.test): + continue + for stmt in node.body: + if _is_bare_pytest_main_call(stmt): + rel = path.relative_to(_REPO_ROOT) + return f"{rel}:{stmt.lineno}" + return None + + +def _is_main_guard(test: ast.expr) -> bool: + """Match `__name__ == "__main__"` (either side).""" + if not isinstance(test, ast.Compare) or len(test.ops) != 1: + return False + if not isinstance(test.ops[0], ast.Eq): + return False + sides = [test.left, *test.comparators] + has_name = any(isinstance(s, ast.Name) and s.id == "__name__" for s in sides) + has_main = any(isinstance(s, ast.Constant) and s.value == "__main__" for s in sides) + return has_name and has_main + + +def _is_bare_pytest_main_call(stmt: ast.stmt) -> bool: + """Match `pytest.main(...)` whose return value is discarded. + `sys.exit(pytest.main(...))` and `code = pytest.main(...)` are fine.""" + if not isinstance(stmt, ast.Expr): + return False + call = stmt.value + if not isinstance(call, ast.Call): + return False + func = call.func + return ( + isinstance(func, ast.Attribute) + and func.attr == "main" + and isinstance(func.value, ast.Name) + and func.value.id == "pytest" + ) + + +if __name__ == "__main__": + unittest.main()