minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Cleanup] Deduplicate kernel tests, diffusion fixtures and benchmark helpers (#40265)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-09-19 19:45:38 +08:00
committed by GitHub
parent 0b0d2c257a
commit cb22f2451e
25 changed files with 512 additions and 1936 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/benchmark/deepseek_utils.py
+65
View File
@@ -0,0 +1,65 @@
"""Shared DeepSeek-V3 benchmark shapes and FP8 input preparation."""
from typing import Tuple
import torch
from triton import cdiv
def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2 and x.size(1) % 128 == 0
m, n = x.shape
x_view = x.view(m, -1, 128)
x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
return (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view(
m, n
), (x_amax / 448.0).view(m, -1)
def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2
m, n = x.shape
x_padded = torch.zeros(
(cdiv(m, 128) * 128, cdiv(n, 128) * 128), dtype=x.dtype, device=x.device
)
x_padded[:m, :n] = x
x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view(
x_view.size(0), x_view.size(2)
)
def get_weight_shapes(tp_size):
"""Return the DeepSeek-V3 (N, K) shapes, including TP-sharded projections."""
# cannot TP
total = [
(512 + 64, 7168),
((128 + 64) * 128, 7168),
(128 * (128 + 128), 512),
(7168, 16384),
(7168, 18432),
]
# N can TP
n_tp = [
(18432 * 2, 7168),
((128 + 64) * 128, 7168),
(128 * (128 + 128), 512),
(24576, 1536),
(4096, 7168),
]
# K can TP
k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)]
weight_shapes = []
for t in total:
weight_shapes.append(t)
for n_t in n_tp:
new_t = (n_t[0] // tp_size, n_t[1])
weight_shapes.append(new_t)
for k_t in k_tp:
new_t = (k_t[0], k_t[1] // tp_size)
weight_shapes.append(new_t)
return weight_shapes
python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py
+2 -21
View File
@@ -15,7 +15,6 @@ import pytest
import torch
from PIL import Image
from sglang.multimodal_gen.runtime import server_args as _sa_mod
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.realtime_chain import (
SanaWMCameraCondStage,
SanaWMNoiseState,
@@ -29,7 +28,6 @@ from sglang.multimodal_gen.runtime.realtime.session import RealtimeSession
from sglang.multimodal_gen.runtime.realtime.states import (
get_realtime_causal_dit_state,
)
from sglang.multimodal_gen.runtime.server_args import set_global_server_args
MC = 8
@@ -39,23 +37,6 @@ class _TestRealtimeStage(SanaWMRealtimeStage):
raise NotImplementedError
@pytest.fixture
def _global_args():
prev = _sa_mod._global_server_args
set_global_server_args(
SimpleNamespace(
comfyui_mode=False,
enable_cfg_parallel=False,
enable_torch_compile=False,
attention_backend=None,
)
)
try:
yield
finally:
set_global_server_args(prev)
def _prep_stage():
return SanaWMRealtimeLatentPrepStage(
use_refiner=True, transformer=None, vae=None, model_path=""
@@ -153,7 +134,7 @@ def test_realtime_camera_conditioning_uses_requested_size():
assert plucker.shape == (1, 48, 3, 15, 26)
def test_latent_prep_plan_and_noise_discipline(_global_args):
def test_latent_prep_plan_and_noise_discipline():
stage = _prep_stage()
session = RealtimeSession()
fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32)
@@ -193,7 +174,7 @@ def test_latent_prep_plan_and_noise_discipline(_global_args):
assert torch.isfinite(batch.latents).all()
def test_latent_prep_open_ended_uniform_chunk0(_global_args):
def test_latent_prep_open_ended_uniform_chunk0():
stage = _prep_stage()
session = RealtimeSession()
fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32)
python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_forward_long.py
+3 -23
View File
@@ -12,14 +12,12 @@ from __future__ import annotations
from types import SimpleNamespace
import pytest
import torch
from sglang.multimodal_gen.configs.models.dits.sana_wm import (
SanaWMArchConfig,
SanaWMConfig,
)
from sglang.multimodal_gen.runtime import server_args as _sa_mod
from sglang.multimodal_gen.runtime.models.dits.sana_wm import (
_CACHE_TYPE_STATE,
_SLOT_CAM_K,
@@ -35,7 +33,6 @@ from sglang.multimodal_gen.runtime.models.dits.sana_wm import (
_slice_rope_to_current_chunk,
process_camera_conditions_ucpe,
)
from sglang.multimodal_gen.runtime.server_args import set_global_server_args
HEAD_DIM = 112
H, W = 2, 3
@@ -236,23 +233,6 @@ def test_forward_long_gdn_reduces_to_dense_with_camera():
# --------------------------------------------------------------------------- #
@pytest.fixture
def _global_args():
prev = _sa_mod._global_server_args
set_global_server_args(
SimpleNamespace(
comfyui_mode=False,
enable_cfg_parallel=False,
enable_torch_compile=False,
attention_backend=None,
)
)
try:
yield
finally:
set_global_server_args(prev)
class _ZeroCross(torch.nn.Module):
def forward(self, x, y, mask=None):
return torch.zeros_like(x)
@@ -284,7 +264,7 @@ def _block():
return b
def test_block_forward_long_reduces_to_dense(_global_args):
def test_block_forward_long_reduces_to_dense():
block = _block()
x = _x()
y = torch.randn(AB, 4, AC, dtype=torch.float64)
@@ -349,7 +329,7 @@ def _model_inputs():
)
def test_model_forward_long_single_chunk_reduces_to_dense(_global_args):
def test_model_forward_long_single_chunk_reduces_to_dense():
m = _tiny_model()
inp = _model_inputs()
with torch.no_grad():
@@ -364,7 +344,7 @@ def test_model_forward_long_single_chunk_reduces_to_dense(_global_args):
assert cache[0][_SLOT_FFN_TCONV] is not None
def test_model_forward_long_two_chunks_runs_and_windows(_global_args):
def test_model_forward_long_two_chunks_runs_and_windows():
m = _tiny_model()
inp = _model_inputs()
split = 2
python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_stage.py
+1 -23
View File
@@ -9,16 +9,12 @@ dim=2 would concat the head axis and silently corrupt every softmax block).
from __future__ import annotations
from types import SimpleNamespace
import pytest
import torch
from sglang.multimodal_gen.configs.models.dits.sana_wm import (
SanaWMArchConfig,
SanaWMConfig,
)
from sglang.multimodal_gen.runtime import server_args as _sa_mod
from sglang.multimodal_gen.runtime.models.dits.sana_wm import (
_CACHE_TYPE_CONCAT,
_CACHE_TYPE_STATE,
@@ -35,7 +31,6 @@ from sglang.multimodal_gen.runtime.models.dits.sana_wm import (
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import (
SanaWMStreamingDenoisingStage as Stage,
)
from sglang.multimodal_gen.runtime.server_args import set_global_server_args
B, Hh, D = 1, 2, 4
@@ -160,23 +155,6 @@ class _ZeroCross(torch.nn.Module):
return torch.zeros_like(x)
@pytest.fixture
def _global_args():
prev = _sa_mod._global_server_args
set_global_server_args(
SimpleNamespace(
comfyui_mode=False,
enable_cfg_parallel=False,
enable_torch_compile=False,
attention_backend=None,
)
)
try:
yield
finally:
set_global_server_args(prev)
def _depth4_model():
arch = SanaWMArchConfig(
in_channels=MC,
@@ -201,7 +179,7 @@ def _depth4_model():
return m
def test_streaming_loop_runs_and_accumulates_concat_block(_global_args):
def test_streaming_loop_runs_and_accumulates_concat_block():
"""Drive forward_long chunk-by-chunk (the stage's core loop) on a depth-4
model: accumulate -> denoise(save=False) -> clean(save=True). Verify finite
output, threaded GDN state, and that the softmax block (idx 3) accumulates a
python/sglang/multimodal_gen/test/unit/test_fp32_layernorm.py
+24 -45
View File
@@ -4,61 +4,40 @@ import torch.nn.functional as F
from sglang.multimodal_gen.runtime.layers.layernorm import FP32LayerNorm
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_fp32_layernorm_cache_matches_reference():
def test_fp32_layernorm_cache_reuse_and_invalidation():
norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16)
inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
actual = norm(inputs)
expected = F.layer_norm(
inputs.float(),
norm.normalized_shape,
norm.weight.float().to(device=inputs.device),
norm.bias.float().to(device=inputs.device),
norm.eps,
).to(inputs.dtype)
for updated in (False, True):
if updated:
previous = norm.__dict__["_weight_fp32_cache"]
norm.weight.add_(1.0)
actual = norm(inputs)
expected = F.layer_norm(
inputs.float(),
norm.normalized_shape,
norm.weight.float(),
norm.bias.float(),
norm.eps,
).to(inputs.dtype)
torch.testing.assert_close(actual, expected)
weight_cache = norm.__dict__["_weight_fp32_cache"]
bias_cache = norm.__dict__["_bias_fp32_cache"]
if updated:
assert weight_cache[0] != previous[0]
assert weight_cache[1] is not previous[1]
norm(inputs)
assert norm.__dict__["_weight_fp32_cache"][1] is weight_cache[1]
assert norm.__dict__["_bias_fp32_cache"][1] is bias_cache[1]
torch.testing.assert_close(actual, expected)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_fp32_layernorm_cache_reuses_converted_params():
norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16)
inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
norm(inputs)
weight_cache = norm.__dict__["_weight_fp32_cache"]
bias_cache = norm.__dict__["_bias_fp32_cache"]
norm(inputs)
assert norm.__dict__["_weight_fp32_cache"][1] is weight_cache[1]
assert norm.__dict__["_bias_fp32_cache"][1] is bias_cache[1]
assert "_weight_fp32_cache" not in norm.state_dict()
assert "_bias_fp32_cache" not in norm.state_dict()
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_fp32_layernorm_cache_invalidates_on_param_update():
norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16)
inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
norm(inputs)
first_key, first_weight = norm.__dict__["_weight_fp32_cache"]
norm.weight.add_(1.0)
norm(inputs)
second_key, second_weight = norm.__dict__["_weight_fp32_cache"]
assert second_key != first_key
assert second_weight is not first_weight
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_fp32_layernorm_grad_mode_preserves_autograd_path():
norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16)
inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True)
python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py
+63 -170
View File
@@ -10,7 +10,6 @@ from sglang.kernels.ops.diffusion import apply_group_norm_silu
from sglang.multimodal_gen.runtime.models.upsampler.latent_upsampler import (
LatentUpsampler,
ResBlock,
SpatialRationalResampler,
)
@@ -21,199 +20,87 @@ def _resblock_eager_reference(block: ResBlock, x: torch.Tensor) -> torch.Tensor:
return block.activation(x + residual)
def _latent_upsampler_eager_reference(
upsampler: LatentUpsampler, latent: torch.Tensor
) -> torch.Tensor:
from einops import rearrange
b, _, f, _, _ = latent.shape
if upsampler.dims == 2:
x = rearrange(latent, "b c f h w -> (b f) c h w")
x = upsampler.initial_activation(
upsampler.initial_norm(upsampler.initial_conv(x))
)
for block in upsampler.res_blocks:
x = _resblock_eager_reference(block, x)
x = upsampler.upsampler(x)
for block in upsampler.post_upsample_res_blocks:
x = _resblock_eager_reference(block, x)
x = upsampler.final_conv(x)
return rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
x = upsampler.initial_activation(
upsampler.initial_norm(upsampler.initial_conv(latent))
)
for block in upsampler.res_blocks:
x = _resblock_eager_reference(block, x)
if upsampler.temporal_upsample:
x = upsampler.upsampler(x)[:, :, 1:, :, :]
elif isinstance(upsampler.upsampler, SpatialRationalResampler):
x = upsampler.upsampler(x)
else:
x = rearrange(x, "b c f h w -> (b f) c h w")
x = upsampler.upsampler(x)
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
for block in upsampler.post_upsample_res_blocks:
x = _resblock_eager_reference(block, x)
return upsampler.final_conv(x)
@pytest.mark.parametrize(
"batch,channels,dims,spatial",
[
(1, 64, 2, (16, 16)),
(2, 64, 2, (8, 24)),
(1, 128, 3, (2, 8, 8)),
],
)
def test_resblock_forward_parity(batch, channels, dims, spatial):
torch.manual_seed(0)
block = ResBlock(channels=channels, dims=dims).eval()
torch.manual_seed(1)
x = torch.randn(batch, channels, *spatial, dtype=torch.float32)
with torch.no_grad():
out = block(x)
ref = _resblock_eager_reference(block, x)
torch.testing.assert_close(out, ref, atol=0.0, rtol=0.0)
@pytest.mark.parametrize(
"dims,latent_shape,mid_channels,num_blocks_per_stage,rational_resampler",
[
(2, (1, 32, 2, 16, 16), 64, 2, False),
(3, (1, 32, 2, 16, 16), 64, 2, False),
(3, (1, 32, 2, 16, 16), 64, 2, True),
],
)
def test_latent_upsampler_forward_parity(
dims, latent_shape, mid_channels, num_blocks_per_stage, rational_resampler
):
torch.manual_seed(2)
upsampler = LatentUpsampler(
in_channels=latent_shape[1],
mid_channels=mid_channels,
num_blocks_per_stage=num_blocks_per_stage,
dims=dims,
spatial_upsample=True,
temporal_upsample=False,
spatial_scale=2.0,
rational_resampler=rational_resampler,
).eval()
torch.manual_seed(3)
latent = torch.randn(*latent_shape, dtype=torch.float32)
with torch.no_grad():
out = upsampler(latent)
ref = _latent_upsampler_eager_reference(upsampler, latent)
torch.testing.assert_close(out, ref, atol=0.0, rtol=0.0)
def test_resblock_fuses_exactly_one_site():
torch.manual_seed(4)
block = ResBlock(channels=64, dims=2).eval()
x = torch.randn(1, 64, 16, 16, dtype=torch.float32)
def _latent_upsampler_eager_reference(upsampler, latent):
with patch.object(
lu_mod, "apply_group_norm_silu", wraps=lu_mod.apply_group_norm_silu
) as spy:
with torch.no_grad():
block(x)
assert spy.call_count == 1
lu_mod, "apply_group_norm_silu", side_effect=lambda x, norm, act: act(norm(x))
):
return upsampler(latent)
@pytest.mark.parametrize(
"dims,num_blocks_per_stage,rational_resampler,expected_calls",
[
(2, 2, False, 1 + 2 * 2),
(2, 4, False, 1 + 4 * 2),
(3, 2, False, 1 + 2 * 2),
(3, 2, True, 1 + 2 * 2),
],
)
def test_latent_upsampler_fuses_expected_sites(
dims, num_blocks_per_stage, rational_resampler, expected_calls
):
torch.manual_seed(5)
upsampler = LatentUpsampler(
in_channels=32,
mid_channels=64,
num_blocks_per_stage=num_blocks_per_stage,
dims=dims,
spatial_upsample=True,
temporal_upsample=False,
spatial_scale=2.0,
rational_resampler=rational_resampler,
).eval()
latent = torch.randn(1, 32, 2, 16, 16, dtype=torch.float32)
with patch.object(
lu_mod, "apply_group_norm_silu", wraps=lu_mod.apply_group_norm_silu
) as spy:
with torch.no_grad():
upsampler(latent)
assert spy.call_count == expected_calls
# CUDA Triton fast path -------------------------------------------------------
requires_cuda = pytest.mark.skipif(
not torch.cuda.is_available(),
reason="Triton fused group_norm_silu requires CUDA",
not torch.cuda.is_available(), reason="CUDA required"
)
# bf16 keeps kernel-level tolerance because its fp32-equivalent exponent range
# absorbs multi-layer conv drift; fp16 needs a looser tolerance on the e2e
# upsampler test where 8+ downstream convs amplify fused-vs-eager rounding.
_RESBLOCK_TOL = {torch.bfloat16: (7e-2, 2e-2), torch.float16: (3e-3, 3e-3)}
_UPSAMPLER_TOL = {torch.bfloat16: (7e-2, 2e-2), torch.float16: (2e-2, 1e-1)}
def _parity_cases(common, cpu_only):
return [
pytest.param(
device, dtype, *case, marks=requires_cuda if device == "cuda" else ()
)
for device, dtype in [
("cpu", torch.float32),
("cuda", torch.bfloat16),
("cuda", torch.float16),
]
for case in common + (cpu_only if device == "cpu" else [])
]
# Downstream convolutions amplify fp16 drift in the complete upsampler.
_RESBLOCK_TOL = {
torch.float32: (0, 0),
torch.bfloat16: (7e-2, 2e-2),
torch.float16: (3e-3, 3e-3),
}
_UPSAMPLER_TOL = {
torch.float32: (0, 0),
torch.bfloat16: (7e-2, 2e-2),
torch.float16: (2e-2, 1e-1),
}
@requires_cuda
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize(
"batch,channels,dims,spatial",
[
(1, 64, 2, (16, 16)),
(1, 128, 3, (2, 8, 8)),
],
"device,dtype,batch,channels,dims,spatial",
_parity_cases(
[(1, 64, 2, (16, 16)), (1, 128, 3, (2, 8, 8))],
cpu_only=[(2, 64, 2, (8, 24))],
),
)
def test_resblock_forward_parity_cuda(dtype, batch, channels, dims, spatial):
def test_resblock_forward_parity(device, dtype, batch, channels, dims, spatial):
torch.manual_seed(0)
device = torch.device("cuda")
block = ResBlock(channels=channels, dims=dims).to(device=device, dtype=dtype).eval()
torch.manual_seed(1)
x = torch.randn(batch, channels, *spatial, device=device, dtype=dtype)
with torch.no_grad():
with (
torch.no_grad(),
patch.object(
lu_mod, "apply_group_norm_silu", wraps=lu_mod.apply_group_norm_silu
) as fused,
):
out = block(x)
assert fused.call_count == 1
with torch.no_grad():
ref = _resblock_eager_reference(block, x)
atol, rtol = _RESBLOCK_TOL[dtype]
torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)
@requires_cuda
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize(
"dims,latent_shape,mid_channels,num_blocks_per_stage,rational_resampler",
[
(2, (1, 32, 2, 16, 16), 64, 2, False),
(3, (1, 32, 2, 16, 16), 64, 2, False),
(3, (1, 32, 2, 16, 16), 64, 2, True),
],
"device,dtype,dims,num_blocks_per_stage,rational_resampler",
_parity_cases(
[(2, 2, False), (3, 2, False), (3, 2, True)],
cpu_only=[(2, 4, False)],
),
)
def test_latent_upsampler_forward_parity_cuda(
dtype, dims, latent_shape, mid_channels, num_blocks_per_stage, rational_resampler
def test_latent_upsampler_forward_parity(
device, dtype, dims, num_blocks_per_stage, rational_resampler
):
torch.manual_seed(2)
device = torch.device("cuda")
upsampler = (
LatentUpsampler(
in_channels=latent_shape[1],
mid_channels=mid_channels,
in_channels=32,
mid_channels=64,
num_blocks_per_stage=num_blocks_per_stage,
dims=dims,
spatial_upsample=True,
@@ -225,12 +112,18 @@ def test_latent_upsampler_forward_parity_cuda(
.eval()
)
torch.manual_seed(3)
latent = torch.randn(*latent_shape, device=device, dtype=dtype)
latent = torch.randn(1, 32, 2, 16, 16, device=device, dtype=dtype)
with torch.no_grad():
with (
torch.no_grad(),
patch.object(
lu_mod, "apply_group_norm_silu", wraps=lu_mod.apply_group_norm_silu
) as fused,
):
out = upsampler(latent)
assert fused.call_count == 1 + 2 * num_blocks_per_stage
with torch.no_grad():
ref = _latent_upsampler_eager_reference(upsampler, latent)
atol, rtol = _UPSAMPLER_TOL[dtype]
torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)
python/sglang/multimodal_gen/test/unit/test_layernorm_cutedsl_dispatch.py
+11 -27
View File
@@ -14,41 +14,25 @@ _CUTEDSL_MODULE = "sglang.kernels.ops.diffusion.norm.scale_residual_norm_cutedsl
@pytest.mark.parametrize("hidden_size", [257, 8448])
def test_norm_scale_shift_cuda_falls_back_for_unsupported_hidden_size(hidden_size):
layer = RMSNormScaleShift(hidden_size)
x = torch.empty(1, 1, hidden_size)
shift = torch.empty(1, 1, hidden_size)
scale = torch.empty(1, 1, hidden_size)
@pytest.mark.parametrize(
"layer_cls,num_inputs",
[(RMSNormScaleShift, 3), (ScaleResidualRMSNormScaleShift, 5)],
)
def test_cuda_falls_back_for_unsupported_hidden_size(
hidden_size, layer_cls, num_inputs
):
layer = layer_cls(hidden_size)
inputs = [torch.empty(1, 1, hidden_size) for _ in range(num_inputs)]
expected = object()
with (
patch.object(layer, "forward_native", return_value=expected) as native,
pytest.warns(UserWarning, match="native fallback"),
):
actual = layer.forward_cuda(x, shift, scale)
actual = layer.forward_cuda(*inputs)
assert actual is expected
native.assert_called_once_with(x, shift, scale)
@pytest.mark.parametrize("hidden_size", [257, 8448])
def test_scale_residual_cuda_falls_back_for_unsupported_hidden_size(hidden_size):
layer = ScaleResidualRMSNormScaleShift(hidden_size)
residual = torch.empty(1, 1, hidden_size)
x = torch.empty(1, 1, hidden_size)
gate = torch.empty(1, 1, hidden_size)
shift = torch.empty(1, 1, hidden_size)
scale = torch.empty(1, 1, hidden_size)
expected = object()
with (
patch.object(layer, "forward_native", return_value=expected) as native,
pytest.warns(UserWarning, match="native fallback"),
):
actual = layer.forward_cuda(residual, x, gate, shift, scale)
assert actual is expected
native.assert_called_once_with(residual, x, gate, shift, scale)
native.assert_called_once_with(*inputs)
def test_norm_scale_shift_cuda_uses_cutedsl_for_supported_hidden_size(monkeypatch):
python/sglang/multimodal_gen/test/unit/test_vae_loader.py
+15 -48
View File
@@ -647,59 +647,26 @@ class TestVAELoader(unittest.TestCase):
self.assertNotIn("latents_mean", loaded)
self.assertNotIn("latents_std", loaded)
def test_channels_last_3d_defaults_true_for_qwen_image_on_cuda(self):
def test_channels_last_3d_cuda_model_defaults(self):
cases = [
(QwenImagePipelineConfig, 1, "vae", True),
(WanT2V480PConfig, 1, "video_vae", True),
(FastWan2_2_TI2V_5B_Config, 1, "video_vae", True),
(Wan2_2_I2V_A14B_Config, 2, "video_vae", False),
(LTX2PipelineConfig, 1, "video_vae", True),
(LTX2PipelineConfig, 2, "video_vae", False),
]
with (
patch.dict("os.environ", {}, clear=True),
patch.object(vae_loader.current_platform, "is_cuda", return_value=True),
patch.object(vae_loader.current_platform, "is_rocm", return_value=False),
):
server_args = _FakeServerArgs(QwenImagePipelineConfig())
self.assertTrue(_should_use_channels_last_3d(server_args, "vae"))
def test_channels_last_3d_defaults_true_for_single_gpu_wan_on_cuda(self):
with (
patch.dict("os.environ", {}, clear=True),
patch.object(vae_loader.current_platform, "is_cuda", return_value=True),
patch.object(vae_loader.current_platform, "is_rocm", return_value=False),
):
server_args = _FakeServerArgs(WanT2V480PConfig(), num_gpus=1)
self.assertTrue(_should_use_channels_last_3d(server_args, "video_vae"))
def test_channels_last_3d_defaults_true_for_single_gpu_fast_wan_on_cuda(self):
with (
patch.dict("os.environ", {}, clear=True),
patch.object(vae_loader.current_platform, "is_cuda", return_value=True),
patch.object(vae_loader.current_platform, "is_rocm", return_value=False),
):
server_args = _FakeServerArgs(FastWan2_2_TI2V_5B_Config(), num_gpus=1)
self.assertTrue(_should_use_channels_last_3d(server_args, "video_vae"))
def test_channels_last_3d_defaults_false_for_multi_gpu_wan_on_cuda(self):
with (
patch.dict("os.environ", {}, clear=True),
patch.object(vae_loader.current_platform, "is_cuda", return_value=True),
patch.object(vae_loader.current_platform, "is_rocm", return_value=False),
):
server_args = _FakeServerArgs(Wan2_2_I2V_A14B_Config(), num_gpus=2)
self.assertFalse(_should_use_channels_last_3d(server_args, "video_vae"))
def test_channels_last_3d_defaults_true_for_single_gpu_ltx_on_cuda(self):
with (
patch.dict("os.environ", {}, clear=True),
patch.object(vae_loader.current_platform, "is_cuda", return_value=True),
patch.object(vae_loader.current_platform, "is_rocm", return_value=False),
):
server_args = _FakeServerArgs(LTX2PipelineConfig(), num_gpus=1)
self.assertTrue(_should_use_channels_last_3d(server_args, "video_vae"))
def test_channels_last_3d_defaults_false_for_multi_gpu_ltx_on_cuda(self):
with (
patch.dict("os.environ", {}, clear=True),
patch.object(vae_loader.current_platform, "is_cuda", return_value=True),
patch.object(vae_loader.current_platform, "is_rocm", return_value=False),
):
server_args = _FakeServerArgs(LTX2PipelineConfig(), num_gpus=2)
self.assertFalse(_should_use_channels_last_3d(server_args, "video_vae"))
for config_cls, num_gpus, component, expected in cases:
with self.subTest(config=config_cls.__name__, num_gpus=num_gpus):
server_args = _FakeServerArgs(config_cls(), num_gpus=num_gpus)
self.assertEqual(
_should_use_channels_last_3d(server_args, component), expected
)
def test_channels_last_3d_can_be_disabled_by_env(self):
with (
python/sglang/test/kernels/paged_mqa.py
+140
View File
@@ -0,0 +1,140 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Shared FP8 paged MQA reference and inputs for the CuTe DSL and DeepGEMM tests."""
import torch
from sglang.srt.layers.attention.dsa.utils import (
fp8_mqa_logits_ceil_to_ue8m0,
fp8_mqa_logits_make_fused_kv,
)
BLOCK_KV = 64
HEAD_DIM = 128
def ref_fp8_paged_mqa_logits(
q_fp8,
kv_fp8,
kv_scales,
weights,
context_lens,
block_table,
max_model_len,
block_kv,
):
B, next_n, H, D = q_fp8.shape
device = q_fp8.device
logits = torch.full(
(B * next_n, max_model_len), float("-inf"), device=device, dtype=torch.float32
)
q_f32 = q_fp8.float()
for b in range(B):
ctx_len = context_lens[b].item()
q_positions = torch.arange(ctx_len - next_n, ctx_len, device=device)
w = weights[b * next_n : (b + 1) * next_n, :]
for blk_idx in range((ctx_len + block_kv - 1) // block_kv):
phys_blk = block_table[b, blk_idx].item()
k_f32 = kv_fp8[phys_blk].float()
scales = kv_scales[phys_blk]
k_positions = torch.arange(
blk_idx * block_kv, (blk_idx + 1) * block_kv, device=device
)
mask = (k_positions[None, :] < ctx_len) & (
k_positions[None, :] <= q_positions[:, None]
)
qk = torch.matmul(q_f32[b].permute(1, 0, 2), k_f32.T)
qk = torch.where(mask[None, :, :], qk, torch.zeros(1, device=device))
qk = torch.relu(qk)
weighted = (w.T[:, :, None] * qk).sum(dim=0)
weighted = weighted * scales[None, :]
start_pos = blk_idx * block_kv
end_pos = start_pos + block_kv
logits[b * next_n : (b + 1) * next_n, start_pos:end_pos] = torch.where(
mask,
weighted,
torch.tensor(float("-inf"), device=device, dtype=torch.float32),
)
return logits
def generate_paged_mqa_test_data(
batch_size,
next_n,
num_heads,
avg_context_len,
max_model_len,
device="cuda",
):
torch.manual_seed(42)
torch.cuda.manual_seed(42)
context_lens = torch.randint(
max(BLOCK_KV, int(0.7 * avg_context_len)),
int(1.3 * avg_context_len) + 1,
(batch_size,),
dtype=torch.int32,
device="cpu",
).clamp(max=max_model_len)
max_blocks_per_seq = (max_model_len + BLOCK_KV - 1) // BLOCK_KV
total_blocks = ((context_lens + BLOCK_KV - 1) // BLOCK_KV).sum().item()
num_phys_blocks = total_blocks + batch_size * 2
block_table = torch.full(
(batch_size, max_blocks_per_seq), 0, dtype=torch.int32, device=device
)
blk_offset = 0
for i in range(batch_size):
n_blks = (context_lens[i].item() + BLOCK_KV - 1) // BLOCK_KV
block_table[i, :n_blks] = torch.arange(
blk_offset, blk_offset + n_blks, dtype=torch.int32, device=device
)
blk_offset += n_blks
q_bf16 = torch.randn(batch_size, next_n, num_heads, HEAD_DIM, device=device)
q_fp8 = q_bf16.to(torch.float8_e4m3fn)
kv_bf16 = torch.randn(num_phys_blocks, BLOCK_KV, HEAD_DIM, device=device)
kv_amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp(1e-4)
kv_scale = fp8_mqa_logits_ceil_to_ue8m0(kv_amax / 448.0).squeeze(-1)
kv_fp8 = (kv_bf16 / kv_scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
weights = torch.randn(
batch_size * next_n, num_heads, device=device, dtype=torch.float32
)
kv_fused = fp8_mqa_logits_make_fused_kv(kv_fp8, kv_scale, BLOCK_KV, HEAD_DIM)
return {
"q_fp8": q_fp8,
"kv_fp8": kv_fp8,
"kv_scales": kv_scale,
"kv_fused": kv_fused,
"weights": weights,
"context_lens": context_lens.to(device),
"block_table": block_table,
}
def assert_paged_mqa_matches_ref(
logits, ref_logits, context_lens, B, next_n, max_model_len
):
device = logits.device
positions = torch.arange(max_model_len, device=device).unsqueeze(0)
row_indices = torch.arange(B * next_n, device=device) // next_n
next_n_offset = torch.arange(B * next_n, device=device) % next_n
end_pos = context_lens[row_indices] - next_n + next_n_offset
mask = positions <= end_pos.unsqueeze(1)
logits_masked = logits.float().masked_fill(~mask, 0)
ref_masked = ref_logits.float().masked_fill(~mask, 0)
torch.testing.assert_close(logits_masked, ref_masked, atol=5e-5, rtol=1e-5)
python/sglang/test/test_marlin_utils.py
+22
View File
@@ -21,11 +21,33 @@ from sglang.srt.layers.quantization.marlin_utils import (
from sglang.srt.layers.quantization.utils import (
get_pack_factor,
gptq_quantize_weights,
pack_cols,
quantize_weights,
sort_weights,
)
def awq_pack(
q_w: torch.Tensor,
num_bits: int,
size_k: int,
size_n: int,
):
assert q_w.shape == (size_k, size_n)
if num_bits == 4:
interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7])
elif num_bits == 8:
interleave = np.array([0, 2, 1, 3])
else:
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel()
q_w = q_w.reshape((-1, size_n)).contiguous()
return pack_cols(q_w, num_bits, size_k, size_n)
class MarlinWorkspace:
def __init__(self, out_features, min_thread_n, max_parallel):
assert out_features % min_thread_n == 0, (