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

[kernels] Reorganize ops/diffusion by operator domain behind a lazy facade (#35114)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-08-18 20:37:43 +08:00
committed by GitHub
co-authored by Claude Opus 5
parent 7605529bdf
commit ae6945e112
167 changed files with 4813 additions and 4340 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/ops/diffusion/test_activation.py
+112
View File
@@ -0,0 +1,112 @@
"""``diffusion.activation``: activation-function fusions.
All of these are bit-exact by construction -- they are elementwise chains with
no reduction, so reproducing aten's per-op fp32-opmath / round-to-bf16
boundaries is enough and ``torch.equal`` is the assertion.
The cublasLt linear+tanh-GELU epilogue is *not* here: it is not bit-exact and
is therefore quality-gated, so it is tested through its mount protocol in
``test_sites.py``.
"""
import sys
import pytest
import torch
import torch.nn.functional as F
from sglang.kernels.ops.diffusion import (
can_use_fused_bias_glu,
can_use_fused_bias_silu,
can_use_fused_silu_mul,
fused_bias_glu,
fused_bias_silu,
fused_packed_silu_mul_bitexact,
fused_silu_mul_bitexact,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=6, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize("channels", [2240, 11200])
def test_sana_bias_silu_is_bit_exact(channels):
torch.manual_seed(0)
x = torch.randn(
(1, channels, 7, 5),
device="cuda",
dtype=torch.bfloat16,
).to(memory_format=torch.channels_last)
bias = torch.randn(channels, device="cuda", dtype=torch.bfloat16)
assert can_use_fused_bias_silu(x, bias)
actual = fused_bias_silu(x, bias)
expected = F.silu(x + bias[None, :, None, None])
assert actual.is_contiguous(memory_format=torch.channels_last)
assert torch.equal(actual, expected)
@pytest.mark.parametrize("channels", [2240, 5600])
def test_sana_bias_glu_is_bit_exact(channels):
torch.manual_seed(1)
x = torch.randn(
(1, 2 * channels, 7, 5),
device="cuda",
dtype=torch.bfloat16,
).to(memory_format=torch.channels_last)
bias = torch.randn(2 * channels, device="cuda", dtype=torch.bfloat16)
assert can_use_fused_bias_glu(x, bias)
actual = fused_bias_glu(x, bias)
biased = x + bias[None, :, None, None]
hidden, gate = torch.chunk(biased, 2, dim=1)
expected = hidden * F.silu(gate)
assert actual.is_contiguous(memory_format=torch.channels_last)
assert torch.equal(actual, expected)
# ---------------------------------------------------------------------------
# silu(a) * b for split-projection SwiGLU
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("shape", [(1, 4096, 3072), (2, 17, 512)])
def test_silu_mul_is_bit_exact(shape):
# Separate gate/up GEMMs, so the concatenated ``silu_and_mul`` kernels do
# not apply without an extra full-width cat -- this kernel replaces the
# eager ``F.silu(a) * b`` pair instead.
torch.manual_seed(0)
a = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
b = torch.randn_like(a)
assert can_use_fused_silu_mul(a, b)
assert torch.equal(fused_silu_mul_bitexact(a, b), F.silu(a) * b)
@pytest.mark.parametrize("hidden", [384, 3072])
@pytest.mark.parametrize("strided", [False, True])
def test_packed_silu_mul_is_bit_exact(hidden, strided):
# The packed form splits one [.., 2 * hidden] projection in-kernel; it must
# accept the strided view a wider projection slice produces.
torch.manual_seed(1)
if strided:
x = torch.randn(1, 19, 3 * hidden, device="cuda", dtype=torch.bfloat16)
x = x[..., : 2 * hidden]
else:
x = torch.randn(1, 19, 2 * hidden, device="cuda", dtype=torch.bfloat16)
expected = F.silu(x[..., :hidden]) * x[..., hidden:]
assert torch.equal(fused_packed_silu_mul_bitexact(x), expected)
def test_silu_mul_rejects_mismatched_operands():
a = torch.randn(1, 8, 64, device="cuda", dtype=torch.bfloat16)
assert not can_use_fused_silu_mul(a, a.float()) # mixed dtypes
assert not can_use_fused_silu_mul(a, a[:, :-1]) # mismatched shapes
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_autoencoder_kl_fastpath.py
-62
View File
@@ -1,62 +0,0 @@
"""Install-path checks for the generic AutoencoderKL CUDA fast path."""
import sys
import pytest
import torch
from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import (
StableDiffusion3VAEConfig,
)
from sglang.multimodal_gen.runtime.models.vaes import flux2_vae_cuda_opt as vae_opt
from sglang.multimodal_gen.runtime.models.vaes.autoencoder import AutoencoderKL
from sglang.multimodal_gen.runtime.models.vaes.fast_path_gate import (
use_vae_fast_path,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def _small_config():
config = StableDiffusion3VAEConfig()
config.arch_config.latent_channels = 2
config.arch_config.block_out_channels = (4, 4)
config.arch_config.down_block_types = ("DownEncoderBlock2D",) * 2
config.arch_config.up_block_types = ("UpDecoderBlock2D",) * 2
config.arch_config.layers_per_block = 1
config.arch_config.norm_num_groups = 1
config.arch_config.sample_size = 8
return config
@torch.no_grad()
def test_autoencoder_kl_fastpath_install():
torch.manual_seed(0)
vae = AutoencoderKL(_small_config()).to("cuda", torch.bfloat16).eval()
ref_names = {n for n, _ in vae.named_parameters()}
ref_sd = {k: v.clone() for k, v in vae.state_dict().items()}
z = torch.randn(1, 2, 8, 8, device="cuda", dtype=torch.bfloat16)
ref = vae.decode(z)
opt = vae_opt.maybe_optimize_autoencoder_kl(vae)
# Wrappers must not change parameter FQNs; strict load must round-trip.
assert {n for n, _ in opt.named_parameters()} == ref_names
opt.load_state_dict(ref_sd, strict=True)
# Gate off: bit-for-bit the original path.
assert torch.equal(opt.decode(z), ref)
# use_vae_fast_path() is a no-op when nothing registered a gate, so check
# the wrappers went in before relying on it to switch paths.
assert any(
isinstance(m, (vae_opt.FusedGroupNormSiLU, vae_opt.FusedUpsample2xConv2d))
for m in opt.modules()
)
# Gate on: fast path runs and stays close; leaving the scope restores exact.
with use_vae_fast_path(opt, True):
torch.testing.assert_close(opt.decode(z).float(), ref.float(), atol=0.1, rtol=0)
assert torch.equal(opt.decode(z), ref)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_bitexact_gate.py
-159
View File
@@ -1,159 +0,0 @@
import sys
from types import ModuleType
from unittest.mock import MagicMock, patch
import pytest
import torch
from sglang.kernels.ops.diffusion.bitexact_gate import (
BitExactFusionGate,
flashinfer_rmsnorm_diagnostic_hint,
tensors_equal,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
def test_bitexact_gate_once_mode_verifies_then_reuses():
gate = BitExactFusionGate("once")
calls = {"fused": 0, "ref": 0}
def fused():
calls["fused"] += 1
return torch.tensor([1.0])
def ref():
calls["ref"] += 1
return torch.tensor([1.0])
assert torch.equal(gate.accept_or_fallback(fused(), ref()), torch.tensor([1.0]))
assert gate.verified and not gate.disabled and calls == {"fused": 1, "ref": 1}
assert torch.equal(fused(), torch.tensor([1.0]))
assert calls == {"fused": 2, "ref": 1}
def test_bitexact_gate_mismatch_disables_permanently():
gate = BitExactFusionGate("mismatch")
out = gate.accept_or_fallback(
torch.tensor([1.0]),
torch.tensor([2.0]),
mismatch_msg="mismatch",
)
assert torch.equal(out, torch.tensor([2.0]))
assert gate.disabled and not gate.verified
def test_bitexact_gate_per_signature_tracks_each_sig():
gate = BitExactFusionGate("sig", per_signature=True)
a = torch.tensor([1.0])
assert torch.equal(gate.accept_or_fallback(a, a, sig=("a",)), a)
assert gate.is_verified(("a",))
assert not gate.is_verified(("b",))
assert torch.equal(gate.accept_or_fallback(a, a, sig=("b",)), a)
assert gate.verified_sigs == {("a",), ("b",)}
def test_bitexact_gate_skips_first_sight_during_graph_capture(monkeypatch):
# Negative-branch contract: an unverified gate must not attempt first-sight
# verification inside CUDA graph capture — the eager-reference host sync
# would abort the capture (and BCG would permanently block the signature).
gate = BitExactFusionGate("capture")
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
assert not gate.can_attempt_once()
# A verified gate replays the fused kernel alone, which is capture-safe.
gate.mark_verified()
assert gate.can_attempt_once()
def test_tensors_equal_supports_sequences():
assert tensors_equal(
(torch.tensor([1.0]), torch.tensor([2.0])),
(torch.tensor([1.0]), torch.tensor([2.0])),
)
assert not tensors_equal(
(torch.tensor([1.0]), torch.tensor([2.0])),
(torch.tensor([1.0]), torch.tensor([3.0])),
)
class TestBitExactFallbackDiagnostics(CustomTestCase):
def test_mismatch_warning_is_actionable_and_diagnostic_is_lazy(self):
logger = MagicMock()
diagnostic = MagicMock(return_value="backend=CuTe DSL")
gate = BitExactFusionGate("diagnostic")
matched = gate.accept_or_fallback(
torch.tensor([1.0]),
torch.tensor([1.0]),
logger=logger,
diagnostic_hint=diagnostic,
)
self.assertTrue(torch.equal(matched, torch.tensor([1.0])))
diagnostic.assert_not_called()
logger.warning_once.assert_not_called()
gate = BitExactFusionGate("diagnostic")
fallback = gate.accept_or_fallback(
torch.tensor([1.0]),
torch.tensor([2.0]),
logger=logger,
diagnostic_hint=diagnostic,
)
self.assertTrue(torch.equal(fallback, torch.tensor([2.0])))
diagnostic.assert_called_once_with()
warning = logger.warning_once.call_args.args[0]
self.assertIn("Correctness is preserved", warning)
self.assertIn("reference kernel or reduction-order change", warning)
self.assertIn("backend=CuTe DSL", warning)
def test_diagnostic_failure_cannot_break_the_eager_fallback(self):
logger = MagicMock()
def broken_diagnostic():
raise RuntimeError("diagnostics unavailable")
gate = BitExactFusionGate("diagnostic")
fallback = gate.accept_or_fallback(
torch.tensor([1.0]),
torch.tensor([2.0]),
logger=logger,
diagnostic_hint=broken_diagnostic,
)
self.assertTrue(torch.equal(fallback, torch.tensor([2.0])))
self.assertTrue(gate.disabled)
self.assertIn("Correctness is preserved", logger.warning_once.call_args.args[0])
def test_flashinfer_rmsnorm_hint_reports_backend_and_versions(self):
flashinfer = ModuleType("flashinfer")
flashinfer_norm = ModuleType("flashinfer.norm")
flashinfer_norm._USE_CUDA_NORM = False
versions = {
"flashinfer-python": "0.6.12",
"flashinfer-cubin": "0.6.12",
"flashinfer-jit-cache": "0.6.12+cu130",
}
with (
patch.dict(
sys.modules,
{"flashinfer": flashinfer, "flashinfer.norm": flashinfer_norm},
),
patch("importlib.metadata.version", side_effect=versions.__getitem__),
patch.dict("os.environ", {"FLASHINFER_USE_CUDA_NORM": "0"}),
):
hint = flashinfer_rmsnorm_diagnostic_hint()
self.assertIn("backend=CuTe DSL", hint)
self.assertIn("FLASHINFER_USE_CUDA_NORM=0", hint)
for package, version in versions.items():
self.assertIn(f"{package}={version}", hint)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_causal_conv3d_cat_pad.py
-89
View File
@@ -1,89 +0,0 @@
import sys
import pytest
import torch
from sglang.kernels.jit.utils import get_ci_test_range
from sglang.kernels.ops.diffusion.causal_conv3d_cat_pad import (
fused_causal_conv3d_cat_pad_cuda,
)
from sglang.kernels.ops.diffusion.triton.causal_conv3d_pad import (
fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton,
)
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
register_amd_ci(est_time=10, stage="jit-kernel-unit", runner_config="amd")
DEVICE = "cuda"
DTYPE = torch.bfloat16
COSMOS3_CASES = get_ci_test_range(
[
(1024, 1, 30, 52, 1),
(1024, 1, 30, 52, 2),
(1024, 2, 60, 104, 1),
(1024, 2, 60, 104, 2),
(512, 4, 120, 208, 1),
(512, 4, 120, 208, 2),
(256, 4, 240, 416, 1),
(256, 4, 240, 416, 2),
],
[(1024, 1, 30, 52, 1), (512, 4, 120, 208, 2)],
)
def _make_inputs(
channels: int,
t_size: int,
h_size: int,
w_size: int,
cache_t: int,
) -> tuple[torch.Tensor, torch.Tensor, tuple[int, ...]]:
generator = torch.Generator(device=DEVICE)
generator.manual_seed(channels * 1009 + t_size * 251 + h_size + cache_t)
x = torch.randn(
(1, channels, t_size, h_size, w_size),
device=DEVICE,
dtype=DTYPE,
generator=generator,
)
cache_x = torch.randn(
(1, channels, cache_t, h_size, w_size),
device=DEVICE,
dtype=DTYPE,
generator=generator,
)
padding = (1, 1, 1, 1, cache_t, 0)
return x, cache_x, padding
@pytest.mark.parametrize("channels,t_size,h_size,w_size,cache_t", COSMOS3_CASES)
def test_causal_conv3d_cat_pad(
channels: int,
t_size: int,
h_size: int,
w_size: int,
cache_t: int,
) -> None:
x, cache_x, padding = _make_inputs(channels, t_size, h_size, w_size, cache_t)
actual = fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding)
expected = fused_causal_conv3d_cat_pad_triton(x, cache_x, padding)
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
def test_causal_conv3d_cat_pad_torch_compile() -> None:
x, cache_x, padding = _make_inputs(1024, 1, 30, 52, 1)
@torch.compile(fullgraph=True)
def fn(x: torch.Tensor, cache_x: torch.Tensor) -> torch.Tensor:
return fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding)
actual = fn(x, cache_x)
expected = fused_causal_conv3d_cat_pad_triton(x, cache_x, padding)
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py
-141
View File
@@ -1,141 +0,0 @@
import sys
import pytest
import torch
from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
ModelOptFp8Config,
ModelOptFp8LinearMethod,
)
from sglang.srt.layers.quantization.fp8_utils import (
cutlass_fp8_supported,
input_to_float8,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
DEVICE = "cuda"
DTYPE = torch.bfloat16
MAX_FP8_DIFF = 5e-4
TEST_CASES = [
pytest.param(19, 150, 80, id="misaligned_projection_shape"),
pytest.param(512, 3072, 4096, id="flux2_added_kv_projection_shape"),
]
def _modelopt_fp8_supported() -> bool:
return torch.cuda.is_available() and cutlass_fp8_supported()
def _calc_diff(x: torch.Tensor, y: torch.Tensor) -> float:
x, y = x.double(), y.double()
denominator = (x * x + y * y).sum()
if denominator == 0:
return 0.0
sim = 2 * (x * y).sum() / denominator
return (1 - sim).item()
def _dequantize_fp8_input(qinput: torch.Tensor, x_scale: torch.Tensor) -> torch.Tensor:
return qinput.to(torch.float32) * x_scale.to(torch.float32)
def _dequantize_fp8_weight(
weight: torch.Tensor, weight_scale: torch.Tensor
) -> torch.Tensor:
if weight_scale.ndim == 0 or weight_scale.numel() == 1:
scale = weight_scale.to(torch.float32)
else:
scale = weight_scale.to(torch.float32).reshape(-1, 1).t()
return weight.to(torch.float32) * scale
def _build_layer(
weight_q: torch.Tensor,
weight_scale: torch.Tensor,
input_scale: torch.Tensor,
) -> tuple[torch.nn.Module, ModelOptFp8LinearMethod]:
output_size, input_size = weight_q.shape
method = ModelOptFp8LinearMethod(
ModelOptFp8Config(is_checkpoint_fp8_serialized=True)
)
layer = torch.nn.Module()
method.create_weights(
layer=layer,
input_size_per_partition=input_size,
output_partition_sizes=[output_size],
input_size=input_size,
output_size=output_size,
params_dtype=DTYPE,
weight_loader=lambda *args, **kwargs: None,
)
layer = layer.to(device=DEVICE)
layer.weight.data.copy_(weight_q)
layer.weight_scale.data.copy_(weight_scale.reshape_as(layer.weight_scale))
layer.input_scale.data.copy_(input_scale.reshape_as(layer.input_scale))
method.process_weights_after_loading(layer)
return layer, method
@pytest.mark.skipif(
not _modelopt_fp8_supported(),
reason="Diffusion ModelOpt FP8 scaled mm correctness requires CUDA FP8 support",
)
@pytest.mark.parametrize("m,n,k", TEST_CASES)
def test_checkpoint_processing(m: int, n: int, k: int) -> None:
generator = torch.Generator(device=DEVICE)
generator.manual_seed(20260410 + m + n + k)
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
weight_q, weight_scale = input_to_float8(weight)
input_scale = torch.tensor(1.0, device=DEVICE, dtype=torch.float32)
layer, _ = _build_layer(weight_q, weight_scale, input_scale)
assert tuple(layer.weight.shape) == (k, n)
assert tuple(layer.weight.stride()) == (1, k)
assert layer.weight.dtype == torch.float8_e4m3fn
assert layer.input_scale.ndim == 0
assert tuple(layer.weight_scale.shape) == (n, 1)
expected_weight = weight_q.t().to(torch.float32) * weight_scale.to(torch.float32)
actual_weight = _dequantize_fp8_weight(layer.weight, layer.weight_scale)
torch.testing.assert_close(actual_weight, expected_weight, atol=0.0, rtol=0.0)
@pytest.mark.skipif(
not _modelopt_fp8_supported(),
reason="Diffusion ModelOpt FP8 scaled mm correctness requires CUDA FP8 support",
)
@pytest.mark.parametrize("m,n,k", TEST_CASES)
def test_shape_correctness(m: int, n: int, k: int) -> None:
generator = torch.Generator(device=DEVICE)
generator.manual_seed(20260410 + m + n + k)
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
weight_q, weight_scale = input_to_float8(weight)
_, input_scale = input_to_float8(x)
layer, method = _build_layer(weight_q, weight_scale, input_scale)
qinput, x_scale = static_quant_fp8(
x.contiguous(),
layer.input_scale,
repeat_scale=method.cutlass_fp8_supported,
)
expected = torch.matmul(
_dequantize_fp8_input(qinput, x_scale),
_dequantize_fp8_weight(layer.weight, layer.weight_scale),
)
actual = method.apply(layer, x)
diff = _calc_diff(actual, expected.to(dtype=DTYPE))
assert diff < MAX_FP8_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_diffusion_nvfp4_scaled_mm.py
-447
View File
@@ -1,447 +0,0 @@
import sys
import flashinfer
import pytest
import torch
from sglang.multimodal_gen.runtime.layers.quantization import (
modelopt_quant as diffusion_modelopt_quant,
)
from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
ModelOptFp4Config,
ModelOptFp4LinearMethod,
)
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.srt.layers.quantization.modelopt_quant import pad_nvfp4_weight
from sglang.test.ci.ci_register import register_cuda_ci
# B200-only correctness coverage for diffusion NVFP4 scaled mm.
register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
DEVICE = "cuda"
DTYPE = torch.bfloat16
BLOCK_SIZE = 16
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
FP4_VALUE_LUT = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
DEEPGEMM_FP4_MAX_DIFF = 0.02
TEST_CASES = [
pytest.param(19, 150, 80, id="padding_regression"),
pytest.param(512, 6144, 128, id="flux2_projection_shape"),
]
FLUX2_PROJECTION_SHAPE = (512, 6144, 128)
def _nvfp4_supported() -> bool:
return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
def _make_global_scale(x: torch.Tensor) -> torch.Tensor:
max_abs = torch.amax(x.abs()).clamp_min_(1e-6)
return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / max_abs).to(torch.float32)
def _calc_diff(x: torch.Tensor, y: torch.Tensor) -> float:
x, y = x.double(), y.double()
denominator = (x * x + y * y).sum()
if denominator == 0:
return 0.0
sim = 2 * (x * y).sum() / denominator
return (1 - sim).item()
def _swap_fp4_nibbles(packed: torch.Tensor) -> torch.Tensor:
return ((packed >> 4) | (packed << 4)).contiguous()
def _fp4_lut(device: torch.device) -> torch.Tensor:
return torch.tensor(FP4_VALUE_LUT, dtype=torch.float32, device=device)
def _unpack_fp4_bytes(packed: torch.Tensor) -> torch.Tensor:
assert packed.dtype == torch.uint8
lut = _fp4_lut(packed.device)
def _decode(nibbles: torch.Tensor) -> torch.Tensor:
values = lut[(nibbles & 0x7).to(torch.long)]
return torch.where((nibbles & 0x8) != 0, -values, values)
low = _decode(packed & 0x0F)
high = _decode((packed & 0xF0) >> 4)
return torch.stack((low, high), dim=-1).reshape(
packed.shape[0], packed.shape[1] * 2
)
def _swizzled_to_linear(
scales_swizzled: torch.Tensor,
rows: int,
cols: int,
) -> torch.Tensor:
scales_swizzled = scales_swizzled.view(torch.float8_e4m3fn)
row_tiles = (rows + 128 - 1) // 128
tile_cols = BLOCK_SIZE * 4
col_tiles = (cols + tile_cols - 1) // tile_cols
tmp = scales_swizzled.reshape(1, row_tiles, col_tiles, 32, 4, 4)
tmp = tmp.permute(0, 1, 4, 3, 2, 5)
linear = tmp.reshape(row_tiles * 128, col_tiles * tile_cols // BLOCK_SIZE)
return linear[:rows, : cols // BLOCK_SIZE]
def _dequantize_nvfp4(
packed: torch.Tensor,
scales_swizzled: torch.Tensor,
global_scale: torch.Tensor,
) -> torch.Tensor:
rows, packed_cols = packed.shape
cols = packed_cols * 2
unpacked = _unpack_fp4_bytes(packed).reshape(rows, cols // BLOCK_SIZE, BLOCK_SIZE)
scales_linear = _swizzled_to_linear(scales_swizzled, rows, cols).to(torch.float32)
return (unpacked * (scales_linear / global_scale).unsqueeze(-1)).reshape(rows, cols)
def _quantize_weight_for_checkpoint(
weight: torch.Tensor, weight_global_scale: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
weight_fp4, weight_scale_linear = flashinfer.fp4_quantize(
weight,
weight_global_scale,
is_sf_swizzled_layout=False,
)
if weight_scale_linear.dtype == torch.uint8:
weight_scale_linear = weight_scale_linear.view(torch.float8_e4m3fn)
return weight_fp4, weight_scale_linear.contiguous()
def _set_diffusion_fp4_backend(
monkeypatch: pytest.MonkeyPatch, backend: str | None
) -> None:
if backend is None:
monkeypatch.delenv(
"SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND", raising=False
)
else:
monkeypatch.setenv("SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND", backend)
current_platform.__class__.get_modelopt_flashinfer_fp4_backend.cache_clear()
current_platform.__class__.get_modelopt_fp4_gemm_op.cache_clear()
diffusion_modelopt_quant._get_fp4_gemm_op.cache_clear()
def _build_layer(
weight_fp4: torch.Tensor,
weight_scale_linear: torch.Tensor,
input_global_scale: torch.Tensor,
weight_global_scale: torch.Tensor,
*,
weight_scale_device: torch.device | str | None = None,
checkpoint_weight_scale_layout: str = "linear",
) -> tuple[ModelOptFp4LinearMethod, torch.nn.Module]:
output_size, input_size_half = weight_fp4.shape
input_size = input_size_half * 2
method = ModelOptFp4LinearMethod(
ModelOptFp4Config(
is_checkpoint_nvfp4_serialized=True,
group_size=BLOCK_SIZE,
swap_weight_nibbles=True,
checkpoint_weight_scale_layout=checkpoint_weight_scale_layout,
)
)
layer = torch.nn.Module()
method.create_weights(
layer,
input_size_per_partition=input_size,
output_partition_sizes=[output_size],
input_size=input_size,
output_size=output_size,
params_dtype=DTYPE,
weight_loader=lambda *args, **kwargs: None,
)
layer = layer.to(device=DEVICE)
checkpoint_weight = _swap_fp4_nibbles(weight_fp4)
layer.weight.data.copy_(checkpoint_weight)
layer.input_scale.data.copy_(
(1.0 / input_global_scale).reshape_as(layer.input_scale)
)
layer.weight_scale_2.data.copy_(
(1.0 / weight_global_scale).reshape_as(layer.weight_scale_2)
)
layer.weight_scale.data.copy_(weight_scale_linear)
if weight_scale_device is not None:
layer.weight_scale = torch.nn.Parameter(
layer.weight_scale.detach().to(weight_scale_device), requires_grad=False
)
method.process_weights_after_loading(layer)
_, flashinfer_backend = current_platform.get_modelopt_fp4_gemm_op()
if flashinfer_backend == "trtllm":
expected_weight, _ = pad_nvfp4_weight(
weight_fp4, n_alignment=128, k_alignment=0
)
expected_scale = (
_swizzled_to_linear(weight_scale_linear, output_size, input_size)
if checkpoint_weight_scale_layout == "swizzled"
else weight_scale_linear
)
if expected_scale.shape[0] != expected_weight.shape[0]:
pad_n = expected_weight.shape[0] - expected_scale.shape[0]
expected_scale = torch.nn.functional.pad(expected_scale, (0, 0, 0, pad_n))
expected_padding_cols = 0
if expected_scale.shape[1] % 4 != 0:
padded_scale_k = ((expected_scale.shape[1] + 4 - 1) // 4) * 4
pad_scale_k = padded_scale_k - expected_scale.shape[1]
expected_scale = torch.nn.functional.pad(
expected_scale, (0, pad_scale_k, 0, 0)
)
pad_weight_k = pad_scale_k * 8
expected_weight = torch.nn.functional.pad(
expected_weight, (0, pad_weight_k, 0, 0)
)
expected_padding_cols = pad_weight_k
expected_weight = flashinfer.shuffle_matrix_a(
expected_weight.view(torch.uint8), 128
)
expected_scale = (
flashinfer.shuffle_matrix_sf_a(expected_scale.view(torch.uint8), 128)
.reshape(expected_scale.shape)
.view(torch.float8_e4m3fn)
)
assert torch.equal(layer.weight, expected_weight)
assert torch.equal(
layer.weight_scale_interleaved.view(torch.uint8),
expected_scale.view(torch.uint8),
)
assert layer.weights_padding_cols == expected_padding_cols
else:
expected_weight, expected_padding_cols = pad_nvfp4_weight(weight_fp4)
expected_scale_shape = (
((output_size + 128 - 1) // 128) * 128,
(((input_size // BLOCK_SIZE) + 4 - 1) // 4) * 4,
)
assert torch.equal(layer.weight, expected_weight)
assert layer.weight_scale_interleaved.shape == expected_scale_shape
assert layer.weight_scale_interleaved.dtype == torch.float8_e4m3fn
assert layer.weights_padding_cols == expected_padding_cols
torch.testing.assert_close(
layer.alpha,
(1.0 / (input_global_scale * weight_global_scale)).to(torch.float32),
)
torch.testing.assert_close(
layer.input_scale_inv,
input_global_scale.to(torch.float32),
)
return method, layer
def _resolve_mode(mode: str):
if mode == "flashinfer2":
return flashinfer.fp4_quantize, flashinfer.mm_fp4, "cudnn"
if mode == "flashinfer_trtllm":
return flashinfer.fp4_quantize, flashinfer.mm_fp4, "trtllm"
raise ValueError(f"Unknown mode: {mode}")
@pytest.mark.skipif(
not _nvfp4_supported(),
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
)
@pytest.mark.parametrize(
"backend", [None, "flashinfer_trtllm"], ids=["default", "flashinfer_trtllm"]
)
@pytest.mark.parametrize("m,n,k", TEST_CASES)
def test_checkpoint_processing(
monkeypatch: pytest.MonkeyPatch, backend: str | None, m: int, n: int, k: int
) -> None:
_set_diffusion_fp4_backend(monkeypatch, backend)
generator = torch.Generator(device=DEVICE)
generator.manual_seed(20260404 + m + n + k)
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
input_global_scale = torch.tensor(512.0, device=DEVICE, dtype=torch.float32)
weight_global_scale = _make_global_scale(weight)
weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
weight, weight_global_scale
)
_build_layer(
weight_fp4, weight_scale_linear, input_global_scale, weight_global_scale
)
@pytest.mark.skipif(
not _nvfp4_supported(),
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
)
@pytest.mark.parametrize("mode", ["flashinfer2"])
def test_flux2_shape_correctness(mode: str) -> None:
m, n, k = FLUX2_PROJECTION_SHAPE
quantize_op, gemm_op, gemm_backend = _resolve_mode(mode)
generator = torch.Generator(device=DEVICE)
generator.manual_seed(20260404 + m + n + k)
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
input_global_scale = _make_global_scale(x)
weight_global_scale = _make_global_scale(weight)
alpha = (1.0 / (input_global_scale * weight_global_scale)).to(torch.float32)
x_fp4, x_scale_swizzled = quantize_op(x, input_global_scale)
weight_fp4, weight_scale_swizzled = quantize_op(weight, weight_global_scale)
if x_scale_swizzled.dtype == torch.uint8:
x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
if weight_scale_swizzled.dtype == torch.uint8:
weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
expected = torch.matmul(
_dequantize_nvfp4(x_fp4, x_scale_swizzled, input_global_scale),
_dequantize_nvfp4(weight_fp4, weight_scale_swizzled, weight_global_scale).t(),
)
actual = gemm_op(
x_fp4,
weight_fp4.t(),
x_scale_swizzled,
weight_scale_swizzled.t(),
alpha,
DTYPE,
backend=gemm_backend,
)
diff = _calc_diff(actual, expected.to(dtype=DTYPE))
assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{mode=}, {m=}, {n=}, {k=}, {diff=:.6f}"
@pytest.mark.skipif(
not _nvfp4_supported(),
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
)
def test_flux2_shape_correctness_flashinfer_trtllm(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
m, n, k = FLUX2_PROJECTION_SHAPE
generator = torch.Generator(device=DEVICE)
generator.manual_seed(20260404 + m + n + k + 17)
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
input_global_scale = _make_global_scale(x)
weight_global_scale = _make_global_scale(weight)
weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
weight, weight_global_scale
)
method, layer = _build_layer(
weight_fp4, weight_scale_linear, input_global_scale, weight_global_scale
)
actual = method.apply(layer, x)
x_fp4, x_scale_swizzled = flashinfer.fp4_quantize(x, input_global_scale)
weight_fp4_ref, weight_scale_swizzled = flashinfer.fp4_quantize(
weight, weight_global_scale
)
if x_scale_swizzled.dtype == torch.uint8:
x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
if weight_scale_swizzled.dtype == torch.uint8:
weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
expected = torch.matmul(
_dequantize_nvfp4(x_fp4, x_scale_swizzled, input_global_scale),
_dequantize_nvfp4(
weight_fp4_ref, weight_scale_swizzled, weight_global_scale
).t(),
)
diff = _calc_diff(actual, expected.to(dtype=DTYPE))
assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
@pytest.mark.skipif(
not _nvfp4_supported(),
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
)
def test_flux2_swizzled_scale_checkpoint_flashinfer_trtllm_matches_cudnn(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
m, n, k = FLUX2_PROJECTION_SHAPE
generator = torch.Generator(device=DEVICE)
generator.manual_seed(20260517 + m + n + k)
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
input_global_scale = _make_global_scale(x)
weight_global_scale = _make_global_scale(weight)
alpha = (1.0 / (input_global_scale * weight_global_scale)).to(torch.float32)
x_fp4, x_scale_swizzled = flashinfer.fp4_quantize(x, input_global_scale)
weight_fp4, weight_scale_swizzled = flashinfer.fp4_quantize(
weight, weight_global_scale
)
if x_scale_swizzled.dtype == torch.uint8:
x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
if weight_scale_swizzled.dtype == torch.uint8:
weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
method, layer = _build_layer(
weight_fp4,
weight_scale_swizzled,
input_global_scale,
weight_global_scale,
checkpoint_weight_scale_layout="swizzled",
)
actual = method.apply(layer, x)
expected = flashinfer.mm_fp4(
x_fp4,
weight_fp4.t(),
x_scale_swizzled,
weight_scale_swizzled.t(),
alpha,
DTYPE,
backend="cudnn",
)
diff = _calc_diff(actual, expected)
assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
@pytest.mark.skipif(
not _nvfp4_supported(),
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
)
def test_checkpoint_processing_flashinfer_trtllm_cpu_weight_scale(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
m, n, k = FLUX2_PROJECTION_SHAPE
generator = torch.Generator(device=DEVICE)
generator.manual_seed(20260413 + m + n + k)
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
input_global_scale = torch.tensor(512.0, device=DEVICE, dtype=torch.float32)
weight_global_scale = _make_global_scale(weight)
weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
weight, weight_global_scale
)
_build_layer(
weight_fp4,
weight_scale_linear,
input_global_scale,
weight_global_scale,
weight_scale_device="cpu",
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_ernie_norm_scale_shift.py
-135
View File
@@ -1,135 +0,0 @@
"""ERNIE fused norm/scale/shift fast paths must stay bit-exact vs eager."""
import sys
from unittest.mock import patch
import pytest
import torch
import sglang.multimodal_gen.runtime.models.dits.ernie_image as ernie_image
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
from sglang.multimodal_gen.runtime.models.dits.ernie_image import (
_ernie_gated_norm_scale_shift,
_ernie_norm_scale_shift,
_ernie_qknorm_rope,
_ernie_qknorm_rope_reference,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize("shape", [(1, 4216, 4096), (2, 1140, 4096), (1, 128, 2048)])
def test_fused_norm_scale_shift_is_bit_exact(shape):
# (1, 4216, 4096) is the real ERNIE-Image shape (1024^2 image + text
# tokens, hidden 4096); 2048 covers the threads_per_row=32 regime.
torch.manual_seed(0)
batch, seq, hidden = shape
norm = RMSNorm(hidden, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
norm.weight.copy_(torch.randn(hidden))
x = torch.randn(batch, seq, hidden, device="cuda", dtype=torch.bfloat16)
residual = torch.randn_like(x)
update = torch.randn_like(x)
scale = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
shift = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
gate = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
out = _ernie_norm_scale_shift(norm, x, scale, shift)
ref = norm(x) * (1 + scale) + shift
assert torch.equal(out, ref)
out2, res = _ernie_gated_norm_scale_shift(
norm, residual, update, gate, scale, shift
)
res_ref = residual + gate * update
ref2 = norm(res_ref) * (1 + scale) + shift
assert torch.equal(res, res_ref)
assert torch.equal(out2, ref2)
# the fast paths must actually be in use (not silently disabled)
assert ernie_image._ERNIE_NORM.verified
assert ernie_image._ERNIE_GATED_NORM.verified
assert not ernie_image._ERNIE_NORM.disabled
assert not ernie_image._ERNIE_GATED_NORM.disabled
def test_fused_qknorm_rope_is_bit_exact():
torch.manual_seed(1)
ernie_image._ERNIE_QKNORM_ROPE.disabled = False
ernie_image._ERNIE_QKNORM_ROPE.verified = False
batch, seq, heads, head_dim = 1, 257, 32, 128
q = torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
q_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
k_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
cos = torch.randn(seq, head_dim, device="cuda", dtype=torch.bfloat16)
sin = torch.randn_like(cos)
cache = torch.cat((cos, sin), dim=-1).contiguous()
positions = torch.arange(seq, device="cuda", dtype=torch.long)
q_ref, k_ref = _ernie_qknorm_rope_reference(
q.clone(), k.clone(), q_norm, k_norm, head_dim, cos, sin
)
q_out, k_out = _ernie_qknorm_rope(
q,
k,
q_norm,
k_norm,
head_dim,
cos,
sin,
cache,
positions,
)
assert torch.equal(q_out, q_ref)
assert torch.equal(k_out, k_ref)
assert ernie_image._ERNIE_QKNORM_ROPE.verified
assert not ernie_image._ERNIE_QKNORM_ROPE.disabled
def test_qknorm_rope_first_attempt_exception_uses_pristine_inputs():
torch.manual_seed(2)
ernie_image._ERNIE_QKNORM_ROPE.disabled = False
ernie_image._ERNIE_QKNORM_ROPE.verified = False
batch, seq, heads, head_dim = 1, 17, 4, 128
q = torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
q_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
k_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
cos = torch.randn(seq, head_dim, device="cuda", dtype=torch.bfloat16)
sin = torch.randn_like(cos)
cache = torch.cat((cos, sin), dim=-1).contiguous()
positions = torch.arange(seq, device="cuda", dtype=torch.long)
q_ref, k_ref = _ernie_qknorm_rope_reference(
q.clone(), k.clone(), q_norm, k_norm, head_dim, cos, sin
)
def mutate_then_raise(**kwargs):
kwargs["q"].zero_()
kwargs["k"].zero_()
raise RuntimeError("synthetic kernel failure")
with patch.object(ernie_image, "apply_qk_norm_rope", mutate_then_raise):
q_out, k_out = _ernie_qknorm_rope(
q,
k,
q_norm,
k_norm,
head_dim,
cos,
sin,
cache,
positions,
)
assert torch.equal(q_out, q_ref)
assert torch.equal(k_out, k_ref)
assert ernie_image._ERNIE_QKNORM_ROPE.disabled
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/diffusion/test_flux2_eager_fusions.py
-101
View File
@@ -1,101 +0,0 @@
"""FLUX.2 eager fusions must be bit-exact for real packed/view layouts."""
import unittest
from unittest.mock import patch
import torch
import torch.nn.functional as F
import sglang.multimodal_gen.runtime.models.dits.flux_2 as flux2
from sglang.multimodal_gen.runtime.models.dits.flux_2 import (
_flux2_norm_modulate,
_flux2_swiglu,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=12, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
class TestFlux2EagerFusions(CustomTestCase):
def setUp(self):
flux2._FLUX2_LN_MOD.disabled = False
flux2._FLUX2_LN_MOD.verified = False
flux2._FLUX2_LN_MOD_SIGS.clear()
flux2._FLUX2_SWIGLU.disabled = False
flux2._FLUX2_SWIGLU.verified = False
flux2._FLUX2_SWIGLU_SIGS.clear()
def test_norm_modulate_is_bit_exact_across_sequence_lengths(self):
torch.manual_seed(0)
hidden = 256
norm = torch.nn.LayerNorm(
hidden, eps=1e-6, elementwise_affine=False, device="cuda"
)
# FLUX.2 modulation values are views of one packed projection.
params = torch.randn(1, 1, 6 * hidden, device="cuda").bfloat16()
shift, scale = params.chunk(6, dim=-1)[:2]
for seq in (17, 65):
x = torch.randn(1, seq, hidden, device="cuda").bfloat16()
expected = norm(x) * (1 + scale) + shift
actual = _flux2_norm_modulate(norm, x, scale, shift)
self.assertTrue(torch.equal(actual, expected))
self.assertFalse(flux2._FLUX2_LN_MOD.disabled)
self.assertEqual(len(flux2._FLUX2_LN_MOD_SIGS), 1)
def test_packed_swiglu_is_bit_exact_for_contiguous_and_strided_views(self):
torch.manual_seed(1)
hidden = 384
inputs = [
torch.randn(1, 19, 2 * hidden, device="cuda").bfloat16(),
torch.randn(1, 19, 3 * hidden, device="cuda").bfloat16()[..., : 2 * hidden],
]
for x in inputs:
expected = F.silu(x[..., :hidden]) * x[..., hidden:]
actual = _flux2_swiglu(x)
self.assertTrue(torch.equal(actual, expected))
self.assertFalse(flux2._FLUX2_SWIGLU.disabled)
self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 2)
def test_fp16_preserves_reference_path(self):
x = torch.randn(1, 17, 512, device="cuda", dtype=torch.float16)
expected = F.silu(x[..., :256]) * x[..., 256:]
actual = _flux2_swiglu(x)
self.assertTrue(torch.equal(actual, expected))
self.assertFalse(flux2._FLUX2_SWIGLU.disabled)
def test_packed_swiglu_rejects_non_dense_outer_stride(self):
base = torch.randn(2, 23, 512, device="cuda", dtype=torch.bfloat16)
x = base[:, :19]
self.assertNotEqual(x.stride(0), x.shape[1] * x.stride(1))
expected = F.silu(x[..., :256]) * x[..., 256:]
actual = _flux2_swiglu(x)
self.assertTrue(torch.equal(actual, expected))
self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 0)
def test_new_swiglu_signature_is_not_verified_during_graph_capture(self):
first = torch.randn(1, 17, 512, device="cuda", dtype=torch.bfloat16)
self.assertTrue(
torch.equal(
_flux2_swiglu(first),
F.silu(first[..., :256]) * first[..., 256:],
)
)
self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1)
second = torch.randn(1, 19, 768, device="cuda", dtype=torch.bfloat16)
with patch("torch.cuda.is_current_stream_capturing", return_value=True):
actual = _flux2_swiglu(second)
expected = F.silu(second[..., :384]) * second[..., 384:]
self.assertTrue(torch.equal(actual, expected))
self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1)
if __name__ == "__main__":
unittest.main()
test/registered/kernels/ops/diffusion/test_flux2_vae_fastpath.py
-64
View File
@@ -1,64 +0,0 @@
"""Focused correctness checks for the FLUX.2 VAE CUDA fast path."""
import sys
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.models.upsampling import Upsample2D
from sglang.kernels.ops.diffusion.triton import group_norm_silu_twopass as gn_kernel
from sglang.multimodal_gen.runtime.models.vaes import flux2_vae_cuda_opt as vae_opt
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@torch.no_grad()
def test_flux2_vae_fastpath():
torch.manual_seed(0)
gate = vae_opt.VaeFastPathGate()
gn = nn.GroupNorm(32, 128, eps=1e-6).to("cuda", torch.bfloat16)
x = torch.randn(1, 128, 64, 64, device="cuda", dtype=torch.bfloat16).to(
memory_format=torch.channels_last
)
ref = F.silu(gn(x))
fused_gn = vae_opt.FusedGroupNormSiLU(gn, gate)
assert set(fused_gn.state_dict()) == {"weight", "bias"}
assert torch.equal(fused_gn(x), ref)
assert (
gn_kernel.group_norm_silu_4d(x.contiguous(), gn.weight, gn.bias, 32, 1e-6)
is None
)
assert gn_kernel.group_norm_silu_4d(x, gn.weight.cpu(), gn.bias, 32, 1e-6) is None
assert (
gn_kernel.group_norm_silu_4d(x[..., :0, :], gn.weight, gn.bias, 32, 1e-6)
is None
)
gate.enabled = True
fast = fused_gn(x)
assert fast.is_contiguous(memory_format=torch.channels_last)
torch.testing.assert_close(fast.float(), ref.float(), atol=0.06, rtol=0)
gate.enabled = False
up = Upsample2D(channels=32, use_conv=True).to("cuda", torch.bfloat16)
fused_up = vae_opt.FusedUpsample2xConv2d(up, gate)
assert set(fused_up.state_dict()) == {"conv.weight", "conv.bias"}
x = torch.randn(2, 32, 33, 29, device="cuda", dtype=torch.bfloat16)
ref = up(x)
assert torch.equal(fused_up(x), ref)
assert fused_up._fused_weight is None
gate.enabled = True
fast = fused_up(x)
assert fused_up._fused_weight is not None
ref_range = ref.float().max() - ref.float().min()
relative_mse = F.mse_loss(fast.float(), ref.float()) / ref_range.square()
assert relative_mse < 3.2e-5
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_flux_ln_modulate.py
-75
View File
@@ -1,75 +0,0 @@
"""FLUX.1 fused LN+modulate fast path must stay bit-exact vs eager."""
import pytest
import torch
import sglang.multimodal_gen.runtime.models.dits.flux as flux
from sglang.kernels.ops.diffusion.fused_ln_modulate import (
mark_fused_ln_modulate_site,
mount_fused_ln_modulate,
)
from sglang.multimodal_gen.runtime.models.dits.flux import (
_flux_fused_ln_modulate,
_flux_norm_modulate,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def _eager(norm, x, scale, shift):
return norm(x) * (1 + scale[:, None]) + shift[:, None]
def _make_site_inputs(shape, chunks, seed):
torch.manual_seed(seed)
batch, seq, hidden = shape
norm = torch.nn.LayerNorm(hidden, eps=1e-6, elementwise_affine=False).cuda()
x = (torch.randn(batch, seq, hidden, device="cuda") * 8).bfloat16()
emb = torch.randn(batch, chunks * hidden, device="cuda").bfloat16()
parts = emb.chunk(chunks, dim=1) # strided adaLN projection views
return norm, x, parts[0], parts[1]
@pytest.mark.parametrize(
"shape,chunks",
[
((1, 4096, 3072), 6), # dual-stream image tokens (1024^2), chunk(6)
((1, 512, 3072), 6), # dual-stream text tokens
((1, 4608, 3072), 3), # single-stream concat, chunk(3)
((2, 300, 3072), 6), # CFG batch, odd seq
],
)
def test_flux_fused_ln_modulate_is_bit_exact(shape, chunks):
# Every distinct (shape, stride, eps) signature the FLUX.1 sites emit
# must verify torch.equal on first sight and stay enabled.
norm, x, shift, scale = _make_site_inputs(shape, chunks, seed=0)
out = _flux_fused_ln_modulate(norm, x, scale, shift)
assert out is not None
assert torch.equal(out, _eager(norm, x, scale, shift))
assert not flux._FLUX_LN_MOD.disabled
assert flux._FLUX_LN_MOD.verified
def test_flux_norm_modulate_bitexact_supersedes_high_fold():
# With the quality="high" affine fold mounted, the bit-exact kernel
# still takes priority, so the site output stays lossless.
site = torch.nn.Module()
mark_fused_ln_modulate_site(site)
assert mount_fused_ln_modulate(site)
norm, x, shift, scale = _make_site_inputs((1, 128, 3072), 6, seed=1)
out = _flux_norm_modulate(site, norm, x, scale, shift)
assert torch.equal(out, _eager(norm, x, scale, shift))
def test_flux_fused_ln_modulate_rejects_unsupported_hidden():
# hidden % 4 != 0 is outside the kernel contract and must bail out.
norm, x, shift, scale = _make_site_inputs((1, 64, 3070), 6, seed=2)
assert _flux_fused_ln_modulate(norm, x, scale, shift) is None
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/diffusion/test_flydsl_fused_norm.py
-133
View File
@@ -1,133 +0,0 @@
import sys
import pytest
import torch
import torch.nn.functional as F
from sglang.test.ci.ci_register import register_amd_ci
register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")
DEVICE = "cuda"
D = 5120
EPS = 1e-6
def _ref_rms_norm(x_f32, weight, eps):
var = x_f32.pow(2).mean(-1, keepdim=True)
return x_f32 * torch.rsqrt(var + eps)
def _ref_fused_residual_norm_ss(
residual, x, gate, weight, bias, scale, shift, norm_type, eps
):
ref_res = residual.float() + x.float() * (gate.float() if gate is not None else 1)
ref_res_bf16 = ref_res.to(torch.bfloat16)
if norm_type == "layer":
normed = F.layer_norm(ref_res_bf16.float(), (D,), weight, bias, eps)
else:
normed = _ref_rms_norm(ref_res_bf16.float(), weight, eps) * weight.float()
y = (normed * (1.0 + scale.float()) + shift.float()).to(torch.bfloat16)
return y, ref_res_bf16
def _ref_norm_ss(x, weight, bias, scale, shift, norm_type, eps):
if norm_type == "layer":
normed = F.layer_norm(x.float(), (D,), weight, bias, eps)
else:
normed = _ref_rms_norm(x.float(), weight, eps) * weight.float()
return (normed * (1.0 + scale.float()) + shift.float()).to(torch.bfloat16)
@pytest.fixture(autouse=True)
def cuda_setup():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
if not hasattr(torch.version, "hip") or not torch.version.hip:
pytest.skip("ROCm/HIP required for FlyDSL kernels")
torch.manual_seed(42)
FUSED_CASES = [
("rms", 1, 16),
("rms", 2, 16),
("layer", 2, 16),
("rms", 1, 90000),
]
@pytest.mark.parametrize("norm_type,B,L", FUSED_CASES)
def test_fused_residual_norm_scale_shift(norm_type, B, L):
from sglang.kernels.ops.diffusion.flydsl.fused_residual_norm import (
flydsl_fused_residual_norm_scale_shift,
)
residual = torch.randn(B, L, D, device=DEVICE, dtype=torch.bfloat16)
x = torch.randn(B, L, D, device=DEVICE, dtype=torch.bfloat16)
gate = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
weight = torch.randn(D, device=DEVICE, dtype=torch.float32)
bias = (
torch.randn(D, device=DEVICE, dtype=torch.float32)
if norm_type == "layer"
else None
)
scale = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
shift = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
y, res_out = flydsl_fused_residual_norm_scale_shift(
residual,
x,
gate,
weight,
bias,
scale,
shift,
norm_type,
EPS,
)
y_ref, res_ref = _ref_fused_residual_norm_ss(
residual,
x,
gate,
weight,
bias,
scale,
shift,
norm_type,
EPS,
)
torch.testing.assert_close(res_out, res_ref, atol=5e-2, rtol=5e-2)
torch.testing.assert_close(y, y_ref, atol=1.0, rtol=5e-2)
NSS_CASES = [
("rms", 2, 16),
("layer", 2, 16),
("rms", 1, 90000),
("layer", 1, 90000),
]
@pytest.mark.parametrize("norm_type,B,L", NSS_CASES)
def test_norm_scale_shift(norm_type, B, L):
from sglang.kernels.ops.diffusion.flydsl.fused_residual_norm import (
flydsl_norm_scale_shift,
)
x = torch.randn(B, L, D, device=DEVICE, dtype=torch.bfloat16)
weight = torch.randn(D, device=DEVICE, dtype=torch.float32)
bias = (
torch.randn(D, device=DEVICE, dtype=torch.float32)
if norm_type == "layer"
else None
)
scale = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
shift = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
y = flydsl_norm_scale_shift(x, weight, bias, scale, shift, norm_type, EPS)
y_ref = _ref_norm_ss(x, weight, bias, scale, shift, norm_type, EPS)
torch.testing.assert_close(y, y_ref, atol=1.0, rtol=5e-2)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_fused_gate_rmsnorm.py
-54
View File
@@ -1,54 +0,0 @@
"""Core checks for the quality-gated fused gate-RMSNorm path."""
import sys
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from sglang.kernels.ops.diffusion import fused_gate_rmsnorm as fgn
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
DIM, EPS = 4608, 1e-5 # Ideogram 4 hidden size / norm_eps
class _Site(nn.Module):
def __init__(self, dtype=torch.bfloat16):
super().__init__()
self.norm = nn.RMSNorm(DIM, eps=EPS, device="cuda", dtype=dtype)
fgn.mark_fused_gate_rmsnorm_site(self, ("norm",))
def test_fused_matches_ideogram_reference():
torch.manual_seed(0)
site = _Site()
w = site.norm.weight.data
x = torch.randn(1, 64, DIM, device="cuda", dtype=torch.bfloat16)
residual = torch.randn_like(x)
# adaln-style strided chunks, as produced by Ideogram's modulation .chunk()
mods = torch.randn(1, 1, 2 * DIM, device="cuda", dtype=torch.bfloat16)
scale, gate = mods.chunk(2, dim=-1)
assert fgn.mount_fused_gate_rmsnorm(site)
got_scale = fgn.fused_rmsnorm_scale(x, w, 1.0 + scale, EPS)
got_gate = fgn.fused_rmsnorm_tanh_residual(x, gate, residual, w, EPS)
ref_scale = F.rms_norm(x, (DIM,), w, EPS) * (1.0 + scale)
ref_gate = residual + torch.tanh(gate) * F.rms_norm(x, (DIM,), w, EPS)
# fused path uses bf16-native norm statistics: close, not bit-exact
torch.testing.assert_close(got_scale, ref_scale, atol=8e-2, rtol=4e-2)
torch.testing.assert_close(got_gate, ref_gate, atol=8e-2, rtol=4e-2)
def test_mount_guards_all_or_nothing():
good, bad = _Site(), _Site(torch.float32)
assert not fgn.mount_fused_gate_rmsnorm(nn.ModuleList([good, bad]))
assert not fgn.fused_gate_rmsnorm_active(good)
assert fgn.mount_fused_gate_rmsnorm(good)
fgn.unmount_fused_gate_rmsnorm(good)
assert not fgn.fused_gate_rmsnorm_active(good)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/diffusion/test_fused_linear_gelu.py
-86
View File
@@ -1,86 +0,0 @@
"""Core checks for the quality-gated linear + tanh-GELU fusion."""
import sys
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from sglang.kernels.ops.diffusion import fused_linear_gelu as gelu
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
class _Site(nn.Module):
def __init__(self, dtype=torch.bfloat16, bias=True):
super().__init__()
self.proj = nn.Linear(64, 256, bias=bias, device="cuda", dtype=dtype)
gelu.mark_fused_gelu_site(self, "proj")
def forward(self, x):
if gelu.fused_gelu_active(self) and gelu.can_fuse_linear_gelu(self.proj, x):
return gelu.fused_linear_gelu_tanh(x, self.proj.weight, self.proj.bias)
return F.gelu(self.proj(x), approximate="tanh")
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
def test_fused_matches_reference(dtype):
torch.manual_seed(0)
site = _Site(dtype)
x = torch.randn(512, 64, device="cuda", dtype=dtype)
ref = site(x)
assert gelu.mount_fused_linear_gelu(site)
atol = 2e-2 if dtype == torch.bfloat16 else 4e-3
torch.testing.assert_close(site(x), ref, atol=atol, rtol=2e-2)
def test_flux_gelu_proj_site():
"""FLUX.1 shared-FF site: gate off is bit-exact, gate on is close."""
from sglang.multimodal_gen.runtime.models.dits.flux import FluxFusedGELUProj
torch.manual_seed(0)
proj = nn.Linear(3072, 12288, device="cuda", dtype=torch.bfloat16)
site = FluxFusedGELUProj(proj)
x = torch.randn(1, 512, 3072, device="cuda", dtype=torch.bfloat16)
ref = F.gelu(proj(x), approximate="tanh")
assert torch.equal(site(x), ref) # unmounted default: bit-exact reference
assert gelu.mount_fused_linear_gelu(site)
torch.testing.assert_close(site(x), ref, atol=2e-2, rtol=2e-2)
gelu.unmount_fused_linear_gelu(site)
assert torch.equal(site(x), ref)
def test_mount_guards_and_lossless_path():
torch.manual_seed(0)
good, bad = _Site(), _Site(torch.float32)
model = nn.ModuleList([good, bad])
assert not gelu.mount_fused_linear_gelu(model)
assert not gelu.fused_gelu_active(good)
x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
ref = good(x)
assert gelu.mount_fused_linear_gelu(good)
gelu.unmount_fused_linear_gelu(good)
assert torch.equal(good(x), ref)
no_bias = nn.Linear(8, 8, bias=False, device="cuda", dtype=torch.bfloat16)
assert not gelu.can_fuse_linear_gelu_static(no_bias)
assert not gelu.can_fuse_linear_gelu(good.proj, x.float())
@torch.no_grad()
def test_mounted_site_torch_compile_fullgraph():
site = _Site()
x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
assert gelu.mount_fused_linear_gelu(site)
expected = site(x)
actual = torch.compile(site, fullgraph=True)(x)
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/diffusion/test_fused_ln_modulate.py
-82
View File
@@ -1,82 +0,0 @@
import pytest
import torch
import torch.nn as nn
from sglang.kernels.ops.diffusion.fused_ln_modulate import (
can_fuse_ln_modulate,
fused_ln_modulate,
fused_ln_modulate_active,
mark_fused_ln_modulate_site,
mount_fused_ln_modulate,
unmount_fused_ln_modulate,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@pytest.fixture(autouse=True)
def cuda_setup():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
torch.cuda.manual_seed(0)
@pytest.mark.parametrize("seq_len", [4096, 512])
def test_fused_ln_modulate_matches_reference(seq_len):
x = torch.randn((1, seq_len, 3072), device="cuda", dtype=torch.bfloat16)
scale = torch.randn((1, 3072), device="cuda", dtype=torch.bfloat16)
shift = torch.randn_like(scale)
assert can_fuse_ln_modulate(x, scale, shift)
out = fused_ln_modulate(x, scale, shift, eps=1e-6)
norm = nn.LayerNorm(3072, eps=1e-6, elementwise_affine=False).cuda()
ref = norm(x) * (1 + scale[:, None]) + shift[:, None]
# Contract: bf16 rounding-order-level difference only, not bit-exact.
torch.testing.assert_close(out, ref, atol=0.0625, rtol=0.05)
def test_fused_ln_modulate_guards_and_mount_protocol():
x = torch.randn((2, 64, 3072), device="cuda", dtype=torch.bfloat16)
row = torch.randn((2, 3072), device="cuda", dtype=torch.bfloat16)
assert not can_fuse_ln_modulate(x, row, row) # folded affine needs B == 1
root = nn.Module()
root.child = nn.Module()
mark_fused_ln_modulate_site(root.child)
assert not fused_ln_modulate_active(root.child)
assert mount_fused_ln_modulate(root)
assert fused_ln_modulate_active(root.child)
unmount_fused_ln_modulate(root)
assert not fused_ln_modulate_active(root.child)
assert not mount_fused_ln_modulate(nn.Module()) # no marked sites
@torch.no_grad()
def test_mounted_ln_modulate_site_torch_compile_fullgraph():
class Site(nn.Module):
def __init__(self):
super().__init__()
mark_fused_ln_modulate_site(self)
def forward(self, x, scale, shift):
if fused_ln_modulate_active(self) and can_fuse_ln_modulate(x, scale, shift):
return fused_ln_modulate(x, scale, shift, eps=1e-6)
return (
nn.functional.layer_norm(x, (x.shape[-1],), eps=1e-6)
* (1 + scale[:, None])
+ shift[:, None]
)
site = Site()
assert mount_fused_ln_modulate(site)
x = torch.randn(1, 64, 128, device="cuda", dtype=torch.bfloat16)
scale = torch.randn(1, 128, device="cuda", dtype=torch.bfloat16)
shift = torch.randn_like(scale)
expected = site(x, scale, shift)
actual = torch.compile(site, fullgraph=True)(x, scale, shift)
torch.testing.assert_close(actual, expected, atol=0.0625, rtol=0.05)
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/diffusion/test_fused_norm_scale_shift.py
-251
View File
@@ -1,251 +0,0 @@
import sys
from typing import Optional, Tuple
import pytest
import torch
from einops import rearrange
from torch import Tensor
from sglang.kernels.ops.diffusion.cutedsl.scale_residual_norm_scale_shift import (
fused_norm_scale_shift,
fused_scale_residual_norm_scale_shift,
validate_scale_shift,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=28, stage="base-b-kernel-unit", runner_config="1-gpu-large")
DEVICE = "cuda"
SHAPE_MAP = {
"1": lambda B, S, F, D: (1,),
"D": lambda B, S, F, D: (D,),
"1D": lambda B, S, F, D: (1, D),
"BD": lambda B, S, F, D: (B, D),
"11D": lambda B, S, F, D: (1, 1, D),
"B1D": lambda B, S, F, D: (B, 1, D),
"1SD": lambda B, S, F, D: (1, S, D),
"BSD": lambda B, S, F, D: (B, S, D),
"BF1D": lambda B, S, F, D: (B, F, 1, D),
}
SHAPES = [
# (B, S, F, D)
(1, 115200, 1, 3072), # Hunyuan
(1, 32760, 1, 1536), # Wan
(1, 6, 1, 3072), # Qwen
(1, 1024, 8, 3072),
(4, 512, 16, 3072),
]
DTYPES = [torch.float16, torch.bfloat16, torch.float32]
NORM_TYPES = ["layer", "rms"]
AFFINE_MODES = ["D", "NAT"]
INDEX_MODES = ["BSD", "1", "1SD", "BD", "B1D", "D", "1D", "11D", "BF1D"]
def _tol(dtype: torch.dtype):
return 1e-5 if dtype == torch.float32 else 5e-2
@pytest.fixture(autouse=True)
def cuda_setup():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
torch.cuda.manual_seed(0)
def _apply_scale_shift(y: Tensor, scale: Tensor, shift: Tensor) -> Tensor:
if scale.ndim == 4:
num_frame = scale.shape[1]
return rearrange(
rearrange(y, "b (f l) d -> b f l d", f=num_frame) * (1 + scale) + shift,
"b f l d -> b (f l) d",
)
else:
scale = rearrange(scale, "b d -> b 1 d") if scale.ndim == 2 else scale
shift = rearrange(shift, "b d -> b 1 d") if shift.ndim == 2 else shift
return y * (1 + scale) + shift
def fused_norm_scale_shift_ref(
x: Tensor,
weight: Optional[Tensor],
bias: Optional[Tensor],
scale: Tensor,
shift: Tensor,
norm_type: str,
eps: float,
) -> Tensor:
original_dtype = x.dtype
x, weight, bias, scale, shift = (
v.float() if v is not None else v for v in [x, weight, bias, scale, shift]
)
if norm_type == "layer":
norm = torch.layer_norm(x, x.shape[-1:], eps=eps, weight=weight, bias=bias)
else:
norm = torch.rms_norm(x, x.shape[-1:], eps=eps, weight=weight)
return _apply_scale_shift(norm, scale, shift).to(original_dtype)
def fused_scale_residual_norm_scale_shift_ref(
residual: Tensor,
x: Tensor,
gate: Optional[Tensor] | int,
weight: Optional[Tensor],
bias: Optional[Tensor],
scale: Tensor,
shift: Tensor,
norm_type: str,
eps: float,
):
original_dtype = x.dtype
residual, x, gate, weight, bias, scale, shift = (
v.float() if isinstance(v, Tensor) else v
for v in [residual, x, gate, weight, bias, scale, shift]
)
if isinstance(gate, int):
x = residual + gate * x
else:
if gate.ndim == 4:
num_frame = gate.shape[1]
x_fld = rearrange(x, "b (f l) d -> b f l d", f=num_frame)
x = residual + rearrange(x_fld * gate, "b f l d -> b (f l) d")
else:
gate = rearrange(gate, "b d -> b 1 d") if gate.ndim == 2 else gate
x = residual + gate * x
if norm_type == "layer":
norm = torch.layer_norm(x, x.shape[-1:], eps=eps, weight=weight, bias=bias)
else:
norm = torch.rms_norm(x, x.shape[-1:], eps=eps, weight=weight)
y_ref = _apply_scale_shift(norm, scale, shift)
return y_ref.to(original_dtype), x.to(original_dtype)
def _make_tensor(index_mode: str, shape: Tuple, dtype: torch.dtype):
if index_mode == "NAT":
return None
return torch.randn(*SHAPE_MAP[index_mode](*shape), device=DEVICE, dtype=dtype)
def test_validate_scale_shift_rejects_non_divisible_frames():
with pytest.raises(ValueError, match=r"S\(10\) must be divisible by F\(4\)"):
validate_scale_shift(
torch.empty((1, 4, 1, 256), device=DEVICE, dtype=torch.float16),
1,
10,
256,
)
@torch.no_grad()
def run_norm_scale_shift(
shape=SHAPES[0],
dtype=DTYPES[0],
affine_dtype=DTYPES[0],
scale_dtype=DTYPES[0],
shift_dtype=DTYPES[0],
norm_type=NORM_TYPES[0],
affine_mode=AFFINE_MODES[0],
scale_mode="BSD",
shift_mode="BSD",
eps=1e-5,
):
x = _make_tensor("BSD", shape, dtype)
weight = _make_tensor(affine_mode, shape, affine_dtype)
bias = _make_tensor(affine_mode, shape, affine_dtype)
scale = _make_tensor(scale_mode, shape, scale_dtype)
shift = _make_tensor(shift_mode, shape, shift_dtype)
y_dev = fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps)
y_ref = fused_norm_scale_shift_ref(x, weight, bias, scale, shift, norm_type, eps)
torch.testing.assert_close(y_dev, y_ref, atol=_tol(dtype), rtol=_tol(dtype))
@torch.no_grad()
def run_scale_resi_norm_scale_shift(
shape=SHAPES[0],
dtype=DTYPES[0],
affine_dtype=DTYPES[0],
scale_dtype=DTYPES[0],
shift_dtype=DTYPES[0],
norm_type=NORM_TYPES[0],
affine_mode=AFFINE_MODES[0],
gate_mode="B1D",
scale_mode="BSD",
shift_mode="BSD",
eps=1e-5,
):
residual = _make_tensor("BSD", shape, dtype)
x = _make_tensor("BSD", shape, dtype)
gate = _make_tensor(gate_mode, shape, dtype)
weight = _make_tensor(affine_mode, shape, affine_dtype)
bias = _make_tensor(affine_mode, shape, affine_dtype)
scale = _make_tensor(scale_mode, shape, scale_dtype)
shift = _make_tensor(shift_mode, shape, shift_dtype)
y_dev, res_dev = fused_scale_residual_norm_scale_shift(
residual, x, gate, weight, bias, scale, shift, norm_type, eps
)
y_ref, res_ref = fused_scale_residual_norm_scale_shift_ref(
residual, x, gate, weight, bias, scale, shift, norm_type, eps
)
torch.testing.assert_close(y_dev, y_ref, atol=_tol(dtype), rtol=_tol(dtype))
torch.testing.assert_close(res_dev, res_ref, atol=_tol(dtype), rtol=_tol(dtype))
@pytest.mark.parametrize("norm_type", NORM_TYPES)
class TestFusedNormScaleShift:
@pytest.mark.parametrize("shape", SHAPES)
@pytest.mark.parametrize("dtype", DTYPES)
def test_shape_dtype(self, shape, dtype, norm_type):
run_norm_scale_shift(shape=shape, dtype=dtype, norm_type=norm_type)
@pytest.mark.parametrize("dtype", DTYPES)
def test_dtype_0(self, dtype, norm_type):
run_norm_scale_shift(affine_dtype=dtype, norm_type=norm_type)
@pytest.mark.parametrize("dtype", DTYPES)
def test_dtype_1(self, dtype, norm_type):
run_norm_scale_shift(scale_dtype=dtype, shift_dtype=dtype, norm_type=norm_type)
@pytest.mark.parametrize("affine_mode", AFFINE_MODES)
def test_normtype_affine(self, affine_mode, norm_type):
run_norm_scale_shift(affine_mode=affine_mode, norm_type=norm_type)
@pytest.mark.parametrize("index_mode", INDEX_MODES)
def test_index_mode(self, index_mode, norm_type):
run_norm_scale_shift(
scale_mode=index_mode, shift_mode=index_mode, norm_type=norm_type
)
@pytest.mark.parametrize("norm_type", NORM_TYPES)
class TestFusedScaleResidualNormScaleShift:
@pytest.mark.parametrize("shape", SHAPES)
@pytest.mark.parametrize("dtype", DTYPES)
def test_shape_dtype(self, shape, dtype, norm_type):
run_scale_resi_norm_scale_shift(shape=shape, dtype=dtype, norm_type=norm_type)
@pytest.mark.parametrize("dtype", DTYPES)
def test_dtype_0(self, dtype, norm_type):
run_scale_resi_norm_scale_shift(affine_dtype=dtype, norm_type=norm_type)
@pytest.mark.parametrize("dtype", DTYPES)
def test_dtype_1(self, dtype, norm_type):
run_scale_resi_norm_scale_shift(
scale_dtype=dtype, shift_dtype=dtype, norm_type=norm_type
)
@pytest.mark.parametrize("affine_mode", AFFINE_MODES)
def test_normtype_affine(self, affine_mode, norm_type):
run_scale_resi_norm_scale_shift(affine_mode=affine_mode, norm_type=norm_type)
@pytest.mark.parametrize("index_mode", INDEX_MODES)
def test_scale_shift_index_mode(self, index_mode, norm_type):
run_scale_resi_norm_scale_shift(
scale_mode=index_mode, shift_mode=index_mode, norm_type=norm_type
)
@pytest.mark.parametrize("index_mode", INDEX_MODES)
def test_gate_index_mode(self, index_mode, norm_type):
run_scale_resi_norm_scale_shift(gate_mode=index_mode, norm_type=norm_type)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py
-55
View File
@@ -1,55 +0,0 @@
"""GLM-Image fused LN+modulate / qk-LN fast paths must stay bit-exact vs eager."""
import pytest
import torch
import sglang.multimodal_gen.runtime.models.dits.glm_image as glm_image
from sglang.multimodal_gen.runtime.models.dits.glm_image import (
_eager_ln_modulate,
_glm_ln_modulate,
_glm_qk_layernorm,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize("shape", [(1, 4096, 4096), (2, 301, 4096), (1, 1, 2560)])
def test_fused_ln_modulate_is_bit_exact(shape):
# (1, 4096, 4096) is the real GLM-Image image-stream shape (1024^2,
# hidden 4096); the others cover the text stream and another hidden.
torch.manual_seed(0)
batch, seq, hidden = shape
norm = torch.nn.LayerNorm(hidden, eps=1e-5, elementwise_affine=False).cuda()
x = (torch.randn(batch, seq, hidden, device="cuda") * 8).bfloat16()
emb = torch.randn(batch, 12 * hidden, device="cuda").bfloat16()
chunks = emb.chunk(12, dim=1) # strided adaLN projection views
shift, scale = chunks[0], chunks[2]
out = _glm_ln_modulate(norm, x, scale, shift, x.dtype)
assert torch.equal(out, _eager_ln_modulate(norm, x, scale, shift, x.dtype))
assert glm_image._GLM_LN_MOD.verified
assert not glm_image._GLM_LN_MOD.disabled
@pytest.mark.parametrize("shape", [(1, 4360, 32, 128), (2, 37, 3, 40), (1, 129, 5, 64)])
def test_fused_qk_head_layernorm_is_bit_exact(shape):
# (1, 4360, 32, 128) is the real GLM-Image q/k shape (text + image
# tokens, 32 heads of dim 128); the others cover partially-filled warps.
torch.manual_seed(1)
batch, seq, heads, head_dim = shape
norm_q = torch.nn.LayerNorm(head_dim, eps=1e-5, elementwise_affine=False).cuda()
norm_k = torch.nn.LayerNorm(head_dim, eps=1e-5, elementwise_affine=False).cuda()
q = (torch.randn(batch, seq, heads, head_dim, device="cuda") * 5).bfloat16()
k = (torch.randn(batch, seq, heads, head_dim, device="cuda") * 5).bfloat16()
q_out, k_out = _glm_qk_layernorm(norm_q, norm_k, q, k, q.dtype)
assert torch.equal(q_out, norm_q(q).to(q.dtype))
assert torch.equal(k_out, norm_k(k).to(k.dtype))
assert glm_image._GLM_QK_LN.verified
assert not glm_image._GLM_QK_LN.disabled
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/diffusion/test_group_norm_silu.py
-104
View File
@@ -1,104 +0,0 @@
import sys
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from sglang.kernels.ops.diffusion.group_norm_silu import apply_group_norm_silu
from sglang.kernels.ops.diffusion.triton.group_norm_silu import triton_group_norm_silu
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_amd_ci(est_time=15, suite="nightly-amd-kernel-1-gpu", nightly=True)
DEVICE = "cuda"
DTYPES = [torch.float16, torch.bfloat16, torch.float32]
TEST_CASES = [
pytest.param((2, 64, 32, 32), 32, id="image_2d"),
pytest.param((1, 64, 4, 16, 16), 32, id="video_3d"),
pytest.param((4, 128), 32, id="token_2d"),
]
LARGE_TILE_CASE = ((1, 128, 20, 256, 256), 32)
def _tol(dtype: torch.dtype) -> tuple[float, float]:
if dtype == torch.float32:
return 1e-5, 1e-5
if dtype == torch.bfloat16:
return 7e-2, 2e-2
return 3e-3, 3e-3
@pytest.fixture(autouse=True)
def cuda_setup():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
torch.cuda.manual_seed(0)
def _reference(
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor,
num_groups: int,
eps: float = 1e-5,
) -> torch.Tensor:
return F.silu(F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps))
@torch.no_grad()
@pytest.mark.parametrize("shape,num_groups", TEST_CASES)
@pytest.mark.parametrize("dtype", DTYPES)
def test_triton_group_norm_silu(
shape: tuple[int, ...], num_groups: int, dtype: torch.dtype
) -> None:
channels = shape[1]
x = torch.randn(shape, device=DEVICE, dtype=dtype)
weight = torch.randn(channels, device=DEVICE, dtype=dtype)
bias = torch.randn(channels, device=DEVICE, dtype=dtype)
actual = triton_group_norm_silu(x, weight, bias, num_groups=num_groups)
expected = _reference(x, weight, bias, num_groups)
atol, rtol = _tol(dtype)
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
@torch.no_grad()
@pytest.mark.parametrize("shape,num_groups", TEST_CASES[:2])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_apply_group_norm_silu(
shape: tuple[int, ...],
num_groups: int,
dtype: torch.dtype,
) -> None:
norm = nn.GroupNorm(num_groups, shape[1], eps=1e-5, affine=True).to(
device=DEVICE, dtype=dtype
)
activation = nn.SiLU()
hidden_states = torch.randn(shape, device=DEVICE, dtype=dtype)
actual = apply_group_norm_silu(hidden_states, norm, activation)
expected = activation(norm(hidden_states))
atol, rtol = _tol(dtype)
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
@torch.no_grad()
def test_triton_group_norm_silu_large_tile_bf16() -> None:
shape, num_groups = LARGE_TILE_CASE
x = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
weight = torch.randn(shape[1], device=DEVICE, dtype=torch.bfloat16)
bias = torch.randn(shape[1], device=DEVICE, dtype=torch.bfloat16)
actual = triton_group_norm_silu(x, weight, bias, num_groups=num_groups)
expected = _reference(x, weight, bias, num_groups)
atol, rtol = _tol(torch.bfloat16)
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_hunyuanvideo_eager_fusions.py
-94
View File
@@ -1,94 +0,0 @@
"""HunyuanVideo eager QKV/RoPE and quality-gated QKNorm tests."""
import sys
from unittest.mock import patch
import pytest
import torch
import sglang.kernels.ops.diffusion.hunyuan_qknorm as hunyuan_qknorm
from sglang.kernels.ops.diffusion.hunyuan_qknorm import (
mark_hunyuan_qknorm_site,
mount_hunyuan_qknorm,
unmount_hunyuan_qknorm,
)
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
from sglang.multimodal_gen.runtime.layers.rotary_embedding.utils import (
_apply_rotary_emb,
)
from sglang.multimodal_gen.runtime.models.dits.hunyuanvideo import (
_hunyuan_pack_qkv,
_hunyuan_qknorm,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize("img_tokens,txt_tokens", [(257, 31), (4096, 256)])
def test_hunyuan_qkv_rope_pack_is_bit_exact(img_tokens, txt_tokens):
torch.manual_seed(0)
shape_img = (1, img_tokens, 24, 128)
shape_txt = (1, txt_tokens, 24, 128)
img_q, img_k, img_v = (
torch.randn(shape_img, device="cuda", dtype=torch.bfloat16) for _ in range(3)
)
txt_q, txt_k, txt_v = (
torch.randn(shape_txt, device="cuda", dtype=torch.bfloat16) for _ in range(3)
)
cos = torch.randn(img_tokens, 64, device="cuda")
sin = torch.randn_like(cos)
q, k, v = _hunyuan_pack_qkv(img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin)
q_ref = torch.cat(
(_apply_rotary_emb(img_q, cos, sin, is_neox_style=False), txt_q), dim=1
)
k_ref = torch.cat(
(_apply_rotary_emb(img_k, cos, sin, is_neox_style=False), txt_k), dim=1
)
v_ref = torch.cat((img_v, txt_v), dim=1)
assert torch.equal(q, q_ref)
assert torch.equal(k, k_ref)
assert torch.equal(v, v_ref)
def test_hunyuan_quality_qknorm_matches_rmsnorm():
torch.manual_seed(1)
site = torch.nn.Module()
mark_hunyuan_qknorm_site(site)
q_norm = RMSNorm(128, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
k_norm = RMSNorm(128, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
packed = torch.randn(1, 257, 3, 24, 128, device="cuda", dtype=torch.bfloat16)
q, k = packed[:, :, 0], packed[:, :, 1]
q_ref = q_norm(q.contiguous()).to(q)
k_ref = k_norm(k.contiguous()).to(k)
q_unmounted, k_unmounted = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
assert torch.equal(q_unmounted, q_ref)
assert torch.equal(k_unmounted, k_ref)
assert mount_hunyuan_qknorm(site)
q_out, k_out = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
torch.testing.assert_close(q_out, q_ref, atol=2e-2, rtol=2e-2)
torch.testing.assert_close(k_out, k_ref, atol=2e-2, rtol=2e-2)
unmount_hunyuan_qknorm(site)
q_unmounted, k_unmounted = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
assert torch.equal(q_unmounted, q_ref)
assert torch.equal(k_unmounted, k_ref)
def test_hunyuan_quality_qknorm_stays_unmounted_without_cute_kernel():
site = torch.nn.Module()
mark_hunyuan_qknorm_site(site)
with patch.object(hunyuan_qknorm, "_get_qk_rmsnorm_cute", return_value=None):
assert not mount_hunyuan_qknorm(site)
assert not hunyuan_qknorm._FUSION.is_enabled(site)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_import_surface.py
+210
View File
@@ -0,0 +1,210 @@
"""Guards that keep the ``diffusion`` package's import surface from eroding.
The reorganization only stays useful if two invariants hold:
1. runtime code imports from ``sglang.kernels.ops.diffusion`` and not from a
submodule, so the internal layout can move without touching call sites;
2. the facade's ``_EXPORTS`` table and the registry's ``_SPECS`` table both
point at symbols that actually exist.
Neither is checkable by the type system, and both fail silently -- a stale
``_EXPORTS`` entry only raises when some model happens to call that kernel, on
a GPU, at serving time. These are pure-CPU tests: they read the tables and
resolve them with ``importlib``/``ast`` without importing torch backends.
"""
import ast
import importlib
import pathlib
import subprocess
import sys
import pytest
from sglang.kernels.ops.diffusion import _EXPORTS, _SPECS
from sglang.kernels.registry import registry
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=4, suite="base-a-test-cpu")
PACKAGE = "sglang.kernels.ops.diffusion"
_PACKAGE_DIR = pathlib.Path(importlib.import_module(PACKAGE).__file__ or "").parent
_REPO_ROOT = _PACKAGE_DIR.parents[4] # <repo>/python/sglang/kernels/ops/diffusion
# Backend-specific test files may name a leaf module on purpose; everything
# else -- all runtime code -- must go through the facade.
_DEEP_IMPORT_ALLOWLIST = {
"python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py",
"test/registered/kernels/ops/diffusion/test_model_fast_paths.py",
"test/registered/kernels/ops/diffusion/test_sites.py",
}
def _module_defines(module_path: str) -> set[str]:
"""Top-level names bound by a submodule, without importing it.
Importing would pull in Triton / CuTe-DSL / FlyDSL, none of which are
installed on the CPU CI lane -- so this reads the source instead.
"""
path = _PACKAGE_DIR / (module_path.replace(".", "/") + ".py")
if not path.exists():
path = _PACKAGE_DIR / module_path.replace(".", "/") / "__init__.py"
assert path.exists(), f"{PACKAGE}.{module_path} does not exist"
names: set[str] = set()
for node in ast.parse(path.read_text(encoding="utf-8")).body:
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
names.add(node.name)
elif isinstance(node, ast.Assign):
names.update(t.id for t in node.targets if isinstance(t, ast.Name))
elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
names.add(node.target.id)
elif isinstance(node, (ast.Import, ast.ImportFrom)):
names.update((a.asname or a.name).split(".")[0] for a in node.names)
elif isinstance(node, (ast.If, ast.Try)):
# Platform-conditional rebinds (``x = select_impl(...)``) and
# guarded defs still bind a public name.
for inner in ast.walk(node):
if isinstance(inner, (ast.FunctionDef, ast.ClassDef)):
names.add(inner.name)
elif isinstance(inner, ast.Assign):
names.update(t.id for t in inner.targets if isinstance(t, ast.Name))
return names
def test_every_export_resolves_to_a_real_symbol():
missing = [
f"{symbol} -> {module}"
for symbol, module in sorted(_EXPORTS.items())
if symbol not in _module_defines(module)
]
assert not missing, f"stale _EXPORTS entries: {missing}"
def test_every_symbol_imported_from_the_facade_is_exported():
"""The reverse of the check above, and the one that actually bites.
A missing ``_EXPORTS`` entry raises ``ImportError`` at module import, so a
module-level ``from ...diffusion import x`` fails loudly. A *function-local*
one -- the pattern used for optional backends -- fails only when that test
or code path runs, on the platform that has the backend. Enumerating the
call sites catches it here instead.
"""
unexported = set()
for root in ("python/sglang", "test", "benchmark"):
root_dir = _REPO_ROOT / root
if not root_dir.exists():
continue
for path in root_dir.rglob("*.py"):
rel = path.relative_to(_REPO_ROOT).as_posix()
if rel.startswith("python/sglang/kernels/ops/diffusion/"):
continue
try:
tree = ast.parse(path.read_text(encoding="utf-8"))
except (SyntaxError, UnicodeDecodeError):
continue
for node in ast.walk(tree):
if isinstance(node, ast.ImportFrom) and node.module == PACKAGE:
unexported.update(
a.name
for a in node.names
if a.name not in _EXPORTS and not a.name.startswith("_")
)
assert not unexported, f"imported but not in _EXPORTS: {sorted(unexported)}"
def test_every_registered_spec_target_resolves():
missing = []
for _op, _backend, target, _caps, _description in _SPECS:
module, _, attr = target.partition(":")
if attr not in _module_defines(module):
missing.append(target)
assert not missing, f"stale _SPECS targets: {missing}"
def test_registry_holds_the_diffusion_ops():
# Registration happens at package import, is metadata-only, and is what
# ``select_kernel`` / the tracing tools read.
registered = {op for op in registry.ops() if op.startswith("diffusion.")}
assert {op for op, *_ in _SPECS} <= registered
def test_facade_rejects_unknown_attributes():
module = sys.modules[PACKAGE]
with pytest.raises(AttributeError):
module.definitely_not_a_kernel
assert set(module.__all__) == set(_EXPORTS)
assert set(_EXPORTS) <= set(dir(module))
def test_importing_the_package_does_not_import_any_leaf_module():
"""The reason ``__getattr__`` is lazy rather than a block of re-exports.
The backends have disjoint, heavy, mutually-exclusive dependencies --
Triton (CUDA/ROCm), CUTLASS/CuTe-DSL, FlyDSL (gfx950), MLX (Apple). If
``_EXPORTS`` ever degrades into eager ``from .norm.x import y`` lines, all
of them become import-time requirements on every platform, which is how a
CPU-only or Apple install starts failing at ``import sglang``.
Asserted on this package's own leaf modules rather than on ``triton`` in
``sys.modules``: sibling operator groups import Triton for their own
reasons, so a global check would not isolate this package's behavior.
Run in a fresh interpreter because this process has already resolved
exports through the facade.
"""
code = (
"import importlib, sys\n"
f"importlib.import_module('{PACKAGE}')\n"
f"prefix = '{PACKAGE}.'\n"
"leaves = [m for m in sys.modules if m.startswith(prefix)"
" and not m.endswith('__init__')]\n"
"print(','.join(sorted(m for m in leaves if '.' in m[len(prefix):]"
" or sys.modules[m].__file__ and not sys.modules[m].__file__"
".endswith('__init__.py'))))\n"
)
result = subprocess.run(
[sys.executable, "-c", code], capture_output=True, text=True, timeout=600
)
assert result.returncode == 0, result.stderr
leaked = [m for m in result.stdout.strip().split(",") if m]
assert not leaked, f"importing {PACKAGE} eagerly imported: {leaked}"
@pytest.mark.parametrize("root", ["python/sglang", "test", "benchmark"])
def test_runtime_code_imports_only_through_the_facade(root):
root_dir = _REPO_ROOT / root
if not root_dir.exists(): # source checkouts only
pytest.skip(f"{root} not present in this install")
offenders = []
for path in root_dir.rglob("*.py"):
rel = path.relative_to(_REPO_ROOT).as_posix()
if rel.startswith("python/sglang/kernels/ops/diffusion/"):
continue # intra-package imports are the point of the subpackages
if rel in _DEEP_IMPORT_ALLOWLIST:
continue
try:
tree = ast.parse(path.read_text(encoding="utf-8"))
except (SyntaxError, UnicodeDecodeError):
continue
for node in ast.walk(tree):
if (
isinstance(node, ast.ImportFrom)
and node.module
and node.module.startswith(f"{PACKAGE}.")
):
offenders.append(f"{rel}:{node.lineno} imports {node.module}")
elif isinstance(node, ast.Import):
offenders.extend(
f"{rel}:{node.lineno} imports {a.name}"
for a in node.names
if a.name.startswith(f"{PACKAGE}.")
)
assert not offenders, (
"import from sglang.kernels.ops.diffusion instead of a submodule:\n "
+ "\n ".join(offenders)
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_layout.py
+597
View File
@@ -0,0 +1,597 @@
"""``diffusion.layout``: data-movement kernels.
Every kernel in this domain only moves values (plus zero fill, plus at most
one same-order add), so each one is *bitwise* identical to the aten chain it
replaces. That makes ``torch.equal`` -- not ``assert_close`` -- the right
assertion throughout this file; a tolerance here would hide a real bug.
Covered: USP output head merge, Ulysses destination-major QKV pack, varlen
pack/scatter, causal Conv3d cat+pad (CUDA and Triton), and the Wan causal-VAE
cache kernels.
"""
import sys
from unittest.mock import patch
import pytest
import torch
import torch.nn.functional as F
from sglang.kernels.jit.utils import get_ci_test_range
from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func
from sglang.kernels.ops.diffusion import (
build_inv_indices,
can_use_usp_merge_heads,
cat_pad_channels_last_3d,
dup_up3d_add,
)
from sglang.kernels.ops.diffusion import (
fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton,
)
from sglang.kernels.ops.diffusion import (
fused_causal_conv3d_cat_pad_cuda,
fused_pack_qkv,
fused_scatter_to_padded,
pack_qkv_destination_major,
usp_merge_heads,
)
from sglang.multimodal_gen.runtime.layers.attention.backends import (
flash_attn as _fa_backend,
)
from sglang.multimodal_gen.runtime.layers.attention.layer import build_varlen_mask_meta
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=110, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
# Nightly is not redundant: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1, which
# expands the get_ci_test_range sweeps below.
register_cuda_ci(est_time=20, stage="nightly", runner_config="1-gpu-large")
register_amd_ci(est_time=10, stage="jit-kernel-unit", runner_config="amd")
register_amd_ci(est_time=15, suite="nightly-amd-kernel-1-gpu", nightly=True)
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
DEVICE = "cuda"
def _cl3d(shape, dtype):
return torch.randn(shape, device=DEVICE, dtype=dtype).contiguous(
memory_format=torch.channels_last_3d
)
# ---------------------------------------------------------------------------
# USP output head merge
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"world,seq,batch,h_local,head_dim",
[
(4, 7936, 1, 14, 128), # H3 768p production shape (Ulysses 4)
(2, 64, 3, 4, 64), # batched
(4, 33, 2, 4, 100), # scalar fallback inside the CUDA kernel
],
)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
@pytest.mark.skipif(
bool(torch.version.hip),
reason="the USP merge-heads JIT fast path is CUDA-only by design -- "
"can_use_usp_merge_heads() returns False under HIP, and the aten fallback "
"it degrades to is covered by the unsupported-inputs test below",
)
def test_usp_merge_heads_bitwise(dtype, world, seq, batch, h_local, head_dim):
generator = torch.Generator(device=DEVICE).manual_seed(4321)
x = torch.randn(
world,
seq,
batch,
h_local,
head_dim,
dtype=dtype,
device=DEVICE,
generator=generator,
)
assert can_use_usp_merge_heads(x)
out = usp_merge_heads(x)
ref = x.permute(2, 1, 0, 3, 4).contiguous()
assert out.shape == ref.shape
assert torch.equal(out, ref)
def test_usp_merge_heads_unsupported_inputs_use_exact_fallback():
# The wrapper degrades to the aten permute for anything the fast path
# rejects -- a wrong rank, a transposed view, an empty leading dim, or a
# ROCm build -- so callers never need their own guard.
x = torch.randn(2, 4, 1, 4, 64, dtype=torch.bfloat16, device=DEVICE)
for value in (x.transpose(0, 1), x[:0], x[0]):
assert not can_use_usp_merge_heads(value)
if value.dim() == 5:
assert torch.equal(
usp_merge_heads(value), value.permute(2, 1, 0, 3, 4).contiguous()
)
with patch.object(torch.version, "hip", "6.3"):
assert not can_use_usp_merge_heads(x)
assert torch.equal(usp_merge_heads(x), x.permute(2, 1, 0, 3, 4).contiguous())
# ---------------------------------------------------------------------------
# Ulysses destination-major QKV pack
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_pack_qkv_destination_major_is_bit_exact(dtype):
torch.manual_seed(0)
rows, world_size, global_heads, head_size = 17, 4, 12, 64
q, k, v = (
torch.randn(rows, global_heads, head_size, device=DEVICE, dtype=dtype)
for _ in range(3)
)
local_heads = global_heads // world_size
expected = torch.empty(
world_size, rows, local_heads, 3 * head_size, device=DEVICE, dtype=dtype
)
for index, tensor in enumerate((q, k, v)):
shards = tensor.view(rows, world_size, local_heads, head_size).permute(
1, 0, 2, 3
)
expected[..., index * head_size : (index + 1) * head_size].copy_(shards)
assert torch.equal(pack_qkv_destination_major(q, k, v, world_size), expected)
def test_pack_qkv_destination_major_validates_inputs():
q = torch.empty(2, 4, 8, device=DEVICE, dtype=torch.bfloat16)
with pytest.raises(ValueError, match="same 3D shape"):
pack_qkv_destination_major(q, q[:, :-1], q, 2)
with pytest.raises(ValueError, match="divide global_heads"):
pack_qkv_destination_major(q, q, q, 3)
with pytest.raises(ValueError, match="expected shape"):
pack_qkv_destination_major(q, q, q, 2, out=torch.empty_like(q))
# ---------------------------------------------------------------------------
# Varlen pack / scatter
# ---------------------------------------------------------------------------
VARLEN_DTYPES = get_ci_test_range([torch.bfloat16, torch.float16], [torch.bfloat16])
# (name, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens)
VARLEN_SHAPES = get_ci_test_range(
[
("small_c2", 2, 64, 128, 4, 64, [32, 48]),
("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
("all_valid_b1", 1, 64, 128, 4, 64, [64]),
("all_valid_b4", 4, 64, 128, 4, 64, [64, 64, 64, 64]),
("c8_prod", 8, 256, 4096, 24, 128, [128, 200, 256, 100, 50, 256, 256, 50]),
# one batch with zero valid text tokens (image side still valid)
("zero_txt_one_batch", 2, 64, 128, 4, 64, [0, 32]),
# bs=1 with no text validity (only image rows packed)
("bs1_zero_txt", 1, 64, 128, 4, 64, [0]),
],
[
("small_c2", 2, 64, 128, 4, 64, [32, 48]),
("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
("all_valid_b4", 4, 64, 128, 4, 64, [64, 64, 64, 64]),
],
)
def _build_mask(bs, s_txt, s_img, valid_txt_lens):
mask = torch.zeros(bs, s_txt + s_img, dtype=torch.bool, device=DEVICE)
for b, vt in enumerate(valid_txt_lens):
mask[b, :vt] = True
mask[b, s_txt:] = True
return mask
def _build_meta(mask):
bs, seq = mask.shape
indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
return indices, build_inv_indices(indices, bs * seq)
@pytest.mark.parametrize("dtype", VARLEN_DTYPES)
@pytest.mark.parametrize("shape", VARLEN_SHAPES, ids=lambda s: s[0])
def test_varlen_pack_matches_index_select(dtype, shape):
_, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
torch.manual_seed(0)
s = s_txt + s_img
indices, _ = _build_meta(_build_mask(bs, s_txt, s_img, valid_txt_lens))
q, k, v = (
torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
for _ in range(3)
)
fused = fused_pack_qkv(q, k, v, indices)
for got, src in zip(fused, (q, k, v), strict=True):
want = src.reshape(bs * s, num_heads, head_dim).index_select(0, indices)
assert torch.equal(got, want)
@pytest.mark.parametrize("dtype", VARLEN_DTYPES)
@pytest.mark.parametrize("shape", VARLEN_SHAPES, ids=lambda s: s[0])
def test_varlen_scatter_matches_index_copy(dtype, shape):
_, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
torch.manual_seed(1)
s = s_txt + s_img
mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
indices, inv_indices = _build_meta(mask)
out_unpad = torch.randn(
indices.shape[0], num_heads, head_dim, dtype=dtype, device=DEVICE
)
flat = torch.zeros(bs * s, num_heads, head_dim, dtype=dtype, device=DEVICE)
flat.index_copy_(0, indices, out_unpad)
out_ref = flat.view(bs, s, num_heads, head_dim)
out_fused = fused_scatter_to_padded(out_unpad, inv_indices, bs, s)
assert torch.equal(out_ref, out_fused)
invalid = ~mask
if invalid.any():
# Padding rows must be exactly zero, not merely small.
assert out_fused[invalid].abs().max().item() == 0.0
def test_varlen_pack_handles_non_contiguous_input():
# Q/K/V arrive as (B, H, S, D) permutes from attention; the helper must
# make them contiguous itself rather than reading the wrong strides.
torch.manual_seed(2)
bs, s_txt, s_img, num_heads, head_dim = 2, 64, 128, 4, 64
indices, _ = _build_meta(_build_mask(bs, s_txt, s_img, [32, 48]))
pre = torch.randn(
bs, num_heads, s_txt + s_img, head_dim, dtype=torch.bfloat16, device=DEVICE
)
q, k, v = (torch.randn_like(pre).permute(0, 2, 1, 3) for _ in range(3))
assert not q.is_contiguous()
fused = fused_pack_qkv(q, k, v, indices)
for got, src in zip(fused, (q, k, v), strict=True):
want = src.contiguous().flatten(0, 1).index_select(0, indices)
assert torch.equal(got, want)
def test_build_inv_indices_matches_manual():
torch.manual_seed(3)
bs, s = 2, 32
mask = torch.bernoulli(torch.full((bs, s), 0.6, device=DEVICE)).to(torch.bool)
indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
manual = torch.full((bs * s,), -1, dtype=torch.int32, device=DEVICE)
if indices.numel():
manual[indices.long()] = torch.arange(
indices.numel(), dtype=torch.int32, device=DEVICE
)
assert torch.equal(build_inv_indices(indices, bs * s), manual)
def test_varlen_empty_valid_set_handled():
# An all-False mask is reachable (a request whose text side is fully
# masked): pack must return empty tensors and scatter an all-zero dense
# output rather than launching a degenerate grid.
bs, s, num_heads, head_dim = 2, 16, 4, 64
indices = torch.zeros(bs, s, dtype=torch.bool, device=DEVICE).reshape(-1).nonzero()
indices = indices.flatten()
inv_indices = build_inv_indices(indices, bs * s)
assert indices.numel() == 0
q = torch.randn(bs, s, num_heads, head_dim, dtype=torch.bfloat16, device=DEVICE)
unpad = fused_pack_qkv(q, q.clone(), q.clone(), indices)
assert all(t.shape == (0, num_heads, head_dim) for t in unpad)
out_padded = fused_scatter_to_padded(unpad[0], inv_indices, bs, s)
assert out_padded.shape == (bs, s, num_heads, head_dim)
assert out_padded.abs().max().item() == 0.0
# The kernels above are unit-tested against index_select/index_copy_; this
# section drives them through the production USPAttention masked branch, where
# a wrong index layout would produce plausible-looking attention output rather
# than an obvious mismatch.
def _sdpa_with_key_mask(q, k, v, key_mask, softmax_scale):
"""Reference: SDPA with a ``[B, S]`` key mask broadcast to ``[B, 1, 1, S]``."""
q_ = q.transpose(1, 2)
k_ = k.transpose(1, 2)
v_ = v.transpose(1, 2)
mask = key_mask.to(dtype=q.dtype)[:, None, None, :]
mask = (mask - 1.0) * torch.finfo(q.dtype).max
out = F.scaled_dot_product_attention(
q_,
k_,
v_,
attn_mask=mask,
dropout_p=0.0,
is_causal=False,
scale=softmax_scale,
)
return out.transpose(1, 2)
def _varlen_path(q, k, v, key_mask, softmax_scale):
"""Production varlen path matching USPAttention.forward masked branch."""
bs, seq = q.shape[0], q.shape[1]
meta = build_varlen_mask_meta(key_mask)
indices = meta["indices"]
if indices.shape[0] == 0:
return torch.zeros_like(q)
q_unpad, k_unpad, v_unpad = fused_pack_qkv(q, k, v, indices)
try:
out_unpad = flash_attn_varlen_func(
q=q_unpad,
k=k_unpad,
v=v_unpad,
cu_seqlens_q=meta["cu_seqlens"],
cu_seqlens_k=meta["cu_seqlens"],
max_seqlen_q=meta["max_seqlen"],
max_seqlen_k=meta["max_seqlen"],
softmax_scale=softmax_scale,
causal=False,
ver=_fa_backend.fa_ver,
)
except ImportError as exc: # pragma: no cover - image-dependent
# ``flash_attn_varlen_func`` resolves its backend lazily, so an image
# without the selected FlashAttention build raises here rather than at
# import. This file also runs on the B200 lane (for the causal-Conv3d
# section), which ships no ``flash_attn`` -- skip only this end-to-end
# comparison there; the pack/scatter kernels themselves are covered
# unit-wise above on every lane.
pytest.skip(f"FlashAttention varlen v{_fa_backend.fa_ver} unavailable: {exc}")
return fused_scatter_to_padded(out_unpad, meta["inv_indices"], bs, seq)
@pytest.mark.parametrize("dtype", VARLEN_DTYPES)
@pytest.mark.parametrize("shape", VARLEN_SHAPES, ids=lambda s: s[0])
def test_varlen_path_matches_sdpa_on_valid_rows(dtype, shape):
"""Valid rows: varlen output ≈ SDPA output within FA tolerance."""
_, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
torch.manual_seed(0)
s = s_txt + s_img
softmax_scale = head_dim**-0.5
mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
out_sdpa = _sdpa_with_key_mask(q, k, v, mask, softmax_scale)
out_varlen = _varlen_path(q, k, v, mask, softmax_scale)
valid = mask[..., None, None].expand_as(out_sdpa)
rtol = 1e-2 if dtype == torch.bfloat16 else 5e-3
atol = 5e-2 if dtype == torch.bfloat16 else 1e-2
torch.testing.assert_close(
out_sdpa[valid],
out_varlen[valid],
rtol=rtol,
atol=atol,
)
@pytest.mark.parametrize("dtype", VARLEN_DTYPES)
@pytest.mark.parametrize("shape", VARLEN_SHAPES, ids=lambda s: s[0])
def test_varlen_path_zeros_masked_rows(dtype, shape):
"""Masked rows: varlen path produces exact zeros (documented contract)."""
_, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
torch.manual_seed(1)
s = s_txt + s_img
softmax_scale = head_dim**-0.5
mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
out_varlen = _varlen_path(q, k, v, mask, softmax_scale)
invalid = ~mask
if invalid.any():
assert (out_varlen[invalid] == 0).all(), "masked rows must be zero-filled"
# ---------------------------------------------------------------------------
# Causal Conv3d cat + pad (CUDA JIT vs Triton)
# ---------------------------------------------------------------------------
CONV3D_CASES = get_ci_test_range(
[
(1024, 1, 30, 52, 1),
(1024, 1, 30, 52, 2),
(1024, 2, 60, 104, 1),
(1024, 2, 60, 104, 2),
(512, 4, 120, 208, 1),
(512, 4, 120, 208, 2),
(256, 4, 240, 416, 1),
(256, 4, 240, 416, 2),
],
[(1024, 1, 30, 52, 1), (512, 4, 120, 208, 2)],
)
def _conv3d_inputs(channels, t_size, h_size, w_size, cache_t):
generator = torch.Generator(device=DEVICE)
generator.manual_seed(channels * 1009 + t_size * 251 + h_size + cache_t)
x = torch.randn(
(1, channels, t_size, h_size, w_size),
device=DEVICE,
dtype=torch.bfloat16,
generator=generator,
)
cache_x = torch.randn(
(1, channels, cache_t, h_size, w_size),
device=DEVICE,
dtype=torch.bfloat16,
generator=generator,
)
return x, cache_x, (1, 1, 1, 1, cache_t, 0)
@pytest.mark.parametrize("channels,t_size,h_size,w_size,cache_t", CONV3D_CASES)
def test_causal_conv3d_cat_pad_cuda_matches_triton(
channels, t_size, h_size, w_size, cache_t
):
x, cache_x, padding = _conv3d_inputs(channels, t_size, h_size, w_size, cache_t)
actual = fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding)
expected = fused_causal_conv3d_cat_pad_triton(x, cache_x, padding)
assert torch.equal(actual, expected)
def test_causal_conv3d_cat_pad_torch_compile():
# The CUDA path is a registered custom op, so a fullgraph compile must not
# graph-break on it.
x, cache_x, padding = _conv3d_inputs(1024, 1, 30, 52, 1)
@torch.compile(fullgraph=True)
def fn(x, cache_x):
return fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding)
assert torch.equal(
fn(x, cache_x), fused_causal_conv3d_cat_pad_triton(x, cache_x, padding)
)
# ---------------------------------------------------------------------------
# Wan causal VAE cache kernels
# ---------------------------------------------------------------------------
def _ref_cat_pad(x, cache, padding):
p = list(padding)
if cache is not None:
x = torch.cat([cache, x], dim=2)
p[4] -= cache.shape[2]
if any(p):
x = F.pad(x, p)
return x.contiguous(memory_format=torch.channels_last_3d)
@torch.no_grad()
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
@pytest.mark.parametrize(
"c,t,h,w,cache_t,pads",
[
(96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)), # first chunk, zero-fill front
(96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)), # legacy 1-frame cache
(96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)), # steady state k3 conv
(64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)), # time_conv (temporal only)
(48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)), # encoder-style T=4 chunk
],
)
def test_cat_pad_channels_last_3d_bitwise(dtype, c, t, h, w, cache_t, pads):
torch.cuda.manual_seed(0)
x = _cl3d((1, c, t, h, w), dtype)
cache = None
if cache_t:
# Strided interior view: caches may arrive as non-contiguous slices.
ph, pw = pads[2], pads[0]
buf = _cl3d((1, c, cache_t, h + 2 * ph, w + 2 * pw), dtype)
cache = buf[:, :, :, ph : ph + h, pw : pw + w]
ref = _ref_cat_pad(x, cache, pads)
out = cat_pad_channels_last_3d(x, cache, pads)
assert out is not None and out.shape == ref.shape
assert out.is_contiguous(memory_format=torch.channels_last_3d)
assert torch.equal(out, ref)
# Dual-output mode: the same pass also emits the compact feature cache
# (unpadded interior of the last frames), bitwise equal to the slice.
pair = cat_pad_channels_last_3d(x, cache, pads, keep_cache_t=2)
assert pair is not None
out2, keep = pair
assert torch.equal(out2, ref)
ph, pw = pads[2], pads[0]
keep_t = min(2, ref.shape[2])
want = ref[:, :, ref.shape[2] - keep_t :, ph : ph + h, pw : pw + w]
assert keep.shape == want.shape
assert keep.is_contiguous(memory_format=torch.channels_last_3d)
assert torch.equal(keep, want)
@torch.no_grad()
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
@pytest.mark.parametrize(
"c_in,c_out,t,h,w,ft,fs,drop",
[
(128, 64, 1, 10, 14, 2, 2, False),
(128, 64, 1, 10, 14, 2, 2, True), # first_chunk slicing
(64, 32, 2, 10, 14, 1, 2, False),
],
)
def test_dup_up3d_add_bitwise(dtype, c_in, c_out, t, h, w, ft, fs, drop):
torch.cuda.manual_seed(0)
repeats = c_out * ft * fs * fs // c_in
src = _cl3d((1, c_in, t, h, w), dtype)
t_out = t * ft - (ft - 1 if drop else 0)
# Main arm as a permuted view, like the WanResample 2D output.
main = torch.randn(
(1, t_out, c_out, h * fs, w * fs), device=DEVICE, dtype=dtype
).permute(0, 2, 1, 3, 4)
dup = src.repeat_interleave(repeats, dim=1)
dup = dup.view(1, c_out, ft, fs, fs, t, h, w)
dup = dup.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
dup = dup.view(1, c_out, t * ft, h * fs, w * fs)
if drop:
dup = dup[:, :, ft - 1 :, :, :]
ref = main + dup
out = dup_up3d_add(main, src, ft, fs, repeats, drop)
assert out is not None and out.shape == ref.shape
# Layout must match the aten add output exactly (downstream reductions
# are layout-sensitive), and every value must be bitwise identical.
assert out.stride() == ref.stride()
assert torch.equal(out, ref)
@torch.no_grad()
@pytest.mark.parametrize("pads_temporal_only", [False, True])
def test_wan_cached_conv_chunk_loop_bitwise(pads_temporal_only):
"""The fused conv-input/compact-cache scheme must reproduce the original
clone/cat bookkeeping bitwise across a chunked decode, including the
first-chunk zero fill and the "Rep" marker start used by WanResample."""
from sglang.multimodal_gen.runtime.models.vaes import wanvae
from sglang.multimodal_gen.runtime.models.vaes.wanvae import (
CACHE_T,
WanCausalConv3d,
_cache_payload,
_run_cached_causal_conv,
)
torch.cuda.manual_seed(0)
c = 64
if pads_temporal_only:
conv = WanCausalConv3d(c, 2 * c, (3, 1, 1), padding=(1, 0, 0))
else:
conv = WanCausalConv3d(c, c, 3, padding=1)
conv = conv.to(device=DEVICE, dtype=torch.float32)
conv.weight.data = conv.weight.data.contiguous(memory_format=torch.channels_last_3d)
chunks = [_cl3d((1, c, 1, 10, 14), torch.float32) for _ in range(4)]
def run(force_fallback, start):
cache = [start]
outs = []
orig = wanvae.cat_pad_channels_last_3d
if force_fallback:
wanvae.cat_pad_channels_last_3d = None
try:
for x in chunks:
outs.append(_run_cached_causal_conv(conv, x, cache, 0))
finally:
wanvae.cat_pad_channels_last_3d = orig
return outs, cache[0]
for start in (None, "Rep"):
fused_outs, fused_cache = run(False, start)
ref_outs, ref_cache = run(True, start)
for got, want in zip(fused_outs, ref_outs, strict=True):
assert torch.equal(got, want)
got_payload = _cache_payload(fused_cache)
assert got_payload is not None and got_payload.shape[2] == CACHE_T
# Reference cache holds the last CACHE_T unpadded frames.
assert torch.equal(got_payload, ref_cache[:, :, -CACHE_T:])
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_ltx2_ada_values.py
-92
View File
@@ -1,92 +0,0 @@
import sys
import pytest
import torch
from sglang.kernels.ops.diffusion.triton.ltx2_ada_values import ltx2_ada_values9
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_amd_ci(est_time=8, suite="nightly-amd-kernel-1-gpu", nightly=True)
DEVICE = "cuda"
@pytest.fixture(autouse=True)
def cuda_setup():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
torch.cuda.manual_seed(0)
def _reference(
scale_shift_table: torch.Tensor,
timestep: torch.Tensor,
) -> tuple[torch.Tensor, ...]:
batch, seq, _ = timestep.shape
hidden = scale_shift_table.shape[1]
return (
scale_shift_table.to(device=timestep.device, dtype=timestep.dtype)
.view(1, 1, 9, hidden)
.add(timestep.reshape(batch, seq, 9, hidden))
.unbind(dim=2)
)
@torch.no_grad()
@pytest.mark.parametrize("batch,seq,hidden", [(1, 1, 4096), (2, 3, 2048)])
@pytest.mark.parametrize("table_dtype", [torch.bfloat16, torch.float32])
def test_ltx2_ada_values9(
batch: int,
seq: int,
hidden: int,
table_dtype: torch.dtype,
) -> None:
scale_shift_table = torch.randn(
9, hidden, device=DEVICE, dtype=table_dtype
).contiguous()
timestep = torch.randn(
batch, seq, 9 * hidden, device=DEVICE, dtype=torch.bfloat16
).contiguous()
actual = ltx2_ada_values9(scale_shift_table, timestep)
expected = _reference(scale_shift_table, timestep)
assert len(actual) == 9
for actual_value, expected_value in zip(actual, expected):
assert actual_value.is_contiguous()
torch.testing.assert_close(actual_value, expected_value, atol=0, rtol=0)
@torch.no_grad()
def test_ltx2_ada_values9_torch_compile_fullgraph() -> None:
hidden = 4096
scale_shift_table = torch.randn(
9, hidden, device=DEVICE, dtype=torch.bfloat16
).contiguous()
timestep = torch.randn(
1, 1, 9 * hidden, device=DEVICE, dtype=torch.bfloat16
).contiguous()
actual = torch.compile(ltx2_ada_values9, fullgraph=True)(
scale_shift_table, timestep
)
expected = _reference(scale_shift_table, timestep)
assert len(actual) == 9
for actual_value, expected_value in zip(actual, expected):
assert actual_value.is_contiguous()
torch.testing.assert_close(actual_value, expected_value, atol=0, rtol=0)
@torch.no_grad()
def test_ltx2_ada_values9_rejects_unsupported_shape() -> None:
scale_shift_table = torch.randn(8, 4096, device=DEVICE, dtype=torch.bfloat16)
timestep = torch.randn(1, 1, 9 * 4096, device=DEVICE, dtype=torch.bfloat16)
with pytest.raises(ValueError, match="scale_shift_table"):
ltx2_ada_values9(scale_shift_table, timestep)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py
-88
View File
@@ -1,88 +0,0 @@
"""LTX-2 quality=high RMSNorm+modulate fusion: gated, close to eager."""
import sys
import pytest
import torch
from torch import nn
import sglang.multimodal_gen.runtime.models.dits.ltx_2 as ltx2_module
from sglang.kernels.ops.diffusion.ltx2_rmsnorm_modulate import (
fused_ltx2_rms_norm_modulate,
mark_ltx2_rms_norm_modulate_site,
mount_ltx2_rms_norm_modulate,
unmount_ltx2_rms_norm_modulate,
)
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNormNoWeight
from sglang.multimodal_gen.runtime.models.dits.ltx_2 import _ltx2_rms_norm_modulate
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_amd_ci(est_time=8, suite="nightly-amd-kernel-1-gpu", nightly=True)
@pytest.fixture(autouse=True)
def _setup():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
torch.cuda.manual_seed(0)
def _eager(rms, x, scale, shift, eps):
return rms(x, eps) * (1 + scale) + shift
def _inputs(hidden, batch=1, seq=4096):
rms = RMSNormNoWeight()
x = torch.randn(batch, seq, hidden, device="cuda", dtype=torch.bfloat16)
scale = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
shift = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
return rms, x, scale, shift
# hidden 4096 = LTX-2 video stream, 2048 = audio stream.
@pytest.mark.parametrize("hidden", [4096, 2048])
def test_lossless_default_is_bitexact(hidden):
# A marked-but-unmounted site uses only the self-verified bit-exact
# modulate fast path after the reference aten RMSNorm.
block = nn.Module()
mark_ltx2_rms_norm_modulate_site(block)
rms, x, scale, shift = _inputs(hidden)
out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
assert torch.equal(out, _eager(rms, x, scale, shift, 1e-6))
def test_lossless_compile_keeps_expression_visible_to_inductor(monkeypatch):
block = nn.Module()
mark_ltx2_rms_norm_modulate_site(block)
rms, x, scale, shift = _inputs(2048, seq=126)
monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
monkeypatch.setattr(
ltx2_module,
"_ltx2_modulate",
lambda *_args: pytest.fail("compiled path must not call the opaque custom op"),
)
out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
assert torch.equal(out, _eager(rms, x, scale, shift, 1e-6))
@pytest.mark.parametrize("hidden", [4096, 2048])
def test_mounted_high_uses_fused_kernel(hidden):
block = nn.Module()
mark_ltx2_rms_norm_modulate_site(block)
assert mount_ltx2_rms_norm_modulate(block)
try:
rms, x, scale, shift = _inputs(hidden)
out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
# The mounted path routes through the fused kernel exactly.
assert torch.equal(out, fused_ltx2_rms_norm_modulate(x, scale, shift, 1e-6))
# And stays within half-precision rounding of the eager reference.
ref = _eager(rms, x, scale, shift, 1e-6)
assert torch.allclose(out.float(), ref.float(), atol=3e-2, rtol=1e-2)
finally:
unmount_ltx2_rms_norm_modulate(block)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_model_fast_paths.py
+712
View File
@@ -0,0 +1,712 @@
"""Per-model fast paths: each model wrapper must reproduce its own reference.
A diffusion kernel is only as good as the wrapper that decides when to use it,
and that decision is model-specific: FLUX.1 and GLM-Image feed different adaLN
layouts to the same LayerNorm+modulate kernel, Sana only engages on non-default
streams, ERNIE runs a bit-exact chain unconditionally. Kernel-level numerics
live in ``test_norm.py`` / ``test_modulate.py`` / ``test_rope.py``; this file
covers the wiring: right kernel, right reference, gate ends in the right state.
Two assertion styles appear, and the difference is the point:
- ``torch.equal`` for the bit-exact default-on paths. These self-verify at
runtime, so a tolerance here would let a real regression through -- the gate
would silently fall back to eager and the fusion would simply stop running.
- a tolerance for the quality-gated paths, which are *documented* as differing
from eager at half-precision rounding-order level.
Each section also asserts the gate ended up ``verified`` / not ``disabled``:
without it a test still passes when the fast path never engaged at all.
"""
import sys
import unittest
from unittest.mock import patch
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.models.upsampling import Upsample2D
import sglang.kernels.ops.diffusion.sites.hunyuan_qknorm_site as hunyuan_qknorm
import sglang.multimodal_gen.runtime.models.dits.ernie_image as ernie_image
import sglang.multimodal_gen.runtime.models.dits.flux as flux
import sglang.multimodal_gen.runtime.models.dits.flux_2 as flux2
import sglang.multimodal_gen.runtime.models.dits.glm_image as glm_image
import sglang.multimodal_gen.runtime.models.dits.ltx_2 as ltx2_module
import sglang.multimodal_gen.runtime.models.dits.sana as sana
from sglang.kernels.ops.diffusion import (
can_use_wan_rmsnorm_silu,
fused_ltx2_rms_norm_modulate,
mark_fused_ln_modulate_site,
mark_hunyuan_qknorm_site,
mark_ltx2_rms_norm_modulate_site,
mount_fused_ln_modulate,
mount_hunyuan_qknorm,
mount_ltx2_rms_norm_modulate,
unmount_hunyuan_qknorm,
unmount_ltx2_rms_norm_modulate,
wan_rmsnorm_silu,
)
from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import (
StableDiffusion3VAEConfig,
)
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, RMSNormNoWeight
from sglang.multimodal_gen.runtime.layers.rotary_embedding.utils import (
_apply_rotary_emb,
)
from sglang.multimodal_gen.runtime.models.dits.ernie_image import (
_ernie_gated_norm_scale_shift,
_ernie_norm_scale_shift,
_ernie_qknorm_rope,
_ernie_qknorm_rope_reference,
)
from sglang.multimodal_gen.runtime.models.dits.flux import (
_flux_fused_ln_modulate,
_flux_norm_modulate,
)
from sglang.multimodal_gen.runtime.models.dits.flux_2 import (
_flux2_norm_modulate,
_flux2_swiglu,
)
from sglang.multimodal_gen.runtime.models.dits.glm_image import (
_eager_ln_modulate as _glm_eager_ln_modulate,
)
from sglang.multimodal_gen.runtime.models.dits.glm_image import (
_glm_ln_modulate,
_glm_qk_layernorm,
)
from sglang.multimodal_gen.runtime.models.dits.hunyuanvideo import (
_hunyuan_pack_qkv,
_hunyuan_qknorm,
)
from sglang.multimodal_gen.runtime.models.dits.ltx_2 import _ltx2_rms_norm_modulate
from sglang.multimodal_gen.runtime.models.dits.sana import (
_eager_ln_modulate as _sana_eager_ln_modulate,
)
from sglang.multimodal_gen.runtime.models.dits.sana import (
_sana_ln_modulate,
)
from sglang.multimodal_gen.runtime.models.vaes import flux2_vae_cuda_opt as vae_opt
from sglang.multimodal_gen.runtime.models.vaes.autoencoder import AutoencoderKL
from sglang.multimodal_gen.runtime.models.vaes.fast_path_gate import use_vae_fast_path
from sglang.multimodal_gen.runtime.models.vaes.wan_vae_cuda_opt import (
FusedWanRMSNormSiLU,
VaeFastPathGate,
)
from sglang.multimodal_gen.runtime.models.vaes.wanvae import WanRMS_norm
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=95, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_amd_ci(est_time=8, suite="nightly-amd-kernel-1-gpu", nightly=True)
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.fixture(autouse=True)
def _seed_cuda():
"""Every wrapper below asserts against a reference computed from the same
random draw, so the seed must be fixed per test, not per module."""
torch.cuda.manual_seed(0)
# -------------------------------------------------------------------------
# FLUX.1 -- LayerNorm + adaLN modulate, and the shared-FF GELU site
# -------------------------------------------------------------------------
def _flux_eager(norm, x, scale, shift):
return norm(x) * (1 + scale[:, None]) + shift[:, None]
def _flux_site_inputs(shape, chunks, seed):
torch.manual_seed(seed)
batch, seq, hidden = shape
norm = torch.nn.LayerNorm(hidden, eps=1e-6, elementwise_affine=False).cuda()
x = (torch.randn(batch, seq, hidden, device="cuda") * 8).bfloat16()
emb = torch.randn(batch, chunks * hidden, device="cuda").bfloat16()
parts = emb.chunk(chunks, dim=1) # strided adaLN projection views
return norm, x, parts[0], parts[1]
@pytest.mark.parametrize(
"shape,chunks",
[
((1, 4096, 3072), 6), # dual-stream image tokens (1024^2), chunk(6)
((1, 512, 3072), 6), # dual-stream text tokens
((1, 4608, 3072), 3), # single-stream concat, chunk(3)
((2, 300, 3072), 6), # CFG batch, odd seq
],
)
def test_flux_fused_ln_modulate_is_bit_exact(shape, chunks):
# Every distinct (shape, stride, eps) signature the FLUX.1 sites emit
# must verify torch.equal on first sight and stay enabled.
norm, x, shift, scale = _flux_site_inputs(shape, chunks, seed=0)
out = _flux_fused_ln_modulate(norm, x, scale, shift)
assert out is not None
assert torch.equal(out, _flux_eager(norm, x, scale, shift))
assert not flux._FLUX_LN_MOD.disabled
assert flux._FLUX_LN_MOD.verified
def test_flux_norm_modulate_bitexact_supersedes_high_fold():
# With the quality="high" affine fold mounted, the bit-exact kernel
# still takes priority, so the site output stays lossless.
site = torch.nn.Module()
mark_fused_ln_modulate_site(site)
assert mount_fused_ln_modulate(site)
norm, x, shift, scale = _flux_site_inputs((1, 128, 3072), 6, seed=1)
out = _flux_norm_modulate(site, norm, x, scale, shift)
assert torch.equal(out, _flux_eager(norm, x, scale, shift))
def test_flux_fused_ln_modulate_rejects_unsupported_hidden():
# hidden % 4 != 0 is outside the kernel contract and must bail out.
norm, x, shift, scale = _flux_site_inputs((1, 64, 3070), 6, seed=2)
assert _flux_fused_ln_modulate(norm, x, scale, shift) is None
# -------------------------------------------------------------------------
# FLUX.2 -- packed norm+modulate and packed SwiGLU views
# -------------------------------------------------------------------------
@unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
class TestFlux2EagerFusions(CustomTestCase):
def setUp(self):
flux2._FLUX2_LN_MOD.disabled = False
flux2._FLUX2_LN_MOD.verified = False
flux2._FLUX2_LN_MOD_SIGS.clear()
flux2._FLUX2_SWIGLU.disabled = False
flux2._FLUX2_SWIGLU.verified = False
flux2._FLUX2_SWIGLU_SIGS.clear()
def test_norm_modulate_is_bit_exact_across_sequence_lengths(self):
torch.manual_seed(0)
hidden = 256
norm = torch.nn.LayerNorm(
hidden, eps=1e-6, elementwise_affine=False, device="cuda"
)
# FLUX.2 modulation values are views of one packed projection.
params = torch.randn(1, 1, 6 * hidden, device="cuda").bfloat16()
shift, scale = params.chunk(6, dim=-1)[:2]
for seq in (17, 65):
x = torch.randn(1, seq, hidden, device="cuda").bfloat16()
expected = norm(x) * (1 + scale) + shift
actual = _flux2_norm_modulate(norm, x, scale, shift)
self.assertTrue(torch.equal(actual, expected))
self.assertFalse(flux2._FLUX2_LN_MOD.disabled)
self.assertEqual(len(flux2._FLUX2_LN_MOD_SIGS), 1)
def test_packed_swiglu_is_bit_exact_for_contiguous_and_strided_views(self):
torch.manual_seed(1)
hidden = 384
inputs = [
torch.randn(1, 19, 2 * hidden, device="cuda").bfloat16(),
torch.randn(1, 19, 3 * hidden, device="cuda").bfloat16()[..., : 2 * hidden],
]
for x in inputs:
expected = F.silu(x[..., :hidden]) * x[..., hidden:]
actual = _flux2_swiglu(x)
self.assertTrue(torch.equal(actual, expected))
self.assertFalse(flux2._FLUX2_SWIGLU.disabled)
self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 2)
def test_fp16_preserves_reference_path(self):
x = torch.randn(1, 17, 512, device="cuda", dtype=torch.float16)
expected = F.silu(x[..., :256]) * x[..., 256:]
actual = _flux2_swiglu(x)
self.assertTrue(torch.equal(actual, expected))
self.assertFalse(flux2._FLUX2_SWIGLU.disabled)
def test_packed_swiglu_rejects_non_dense_outer_stride(self):
base = torch.randn(2, 23, 512, device="cuda", dtype=torch.bfloat16)
x = base[:, :19]
self.assertNotEqual(x.stride(0), x.shape[1] * x.stride(1))
expected = F.silu(x[..., :256]) * x[..., 256:]
actual = _flux2_swiglu(x)
self.assertTrue(torch.equal(actual, expected))
self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 0)
def test_new_swiglu_signature_is_not_verified_during_graph_capture(self):
first = torch.randn(1, 17, 512, device="cuda", dtype=torch.bfloat16)
self.assertTrue(
torch.equal(
_flux2_swiglu(first),
F.silu(first[..., :256]) * first[..., 256:],
)
)
self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1)
second = torch.randn(1, 19, 768, device="cuda", dtype=torch.bfloat16)
with patch("torch.cuda.is_current_stream_capturing", return_value=True):
actual = _flux2_swiglu(second)
expected = F.silu(second[..., :384]) * second[..., 384:]
self.assertTrue(torch.equal(actual, expected))
self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1)
# -------------------------------------------------------------------------
# GLM-Image -- LayerNorm + modulate and per-head qk LayerNorm
# -------------------------------------------------------------------------
@pytest.mark.parametrize("shape", [(1, 4096, 4096), (2, 301, 4096), (1, 1, 2560)])
def test_glm_ln_modulate_is_bit_exact(shape):
# (1, 4096, 4096) is the real GLM-Image image-stream shape (1024^2,
# hidden 4096); the others cover the text stream and another hidden.
torch.manual_seed(0)
batch, seq, hidden = shape
norm = torch.nn.LayerNorm(hidden, eps=1e-5, elementwise_affine=False).cuda()
x = (torch.randn(batch, seq, hidden, device="cuda") * 8).bfloat16()
emb = torch.randn(batch, 12 * hidden, device="cuda").bfloat16()
chunks = emb.chunk(12, dim=1) # strided adaLN projection views
shift, scale = chunks[0], chunks[2]
out = _glm_ln_modulate(norm, x, scale, shift, x.dtype)
assert torch.equal(out, _glm_eager_ln_modulate(norm, x, scale, shift, x.dtype))
assert glm_image._GLM_LN_MOD.verified
assert not glm_image._GLM_LN_MOD.disabled
@pytest.mark.parametrize("shape", [(1, 4360, 32, 128), (2, 37, 3, 40), (1, 129, 5, 64)])
def test_glm_qk_head_layernorm_is_bit_exact(shape):
# (1, 4360, 32, 128) is the real GLM-Image q/k shape (text + image
# tokens, 32 heads of dim 128); the others cover partially-filled warps.
torch.manual_seed(1)
batch, seq, heads, head_dim = shape
norm_q = torch.nn.LayerNorm(head_dim, eps=1e-5, elementwise_affine=False).cuda()
norm_k = torch.nn.LayerNorm(head_dim, eps=1e-5, elementwise_affine=False).cuda()
q = (torch.randn(batch, seq, heads, head_dim, device="cuda") * 5).bfloat16()
k = (torch.randn(batch, seq, heads, head_dim, device="cuda") * 5).bfloat16()
q_out, k_out = _glm_qk_layernorm(norm_q, norm_k, q, k, q.dtype)
assert torch.equal(q_out, norm_q(q).to(q.dtype))
assert torch.equal(k_out, norm_k(k).to(k.dtype))
assert glm_image._GLM_QK_LN.verified
assert not glm_image._GLM_QK_LN.disabled
# -------------------------------------------------------------------------
# Sana -- stream-conditional LayerNorm + modulate
# -------------------------------------------------------------------------
@pytest.mark.parametrize(
"shape,nmod,transposed",
[
((2, 1024, 2240), 6, False),
((2, 1024, 2240), 2, False),
((1, 1024, 2240), 6, True),
((1, 37, 2240), 6, False),
],
)
def test_sana_fused_ln_modulate_is_bit_exact(shape, nmod, transposed):
# (., 1024, 2240) is the real Sana 1024px shape; hidden 2240 % 512 != 0
# exercises the kernel's partial tail chunk. nmod mirrors the two adaLN
# chunk layouts, transposed the permuted layout the Sana DiT serves.
torch.manual_seed(0)
batch, seq, hidden = shape
norm = torch.nn.LayerNorm(hidden, eps=1e-6, elementwise_affine=False).cuda()
x = (torch.randn(batch, seq, hidden, device="cuda") * 4).bfloat16()
if transposed:
x = x.permute(0, 2, 1).contiguous().permute(0, 2, 1)
emb = torch.randn(batch, nmod, hidden, device="cuda").bfloat16()
shift, scale = emb.chunk(nmod, dim=1)[0], emb.chunk(nmod, dim=1)[-1]
# default-stream eager serving must stay on the untouched eager chain
n_sigs = len(sana._SANA_LN_MOD.verified_sigs)
_sana_ln_modulate(norm, x, scale, shift)
assert len(sana._SANA_LN_MOD.verified_sigs) == n_sigs
# The fusion engages on non-default streams (the BCG warmup/capture path).
# x/scale/shift were filled on the default stream, so the side stream must
# wait for that work before reading them -- without this the fused kernel
# can read a half-written tensor, the first-sight torch.equal check fails,
# and the gate disables itself *permanently*, which then breaks every later
# parametrization too. It only loses the race when the GPU is contended,
# which is why it shows up on shared CI runners and not on an idle box.
side = torch.cuda.Stream()
side.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(side):
out = _sana_ln_modulate(norm, x, scale, shift)
assert len(sana._SANA_LN_MOD.verified_sigs) == n_sigs + 1 # verified
out2 = _sana_ln_modulate(norm, x, scale, shift) # verified-sig lane
torch.cuda.current_stream().wait_stream(side)
torch.cuda.synchronize()
assert torch.equal(out, _sana_eager_ln_modulate(norm, x, scale, shift))
assert torch.equal(out2, out) and not sana._SANA_LN_MOD.disabled
# -------------------------------------------------------------------------
# ERNIE-Image -- bit-exact RMSNorm scale/shift and rotate-half RoPE
# -------------------------------------------------------------------------
@pytest.mark.parametrize("shape", [(1, 4216, 4096), (2, 1140, 4096), (1, 128, 2048)])
def test_ernie_norm_scale_shift_is_bit_exact(shape):
# (1, 4216, 4096) is the real ERNIE-Image shape (1024^2 image + text
# tokens, hidden 4096); 2048 covers the threads_per_row=32 regime.
torch.manual_seed(0)
batch, seq, hidden = shape
norm = RMSNorm(hidden, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
norm.weight.copy_(torch.randn(hidden))
x = torch.randn(batch, seq, hidden, device="cuda", dtype=torch.bfloat16)
residual = torch.randn_like(x)
update = torch.randn_like(x)
scale = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
shift = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
gate = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
out = _ernie_norm_scale_shift(norm, x, scale, shift)
ref = norm(x) * (1 + scale) + shift
assert torch.equal(out, ref)
out2, res = _ernie_gated_norm_scale_shift(
norm, residual, update, gate, scale, shift
)
res_ref = residual + gate * update
ref2 = norm(res_ref) * (1 + scale) + shift
assert torch.equal(res, res_ref)
assert torch.equal(out2, ref2)
# the fast paths must actually be in use (not silently disabled)
assert ernie_image._ERNIE_NORM.verified
assert ernie_image._ERNIE_GATED_NORM.verified
assert not ernie_image._ERNIE_NORM.disabled
assert not ernie_image._ERNIE_GATED_NORM.disabled
def test_ernie_qknorm_rope_is_bit_exact():
torch.manual_seed(1)
ernie_image._ERNIE_QKNORM_ROPE.disabled = False
ernie_image._ERNIE_QKNORM_ROPE.verified = False
batch, seq, heads, head_dim = 1, 257, 32, 128
q = torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
q_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
k_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
cos = torch.randn(seq, head_dim, device="cuda", dtype=torch.bfloat16)
sin = torch.randn_like(cos)
cache = torch.cat((cos, sin), dim=-1).contiguous()
positions = torch.arange(seq, device="cuda", dtype=torch.long)
q_ref, k_ref = _ernie_qknorm_rope_reference(
q.clone(), k.clone(), q_norm, k_norm, head_dim, cos, sin
)
q_out, k_out = _ernie_qknorm_rope(
q,
k,
q_norm,
k_norm,
head_dim,
cos,
sin,
cache,
positions,
)
assert torch.equal(q_out, q_ref)
assert torch.equal(k_out, k_ref)
assert ernie_image._ERNIE_QKNORM_ROPE.verified
assert not ernie_image._ERNIE_QKNORM_ROPE.disabled
def test_ernie_qknorm_rope_first_attempt_exception_uses_pristine_inputs():
torch.manual_seed(2)
ernie_image._ERNIE_QKNORM_ROPE.disabled = False
ernie_image._ERNIE_QKNORM_ROPE.verified = False
batch, seq, heads, head_dim = 1, 17, 4, 128
q = torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
q_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
k_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
cos = torch.randn(seq, head_dim, device="cuda", dtype=torch.bfloat16)
sin = torch.randn_like(cos)
cache = torch.cat((cos, sin), dim=-1).contiguous()
positions = torch.arange(seq, device="cuda", dtype=torch.long)
q_ref, k_ref = _ernie_qknorm_rope_reference(
q.clone(), k.clone(), q_norm, k_norm, head_dim, cos, sin
)
def mutate_then_raise(**kwargs):
kwargs["q"].zero_()
kwargs["k"].zero_()
raise RuntimeError("synthetic kernel failure")
with patch.object(ernie_image, "apply_qk_norm_rope", mutate_then_raise):
q_out, k_out = _ernie_qknorm_rope(
q,
k,
q_norm,
k_norm,
head_dim,
cos,
sin,
cache,
positions,
)
assert torch.equal(q_out, q_ref)
assert torch.equal(k_out, k_ref)
assert ernie_image._ERNIE_QKNORM_ROPE.disabled
# -------------------------------------------------------------------------
# LTX-2 -- weightless RMSNorm + modulate (quality-gated)
# -------------------------------------------------------------------------
def _ltx2_eager(rms, x, scale, shift, eps):
return rms(x, eps) * (1 + scale) + shift
def _ltx2_inputs(hidden, batch=1, seq=4096):
rms = RMSNormNoWeight()
x = torch.randn(batch, seq, hidden, device="cuda", dtype=torch.bfloat16)
scale = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
shift = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
return rms, x, scale, shift
# hidden 4096 = LTX-2 video stream, 2048 = audio stream.
@pytest.mark.parametrize("hidden", [4096, 2048])
def test_ltx2_lossless_default_is_bitexact(hidden):
# A marked-but-unmounted site uses only the self-verified bit-exact
# modulate fast path after the reference aten RMSNorm.
block = nn.Module()
mark_ltx2_rms_norm_modulate_site(block)
rms, x, scale, shift = _ltx2_inputs(hidden)
out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
assert torch.equal(out, _ltx2_eager(rms, x, scale, shift, 1e-6))
def test_ltx2_lossless_compile_keeps_expression_visible_to_inductor(monkeypatch):
block = nn.Module()
mark_ltx2_rms_norm_modulate_site(block)
rms, x, scale, shift = _ltx2_inputs(2048, seq=126)
monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
monkeypatch.setattr(
ltx2_module,
"_ltx2_modulate",
lambda *_args: pytest.fail("compiled path must not call the opaque custom op"),
)
out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
assert torch.equal(out, _ltx2_eager(rms, x, scale, shift, 1e-6))
@pytest.mark.parametrize("hidden", [4096, 2048])
def test_ltx2_mounted_high_uses_fused_kernel(hidden):
block = nn.Module()
mark_ltx2_rms_norm_modulate_site(block)
assert mount_ltx2_rms_norm_modulate(block)
try:
rms, x, scale, shift = _ltx2_inputs(hidden)
out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
# The mounted path routes through the fused kernel exactly.
assert torch.equal(out, fused_ltx2_rms_norm_modulate(x, scale, shift, 1e-6))
# And stays within half-precision rounding of the eager reference.
ref = _ltx2_eager(rms, x, scale, shift, 1e-6)
assert torch.allclose(out.float(), ref.float(), atol=3e-2, rtol=1e-2)
finally:
unmount_ltx2_rms_norm_modulate(block)
# -------------------------------------------------------------------------
# HunyuanVideo -- QKV/RoPE pack and quality-gated QK RMSNorm
# -------------------------------------------------------------------------
@pytest.mark.parametrize("img_tokens,txt_tokens", [(257, 31), (4096, 256)])
def test_hunyuan_qkv_rope_pack_is_bit_exact(img_tokens, txt_tokens):
torch.manual_seed(0)
shape_img = (1, img_tokens, 24, 128)
shape_txt = (1, txt_tokens, 24, 128)
img_q, img_k, img_v = (
torch.randn(shape_img, device="cuda", dtype=torch.bfloat16) for _ in range(3)
)
txt_q, txt_k, txt_v = (
torch.randn(shape_txt, device="cuda", dtype=torch.bfloat16) for _ in range(3)
)
cos = torch.randn(img_tokens, 64, device="cuda")
sin = torch.randn_like(cos)
q, k, v = _hunyuan_pack_qkv(img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin)
q_ref = torch.cat(
(_apply_rotary_emb(img_q, cos, sin, is_neox_style=False), txt_q), dim=1
)
k_ref = torch.cat(
(_apply_rotary_emb(img_k, cos, sin, is_neox_style=False), txt_k), dim=1
)
v_ref = torch.cat((img_v, txt_v), dim=1)
assert torch.equal(q, q_ref)
assert torch.equal(k, k_ref)
assert torch.equal(v, v_ref)
def test_hunyuan_quality_qknorm_matches_rmsnorm():
torch.manual_seed(1)
site = torch.nn.Module()
mark_hunyuan_qknorm_site(site)
q_norm = RMSNorm(128, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
k_norm = RMSNorm(128, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
packed = torch.randn(1, 257, 3, 24, 128, device="cuda", dtype=torch.bfloat16)
q, k = packed[:, :, 0], packed[:, :, 1]
q_ref = q_norm(q.contiguous()).to(q)
k_ref = k_norm(k.contiguous()).to(k)
q_unmounted, k_unmounted = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
assert torch.equal(q_unmounted, q_ref)
assert torch.equal(k_unmounted, k_ref)
assert mount_hunyuan_qknorm(site)
q_out, k_out = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
torch.testing.assert_close(q_out, q_ref, atol=2e-2, rtol=2e-2)
torch.testing.assert_close(k_out, k_ref, atol=2e-2, rtol=2e-2)
unmount_hunyuan_qknorm(site)
q_unmounted, k_unmounted = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
assert torch.equal(q_unmounted, q_ref)
assert torch.equal(k_unmounted, k_ref)
def test_hunyuan_quality_qknorm_stays_unmounted_without_cute_kernel():
site = torch.nn.Module()
mark_hunyuan_qknorm_site(site)
with patch.object(hunyuan_qknorm, "_get_qk_rmsnorm_cute", return_value=None):
assert not mount_hunyuan_qknorm(site)
assert not hunyuan_qknorm._FUSION.is_enabled(site)
# -------------------------------------------------------------------------
# Wan VAE -- fused RMSNorm+SiLU module gate
# -------------------------------------------------------------------------
def _wan_cl3d(shape, dtype):
return torch.randn(shape, device="cuda", dtype=dtype).contiguous(
memory_format=torch.channels_last_3d
)
@torch.no_grad()
def test_wan_vae_gate_dispatch() -> None:
# Gate off must stay bit-exact; gate on must route to the fused kernel.
torch.cuda.manual_seed(0)
norm = WanRMS_norm(96, images=False).to(device="cuda", dtype=torch.bfloat16)
norm.gamma.add_(torch.randn_like(norm.gamma))
gate = VaeFastPathGate()
fused = FusedWanRMSNormSiLU(norm, gate)
# Parameter names must not change (weight transfer matches by name).
assert [n for n, _ in fused.named_parameters()] == ["gamma"]
x = _wan_cl3d((1, 96, 3, 10, 14), torch.bfloat16)
assert torch.equal(fused(x), nn.SiLU()(norm(x)))
gate.enabled = True
expected = wan_rmsnorm_silu(x, norm.gamma, rms_scale=float(norm.scale))
assert torch.equal(fused(x), expected)
@torch.no_grad()
def test_wan_vae_rejects_empty_input() -> None:
x = torch.empty(1, 96, 0, 2, 2, device="cuda", dtype=torch.bfloat16).to(
memory_format=torch.channels_last_3d
)
gamma = torch.ones(96, 1, 1, 1, device="cuda", dtype=torch.bfloat16)
assert not can_use_wan_rmsnorm_silu(x, gamma, None)
# -------------------------------------------------------------------------
# FLUX.2 VAE -- fused GroupNorm+SiLU and folded 2x upsample conv
# -------------------------------------------------------------------------
@torch.no_grad()
def test_flux2_vae_fast_path():
torch.manual_seed(0)
gate = vae_opt.VaeFastPathGate()
gn = nn.GroupNorm(32, 128, eps=1e-6).to("cuda", torch.bfloat16)
x = torch.randn(1, 128, 64, 64, device="cuda", dtype=torch.bfloat16).to(
memory_format=torch.channels_last
)
ref = F.silu(gn(x))
fused_gn = vae_opt.FusedGroupNormSiLU(gn, gate)
assert set(fused_gn.state_dict()) == {"weight", "bias"}
assert torch.equal(fused_gn(x), ref) # gate off: bit-exact reference
gate.enabled = True
fast = fused_gn(x)
assert fast.is_contiguous(memory_format=torch.channels_last)
torch.testing.assert_close(fast.float(), ref.float(), atol=0.06, rtol=0)
gate.enabled = False
up = Upsample2D(channels=32, use_conv=True).to("cuda", torch.bfloat16)
fused_up = vae_opt.FusedUpsample2xConv2d(up, gate)
assert set(fused_up.state_dict()) == {"conv.weight", "conv.bias"}
x = torch.randn(2, 32, 33, 29, device="cuda", dtype=torch.bfloat16)
ref = up(x)
assert torch.equal(fused_up(x), ref)
assert fused_up._fused_weight is None
gate.enabled = True
fast = fused_up(x)
assert fused_up._fused_weight is not None
ref_range = ref.float().max() - ref.float().min()
relative_mse = F.mse_loss(fast.float(), ref.float()) / ref_range.square()
assert relative_mse < 3.2e-5
# ---------------------------------------------------------------------------
# AutoencoderKL (generic) -- fast-path install must not disturb the checkpoint
# ---------------------------------------------------------------------------
def _small_config():
config = StableDiffusion3VAEConfig()
config.arch_config.latent_channels = 2
config.arch_config.block_out_channels = (4, 4)
config.arch_config.down_block_types = ("DownEncoderBlock2D",) * 2
config.arch_config.up_block_types = ("UpDecoderBlock2D",) * 2
config.arch_config.layers_per_block = 1
config.arch_config.norm_num_groups = 1
config.arch_config.sample_size = 8
return config
@torch.no_grad()
def test_autoencoder_kl_fastpath_install():
torch.manual_seed(0)
vae = AutoencoderKL(_small_config()).to("cuda", torch.bfloat16).eval()
ref_names = {n for n, _ in vae.named_parameters()}
ref_sd = {k: v.clone() for k, v in vae.state_dict().items()}
z = torch.randn(1, 2, 8, 8, device="cuda", dtype=torch.bfloat16)
ref = vae.decode(z)
opt = vae_opt.maybe_optimize_autoencoder_kl(vae)
# Wrappers must not change parameter FQNs; strict load must round-trip.
assert {n for n, _ in opt.named_parameters()} == ref_names
opt.load_state_dict(ref_sd, strict=True)
# Gate off: bit-for-bit the original path.
assert torch.equal(opt.decode(z), ref)
# use_vae_fast_path() is a no-op when nothing registered a gate, so check
# the wrappers went in before relying on it to switch paths.
assert any(
isinstance(m, (vae_opt.FusedGroupNormSiLU, vae_opt.FusedUpsample2xConv2d))
for m in opt.modules()
)
# Gate on: fast path runs and stays close; leaving the scope restores exact.
with use_vae_fast_path(opt, True):
torch.testing.assert_close(opt.decode(z).float(), ref.float(), atol=0.1, rtol=0)
assert torch.equal(opt.decode(z), ref)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_modulate.py
+424
View File
@@ -0,0 +1,424 @@
"""``diffusion.modulate``: adaLN modulation, gating and timestep conditioning.
The bit-exact kernels here (``modulate_scale_shift``, ``residual_gate_add``,
``ltx2_ada_values9``, ``try_fused_scaled_residual_add_exact``) reproduce every
aten rounding boundary, so they are asserted with ``torch.equal``. The
select-0/1 LayerNorm fusions compute their statistics differently from the
reference chain and are asserted with a tolerance.
"""
import sys
import pytest
import torch
from sglang.kernels.jit.utils import get_ci_test_range
from sglang.kernels.ops.diffusion import (
can_use_modulate_scale_shift_cuda,
can_use_residual_gate_add_cuda,
fuse_layernorm_scale_shift_gate_select01_kernel,
fuse_residual_layernorm_scale_shift_gate_select01_kernel,
ltx2_ada_values9,
modulate_scale_shift,
modulate_scale_shift_cuda,
norm_infer,
residual_gate_add,
residual_gate_add_cuda,
timestep_embedding,
try_fused_scaled_residual_add_exact,
)
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=75, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
# Nightly is not redundant: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1, which
# expands the get_ci_test_range sweeps below.
register_cuda_ci(est_time=50, stage="nightly", runner_config="1-gpu-large")
register_amd_ci(est_time=38, suite="nightly-amd-kernel-1-gpu", nightly=True)
DEVICE = "cuda"
@pytest.fixture(autouse=True)
def cuda_setup():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
torch.cuda.manual_seed(0)
# ---------------------------------------------------------------------------
# modulate: x * (1 + scale) + shift
# ---------------------------------------------------------------------------
# FLUX.1 1024^2 adaLN shapes (D=3072) plus batched and odd-length coverage.
MODULATE_CASES = [
(1, 4096, 3072),
(1, 512, 3072),
(1, 4608, 3072),
(2, 1024, 3072),
(1, 17, 64),
]
def _eager_modulate(x, scale, shift):
return x * (1 + scale[:, None]) + shift[:, None]
@pytest.mark.parametrize("shape", MODULATE_CASES)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
def test_modulate_scale_shift_matches_eager(shape, dtype):
x = torch.randn(shape, device=DEVICE, dtype=dtype)
scale = torch.randn((shape[0], shape[-1]), device=DEVICE, dtype=dtype)
shift = torch.randn_like(scale)
assert torch.equal(
modulate_scale_shift_cuda(x, scale, shift), _eager_modulate(x, scale, shift)
)
def test_modulate_scale_shift_accepts_adaln_chunk_views():
# Production feeds strided ``emb.chunk(6)`` views, not fresh tensors.
x = torch.randn((1, 4096, 3072), device=DEVICE, dtype=torch.bfloat16)
emb = torch.randn((1, 6 * 3072), device=DEVICE, dtype=torch.bfloat16)
shift, scale = emb.chunk(6, dim=1)[:2]
assert can_use_modulate_scale_shift_cuda(x, scale, shift)
assert torch.equal(
modulate_scale_shift_cuda(x, scale, shift), _eager_modulate(x, scale, shift)
)
def test_modulate_scale_shift_guards_reject_fp32():
x = torch.randn((1, 64, 64), device=DEVICE, dtype=torch.float32)
row = torch.randn((1, 64), device=DEVICE, dtype=torch.float32)
assert not can_use_modulate_scale_shift_cuda(x, row, row)
# The public wrapper still returns the eager result on a rejected input.
assert torch.equal(modulate_scale_shift(x, row, row), _eager_modulate(x, row, row))
# ---------------------------------------------------------------------------
# residual + gate * update
# ---------------------------------------------------------------------------
GATE_CASES = [
((1, 1024, 4096), (1, 1, 4096)),
((1, 512, 4096), (1, 512, 4096)),
((1, 17, 65), (1, 1, 65)),
((1, 17, 65), (1, 17, 65)),
# FLUX.1 / FLUX.2-klein 1024^2 shapes (D=3072): dual-stream image/text
# and single-stream/joint concat; gates are [1, 1, D] modulation rows.
((1, 4096, 3072), (1, 1, 3072)),
((1, 512, 3072), (1, 1, 3072)),
((1, 4608, 3072), (1, 1, 3072)),
# FLUX.2-dev (D=6144) joint sequence.
((1, 4608, 6144), (1, 1, 6144)),
# ERNIE-4.5-VL 1024^2 image tokens plus text tokens.
((1, 4216, 4096), (1, 1, 4096)),
]
def _assert_gate_add(out, ref):
if ref.dtype == torch.float32:
# fp32 has no rounding boundary to reproduce; the kernel keeps the
# accumulation in fp32 and only order may differ.
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
else:
assert torch.equal(out, ref)
@pytest.mark.parametrize("residual_shape,gate_shape", GATE_CASES)
def test_residual_gate_add_matches_torch(residual_shape, gate_shape):
residual = torch.randn(residual_shape, device=DEVICE, dtype=torch.bfloat16)
update = torch.randn_like(residual)
gate = torch.randn(gate_shape, device=DEVICE, dtype=torch.bfloat16)
ref = residual + update * gate
_assert_gate_add(residual_gate_add_cuda(residual, update, gate), ref)
assert torch.equal(residual_gate_add(residual, update, gate), ref)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@pytest.mark.parametrize("gate_shape", [(1, 1, 64), (1, 9, 64)])
def test_residual_gate_add_dtypes(dtype, gate_shape):
residual = torch.randn((1, 9, 64), device=DEVICE, dtype=dtype)
update = torch.randn_like(residual)
gate = torch.randn(gate_shape, device=DEVICE, dtype=dtype)
_assert_gate_add(
residual_gate_add_cuda(residual, update, gate), residual + update * gate
)
def test_residual_gate_add_guards_and_eager_fallback():
residual = torch.randn((1, 8, 64), device=DEVICE, dtype=torch.bfloat16)
update = torch.randn_like(residual)
gate = torch.randn((1, 1, 64), device=DEVICE, dtype=torch.bfloat16)
assert can_use_residual_gate_add_cuda(residual, update, gate)
rejected = [
(residual.cpu(), update, gate), # not on device
(residual, update.float(), gate), # mixed dtypes
(residual, update[:, ::2], gate), # strided update
(residual, update, gate[:, :, ::2]), # strided gate
(residual[:, :0], update[:, :0], gate), # empty token dim
]
for args in rejected:
assert not can_use_residual_gate_add_cuda(*args)
# Only [1, ..., 1, D] row-broadcast gates are supported; a batched
# [B>1, 1, D] gate is not row-broadcast here and must fall back.
batched = torch.randn((2, 8, 64), device=DEVICE, dtype=torch.bfloat16)
batched_update = torch.randn_like(batched)
batched_gate = torch.randn((2, 1, 64), device=DEVICE, dtype=torch.bfloat16)
assert not can_use_residual_gate_add_cuda(batched, batched_update, batched_gate)
assert torch.equal(
residual_gate_add(batched, batched_update, batched_gate),
batched + batched_update * batched_gate,
)
def test_residual_gate_add_torch_compile_fullgraph():
residual = torch.randn((1, 32, 128), device=DEVICE, dtype=torch.bfloat16)
update = torch.randn_like(residual)
gate = torch.randn((1, 1, 128), device=DEVICE, dtype=torch.bfloat16)
compiled = torch.compile(residual_gate_add, fullgraph=True)
assert torch.equal(compiled(residual, update, gate), residual + update * gate)
@torch.no_grad()
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_scaled_residual_add_is_bit_exact(dtype):
# fp32 residual accumulator + half-precision update, as the DiT blocks
# that keep their residual stream in fp32 emit it.
residual = torch.randn(2, 17, 64, device=DEVICE, dtype=torch.float32)
x = torch.randn(2, 17, 64, device=DEVICE, dtype=dtype)
scale = torch.randn(64, device=DEVICE, dtype=torch.float32)
actual = try_fused_scaled_residual_add_exact(residual, x, scale)
assert actual is not None
assert torch.equal(actual, residual + x * scale)
@torch.no_grad()
def test_scaled_residual_add_rejects_unsupported_inputs():
residual = torch.empty(2, 3, 8, device=DEVICE, dtype=torch.float32)
x = torch.empty_like(residual)
scale = torch.empty(8, device=DEVICE, dtype=torch.float32)
# A too-small hidden dim and a mismatched scale length both bail out;
# ``try_`` returning None is this helper's documented contract.
assert try_fused_scaled_residual_add_exact(residual, x, scale) is None
assert try_fused_scaled_residual_add_exact(residual, x.half(), scale[:-1]) is None
# ---------------------------------------------------------------------------
# LTX-2 nine-way adaLN value split
# ---------------------------------------------------------------------------
def _ltx2_reference(scale_shift_table, timestep):
batch, seq, _ = timestep.shape
hidden = scale_shift_table.shape[1]
return (
scale_shift_table.to(device=timestep.device, dtype=timestep.dtype)
.view(1, 1, 9, hidden)
.add(timestep.reshape(batch, seq, 9, hidden))
.unbind(dim=2)
)
@torch.no_grad()
@pytest.mark.parametrize("batch,seq,hidden", [(1, 1, 4096), (2, 3, 2048)])
@pytest.mark.parametrize("table_dtype", [torch.bfloat16, torch.float32])
@pytest.mark.parametrize("compiled", [False, True])
def test_ltx2_ada_values9(batch, seq, hidden, table_dtype, compiled):
scale_shift_table = torch.randn(9, hidden, device=DEVICE, dtype=table_dtype)
timestep = torch.randn(batch, seq, 9 * hidden, device=DEVICE, dtype=torch.bfloat16)
fn = (
torch.compile(ltx2_ada_values9, fullgraph=True)
if compiled
else ltx2_ada_values9
)
actual = fn(scale_shift_table, timestep)
expected = _ltx2_reference(scale_shift_table, timestep)
assert len(actual) == 9
for got, want in zip(actual, expected, strict=True):
# Each slice must come out naturally contiguous -- that is the point
# of the kernel; a strided slice would re-add the downstream copy.
assert got.is_contiguous()
assert torch.equal(got, want)
@torch.no_grad()
def test_ltx2_ada_values9_rejects_unsupported_shape():
scale_shift_table = torch.randn(8, 4096, device=DEVICE, dtype=torch.bfloat16)
timestep = torch.randn(1, 1, 9 * 4096, device=DEVICE, dtype=torch.bfloat16)
with pytest.raises(ValueError, match="scale_shift_table"):
ltx2_ada_values9(scale_shift_table, timestep)
# ---------------------------------------------------------------------------
# select-0/1 LayerNorm modulation (Qwen-Image)
# ---------------------------------------------------------------------------
SELECT01_DTYPES = get_ci_test_range(
[torch.float16, torch.bfloat16, torch.float32], [torch.float16, torch.bfloat16]
)
SELECT01_SHAPES = get_ci_test_range(
[(b, s, h) for b in (1, 2, 4) for s in (6, 33, 128, 257) for h in (512, 3072)],
[(1, 6, 512), (2, 128, 3072)],
)
EPS = 1e-6
def _select01_reference(normalized, mods, index):
scale0, shift0, gate0, scale1, shift1, gate1 = mods
idx = index.bool().unsqueeze(-1)
scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
return normalized * (1 + scale) + shift, gate
@pytest.mark.parametrize("dtype", SELECT01_DTYPES)
@pytest.mark.parametrize("shape", SELECT01_SHAPES)
@pytest.mark.parametrize("with_residual", [False, True])
def test_layernorm_scale_shift_gate_select01(dtype, shape, with_residual):
batch_size, seq_len, hidden_size = shape
x = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE, dtype=dtype)
weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
bias = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
index = torch.randint(0, 2, (batch_size, seq_len), device=DEVICE, dtype=torch.int32)
mods = tuple(
torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
for _ in range(6)
)
scale0, shift0, gate0, scale1, shift1, gate1 = mods
if with_residual:
residual = torch.randn_like(x)
residual_gate = torch.randn_like(x)
residual_ref = residual + residual_gate * x
normalized = norm_infer(
residual_ref.flatten(0, 1), weight, bias, eps=EPS, is_rms_norm=False
).view_as(residual_ref)
out_ref, gate_ref = _select01_reference(normalized, mods, index)
out, residual_out, gate = (
fuse_residual_layernorm_scale_shift_gate_select01_kernel(
x.contiguous(),
residual=residual.contiguous(),
residual_gate=residual_gate.contiguous(),
weight=weight,
bias=bias,
scale0=scale0,
shift0=shift0,
gate0=gate0,
scale1=scale1,
shift1=shift1,
gate1=gate1,
index=index,
eps=EPS,
)
)
else:
normalized = norm_infer(
x.flatten(0, 1), weight, bias, eps=EPS, is_rms_norm=False
).view_as(x)
out_ref, gate_ref = _select01_reference(normalized, mods, index)
residual_ref = residual_out = None
out, gate = fuse_layernorm_scale_shift_gate_select01_kernel(
x.contiguous(),
weight=weight,
bias=bias,
scale0=scale0,
shift0=shift0,
gate0=gate0,
scale1=scale1,
shift1=shift1,
gate1=gate1,
index=index,
eps=EPS,
)
tol = 1e-5 if dtype == torch.float32 else 5e-2
torch.testing.assert_close(out, out_ref, atol=tol, rtol=tol)
torch.testing.assert_close(gate, gate_ref, atol=tol, rtol=tol)
if with_residual:
torch.testing.assert_close(residual_out, residual_ref, atol=tol, rtol=tol)
# ---------------------------------------------------------------------------
# Sinusoidal timestep embedding
# ---------------------------------------------------------------------------
TIMESTEP_BATCHES = get_ci_test_range(
[1, 2, 8, 128, 256, 512, 1536, 2048, 4096, 11008, 16384], [1, 128, 2048, 16384]
)
TIMESTEP_DIMS = get_ci_test_range(
[32, 128, 256, 512, 1536, 2048, 4096, 8192], [32, 512, 8192]
)
TIMESTEP_DTYPES = get_ci_test_range(
[torch.float16, torch.bfloat16, torch.float32], [torch.float16, torch.bfloat16]
)
def timestep_embedding_reference(
timesteps,
dim,
*,
flip_sin_to_cos=False,
downscale_freq_shift=1,
scale=1,
max_period=10000,
):
"""diffusers' ``get_timestep_embedding``, kept verbatim as the oracle."""
assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
timesteps = timesteps.to(torch.float32)
half_dim = dim // 2
exponent = -torch.log(
torch.tensor(max_period, dtype=torch.float32, device=timesteps.device)
) * torch.arange(
start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
)
exponent = exponent / (half_dim - downscale_freq_shift)
emb = torch.exp(exponent)
emb = timesteps[:, None].float() * emb[None, :]
emb = scale * emb
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
if flip_sin_to_cos:
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
if dim % 2 == 1:
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
return emb
@pytest.mark.parametrize("batch_size", TIMESTEP_BATCHES)
@pytest.mark.parametrize("dim", TIMESTEP_DIMS)
@pytest.mark.parametrize("dtype", TIMESTEP_DTYPES)
@pytest.mark.parametrize(
"flip_sin_to_cos,downscale_freq_shift,scale",
[
(True, 0, 1), # the sgl-diffusion default
(False, 1, 1), # the diffusers default
(True, 1, 0.01), # scaled variant used by the SD-style embedders
],
)
def test_timestep_embedding_matches_diffusers(
batch_size, dim, dtype, flip_sin_to_cos, downscale_freq_shift, scale
):
t = torch.randint(low=0, high=1000, size=(batch_size,), device=DEVICE).to(dtype)
kwargs = dict(
flip_sin_to_cos=flip_sin_to_cos,
downscale_freq_shift=downscale_freq_shift,
scale=scale,
max_period=10000,
)
torch.testing.assert_close(
timestep_embedding(t, dim, **kwargs),
timestep_embedding_reference(t, dim, **kwargs),
atol=1e-3,
rtol=1e-3,
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_modulate_scale_shift.py
-58
View File
@@ -1,58 +0,0 @@
import pytest
import torch
from sglang.kernels.ops.diffusion.modulate_scale_shift import (
can_use_modulate_scale_shift_cuda,
modulate_scale_shift,
modulate_scale_shift_cuda,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
# FLUX.1 1024^2 adaLN shapes (D=3072) plus batched and odd-length coverage.
CASES = [(1, 4096, 3072), (1, 512, 3072), (1, 4608, 3072), (2, 1024, 3072), (1, 17, 64)]
@pytest.fixture(autouse=True)
def cuda_setup():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
torch.cuda.manual_seed(0)
def _eager(x, scale, shift):
return x * (1 + scale[:, None]) + shift[:, None]
@pytest.mark.parametrize("shape", CASES)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
def test_modulate_scale_shift_matches_eager(shape, dtype):
x = torch.randn(shape, device="cuda", dtype=dtype)
scale = torch.randn((shape[0], shape[-1]), device="cuda", dtype=dtype)
shift = torch.randn_like(scale)
out = modulate_scale_shift_cuda(x, scale, shift)
assert torch.equal(out, _eager(x, scale, shift)) # bitwise contract
def test_modulate_scale_shift_adaln_chunk_views():
x = torch.randn((1, 4096, 3072), device="cuda", dtype=torch.bfloat16)
emb = torch.randn((1, 6 * 3072), device="cuda", dtype=torch.bfloat16)
shift, scale = emb.chunk(6, dim=1)[:2]
assert can_use_modulate_scale_shift_cuda(x, scale, shift)
out = modulate_scale_shift_cuda(x, scale, shift)
assert torch.equal(out, _eager(x, scale, shift))
def test_modulate_scale_shift_guards_reject_fp32():
x = torch.randn((1, 64, 64), device="cuda", dtype=torch.float32)
row = torch.randn((1, 64), device="cuda", dtype=torch.float32)
assert not can_use_modulate_scale_shift_cuda(x, row, row)
assert torch.equal(modulate_scale_shift(x, row, row), _eager(x, row, row))
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/diffusion/test_native_bf16_rmsnorm.py
-87
View File
@@ -1,87 +0,0 @@
import pytest
import torch
from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import (
rmsnorm_scale,
rmsnorm_tanh_residual,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
EPS = 1e-5
def _native_bf16_rmsnorm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
square = (x * x).to(torch.bfloat16)
mean_square = square.mean(dim=-1, keepdim=True).to(torch.bfloat16)
rstd = torch.rsqrt((mean_square + EPS).to(torch.bfloat16).float()).to(
torch.bfloat16
)
return ((x * rstd).to(torch.bfloat16) * weight).to(torch.bfloat16)
def test_native_bf16_rmsnorm_rejects_unsupported_inputs():
x = torch.randn(2, 3, 16, dtype=torch.bfloat16)
weight = torch.randn(16, dtype=torch.bfloat16)
modulation = torch.randn(2, 1, 16, dtype=torch.bfloat16)
residual = torch.randn_like(x)
assert rmsnorm_scale(x, weight, modulation, EPS) is None
assert rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
assert rmsnorm_scale(x, weight[:-1], modulation, EPS) is None
assert rmsnorm_tanh_residual(x, modulation, residual[..., :-1], weight, EPS) is None
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize("shape", [(1, 32, 2560), (2, 17, 256)])
def test_rmsnorm_scale_matches_native_bf16(shape):
torch.manual_seed(0)
batch, _, dim = shape
x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(dim, device="cuda", dtype=torch.bfloat16)
scale = torch.randn(batch, 1, dim, device="cuda", dtype=torch.bfloat16)
actual = rmsnorm_scale(x, weight, scale, EPS)
expected = (_native_bf16_rmsnorm(x, weight) * scale).to(torch.bfloat16)
assert actual is not None
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize("shape", [(1, 32, 2560), (2, 17, 256)])
def test_rmsnorm_tanh_residual_matches_native_bf16(shape):
torch.manual_seed(0)
batch, _, dim = shape
x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
gate = torch.randn(batch, 1, dim, device="cuda", dtype=torch.bfloat16)
residual = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(dim, device="cuda", dtype=torch.bfloat16)
actual = rmsnorm_tanh_residual(x, gate, residual, weight, EPS)
norm = _native_bf16_rmsnorm(x, weight)
gated = (torch.tanh(gate.float()).to(torch.bfloat16) * norm).to(torch.bfloat16)
expected = (residual + gated).to(torch.bfloat16)
assert actual is not None
# Triton's exp-based tanh can differ slightly from torch.tanh in BF16.
torch.testing.assert_close(actual, expected, atol=4e-2, rtol=2e-2)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_native_bf16_rmsnorm_rejects_hidden_size_above_limit():
dim = 8448
x = torch.empty(1, 1, dim, device="cuda", dtype=torch.bfloat16)
weight = torch.empty(dim, device="cuda", dtype=torch.bfloat16)
modulation = torch.empty(1, 1, dim, device="cuda", dtype=torch.bfloat16)
residual = torch.empty_like(x)
assert rmsnorm_scale(x, weight, modulation, EPS) is None
assert rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/diffusion/test_norm.py
+466
View File
@@ -0,0 +1,466 @@
"""``diffusion.norm``: GroupNorm / RMSNorm / LayerNorm and their fused epilogues.
This domain has the most implementations of any in the package (see the
selection matrix in ``sglang/kernels/ops/diffusion/README.md``), so the suite
is organized by *kernel*, and each section states which oracle it is held to:
- ``triton_group_norm_silu`` / ``apply_group_norm_silu`` -> ``F.group_norm`` +
``F.silu`` with a per-dtype tolerance (fp32 statistics, different reduction).
- the two-pass channels-last GroupNorm -> same oracle, plus its support
predicates (the kernels raise on an unsupported input rather than returning
``None``).
- ``rmsnorm_scale`` / ``rmsnorm_tanh_residual`` -> a bf16-native reference that
reproduces Z-Image's own norm, with a tolerance for Triton's exp-based tanh.
- the CuTe-DSL fused norm+scale/shift -> an fp32 reference chain.
The FlyDSL norms live in ``test_norm_flydsl.py``: they are ROCm gfx950-only,
so they run on a CI lane this file does not, and keeping them here dragged the
CUDA-only CuTe-DSL cases onto the AMD runner.
The *bit-exact* norms (``fused_rmsnorm_scale_shift_bitexact``,
``fused_layernorm_modulate``, ``zimage_qk_rmsnorm_native``) are exercised
through their model wrappers in ``test_model_fast_paths.py``, where the live
eager chain they must reproduce is available.
"""
import sys
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from sglang.kernels.ops.diffusion import (
apply_group_norm_silu,
can_use_group_norm_silu_4d,
can_use_wan_rmsnorm_silu,
group_norm_silu_4d,
rmsnorm_scale,
rmsnorm_tanh_residual,
triton_group_norm_silu,
wan_rmsnorm_silu,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=70, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
DEVICE = "cuda"
DTYPES = [torch.float16, torch.bfloat16, torch.float32]
EPS = 1e-5
def _tol(dtype: torch.dtype) -> tuple[float, float]:
if dtype == torch.float32:
return 1e-5, 1e-5
if dtype == torch.bfloat16:
return 7e-2, 2e-2
return 3e-3, 3e-3
@pytest.fixture(autouse=True)
def cuda_setup():
torch.cuda.manual_seed(0)
def _cl3d(shape, dtype):
return torch.randn(shape, device=DEVICE, dtype=dtype).contiguous(
memory_format=torch.channels_last_3d
)
# ---------------------------------------------------------------------------
# GroupNorm + SiLU
# ---------------------------------------------------------------------------
GN_CASES = [
pytest.param((2, 64, 32, 32), 32, id="image_2d"),
pytest.param((1, 64, 4, 16, 16), 32, id="video_3d"),
pytest.param((4, 128), 32, id="token_2d"),
]
def _gn_silu_reference(x, weight, bias, num_groups, eps=EPS):
return F.silu(F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps))
@torch.no_grad()
@pytest.mark.parametrize("shape,num_groups", GN_CASES)
@pytest.mark.parametrize("dtype", DTYPES)
def test_triton_group_norm_silu(shape, num_groups, dtype):
channels = shape[1]
x = torch.randn(shape, device=DEVICE, dtype=dtype)
weight = torch.randn(channels, device=DEVICE, dtype=dtype)
bias = torch.randn(channels, device=DEVICE, dtype=dtype)
atol, rtol = _tol(dtype)
torch.testing.assert_close(
triton_group_norm_silu(x, weight, bias, num_groups=num_groups),
_gn_silu_reference(x, weight, bias, num_groups),
atol=atol,
rtol=rtol,
)
@torch.no_grad()
def test_triton_group_norm_silu_large_tile_bf16():
# A tile large enough to force the chunked launch path (128 channels over
# 20x256x256), which the small cases above never reach.
shape, num_groups = (1, 128, 20, 256, 256), 32
x = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
weight = torch.randn(shape[1], device=DEVICE, dtype=torch.bfloat16)
bias = torch.randn(shape[1], device=DEVICE, dtype=torch.bfloat16)
atol, rtol = _tol(torch.bfloat16)
torch.testing.assert_close(
triton_group_norm_silu(x, weight, bias, num_groups=num_groups),
_gn_silu_reference(x, weight, bias, num_groups),
atol=atol,
rtol=rtol,
)
@torch.no_grad()
@pytest.mark.parametrize("shape,num_groups", GN_CASES[:2])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_apply_group_norm_silu_module_wrapper(shape, num_groups, dtype):
# The nn.Module-taking wrapper must match the eager module pair it stands
# in for, including its own guard set (affine=True, non-inplace SiLU).
norm = nn.GroupNorm(num_groups, shape[1], eps=EPS, affine=True).to(
device=DEVICE, dtype=dtype
)
activation = nn.SiLU()
x = torch.randn(shape, device=DEVICE, dtype=dtype)
atol, rtol = _tol(dtype)
torch.testing.assert_close(
apply_group_norm_silu(x, norm, activation),
activation(norm(x)),
atol=atol,
rtol=rtol,
)
@torch.no_grad()
def test_group_norm_silu_4d_channels_last_and_guards():
gn = nn.GroupNorm(32, 128, eps=1e-6).to(DEVICE, torch.bfloat16)
x = torch.randn(1, 128, 64, 64, device=DEVICE, dtype=torch.bfloat16).to(
memory_format=torch.channels_last
)
assert can_use_group_norm_silu_4d(x, gn.weight, gn.bias, 32)
out = group_norm_silu_4d(x, gn.weight, gn.bias, 32, 1e-6)
assert out.is_contiguous(memory_format=torch.channels_last)
torch.testing.assert_close(out.float(), F.silu(gn(x)).float(), atol=0.06, rtol=0)
# Guards: the kernel exists only for channels_last inputs with device-side
# affine params and a non-empty spatial extent. Each rejected case must
# fail the predicate *and* raise if called anyway -- silently returning
# ``None`` is what this protocol replaced.
rejected = [
(x.contiguous(), gn.weight, gn.bias), # contiguous (NCHW) layout
(x, gn.weight.cpu(), gn.bias), # host-side affine
(x[..., :0, :], gn.weight, gn.bias), # empty spatial extent
]
for args in rejected:
assert not can_use_group_norm_silu_4d(*args, 32)
with pytest.raises(ValueError):
group_norm_silu_4d(*args, 32, 1e-6)
# ---------------------------------------------------------------------------
# BF16-native RMSNorm fusions (Z-Image / Ideogram)
# ---------------------------------------------------------------------------
def _native_bf16_rmsnorm(x, weight):
"""Z-Image's own norm: every step materialized in bf16, no fp32 carry."""
square = (x * x).to(torch.bfloat16)
mean_square = square.mean(dim=-1, keepdim=True).to(torch.bfloat16)
rstd = torch.rsqrt((mean_square + EPS).to(torch.bfloat16).float()).to(
torch.bfloat16
)
return ((x * rstd).to(torch.bfloat16) * weight).to(torch.bfloat16)
@pytest.mark.parametrize("shape", [(1, 32, 2560), (2, 17, 256)])
def test_rmsnorm_scale_matches_native_bf16(shape):
batch, _, dim = shape
x = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
weight = torch.randn(dim, device=DEVICE, dtype=torch.bfloat16)
scale = torch.randn(batch, 1, dim, device=DEVICE, dtype=torch.bfloat16)
actual = rmsnorm_scale(x, weight, scale, EPS)
assert actual is not None
torch.testing.assert_close(
actual,
(_native_bf16_rmsnorm(x, weight) * scale).to(torch.bfloat16),
atol=2e-2,
rtol=2e-2,
)
@pytest.mark.parametrize("shape", [(1, 32, 2560), (2, 17, 256)])
def test_rmsnorm_tanh_residual_matches_native_bf16(shape):
batch, _, dim = shape
x = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
gate = torch.randn(batch, 1, dim, device=DEVICE, dtype=torch.bfloat16)
residual = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
weight = torch.randn(dim, device=DEVICE, dtype=torch.bfloat16)
actual = rmsnorm_tanh_residual(x, gate, residual, weight, EPS)
norm = _native_bf16_rmsnorm(x, weight)
gated = (torch.tanh(gate.float()).to(torch.bfloat16) * norm).to(torch.bfloat16)
assert actual is not None
# Triton's exp-based tanh can differ slightly from torch.tanh in bf16.
torch.testing.assert_close(
actual, (residual + gated).to(torch.bfloat16), atol=4e-2, rtol=2e-2
)
@pytest.mark.parametrize("on_host", [True, False])
def test_native_bf16_rmsnorm_rejects_unsupported_inputs(on_host):
# Host tensors and a hidden size past the kernel limit are both outside
# the contract; these entry points signal that by returning None (they
# are internal fast-path probes, not public predicate+kernel pairs).
device = "cpu" if on_host else DEVICE
dim = 16 if on_host else 8448
x = torch.randn(2, 3, dim, dtype=torch.bfloat16, device=device)
weight = torch.randn(dim, dtype=torch.bfloat16, device=device)
modulation = torch.randn(2, 1, dim, dtype=torch.bfloat16, device=device)
residual = torch.randn_like(x)
assert rmsnorm_scale(x, weight, modulation, EPS) is None
assert rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
if on_host:
# Mismatched trailing dims are rejected too.
assert rmsnorm_scale(x, weight[:-1], modulation, EPS) is None
assert (
rmsnorm_tanh_residual(x, modulation, residual[..., :-1], weight, EPS)
is None
)
# ---------------------------------------------------------------------------
# Wan VAE channels_last_3d RMSNorm + SiLU
# ---------------------------------------------------------------------------
@torch.no_grad()
@pytest.mark.parametrize(
"x_dtype,affine_dtype,atol,rtol",
[
(torch.float32, torch.float32, 1e-5, 1e-5), # FastWan2.2 fp32 decode
(torch.bfloat16, torch.float32, 1.5e-1, 3e-2), # Wan2.1 bf16 autocast
],
)
def test_wan_rmsnorm_silu_numerics(x_dtype, affine_dtype, atol, rtol):
x = _cl3d((1, 96, 3, 10, 14), x_dtype)
gamma = torch.randn((96, 1, 1, 1), device=DEVICE, dtype=affine_dtype)
for bias in (None, torch.randn_like(gamma)):
expected = F.silu(
F.normalize(x, dim=1) * 96**0.5 * gamma + (0 if bias is None else bias)
)
actual = wan_rmsnorm_silu(x, gamma, bias)
assert actual.dtype == expected.dtype
# The kernel must preserve the channels_last_3d layout; a relayout
# here would undo the reason the decoder runs in that format.
assert actual.stride() == x.stride()
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
@torch.no_grad()
def test_wan_rmsnorm_silu_rejects_empty_input():
x = torch.empty(1, 96, 0, 2, 2, device=DEVICE, dtype=torch.bfloat16).to(
memory_format=torch.channels_last_3d
)
gamma = torch.ones(96, 1, 1, 1, device=DEVICE, dtype=torch.bfloat16)
assert not can_use_wan_rmsnorm_silu(x, gamma, None)
with pytest.raises(ValueError):
wan_rmsnorm_silu(x, gamma)
# ---------------------------------------------------------------------------
# CuTe-DSL fused (residual +) norm + scale/shift
# ---------------------------------------------------------------------------
SHAPE_MAP = {
"1": lambda B, S, F_, D: (1,),
"D": lambda B, S, F_, D: (D,),
"1D": lambda B, S, F_, D: (1, D),
"BD": lambda B, S, F_, D: (B, D),
"11D": lambda B, S, F_, D: (1, 1, D),
"B1D": lambda B, S, F_, D: (B, 1, D),
"1SD": lambda B, S, F_, D: (1, S, D),
"BSD": lambda B, S, F_, D: (B, S, D),
"BF1D": lambda B, S, F_, D: (B, F_, 1, D),
}
# (B, S, F, D)
CUTE_SHAPES = [
(1, 115200, 1, 3072), # HunyuanVideo
(1, 32760, 1, 1536), # Wan
(1, 6, 1, 3072), # Qwen-Image
(1, 1024, 8, 3072),
(4, 512, 16, 3072),
]
NORM_TYPES = ["layer", "rms"]
AFFINE_MODES = ["D", "NAT"]
INDEX_MODES = ["BSD", "1", "1SD", "BD", "B1D", "D", "1D", "11D", "BF1D"]
def _import_cutedsl():
"""Import the CuTe-DSL entry points, skipping when the backend is absent.
This file is registered on the AMD lane for its FlyDSL section, but the
CuTe-DSL norms need cuda-python and CUTLASS, which the ROCm image does not
ship. Guarded per test rather than by dropping this file from the AMD
lane, so the Triton and FlyDSL sections keep running there.
"""
try:
from sglang.kernels.ops.diffusion import (
fused_norm_scale_shift,
fused_scale_residual_norm_scale_shift,
)
except ImportError as exc: # pragma: no cover - platform-dependent
pytest.skip(f"CuTe-DSL backend unavailable: {exc}")
return fused_norm_scale_shift, fused_scale_residual_norm_scale_shift
def _make_tensor(index_mode, shape, dtype):
if index_mode == "NAT":
return None
return torch.randn(*SHAPE_MAP[index_mode](*shape), device=DEVICE, dtype=dtype)
def _apply_scale_shift(y, scale, shift):
if scale.ndim == 4:
num_frame = scale.shape[1]
return rearrange(
rearrange(y, "b (f l) d -> b f l d", f=num_frame) * (1 + scale) + shift,
"b f l d -> b (f l) d",
)
scale = rearrange(scale, "b d -> b 1 d") if scale.ndim == 2 else scale
shift = rearrange(shift, "b d -> b 1 d") if shift.ndim == 2 else shift
return y * (1 + scale) + shift
def _cute_reference(residual, x, gate, weight, bias, scale, shift, norm_type, eps):
"""fp32 oracle for both variants; ``residual is None`` = no-residual form."""
original_dtype = x.dtype
residual, x, gate, weight, bias, scale, shift = (
v.float() if isinstance(v, torch.Tensor) else v
for v in (residual, x, gate, weight, bias, scale, shift)
)
residual_out = None
if residual is not None:
if isinstance(gate, int):
x = residual + gate * x
elif gate.ndim == 4:
folded = rearrange(x, "b (f l) d -> b f l d", f=gate.shape[1])
x = residual + rearrange(folded * gate, "b f l d -> b (f l) d")
else:
g = rearrange(gate, "b d -> b 1 d") if gate.ndim == 2 else gate
x = residual + g * x
residual_out = x.to(original_dtype)
if norm_type == "layer":
norm = torch.layer_norm(x, x.shape[-1:], eps=eps, weight=weight, bias=bias)
else:
norm = torch.rms_norm(x, x.shape[-1:], eps=eps, weight=weight)
return _apply_scale_shift(norm, scale, shift).to(original_dtype), residual_out
@torch.no_grad()
def _run_cute(
with_residual,
shape=CUTE_SHAPES[0],
dtype=DTYPES[0],
affine_dtype=DTYPES[0],
mod_dtype=DTYPES[0],
norm_type=NORM_TYPES[0],
affine_mode=AFFINE_MODES[0],
gate_mode="B1D",
index_mode="BSD",
eps=EPS,
):
fused_norm_scale_shift, fused_scale_residual_norm_scale_shift = _import_cutedsl()
x = _make_tensor("BSD", shape, dtype)
weight = _make_tensor(affine_mode, shape, affine_dtype)
bias = _make_tensor(affine_mode, shape, affine_dtype)
scale = _make_tensor(index_mode, shape, mod_dtype)
shift = _make_tensor(index_mode, shape, mod_dtype)
tol = 1e-5 if dtype == torch.float32 else 5e-2
if with_residual:
residual = _make_tensor("BSD", shape, dtype)
gate = _make_tensor(gate_mode, shape, dtype)
y, res = fused_scale_residual_norm_scale_shift(
residual, x, gate, weight, bias, scale, shift, norm_type, eps
)
y_ref, res_ref = _cute_reference(
residual, x, gate, weight, bias, scale, shift, norm_type, eps
)
torch.testing.assert_close(res, res_ref, atol=tol, rtol=tol)
else:
y = fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps)
y_ref, _ = _cute_reference(
None, x, None, weight, bias, scale, shift, norm_type, eps
)
torch.testing.assert_close(y, y_ref, atol=tol, rtol=tol)
@pytest.mark.parametrize("with_residual", [False, True])
@pytest.mark.parametrize("norm_type", NORM_TYPES)
@pytest.mark.parametrize("shape", CUTE_SHAPES)
@pytest.mark.parametrize("dtype", DTYPES)
def test_cutedsl_norm_scale_shift_shapes(with_residual, norm_type, shape, dtype):
_run_cute(with_residual, shape=shape, dtype=dtype, norm_type=norm_type)
@pytest.mark.parametrize("with_residual", [False, True])
@pytest.mark.parametrize("norm_type", NORM_TYPES)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("operand", ["affine", "modulation"])
def test_cutedsl_norm_scale_shift_mixed_operand_dtypes(
with_residual, norm_type, dtype, operand
):
# The affine params and the modulation rows may each arrive in a dtype
# other than the activation's; both combinations must dispatch.
kwargs = {"affine_dtype" if operand == "affine" else "mod_dtype": dtype}
_run_cute(with_residual, norm_type=norm_type, **kwargs)
@pytest.mark.parametrize("with_residual", [False, True])
@pytest.mark.parametrize("norm_type", NORM_TYPES)
@pytest.mark.parametrize("affine_mode", AFFINE_MODES)
def test_cutedsl_norm_scale_shift_affine_modes(with_residual, norm_type, affine_mode):
_run_cute(with_residual, norm_type=norm_type, affine_mode=affine_mode)
@pytest.mark.parametrize("with_residual", [False, True])
@pytest.mark.parametrize("norm_type", NORM_TYPES)
@pytest.mark.parametrize("index_mode", INDEX_MODES)
def test_cutedsl_norm_scale_shift_index_modes(with_residual, norm_type, index_mode):
_run_cute(with_residual, norm_type=norm_type, index_mode=index_mode)
@pytest.mark.parametrize("norm_type", NORM_TYPES)
@pytest.mark.parametrize("index_mode", INDEX_MODES)
def test_cutedsl_scale_residual_gate_index_modes(norm_type, index_mode):
_run_cute(True, norm_type=norm_type, gate_mode=index_mode)
def test_validate_scale_shift_rejects_non_divisible_frames():
_import_cutedsl()
from sglang.kernels.ops.diffusion import validate_scale_shift
with pytest.raises(ValueError, match=r"S\(10\) must be divisible by F\(4\)"):
validate_scale_shift(
torch.empty((1, 4, 1, 256), device=DEVICE, dtype=torch.float16), 1, 10, 256
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_norm_flydsl.py
+98
View File
@@ -0,0 +1,98 @@
"""``diffusion.norm``: the FlyDSL fused norm + scale/shift kernels (ROCm).
Split out of ``test_norm.py`` rather than merged with the other norm backends:
FlyDSL is an AMD gfx950-only compiler, so these run on the AMD CI lane and
nothing else in that file does. Keeping them together forced the CUDA-only
CuTe-DSL cases onto the ROCm runner, where cuda-python does not exist.
Oracle: an fp32 reference chain, with a tolerance -- the kernel keeps fp32
statistics but reorders the reduction.
"""
import sys
import pytest
import torch
import torch.nn.functional as F
from sglang.test.ci.ci_register import register_amd_ci
register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU required")
DEVICE = "cuda"
FLYDSL_D = 5120
FLYDSL_EPS = 1e-6
def _require_rocm():
if not torch.version.hip:
pytest.skip("ROCm/HIP required for FlyDSL kernels")
def _flydsl_reference(residual, x, gate, weight, bias, scale, shift, norm_type, eps):
if residual is not None:
x = (residual.float() + x.float() * gate.float()).to(torch.bfloat16)
residual_out = x
else:
residual_out = None
if norm_type == "layer":
normed = F.layer_norm(x.float(), (FLYDSL_D,), weight, bias, eps)
else:
var = x.float().pow(2).mean(-1, keepdim=True)
normed = x.float() * torch.rsqrt(var + eps) * weight.float()
y = (normed * (1.0 + scale.float()) + shift.float()).to(torch.bfloat16)
return y, residual_out
@pytest.mark.parametrize("with_residual", [False, True])
@pytest.mark.parametrize(
"norm_type,B,L",
[("rms", 1, 16), ("rms", 2, 16), ("layer", 2, 16), ("rms", 1, 90000)],
)
def test_flydsl_norm_scale_shift(with_residual, norm_type, B, L):
_require_rocm()
# Imported inside the test: the FlyDSL compiler only exists on ROCm, and
# the facade resolves an export the moment it is named -- a module-level
# import here would fail collection of this whole file on CUDA.
from sglang.kernels.ops.diffusion import (
flydsl_fused_residual_norm_scale_shift,
flydsl_norm_scale_shift,
)
torch.manual_seed(42)
shape = (B, L, FLYDSL_D)
x = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
weight = torch.randn(FLYDSL_D, device=DEVICE, dtype=torch.float32)
bias = (
torch.randn(FLYDSL_D, device=DEVICE, dtype=torch.float32)
if norm_type == "layer"
else None
)
scale = torch.randn(B, 1, FLYDSL_D, device=DEVICE, dtype=torch.bfloat16)
shift = torch.randn(B, 1, FLYDSL_D, device=DEVICE, dtype=torch.bfloat16)
if with_residual:
residual = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
gate = torch.randn(B, 1, FLYDSL_D, device=DEVICE, dtype=torch.bfloat16)
y, res = flydsl_fused_residual_norm_scale_shift(
residual, x, gate, weight, bias, scale, shift, norm_type, FLYDSL_EPS
)
y_ref, res_ref = _flydsl_reference(
residual, x, gate, weight, bias, scale, shift, norm_type, FLYDSL_EPS
)
torch.testing.assert_close(res, res_ref, atol=5e-2, rtol=5e-2)
else:
y = flydsl_norm_scale_shift(
x, weight, bias, scale, shift, norm_type, FLYDSL_EPS
)
y_ref, _ = _flydsl_reference(
None, x, None, weight, bias, scale, shift, norm_type, FLYDSL_EPS
)
torch.testing.assert_close(y, y_ref, atol=1.0, rtol=5e-2)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_quality_gate.py
-47
View File
@@ -1,47 +0,0 @@
import sys
import pytest
import torch.nn as nn
from sglang.kernels.ops.diffusion.quality_gate import QualityGatedFusion
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
def test_quality_gate_mounts_and_unmounts_all_sites():
fusion = QualityGatedFusion(
name="test fusion",
marker_attr="_test_fusion_site",
enabled_attr="_test_fusion_enabled",
)
root = nn.ModuleList([nn.Module(), nn.Module()])
for index, site in enumerate(root):
fusion.mark(site, index)
assert [fusion.metadata(site) for site in root] == [0, 1]
assert fusion.mount(root)
assert all(fusion.is_enabled(site) for site in root)
fusion.unmount(root)
assert not any(fusion.is_enabled(site) for site in root)
def test_quality_gate_rejection_is_all_or_nothing():
fusion = QualityGatedFusion(
name="test fusion",
marker_attr="_test_fusion_site",
enabled_attr="_test_fusion_enabled",
)
root = nn.ModuleList([nn.Module(), nn.Module()])
for index, site in enumerate(root):
fusion.mark(site, index)
assert not fusion.mount(
root, reject_reason=lambda site: "rejected" if fusion.metadata(site) else None
)
assert not any(fusion.is_enabled(site) for site in root)
assert not fusion.mount(nn.Module())
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_qwen_image_modulation.py
-228
View File
@@ -1,228 +0,0 @@
import sys
import pytest
import torch
import triton
from sglang.kernels.jit.utils import get_ci_test_range
from sglang.kernels.ops.diffusion.triton.norm import norm_infer
from sglang.kernels.ops.diffusion.triton.scale_shift import (
fuse_layernorm_scale_shift_gate_select01_kernel,
fuse_residual_layernorm_scale_shift_gate_select01_kernel,
)
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
register_cuda_ci(est_time=30, stage="nightly", runner_config="1-gpu-large")
register_amd_ci(est_time=30, suite="nightly-amd-kernel-1-gpu", nightly=True)
DEVICE = "cuda"
DTYPES = get_ci_test_range(
[torch.float16, torch.bfloat16, torch.float32], [torch.float16, torch.bfloat16]
)
BATCH_SIZES = get_ci_test_range([1, 2, 4], [1, 2])
SEQ_LENS = get_ci_test_range([6, 33, 128, 257], [6, 128])
HIDDEN_SIZES = get_ci_test_range([512, 1024, 1536, 3072], [512, 3072])
EPS = 1e-6
def _tol(dtype: torch.dtype) -> tuple[float, float]:
if dtype == torch.float32:
return 1e-5, 1e-5
return 5e-2, 5e-2
def _make_modulation_tensors(batch_size: int, hidden_size: int, dtype: torch.dtype):
scale0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
shift0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
gate0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
scale1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
shift1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
gate1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
return scale0, shift0, gate0, scale1, shift1, gate1
def _baseline_select01_modulation(
x: torch.Tensor,
weight: torch.Tensor | None,
bias: torch.Tensor | None,
scale0: torch.Tensor,
shift0: torch.Tensor,
gate0: torch.Tensor,
scale1: torch.Tensor,
shift1: torch.Tensor,
gate1: torch.Tensor,
index: torch.Tensor,
eps: float,
):
normalized = norm_infer(
x.view(-1, x.shape[-1]),
weight,
bias,
eps=eps,
is_rms_norm=False,
).view_as(x)
return _apply_select01_modulation(
normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
)
def _baseline_residual_select01_modulation(
x: torch.Tensor,
residual: torch.Tensor,
residual_gate: torch.Tensor,
weight: torch.Tensor | None,
bias: torch.Tensor | None,
scale0: torch.Tensor,
shift0: torch.Tensor,
gate0: torch.Tensor,
scale1: torch.Tensor,
shift1: torch.Tensor,
gate1: torch.Tensor,
index: torch.Tensor,
eps: float,
):
residual_out = residual + residual_gate * x
normalized = norm_infer(
residual_out.view(-1, residual_out.shape[-1]),
weight,
bias,
eps=eps,
is_rms_norm=False,
).view_as(residual_out)
output, gate_out = _apply_select01_modulation(
normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
)
return output, residual_out, gate_out
def _apply_select01_modulation(
x: torch.Tensor,
scale0: torch.Tensor,
shift0: torch.Tensor,
gate0: torch.Tensor,
scale1: torch.Tensor,
shift1: torch.Tensor,
gate1: torch.Tensor,
index: torch.Tensor,
):
idx = index.bool().unsqueeze(-1)
scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
return x * (1 + scale) + shift, gate
@pytest.fixture(autouse=True)
def cuda_setup():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
torch.cuda.manual_seed(0)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
@pytest.mark.parametrize("seq_len", SEQ_LENS)
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
def test_fused_layernorm_scale_shift_gate_select01(
dtype, batch_size, seq_len, hidden_size
):
x = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE, dtype=dtype)
weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
bias = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
index = torch.randint(0, 2, (batch_size, seq_len), device=DEVICE, dtype=torch.int32)
scale0, shift0, gate0, scale1, shift1, gate1 = _make_modulation_tensors(
batch_size, hidden_size, dtype
)
out_ref, gate_ref = _baseline_select01_modulation(
x,
weight,
bias,
scale0,
shift0,
gate0,
scale1,
shift1,
gate1,
index,
EPS,
)
out_fused, gate_fused = fuse_layernorm_scale_shift_gate_select01_kernel(
x.contiguous(),
weight=weight,
bias=bias,
scale0=scale0,
shift0=shift0,
gate0=gate0,
scale1=scale1,
shift1=shift1,
gate1=gate1,
index=index,
eps=EPS,
)
atol, rtol = _tol(dtype)
triton.testing.assert_close(out_ref, out_fused, atol=atol, rtol=rtol)
triton.testing.assert_close(gate_ref, gate_fused, atol=atol, rtol=rtol)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
@pytest.mark.parametrize("seq_len", SEQ_LENS)
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
def test_fused_residual_layernorm_scale_shift_gate_select01(
dtype, batch_size, seq_len, hidden_size
):
x = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE, dtype=dtype)
residual = torch.randn_like(x)
residual_gate = torch.randn_like(x)
weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
bias = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
index = torch.randint(0, 2, (batch_size, seq_len), device=DEVICE, dtype=torch.int32)
scale0, shift0, gate0, scale1, shift1, gate1 = _make_modulation_tensors(
batch_size, hidden_size, dtype
)
out_ref, residual_ref, gate_ref = _baseline_residual_select01_modulation(
x,
residual,
residual_gate,
weight,
bias,
scale0,
shift0,
gate0,
scale1,
shift1,
gate1,
index,
EPS,
)
out_fused, residual_fused, gate_fused = (
fuse_residual_layernorm_scale_shift_gate_select01_kernel(
x.contiguous(),
residual=residual.contiguous(),
residual_gate=residual_gate.contiguous(),
weight=weight,
bias=bias,
scale0=scale0,
shift0=shift0,
gate0=gate0,
scale1=scale1,
shift1=shift1,
gate1=gate1,
index=index,
eps=EPS,
)
)
atol, rtol = _tol(dtype)
triton.testing.assert_close(out_ref, out_fused, atol=atol, rtol=rtol)
triton.testing.assert_close(residual_ref, residual_fused, atol=atol, rtol=rtol)
triton.testing.assert_close(gate_ref, gate_fused, atol=atol, rtol=rtol)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_residual_gate_add.py
-123
View File
@@ -1,123 +0,0 @@
import sys
import pytest
import torch
from sglang.kernels.ops.diffusion.residual_gate_add import (
can_use_residual_gate_add_cuda,
residual_gate_add,
residual_gate_add_cuda,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
CASES = [
((1, 1024, 4096), (1, 1, 4096)),
((1, 512, 4096), (1, 512, 4096)),
((1, 17, 65), (1, 1, 65)),
((1, 17, 65), (1, 17, 65)),
# FLUX.1 / FLUX.2-klein 1024^2 shapes (D=3072): dual-stream image/text
# and single-stream/joint concat; gates are [1, 1, D] modulation rows.
((1, 4096, 3072), (1, 1, 3072)),
((1, 512, 3072), (1, 1, 3072)),
((1, 4608, 3072), (1, 1, 3072)),
# FLUX.2-dev (D=6144) joint sequence.
((1, 4608, 6144), (1, 1, 6144)),
# ERNIE-4.5-VL 1024^2 image tokens plus text tokens.
((1, 4216, 4096), (1, 1, 4096)),
]
def _tol(dtype: torch.dtype) -> float:
return 1e-5 if dtype == torch.float32 else 5e-2
def _assert_matches_torch(out: torch.Tensor, ref: torch.Tensor) -> None:
if ref.dtype == torch.float32:
torch.testing.assert_close(out, ref, atol=_tol(ref.dtype), rtol=_tol(ref.dtype))
else:
torch.testing.assert_close(out, ref, atol=0, rtol=0)
@pytest.fixture(autouse=True)
def cuda_setup():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
torch.cuda.manual_seed(0)
@pytest.mark.parametrize("residual_shape,gate_shape", CASES)
def test_residual_gate_add_matches_torch(residual_shape, gate_shape):
residual = torch.randn(residual_shape, device="cuda", dtype=torch.bfloat16)
update = torch.randn_like(residual)
gate = torch.randn(gate_shape, device="cuda", dtype=torch.bfloat16)
out = residual_gate_add_cuda(residual, update, gate)
ref = residual + update * gate
_assert_matches_torch(out, ref)
assert torch.equal(residual_gate_add(residual, update, gate), ref)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@pytest.mark.parametrize("gate_shape", [(1, 1, 64), (1, 9, 64)])
def test_residual_gate_add_dtypes(dtype, gate_shape):
residual = torch.randn((1, 9, 64), device="cuda", dtype=dtype)
update = torch.randn_like(residual)
gate = torch.randn(gate_shape, device="cuda", dtype=dtype)
out = residual_gate_add_cuda(residual, update, gate)
ref = residual + update * gate
_assert_matches_torch(out, ref)
def test_can_use_residual_gate_add_cuda_rejects_unsupported_inputs():
residual = torch.randn((1, 8, 64), device="cuda", dtype=torch.bfloat16)
update = torch.randn_like(residual)
gate = torch.randn((1, 1, 64), device="cuda", dtype=torch.bfloat16)
assert can_use_residual_gate_add_cuda(residual, update, gate)
assert not can_use_residual_gate_add_cuda(residual.cpu(), update, gate)
assert not can_use_residual_gate_add_cuda(residual, update.float(), gate)
assert not can_use_residual_gate_add_cuda(residual, update[:, ::2], gate)
assert not can_use_residual_gate_add_cuda(residual, update, gate[:, :, ::2])
empty_residual = residual[:, :0]
empty_update = update[:, :0]
assert not can_use_residual_gate_add_cuda(empty_residual, empty_update, gate)
assert torch.equal(
residual_gate_add(empty_residual, empty_update, gate),
empty_residual + empty_update * gate,
)
# Only [1, ..., 1, D] row-broadcast gates are supported; a batched
# [B>1, 1, D] gate is not row-broadcast here and must fall back.
batched_residual = torch.randn((2, 8, 64), device="cuda", dtype=torch.bfloat16)
batched_update = torch.randn_like(batched_residual)
batched_gate = torch.randn((2, 1, 64), device="cuda", dtype=torch.bfloat16)
assert not can_use_residual_gate_add_cuda(
batched_residual, batched_update, batched_gate
)
assert torch.equal(
residual_gate_add(batched_residual, batched_update, batched_gate),
batched_residual + batched_update * batched_gate,
)
def test_residual_gate_add_custom_op_torch_compile_fullgraph():
residual = torch.randn((1, 32, 128), device="cuda", dtype=torch.bfloat16)
update = torch.randn_like(residual)
gate = torch.randn((1, 1, 128), device="cuda", dtype=torch.bfloat16)
def fn(residual, update, gate):
return residual_gate_add(residual, update, gate)
compiled = torch.compile(fn, fullgraph=True)
out = compiled(residual, update, gate)
ref = residual + update * gate
_assert_matches_torch(out, ref)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_qknorm_rope.py → test/registered/kernels/ops/diffusion/test_rope.py
+27 -23
View File
@@ -1,3 +1,20 @@
"""``diffusion.rope``: rotary embeddings and the QK-norm chains fused into them.
Two families with different oracles:
- ``fused_inplace_qknorm_rope`` / ``fused_qknorm_rope_pack_kv`` are compared
against the *split* baseline (a separate qknorm kernel plus FlashInfer or
sgl_kernel RoPE). In the default mode the two differ by about one bf16
rounding step, so those cases use a tolerance; with
``round_norm_before_rope=True`` the fused kernel reproduces the split
rounding exactly and ``torch.equal`` applies.
The LTX-2 split-RoPE kernel lives in ``test_rope_ltx2.py``: it is validated on
B200 and registered on that lane alone, which the cases here cannot share --
their oracle is the *split* baseline (a separate qknorm kernel plus sgl_kernel
or FlashInfer RoPE), whose dispatch differs on Blackwell, so the bit-exact
assertions below do not hold there.
"""
import itertools
import sys
@@ -6,12 +23,20 @@ import torch
import triton
from sglang.kernels.jit.utils import get_ci_test_range
from sglang.kernels.ops.diffusion import (
can_use_fused_inplace_qknorm_rope,
fused_inplace_qknorm_rope,
fused_qknorm_rope_pack_kv,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=44, stage="base-b-kernel-unit", runner_config="1-gpu-large")
# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
# Nightly is not redundant: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1, which
# expands the get_ci_test_range sweeps below.
register_cuda_ci(est_time=220, stage="nightly", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
DEVICE = "cuda"
DTYPE = torch.bfloat16
MAX_SEQ_LEN = 131072
@@ -70,8 +95,6 @@ def fused_qknorm_rope(
positions: torch.Tensor,
is_neox: bool,
) -> None:
from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
fused_inplace_qknorm_rope(
q,
k,
@@ -85,10 +108,6 @@ def fused_qknorm_rope(
def test_qknorm_rope_rejects_unsupported_dtypes() -> None:
from sglang.kernels.ops.diffusion.qknorm_rope import (
can_use_fused_inplace_qknorm_rope,
)
assert not can_use_fused_inplace_qknorm_rope(128, 128, False, torch.float32)
assert not can_use_fused_inplace_qknorm_rope(
128, 128, False, torch.bfloat16, torch.float64
@@ -164,9 +183,6 @@ def test_qknorm_rope(
def test_qknorm_rope_preserves_split_bf16_rounding() -> None:
from sgl_kernel import rotary_embedding
from sglang.kernels.ops.diffusion.qknorm_rope import (
fused_inplace_qknorm_rope,
)
from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
num_tokens, num_heads, head_dim, rope_dim = 257, 28, 128, 96
@@ -217,7 +233,6 @@ def test_qknorm_rope_preserves_split_bf16_rounding() -> None:
def test_qknorm_rope_preserves_full_width_neox_cache() -> None:
from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
num_tokens, num_heads, head_dim = 257, 32, 128
@@ -256,9 +271,6 @@ def test_qknorm_rope_preserves_full_width_neox_cache() -> None:
def test_qknorm_rope_requires_opt_in_for_strided_packed_gqa() -> None:
from sglang.kernels.ops.diffusion.qknorm_rope import (
fused_inplace_qknorm_rope,
)
from sglang.multimodal_gen.runtime.layers.layernorm import (
RMSNorm,
apply_qk_norm_rope,
@@ -329,10 +341,6 @@ def test_qknorm_rope_requires_opt_in_for_strided_packed_gqa() -> None:
def test_qknorm_rope_pack_kv_matches_separate_ops() -> None:
from sglang.kernels.ops.diffusion.qknorm_rope import (
fused_inplace_qknorm_rope,
fused_qknorm_rope_pack_kv,
)
batch_size = 2
prefix_tokens, suffix_tokens = 17, 257
@@ -416,9 +424,6 @@ def test_qknorm_rope_pack_kv_matches_separate_ops() -> None:
def test_qknorm_rope_pack_kv_preserves_split_bf16_rounding() -> None:
from sgl_kernel import rotary_embedding
from sglang.kernels.ops.diffusion.qknorm_rope import (
fused_qknorm_rope_pack_kv,
)
from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
batch_size = 1
@@ -507,7 +512,6 @@ def test_qknorm_rope_pack_kv_preserves_split_bf16_rounding() -> None:
def test_qknorm_rope_accepts_empty_token_dimension() -> None:
from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
num_heads, head_dim = 8, 128
q = torch.empty(0, num_heads, head_dim, device=DEVICE, dtype=DTYPE)
@@ -530,4 +534,4 @@ def test_qknorm_rope_accepts_empty_token_dimension() -> None:
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_ltx2_qknorm_split_rope.py → test/registered/kernels/ops/diffusion/test_rope_ltx2.py
+30 -18
View File
@@ -1,10 +1,19 @@
"""``diffusion.rope``: the LTX-2 QK-norm + split-RoPE CUDA kernel.
Split out of ``test_rope.py`` rather than merged with the other RoPE kernels:
this one is validated on B200 and registered on that lane alone, while the
``fused_inplace_qknorm_rope`` cases there are held to the *split* baseline,
whose sgl_kernel / FlashInfer dispatch differs on Blackwell -- their bit-exact
assertions fail on B200. One file cannot carry both lane sets.
"""
import sys
import pytest
import torch
import torch.nn.functional as F
from sglang.kernels.ops.diffusion.ltx2_qknorm_split_rope import (
from sglang.kernels.ops.diffusion import (
can_use_ltx2_qknorm_split_rope_cuda,
ltx2_qknorm_split_rope_cuda,
)
@@ -12,23 +21,18 @@ from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
DEVICE = "cuda"
BF16_FUSED_ATOL = 1.6e-1
def _require_cuda_b200() -> None:
def _require_b200() -> None:
if not torch.cuda.is_available():
pytest.skip("CUDA required")
if torch.cuda.get_device_capability()[0] < 10:
pytest.skip("LTX2 QKNorm split-RoPE CUDA path is validated on B200")
@pytest.fixture(autouse=True)
def cuda_setup():
_require_cuda_b200()
torch.cuda.manual_seed(20260630)
def _make_cos_sin(
def _ltx2_make_cos_sin(
batch: int, seq_len: int, num_heads: int, head_dim: int
) -> tuple[torch.Tensor, torch.Tensor]:
half_dim = head_dim // 2
@@ -64,7 +68,7 @@ def _apply_split_rotary_ref(
return out.swapaxes(1, 2).reshape(batch, seq_len, -1).to(dtype=x_dtype)
def _reference(
def _ltx2_reference(
q: torch.Tensor,
k: torch.Tensor,
q_cos: torch.Tensor,
@@ -95,12 +99,14 @@ def _reference(
def test_ltx2_qknorm_split_rope_matches_torch_exactly(
batch: int, q_seq: int, k_seq: int, num_heads: int, head_dim: int
) -> None:
_require_b200()
torch.cuda.manual_seed(20260630)
hidden = num_heads * head_dim
eps = 1e-6
q = torch.randn(batch, q_seq, hidden, device="cuda", dtype=torch.bfloat16)
k = torch.randn(batch, k_seq, hidden, device="cuda", dtype=torch.bfloat16)
q_cos, q_sin = _make_cos_sin(batch, q_seq, num_heads, head_dim)
k_cos, k_sin = _make_cos_sin(batch, k_seq, num_heads, head_dim)
q_cos, q_sin = _ltx2_make_cos_sin(batch, q_seq, num_heads, head_dim)
k_cos, k_sin = _ltx2_make_cos_sin(batch, k_seq, num_heads, head_dim)
q_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
k_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
@@ -117,7 +123,9 @@ def test_ltx2_qknorm_split_rope_matches_torch_exactly(
head_dim=head_dim,
)
q_ref, k_ref = _reference(q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight, eps)
q_ref, k_ref = _ltx2_reference(
q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight, eps
)
q_out, k_out = ltx2_qknorm_split_rope_cuda(
q,
q_cos,
@@ -138,9 +146,11 @@ def test_ltx2_qknorm_split_rope_matches_torch_exactly(
def test_ltx2_qknorm_split_rope_rejects_unsupported_inputs() -> None:
_require_b200()
torch.cuda.manual_seed(20260630)
q = torch.randn((1, 3, 4096), device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
q_cos, q_sin = _make_cos_sin(1, 3, 32, 128)
q_cos, q_sin = _ltx2_make_cos_sin(1, 3, 32, 128)
q_weight = torch.randn(4096, device="cuda", dtype=torch.bfloat16)
k_weight = torch.randn(4096, device="cuda", dtype=torch.bfloat16)
@@ -183,12 +193,14 @@ def test_ltx2_qknorm_split_rope_rejects_unsupported_inputs() -> None:
def test_ltx2_qknorm_split_rope_custom_op_torch_compile_fullgraph() -> None:
_require_b200()
torch.cuda.manual_seed(20260630)
batch, q_seq, k_seq, num_heads, head_dim = 1, 3, 2, 32, 64
hidden = num_heads * head_dim
q = torch.randn(batch, q_seq, hidden, device="cuda", dtype=torch.bfloat16)
k = torch.randn(batch, k_seq, hidden, device="cuda", dtype=torch.bfloat16)
q_cos, q_sin = _make_cos_sin(batch, q_seq, num_heads, head_dim)
k_cos, k_sin = _make_cos_sin(batch, k_seq, num_heads, head_dim)
q_cos, q_sin = _ltx2_make_cos_sin(batch, q_seq, num_heads, head_dim)
k_cos, k_sin = _ltx2_make_cos_sin(batch, k_seq, num_heads, head_dim)
q_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
k_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
@@ -209,7 +221,7 @@ def test_ltx2_qknorm_split_rope_custom_op_torch_compile_fullgraph() -> None:
compiled = torch.compile(fn, fullgraph=True)
q_out, k_out = compiled(q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight)
q_ref, k_ref = _reference(
q_ref, k_ref = _ltx2_reference(
q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight, 1e-6
)
torch.cuda.synchronize()
@@ -218,4 +230,4 @@ def test_ltx2_qknorm_split_rope_custom_op_torch_compile_fullgraph() -> None:
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_sana_conv_post.py
-60
View File
@@ -1,60 +0,0 @@
"""Sana GLUMB post-processing fusions must match the eager bf16 chain."""
import pytest
import torch
import torch.nn.functional as F
from sglang.kernels.ops.diffusion.triton.sana_conv_post import (
can_use_fused_bias_glu,
can_use_fused_bias_silu,
fused_bias_glu,
fused_bias_silu,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=3, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize("channels", [2240, 11200])
def test_sana_bias_silu_is_bit_exact(channels):
torch.manual_seed(0)
x = torch.randn(
(1, channels, 7, 5),
device="cuda",
dtype=torch.bfloat16,
).to(memory_format=torch.channels_last)
bias = torch.randn(channels, device="cuda", dtype=torch.bfloat16)
assert can_use_fused_bias_silu(x, bias)
actual = fused_bias_silu(x, bias)
expected = F.silu(x + bias[None, :, None, None])
assert actual.is_contiguous(memory_format=torch.channels_last)
assert torch.equal(actual, expected)
@pytest.mark.parametrize("channels", [2240, 5600])
def test_sana_bias_glu_is_bit_exact(channels):
torch.manual_seed(1)
x = torch.randn(
(1, 2 * channels, 7, 5),
device="cuda",
dtype=torch.bfloat16,
).to(memory_format=torch.channels_last)
bias = torch.randn(2 * channels, device="cuda", dtype=torch.bfloat16)
assert can_use_fused_bias_glu(x, bias)
actual = fused_bias_glu(x, bias)
biased = x + bias[None, :, None, None]
hidden, gate = torch.chunk(biased, 2, dim=1)
expected = hidden * F.silu(gate)
assert actual.is_contiguous(memory_format=torch.channels_last)
assert torch.equal(actual, expected)
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/diffusion/test_sana_ln_modulate.py
-55
View File
@@ -1,55 +0,0 @@
"""Sana fused LN+modulate fast path must stay bit-exact vs eager."""
import pytest
import torch
import sglang.multimodal_gen.runtime.models.dits.sana as sana
from sglang.multimodal_gen.runtime.models.dits.sana import (
_eager_ln_modulate,
_sana_ln_modulate,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=3, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize(
"shape,nmod,transposed",
[
((2, 1024, 2240), 6, False),
((2, 1024, 2240), 2, False),
((1, 1024, 2240), 6, True),
((1, 37, 2240), 6, False),
],
)
def test_sana_fused_ln_modulate_is_bit_exact(shape, nmod, transposed):
# (., 1024, 2240) is the real Sana 1024px shape; hidden 2240 % 512 != 0
# exercises the kernel's partial tail chunk. nmod mirrors the two adaLN
# chunk layouts, transposed the permuted layout the Sana DiT serves.
torch.manual_seed(0)
batch, seq, hidden = shape
norm = torch.nn.LayerNorm(hidden, eps=1e-6, elementwise_affine=False).cuda()
x = (torch.randn(batch, seq, hidden, device="cuda") * 4).bfloat16()
if transposed:
x = x.permute(0, 2, 1).contiguous().permute(0, 2, 1)
emb = torch.randn(batch, nmod, hidden, device="cuda").bfloat16()
shift, scale = emb.chunk(nmod, dim=1)[0], emb.chunk(nmod, dim=1)[-1]
# default-stream eager serving must stay on the untouched eager chain
n_sigs = len(sana._SANA_LN_MOD.verified_sigs)
_sana_ln_modulate(norm, x, scale, shift)
assert len(sana._SANA_LN_MOD.verified_sigs) == n_sigs
# the fusion engages on non-default streams (the BCG warmup/capture path)
with torch.cuda.stream(torch.cuda.Stream()):
out = _sana_ln_modulate(norm, x, scale, shift)
assert len(sana._SANA_LN_MOD.verified_sigs) == n_sigs + 1 # verified
out2 = _sana_ln_modulate(norm, x, scale, shift) # verified-sig lane
torch.cuda.synchronize()
assert torch.equal(out, _eager_ln_modulate(norm, x, scale, shift))
assert torch.equal(out2, out) and not sana._SANA_LN_MOD.disabled
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/diffusion/test_scale_shift.py
-40
View File
@@ -1,40 +0,0 @@
import sys
import pytest
import torch
from sglang.kernels.ops.diffusion.triton.scale_shift import (
try_fused_scaled_residual_add_exact,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@torch.no_grad()
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_scaled_residual_add_is_bit_exact(dtype):
torch.manual_seed(0)
residual = torch.randn(2, 17, 64, device="cuda", dtype=torch.float32)
x = torch.randn(2, 17, 64, device="cuda", dtype=dtype)
scale = torch.randn(64, device="cuda", dtype=torch.float32)
actual = try_fused_scaled_residual_add_exact(residual, x, scale)
expected = residual + x * scale
assert actual is not None
assert torch.equal(actual, expected)
@torch.no_grad()
def test_scaled_residual_add_rejects_unsupported_inputs():
residual = torch.empty(2, 3, 8, device="cuda", dtype=torch.float32)
x = torch.empty_like(residual)
scale = torch.empty(8, device="cuda", dtype=torch.float32)
assert try_fused_scaled_residual_add_exact(residual, x, scale) is None
assert try_fused_scaled_residual_add_exact(residual, x.half(), scale[:-1]) is None
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_sites.py
+416
View File
@@ -0,0 +1,416 @@
"""``diffusion.sites``: the two gate protocols and the mount lifecycle.
Nothing here is a kernel. ``sites`` decides *whether* a fusion is allowed to
run for a given request and model:
- :class:`QualityGatedFusion` -- for fusions that are **not** bit-exact. They
mount onto marked ``nn.Module`` sites only for ``quality="high"``, at batch
boundaries, all-or-nothing per transformer.
- :class:`BitExactFusionGate` -- for fusions that **are** bit-exact. They run
by default but verify themselves against the live eager chain on first
sight and disable permanently on any mismatch.
The protocol tests are pure-CPU. The mount-lifecycle tests below use
synthetic sites (a bare ``nn.Module`` with the marker attribute) so they test
the protocol rather than any one model; the real model wrappers live in
``test_model_fast_paths.py``.
"""
import sys
from types import ModuleType
from unittest.mock import MagicMock, patch
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
import sglang.kernels.ops.diffusion.sites.fused_gate_rmsnorm_site as gate_rmsnorm
import sglang.kernels.ops.diffusion.sites.fused_linear_gelu_site as linear_gelu
from sglang.kernels.ops.diffusion import (
BitExactFusionGate,
QualityGatedFusion,
can_use_ln_modulate,
flashinfer_rmsnorm_diagnostic_hint,
fused_ln_modulate,
fused_ln_modulate_active,
mark_fused_ln_modulate_site,
mount_fused_ln_modulate,
tensors_equal,
unmount_fused_ln_modulate,
)
from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
register_cuda_ci(est_time=38, stage="base-b-kernel-unit", runner_config="1-gpu-large")
requires_cuda = pytest.mark.skipif(
not torch.cuda.is_available(), reason="CUDA required"
)
# ---------------------------------------------------------------------------
# QualityGatedFusion protocol (CPU)
# ---------------------------------------------------------------------------
def test_quality_gate_mounts_and_unmounts_all_sites():
fusion = QualityGatedFusion(
name="test fusion",
marker_attr="_test_fusion_site",
enabled_attr="_test_fusion_enabled",
)
root = nn.ModuleList([nn.Module(), nn.Module()])
for index, site in enumerate(root):
fusion.mark(site, index)
assert [fusion.metadata(site) for site in root] == [0, 1]
assert fusion.mount(root)
assert all(fusion.is_enabled(site) for site in root)
fusion.unmount(root)
assert not any(fusion.is_enabled(site) for site in root)
def test_quality_gate_rejection_is_all_or_nothing():
fusion = QualityGatedFusion(
name="test fusion",
marker_attr="_test_fusion_site",
enabled_attr="_test_fusion_enabled",
)
root = nn.ModuleList([nn.Module(), nn.Module()])
for index, site in enumerate(root):
fusion.mark(site, index)
assert not fusion.mount(
root, reject_reason=lambda site: "rejected" if fusion.metadata(site) else None
)
assert not any(fusion.is_enabled(site) for site in root)
assert not fusion.mount(nn.Module())
# -------------------------------------------------------------------------
# BitExactFusionGate protocol (CPU)
# -------------------------------------------------------------------------
def test_bitexact_gate_once_mode_verifies_then_reuses():
gate = BitExactFusionGate("once")
calls = {"fused": 0, "ref": 0}
def fused():
calls["fused"] += 1
return torch.tensor([1.0])
def ref():
calls["ref"] += 1
return torch.tensor([1.0])
assert torch.equal(gate.accept_or_fallback(fused(), ref()), torch.tensor([1.0]))
assert gate.verified and not gate.disabled and calls == {"fused": 1, "ref": 1}
assert torch.equal(fused(), torch.tensor([1.0]))
assert calls == {"fused": 2, "ref": 1}
def test_bitexact_gate_mismatch_disables_permanently():
gate = BitExactFusionGate("mismatch")
out = gate.accept_or_fallback(
torch.tensor([1.0]),
torch.tensor([2.0]),
mismatch_msg="mismatch",
)
assert torch.equal(out, torch.tensor([2.0]))
assert gate.disabled and not gate.verified
def test_bitexact_gate_per_signature_tracks_each_sig():
gate = BitExactFusionGate("sig", per_signature=True)
a = torch.tensor([1.0])
assert torch.equal(gate.accept_or_fallback(a, a, sig=("a",)), a)
assert gate.is_verified(("a",))
assert not gate.is_verified(("b",))
assert torch.equal(gate.accept_or_fallback(a, a, sig=("b",)), a)
assert gate.verified_sigs == {("a",), ("b",)}
def test_bitexact_gate_skips_first_sight_during_graph_capture(monkeypatch):
# Negative-branch contract: an unverified gate must not attempt first-sight
# verification inside CUDA graph capture — the eager-reference host sync
# would abort the capture (and BCG would permanently block the signature).
gate = BitExactFusionGate("capture")
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
assert not gate.can_attempt_once()
# A verified gate replays the fused kernel alone, which is capture-safe.
gate.mark_verified()
assert gate.can_attempt_once()
def test_tensors_equal_supports_sequences():
assert tensors_equal(
(torch.tensor([1.0]), torch.tensor([2.0])),
(torch.tensor([1.0]), torch.tensor([2.0])),
)
assert not tensors_equal(
(torch.tensor([1.0]), torch.tensor([2.0])),
(torch.tensor([1.0]), torch.tensor([3.0])),
)
class TestBitExactFallbackDiagnostics(CustomTestCase):
def test_mismatch_warning_is_actionable_and_diagnostic_is_lazy(self):
logger = MagicMock()
diagnostic = MagicMock(return_value="backend=CuTe DSL")
gate = BitExactFusionGate("diagnostic")
matched = gate.accept_or_fallback(
torch.tensor([1.0]),
torch.tensor([1.0]),
logger=logger,
diagnostic_hint=diagnostic,
)
self.assertTrue(torch.equal(matched, torch.tensor([1.0])))
diagnostic.assert_not_called()
logger.warning_once.assert_not_called()
gate = BitExactFusionGate("diagnostic")
fallback = gate.accept_or_fallback(
torch.tensor([1.0]),
torch.tensor([2.0]),
logger=logger,
diagnostic_hint=diagnostic,
)
self.assertTrue(torch.equal(fallback, torch.tensor([2.0])))
diagnostic.assert_called_once_with()
warning = logger.warning_once.call_args.args[0]
self.assertIn("Correctness is preserved", warning)
self.assertIn("reference kernel or reduction-order change", warning)
self.assertIn("backend=CuTe DSL", warning)
def test_diagnostic_failure_cannot_break_the_eager_fallback(self):
logger = MagicMock()
def broken_diagnostic():
raise RuntimeError("diagnostics unavailable")
gate = BitExactFusionGate("diagnostic")
fallback = gate.accept_or_fallback(
torch.tensor([1.0]),
torch.tensor([2.0]),
logger=logger,
diagnostic_hint=broken_diagnostic,
)
self.assertTrue(torch.equal(fallback, torch.tensor([2.0])))
self.assertTrue(gate.disabled)
self.assertIn("Correctness is preserved", logger.warning_once.call_args.args[0])
def test_flashinfer_rmsnorm_hint_reports_backend_and_versions(self):
flashinfer = ModuleType("flashinfer")
flashinfer_norm = ModuleType("flashinfer.norm")
flashinfer_norm._USE_CUDA_NORM = False
versions = {
"flashinfer-python": "0.6.12",
"flashinfer-cubin": "0.6.12",
"flashinfer-jit-cache": "0.6.12+cu130",
}
with (
patch.dict(
sys.modules,
{"flashinfer": flashinfer, "flashinfer.norm": flashinfer_norm},
),
patch("importlib.metadata.version", side_effect=versions.__getitem__),
patch.dict("os.environ", {"FLASHINFER_USE_CUDA_NORM": "0"}),
):
hint = flashinfer_rmsnorm_diagnostic_hint()
self.assertIn("backend=CuTe DSL", hint)
self.assertIn("FLASHINFER_USE_CUDA_NORM=0", hint)
for package, version in versions.items():
self.assertIn(f"{package}={version}", hint)
# ---------------------------------------------------------------------------
# Mount lifecycle on synthetic sites (CUDA)
# ---------------------------------------------------------------------------
LN_DIM = 3072 # FLUX.1 hidden size
@requires_cuda
@pytest.mark.parametrize("seq_len", [4096, 512])
def test_fused_ln_modulate_matches_reference(seq_len):
torch.cuda.manual_seed(0)
x = torch.randn((1, seq_len, LN_DIM), device="cuda", dtype=torch.bfloat16)
scale = torch.randn((1, LN_DIM), device="cuda", dtype=torch.bfloat16)
shift = torch.randn_like(scale)
assert can_use_ln_modulate(x, scale, shift)
norm = nn.LayerNorm(LN_DIM, eps=1e-6, elementwise_affine=False).cuda()
ref = norm(x) * (1 + scale[:, None]) + shift[:, None]
# Contract: bf16 rounding-order-level difference only, not bit-exact --
# which is exactly why this fusion is quality-gated.
torch.testing.assert_close(
fused_ln_modulate(x, scale, shift, eps=1e-6), ref, atol=0.0625, rtol=0.05
)
@requires_cuda
def test_fused_ln_modulate_guards_and_mount_protocol():
x = torch.randn((2, 64, LN_DIM), device="cuda", dtype=torch.bfloat16)
row = torch.randn((2, LN_DIM), device="cuda", dtype=torch.bfloat16)
assert not can_use_ln_modulate(x, row, row) # folded affine needs B == 1
root = nn.Module()
root.child = nn.Module()
mark_fused_ln_modulate_site(root.child)
assert not fused_ln_modulate_active(root.child)
assert mount_fused_ln_modulate(root)
assert fused_ln_modulate_active(root.child)
unmount_fused_ln_modulate(root)
assert not fused_ln_modulate_active(root.child)
assert not mount_fused_ln_modulate(nn.Module()) # no marked sites
@requires_cuda
@torch.no_grad()
def test_mounted_ln_modulate_site_compiles_fullgraph():
class Site(nn.Module):
def __init__(self):
super().__init__()
mark_fused_ln_modulate_site(self)
def forward(self, x, scale, shift):
if fused_ln_modulate_active(self) and can_use_ln_modulate(x, scale, shift):
return fused_ln_modulate(x, scale, shift, eps=1e-6)
return (
F.layer_norm(x, (x.shape[-1],), eps=1e-6) * (1 + scale[:, None])
+ shift[:, None]
)
site = Site()
assert mount_fused_ln_modulate(site)
x = torch.randn(1, 64, 128, device="cuda", dtype=torch.bfloat16)
scale = torch.randn(1, 128, device="cuda", dtype=torch.bfloat16)
shift = torch.randn_like(scale)
expected = site(x, scale, shift)
torch.testing.assert_close(
torch.compile(site, fullgraph=True)(x, scale, shift),
expected,
atol=0.0625,
rtol=0.05,
)
GATE_RMSNORM_DIM, GATE_RMSNORM_EPS = 4608, 1e-5 # Ideogram 4 hidden / norm_eps
class _GateRMSNormSite(nn.Module):
def __init__(self, dtype=torch.bfloat16):
super().__init__()
self.norm = nn.RMSNorm(
GATE_RMSNORM_DIM, eps=GATE_RMSNORM_EPS, device="cuda", dtype=dtype
)
gate_rmsnorm.mark_fused_gate_rmsnorm_site(self, ("norm",))
@requires_cuda
def test_fused_gate_rmsnorm_matches_ideogram_reference():
torch.manual_seed(0)
site = _GateRMSNormSite()
w = site.norm.weight.data
dim = GATE_RMSNORM_DIM
x = torch.randn(1, 64, dim, device="cuda", dtype=torch.bfloat16)
residual = torch.randn_like(x)
# adaln-style strided chunks, as produced by Ideogram's modulation .chunk()
mods = torch.randn(1, 1, 2 * dim, device="cuda", dtype=torch.bfloat16)
scale, gate = mods.chunk(2, dim=-1)
assert gate_rmsnorm.mount_fused_gate_rmsnorm(site)
got_scale = gate_rmsnorm.fused_rmsnorm_scale(x, w, 1.0 + scale, GATE_RMSNORM_EPS)
got_gate = gate_rmsnorm.fused_rmsnorm_tanh_residual(
x, gate, residual, w, GATE_RMSNORM_EPS
)
norm = F.rms_norm(x, (dim,), w, GATE_RMSNORM_EPS)
# The fused path uses bf16-native norm statistics: close, not bit-exact.
torch.testing.assert_close(got_scale, norm * (1.0 + scale), atol=8e-2, rtol=4e-2)
torch.testing.assert_close(
got_gate, residual + torch.tanh(gate) * norm, atol=8e-2, rtol=4e-2
)
@requires_cuda
def test_fused_gate_rmsnorm_mount_is_all_or_nothing():
good, bad = _GateRMSNormSite(), _GateRMSNormSite(torch.float32)
# One fp32 norm anywhere in the tree keeps *every* site on the reference.
assert not gate_rmsnorm.mount_fused_gate_rmsnorm(nn.ModuleList([good, bad]))
assert not gate_rmsnorm.fused_gate_rmsnorm_active(good)
assert gate_rmsnorm.mount_fused_gate_rmsnorm(good)
gate_rmsnorm.unmount_fused_gate_rmsnorm(good)
assert not gate_rmsnorm.fused_gate_rmsnorm_active(good)
class _GeluSite(nn.Module):
def __init__(self, dtype=torch.bfloat16, bias=True):
super().__init__()
self.proj = nn.Linear(64, 256, bias=bias, device="cuda", dtype=dtype)
linear_gelu.mark_fused_gelu_site(self, "proj")
def forward(self, x):
if linear_gelu.fused_gelu_active(self) and linear_gelu.can_use_linear_gelu(
self.proj, x
):
return linear_gelu.fused_linear_gelu_tanh(
x, self.proj.weight, self.proj.bias
)
return F.gelu(self.proj(x), approximate="tanh")
@requires_cuda
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
def test_fused_linear_gelu_matches_reference(dtype):
torch.manual_seed(0)
site = _GeluSite(dtype)
x = torch.randn(512, 64, device="cuda", dtype=dtype)
ref = site(x) # unmounted: the reference chain
assert linear_gelu.mount_fused_linear_gelu(site)
atol = 2e-2 if dtype == torch.bfloat16 else 4e-3
torch.testing.assert_close(site(x), ref, atol=atol, rtol=2e-2)
@requires_cuda
def test_fused_linear_gelu_guards_and_lossless_path():
torch.manual_seed(0)
good, bad = _GeluSite(), _GeluSite(torch.float32)
assert not linear_gelu.mount_fused_linear_gelu(nn.ModuleList([good, bad]))
assert not linear_gelu.fused_gelu_active(good)
x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
ref = good(x)
assert linear_gelu.mount_fused_linear_gelu(good)
linear_gelu.unmount_fused_linear_gelu(good)
# Unmounting must restore the reference path bit-for-bit.
assert torch.equal(good(x), ref)
no_bias = nn.Linear(8, 8, bias=False, device="cuda", dtype=torch.bfloat16)
assert not linear_gelu.can_use_linear_gelu_static(no_bias)
assert not linear_gelu.can_use_linear_gelu(good.proj, x.float())
@requires_cuda
@torch.no_grad()
def test_mounted_gelu_site_compiles_fullgraph():
site = _GeluSite()
x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
assert linear_gelu.mount_fused_linear_gelu(site)
expected = site(x)
torch.testing.assert_close(
torch.compile(site, fullgraph=True)(x), expected, atol=2e-2, rtol=2e-2
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_stage_profiler_sync.py
-50
View File
@@ -1,50 +0,0 @@
"""SGLANG_DIFFUSION_SYNC_STAGE_PROFILING must drain the GPU queue at the
timing start of *stage* records too — otherwise a stage that only launches
kernels (DenoisingStage's tail) leaks its queued work into whichever later
stage blocks first, inflating e.g. DecodingStage readings 2-3x."""
import sys
import time
import pytest
import torch
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.perf_logger import (
RequestMetrics,
StageProfiler,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
def test_stage_entry_sync_excludes_previous_stage_tail(monkeypatch):
monkeypatch.setenv("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "1")
logger = init_logger(__name__)
metrics = RequestMetrics("stage-sync-test")
# Calibrate ~0.5 s of queued GPU work.
torch.cuda.synchronize()
t0 = time.perf_counter()
torch.cuda._sleep(10_000_000)
torch.cuda.synchronize()
cycles = int(10_000_000 / max(time.perf_counter() - t0, 1e-9) * 0.5)
# Producer stage queues work without awaiting it (a denoise tail).
with StageProfiler("producer", logger, metrics, perf_dump_path_provided=True):
torch.cuda._sleep(cycles)
# Consumer stage's first blocking op used to absorb the producer's tail.
with StageProfiler("consumer", logger, metrics, perf_dump_path_provided=True):
torch.ones(8, device="cuda").sum().cpu()
producer_ms, consumer_ms = metrics.stages["producer"], metrics.stages["consumer"]
assert (
producer_ms > 250
), f"queued work not attributed to producer: {metrics.stages}"
assert consumer_ms < 100, f"producer tail leaked into consumer: {metrics.stages}"
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_timestep_embedding.py
-184
View File
@@ -1,184 +0,0 @@
import os
import sys
import numpy as np
import pytest
import torch
try:
import tabulate
except Exception:
tabulate = None
from sglang.kernels.jit.utils import get_ci_test_range
from sglang.kernels.ops.diffusion.timestep_embedding import (
timestep_embedding as timestep_embedding_cuda,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=16, stage="base-b-kernel-unit", runner_config="1-gpu-large")
# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
register_cuda_ci(est_time=20, stage="nightly", runner_config="1-gpu-large")
CORRECTNESS_BATCH_SIZES = get_ci_test_range(
[1, 2, 8, 128, 256, 512, 1536, 2048, 4096, 11008, 16384],
[1, 128, 2048, 16384],
)
CORRECTNESS_DIMS = get_ci_test_range(
[32, 128, 256, 512, 1536, 2048, 4096, 8192],
[32, 512, 8192],
)
DIFFUSERS_BATCH_SIZES = get_ci_test_range(
[1, 2, 8, 128, 256, 512, 1536, 2048, 16384],
[1, 512, 16384],
)
DIFFUSERS_DIMS = get_ci_test_range([32, 256, 512, 1536, 8192], [32, 512, 8192])
DTYPES = get_ci_test_range(
[torch.float16, torch.bfloat16, torch.float32],
[torch.float16, torch.bfloat16],
)
SCALES = get_ci_test_range([1, 0.01], [1, 0.01])
def get_timestep_embedding_reference(
timesteps: torch.Tensor,
dim: int,
*,
flip_sin_to_cos: bool = False,
downscale_freq_shift: float = 1,
scale: float = 1,
max_period: int = 10000,
):
assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
timesteps = timesteps.to(torch.float32)
half_dim = dim // 2
exponent = -torch.log(
torch.tensor(max_period, dtype=torch.float32, device=timesteps.device)
) * torch.arange(
start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
)
exponent = exponent / (half_dim - downscale_freq_shift)
emb = torch.exp(exponent)
emb = timesteps[:, None].float() * emb[None, :]
emb = scale * emb
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
if flip_sin_to_cos:
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
if dim % 2 == 1:
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
return emb
@pytest.mark.parametrize("batch_size", CORRECTNESS_BATCH_SIZES)
@pytest.mark.parametrize("dim", CORRECTNESS_DIMS)
@pytest.mark.parametrize("dtype", DTYPES)
def test_timestep_embedding_correctness_with_sgld(batch_size, dim, dtype):
device = "cuda"
t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype)
torch_output = get_timestep_embedding_reference(
t, dim, flip_sin_to_cos=True, downscale_freq_shift=0
)
cuda_output = timestep_embedding_cuda(
t, dim, flip_sin_to_cos=True, downscale_freq_shift=0
)
torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3)
@pytest.mark.parametrize("batch_size", DIFFUSERS_BATCH_SIZES)
@pytest.mark.parametrize("dim", DIFFUSERS_DIMS)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("flip_sin_to_cos", [False, True])
@pytest.mark.parametrize("downscale_freq_shift", [0, 1])
@pytest.mark.parametrize("scale", SCALES)
def test_timestep_embedding_correctness_with_diffusers(
batch_size, dim, flip_sin_to_cos, downscale_freq_shift, scale, dtype
):
device = "cuda"
t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype)
torch_output = get_timestep_embedding_reference(
t,
dim,
flip_sin_to_cos=flip_sin_to_cos,
downscale_freq_shift=downscale_freq_shift,
scale=scale,
max_period=10000,
)
cuda_output = timestep_embedding_cuda(
t,
dim,
flip_sin_to_cos=flip_sin_to_cos,
downscale_freq_shift=downscale_freq_shift,
scale=scale,
max_period=10000,
)
torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3)
def test_timestep_embedding_perf():
if os.environ.get("SGLANG_RUN_JIT_KERNEL_PERF_TESTS") != "1":
pytest.skip("Perf test disabled by default")
if tabulate is None:
pytest.skip("Optional dependency 'tabulate' is not installed")
NUM_BATCH = [1, 2, 8, 63, 256, 512, 613, 1024, 1536]
NUM_DIM = [32, 64, 128, 256, 512, 1024, 2048, 4096]
def perf_kernel_fn(kernel_fn: callable, *args, **kwargs):
warmup_times = 4
repeat_times = 20
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
for _ in range(warmup_times):
kernel_fn(*args, **kwargs)
torch.cuda.synchronize()
start.record()
for _ in range(repeat_times):
kernel_fn(*args, **kwargs)
end.record()
end.synchronize()
return start.elapsed_time(end) / repeat_times
device = "cuda"
results = []
cuda_speedups = []
for B in NUM_BATCH:
for dim in NUM_DIM:
t = torch.linspace(0, max(100000, B), steps=B, device=device).to(
torch.float32
)
time_torch = perf_kernel_fn(get_timestep_embedding_reference, t, dim)
time_cuda = perf_kernel_fn(timestep_embedding_cuda, t, dim)
speedup_cuda = time_torch / time_cuda
results.append(
{
"Batch Size": B,
"Dimension": dim,
"Torch Time (ms)": time_torch,
"CUDA Time (ms)": time_cuda,
"Speedup (CUDA)": speedup_cuda,
}
)
cuda_speedups.append(speedup_cuda)
print("=== Timestep Embedding Benchmark Results ===")
print(
tabulate.tabulate(
results,
headers="keys",
tablefmt="fancy_grid",
floatfmt=(".0f", ".0f", ".6f", ".6f", ".5f"),
)
)
print(f"Average Speedup(cuda): {np.mean(cuda_speedups):.4f}")
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_ulysses_qkv.py
-54
View File
@@ -1,54 +0,0 @@
import sys
import pytest
import torch
from sglang.kernels.ops.diffusion.triton.ulysses_qkv import (
pack_qkv_destination_major,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_pack_qkv_destination_major_is_bit_exact(dtype):
torch.manual_seed(0)
rows, world_size, global_heads, head_size = 17, 4, 12, 64
q, k, v = (
torch.randn(rows, global_heads, head_size, device="cuda", dtype=dtype)
for _ in range(3)
)
local_heads = global_heads // world_size
expected = torch.empty(
world_size,
rows,
local_heads,
3 * head_size,
device="cuda",
dtype=dtype,
)
for index, tensor in enumerate((q, k, v)):
shards = tensor.view(rows, world_size, local_heads, head_size).permute(
1, 0, 2, 3
)
expected[..., index * head_size : (index + 1) * head_size].copy_(shards)
actual = pack_qkv_destination_major(q, k, v, world_size)
assert torch.equal(actual, expected)
def test_pack_qkv_destination_major_validates_inputs():
q = torch.empty(2, 4, 8, device="cuda", dtype=torch.bfloat16)
with pytest.raises(ValueError, match="same 3D shape"):
pack_qkv_destination_major(q, q[:, :-1], q, 2)
with pytest.raises(ValueError, match="divide global_heads"):
pack_qkv_destination_major(q, q, q, 3)
with pytest.raises(ValueError, match="expected shape"):
pack_qkv_destination_major(q, q, q, 2, out=torch.empty_like(q))
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_usp_relayout.py
-69
View File
@@ -1,69 +0,0 @@
"""Bitwise tests for the generic Ulysses output head-merge fast path."""
import sys
from unittest.mock import patch
import pytest
import torch
from sglang.kernels.ops.diffusion.usp_relayout import (
can_use_usp_merge_heads,
usp_merge_heads,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
DEVICE = "cuda"
@pytest.mark.parametrize(
"world,seq,batch,h_local,head_dim",
[
(4, 7936, 1, 14, 128), # H3 768p production shape (Ulysses 4)
(2, 64, 3, 4, 64), # batched
(4, 33, 2, 4, 100), # scalar fallback inside the CUDA kernel
],
)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
def test_usp_merge_heads_bitwise(dtype, world, seq, batch, h_local, head_dim):
generator = torch.Generator(device=DEVICE).manual_seed(4321)
x = torch.randn(
world,
seq,
batch,
h_local,
head_dim,
dtype=dtype,
device=DEVICE,
generator=generator,
)
assert can_use_usp_merge_heads(x)
out = usp_merge_heads(x)
ref = x.permute(2, 1, 0, 3, 4).contiguous()
assert out.shape == ref.shape
assert torch.equal(out, ref)
def test_usp_merge_heads_unsupported_inputs_use_exact_fallback():
x = torch.randn(2, 4, 1, 4, 64, dtype=torch.bfloat16, device=DEVICE)
unsupported = [x.transpose(0, 1), x[:0]]
for value in unsupported:
assert not can_use_usp_merge_heads(value)
assert torch.equal(
usp_merge_heads(value), value.permute(2, 1, 0, 3, 4).contiguous()
)
with patch.object(torch.version, "hip", "6.3"):
assert not can_use_usp_merge_heads(x)
assert torch.equal(usp_merge_heads(x), x.permute(2, 1, 0, 3, 4).contiguous())
def test_usp_merge_heads_fast_path_rejects_wrong_rank():
x = torch.randn(2, 4, 1, 4, 64, dtype=torch.bfloat16, device=DEVICE)
assert not can_use_usp_merge_heads(x[0])
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py
-197
View File
@@ -1,197 +0,0 @@
"""Numerical correctness for fused varlen pack/scatter Triton kernels.
Bit-exact comparison against the equivalent PyTorch ops (index_select,
zeros + index_copy_) across bf16/fp16 and several shape/mask cases.
"""
import pytest
import torch
from sglang.kernels.jit.utils import get_ci_test_range
from sglang.kernels.ops.diffusion.triton.varlen_pack_pad import (
build_inv_indices,
fused_pack_qkv,
fused_scatter_to_padded,
)
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
register_cuda_ci(est_time=20, stage="nightly", runner_config="1-gpu-large")
register_amd_ci(est_time=15, suite="nightly-amd-kernel-1-gpu", nightly=True)
DEVICE = "cuda"
DTYPES = get_ci_test_range([torch.bfloat16, torch.float16], [torch.bfloat16])
# (bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens) tuples
SHAPES = get_ci_test_range(
[
# name, bs, s_txt, s_img, H, D, valid_txt_lens
("small_c2", 2, 64, 128, 4, 64, [32, 48]),
("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
("all_valid_b1", 1, 64, 128, 4, 64, [64]),
("all_valid_b4", 4, 64, 128, 4, 64, [64, 64, 64, 64]),
("c8_prod", 8, 256, 4096, 24, 128, [128, 200, 256, 100, 50, 256, 256, 50]),
# one batch with zero valid text tokens (image side still valid)
("zero_txt_one_batch", 2, 64, 128, 4, 64, [0, 32]),
# bs=1 with no text validity (only image rows packed)
("bs1_zero_txt", 1, 64, 128, 4, 64, [0]),
],
[
("small_c2", 2, 64, 128, 4, 64, [32, 48]),
("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
("all_valid_b4", 4, 64, 128, 4, 64, [64, 64, 64, 64]),
],
)
def _build_mask(bs, s_txt, s_img, valid_txt_lens):
s = s_txt + s_img
mask = torch.zeros(bs, s, dtype=torch.bool, device=DEVICE)
for b, vt in enumerate(valid_txt_lens):
mask[b, :vt] = True
mask[b, s_txt:] = True
return mask
def _ref_pack(q, k, v, indices):
bs, seq = q.shape[:2]
def flat(t):
return t.reshape(bs * seq, *t.shape[2:])
return (
flat(q).index_select(0, indices),
flat(k).index_select(0, indices),
flat(v).index_select(0, indices),
)
def _ref_scatter(out_unpad, indices, bs, seq):
_, num_heads, head_dim = out_unpad.shape
flat = torch.zeros(
bs * seq, num_heads, head_dim, dtype=out_unpad.dtype, device=DEVICE
)
flat.index_copy_(0, indices, out_unpad)
return flat.view(bs, seq, num_heads, head_dim)
def _build_meta(mask):
bs, seq = mask.shape
indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
inv_indices = build_inv_indices(indices, bs * seq)
return indices, inv_indices
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize(
"shape", SHAPES, ids=lambda s: s[0] if isinstance(s, tuple) else str(s)
)
def test_pack_matches_index_select(dtype, shape):
_, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
torch.manual_seed(0)
s = s_txt + s_img
mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
indices, _ = _build_meta(mask)
q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
q_ref, k_ref, v_ref = _ref_pack(q, k, v, indices)
q_fused, k_fused, v_fused = fused_pack_qkv(q, k, v, indices)
# bit-exact: pack is pure gather, no math
assert torch.equal(q_ref, q_fused)
assert torch.equal(k_ref, k_fused)
assert torch.equal(v_ref, v_fused)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize(
"shape", SHAPES, ids=lambda s: s[0] if isinstance(s, tuple) else str(s)
)
def test_scatter_matches_index_copy(dtype, shape):
_, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
torch.manual_seed(1)
s = s_txt + s_img
mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
indices, inv_indices = _build_meta(mask)
n_valid = indices.shape[0]
out_unpad = torch.randn(n_valid, num_heads, head_dim, dtype=dtype, device=DEVICE)
out_ref = _ref_scatter(out_unpad, indices, bs, s)
out_fused = fused_scatter_to_padded(out_unpad, inv_indices, bs, s)
# bit-exact: scatter is pure copy + zero-fill
assert torch.equal(out_ref, out_fused)
# Padding rows must be exactly zero
invalid = ~mask
if invalid.any():
assert out_fused[invalid].abs().max().item() == 0.0
def test_pack_handles_non_contiguous_input():
"""Helper must accept non-contiguous Q/K/V (auto .contiguous() inside)."""
torch.manual_seed(2)
bs, s_txt, s_img, num_heads, head_dim = 2, 64, 128, 4, 64
s = s_txt + s_img
mask = _build_mask(bs, s_txt, s_img, [32, 48])
indices, _ = _build_meta(mask)
# Build non-contiguous tensors via permute
qkv_pre = torch.randn(
bs, num_heads, s, head_dim, dtype=torch.bfloat16, device=DEVICE
)
q = qkv_pre.permute(0, 2, 1, 3)
k = torch.randn_like(qkv_pre).permute(0, 2, 1, 3)
v = torch.randn_like(qkv_pre).permute(0, 2, 1, 3)
assert not q.is_contiguous()
q_ref, k_ref, v_ref = _ref_pack(
q.contiguous(), k.contiguous(), v.contiguous(), indices
)
q_fused, k_fused, v_fused = fused_pack_qkv(q, k, v, indices)
assert torch.equal(q_ref, q_fused)
assert torch.equal(k_ref, k_fused)
assert torch.equal(v_ref, v_fused)
def test_build_inv_indices_matches_manual():
"""build_inv_indices output should match the manual full+scatter form."""
torch.manual_seed(3)
bs, s = 2, 32
mask = torch.bernoulli(torch.full((bs, s), 0.6, device=DEVICE)).to(torch.bool)
indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
n_valid = indices.shape[0]
manual = torch.full((bs * s,), -1, dtype=torch.int32, device=DEVICE)
if n_valid > 0:
manual[indices.long()] = torch.arange(n_valid, dtype=torch.int32, device=DEVICE)
built = build_inv_indices(indices, bs * s)
assert torch.equal(built, manual)
def test_empty_valid_set_handled():
"""All-False mask: pack returns empty tensors; scatter writes all zeros."""
bs, s, num_heads, head_dim = 2, 16, 4, 64
mask = torch.zeros(bs, s, dtype=torch.bool, device=DEVICE)
indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
inv_indices = build_inv_indices(indices, bs * s)
assert indices.numel() == 0
q = torch.randn(bs, s, num_heads, head_dim, dtype=torch.bfloat16, device=DEVICE)
q_unpad, k_unpad, v_unpad = fused_pack_qkv(q, q.clone(), q.clone(), indices)
assert q_unpad.shape == (0, num_heads, head_dim)
assert k_unpad.shape == (0, num_heads, head_dim)
assert v_unpad.shape == (0, num_heads, head_dim)
out_padded = fused_scatter_to_padded(q_unpad, inv_indices, bs, s)
assert out_padded.shape == (bs, s, num_heads, head_dim)
assert out_padded.abs().max().item() == 0.0
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_varlen_uspattn_equivalence.py
-157
View File
@@ -1,157 +0,0 @@
"""End-to-end equivalence between USPAttention varlen path and SDPA reference.
Compares the production varlen path (``build_varlen_mask_meta`` +
``fused_pack_qkv`` + ``flash_attn_varlen_func`` + ``fused_scatter_to_padded``)
against ``torch.nn.functional.scaled_dot_product_attention`` with a broadcast
key mask, for inputs the gating in ``USPAttention.forward`` would accept.
Verifies the documented contract:
* Valid (non-masked) query rows match SDPA within FA-vs-SDPA tolerance.
* Masked query rows are exactly zero in the varlen path (differs from
SDPA, which produces deterministic attention output at those rows).
"""
import pytest
import torch
import torch.nn.functional as F
from sglang.kernels.jit.utils import get_ci_test_range
from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func
from sglang.kernels.ops.diffusion.triton.varlen_pack_pad import (
fused_pack_qkv,
fused_scatter_to_padded,
)
from sglang.multimodal_gen.runtime.layers.attention.backends import (
flash_attn as _fa_backend,
)
from sglang.multimodal_gen.runtime.layers.attention.layer import (
build_varlen_mask_meta,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
register_cuda_ci(est_time=20, stage="nightly", runner_config="1-gpu-large")
DEVICE = "cuda"
DTYPES = get_ci_test_range([torch.bfloat16, torch.float16], [torch.bfloat16])
# (name, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens)
SHAPES = get_ci_test_range(
[
("small_c2", 2, 64, 128, 4, 64, [32, 48]),
("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
("all_valid_b1", 1, 64, 128, 4, 64, [64]),
("zero_txt_one_batch", 2, 64, 128, 4, 64, [0, 32]),
],
[
("small_c2", 2, 64, 128, 4, 64, [32, 48]),
],
)
def _build_mask(bs, s_txt, s_img, valid_txt_lens):
s = s_txt + s_img
mask = torch.zeros(bs, s, dtype=torch.bool, device=DEVICE)
for b, vt in enumerate(valid_txt_lens):
mask[b, :vt] = True
mask[b, s_txt:] = True
return mask
def _sdpa_with_key_mask(q, k, v, key_mask, softmax_scale):
"""Reference: SDPA with a ``[B, S]`` key mask broadcast to ``[B, 1, 1, S]``."""
q_ = q.transpose(1, 2)
k_ = k.transpose(1, 2)
v_ = v.transpose(1, 2)
mask = key_mask.to(dtype=q.dtype)[:, None, None, :]
mask = (mask - 1.0) * torch.finfo(q.dtype).max
out = F.scaled_dot_product_attention(
q_,
k_,
v_,
attn_mask=mask,
dropout_p=0.0,
is_causal=False,
scale=softmax_scale,
)
return out.transpose(1, 2)
def _varlen_path(q, k, v, key_mask, softmax_scale):
"""Production varlen path matching USPAttention.forward masked branch."""
bs, seq = q.shape[0], q.shape[1]
meta = build_varlen_mask_meta(key_mask)
indices = meta["indices"]
if indices.shape[0] == 0:
return torch.zeros_like(q)
q_unpad, k_unpad, v_unpad = fused_pack_qkv(q, k, v, indices)
out_unpad = flash_attn_varlen_func(
q=q_unpad,
k=k_unpad,
v=v_unpad,
cu_seqlens_q=meta["cu_seqlens"],
cu_seqlens_k=meta["cu_seqlens"],
max_seqlen_q=meta["max_seqlen"],
max_seqlen_k=meta["max_seqlen"],
softmax_scale=softmax_scale,
causal=False,
ver=_fa_backend.fa_ver,
)
return fused_scatter_to_padded(out_unpad, meta["inv_indices"], bs, seq)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize(
"shape", SHAPES, ids=lambda s: s[0] if isinstance(s, tuple) else str(s)
)
def test_varlen_path_matches_sdpa_on_valid_rows(dtype, shape):
"""Valid rows: varlen output ≈ SDPA output within FA tolerance."""
_, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
torch.manual_seed(0)
s = s_txt + s_img
softmax_scale = head_dim**-0.5
mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
out_sdpa = _sdpa_with_key_mask(q, k, v, mask, softmax_scale)
out_varlen = _varlen_path(q, k, v, mask, softmax_scale)
valid = mask[..., None, None].expand_as(out_sdpa)
rtol = 1e-2 if dtype == torch.bfloat16 else 5e-3
atol = 5e-2 if dtype == torch.bfloat16 else 1e-2
torch.testing.assert_close(
out_sdpa[valid],
out_varlen[valid],
rtol=rtol,
atol=atol,
)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize(
"shape", SHAPES, ids=lambda s: s[0] if isinstance(s, tuple) else str(s)
)
def test_varlen_path_zeros_masked_rows(dtype, shape):
"""Masked rows: varlen path produces exact zeros (documented contract)."""
_, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
torch.manual_seed(1)
s = s_txt + s_img
softmax_scale = head_dim**-0.5
mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
out_varlen = _varlen_path(q, k, v, mask, softmax_scale)
invalid = ~mask
if invalid.any():
assert (out_varlen[invalid] == 0).all(), "masked rows must be zero-filled"
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_wan_causal_cache.py
-164
View File
@@ -1,164 +0,0 @@
"""Wan causal VAE data-movement kernels: the fused conv-input builder and the
fused DupUp3D shortcut add must be bitwise identical to the aten op chains
they replace (they are pure data movement plus zero fill / one fp32 add)."""
import sys
import pytest
import torch
import torch.nn.functional as F
from sglang.kernels.ops.diffusion.triton.wan_causal_cache import (
cat_pad_channels_last_3d,
dup_up3d_add,
)
from sglang.multimodal_gen.runtime.models.vaes import wanvae
from sglang.multimodal_gen.runtime.models.vaes.wanvae import (
CACHE_T,
WanCausalConv3d,
_cache_payload,
_run_cached_causal_conv,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def _cl3d(shape, dtype):
return torch.randn(shape, device="cuda", dtype=dtype).contiguous(
memory_format=torch.channels_last_3d
)
def _ref_cat_pad(x, cache, padding):
p = list(padding)
if cache is not None:
x = torch.cat([cache, x], dim=2)
p[4] -= cache.shape[2]
if any(p):
x = F.pad(x, p)
return x.contiguous(memory_format=torch.channels_last_3d)
@torch.no_grad()
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
@pytest.mark.parametrize(
"c,t,h,w,cache_t,pads",
[
(96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)), # first chunk, zero-fill front
(96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)), # legacy 1-frame cache
(96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)), # steady state k3 conv
(64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)), # time_conv (temporal only)
(48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)), # encoder-style T=4 chunk
],
)
def test_cat_pad_bitwise(dtype, c, t, h, w, cache_t, pads) -> None:
torch.cuda.manual_seed(0)
x = _cl3d((1, c, t, h, w), dtype)
cache = None
if cache_t:
# Strided interior view: caches may arrive as non-contiguous slices.
ph, pw = pads[2], pads[0]
buf = _cl3d((1, c, cache_t, h + 2 * ph, w + 2 * pw), dtype)
cache = buf[:, :, :, ph : ph + h, pw : pw + w]
out = cat_pad_channels_last_3d(x, cache, pads)
ref = _ref_cat_pad(x, cache, pads)
assert out is not None and out.shape == ref.shape
assert out.is_contiguous(memory_format=torch.channels_last_3d)
assert torch.equal(out, ref)
# Dual-output mode: the same pass also emits the compact feature cache
# (unpadded interior of the last frames), bitwise equal to the slice.
pair = cat_pad_channels_last_3d(x, cache, pads, keep_cache_t=2)
assert pair is not None
out2, keep = pair
assert torch.equal(out2, ref)
ph, pw = pads[2], pads[0]
keep_t = min(2, ref.shape[2])
want = ref[:, :, ref.shape[2] - keep_t :, ph : ph + h, pw : pw + w]
assert keep.shape == want.shape
assert keep.is_contiguous(memory_format=torch.channels_last_3d)
assert torch.equal(keep, want)
@torch.no_grad()
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
@pytest.mark.parametrize(
"c_in,c_out,t,h,w,ft,fs,drop",
[
(128, 64, 1, 10, 14, 2, 2, False),
(128, 64, 1, 10, 14, 2, 2, True), # first_chunk slicing
(64, 32, 2, 10, 14, 1, 2, False),
],
)
def test_dup_up3d_add_bitwise(dtype, c_in, c_out, t, h, w, ft, fs, drop) -> None:
torch.cuda.manual_seed(0)
repeats = c_out * ft * fs * fs // c_in
src = _cl3d((1, c_in, t, h, w), dtype)
t_out = t * ft - (ft - 1 if drop else 0)
# Main arm as a permuted view, like the WanResample 2D output.
main = torch.randn(
(1, t_out, c_out, h * fs, w * fs), device="cuda", dtype=dtype
).permute(0, 2, 1, 3, 4)
dup = src.repeat_interleave(repeats, dim=1)
dup = dup.view(1, c_out, ft, fs, fs, t, h, w)
dup = dup.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
dup = dup.view(1, c_out, t * ft, h * fs, w * fs)
if drop:
dup = dup[:, :, ft - 1 :, :, :]
ref = main + dup
out = dup_up3d_add(main, src, ft, fs, repeats, drop)
assert out is not None and out.shape == ref.shape
# Layout must match the aten add output exactly (downstream reductions
# are layout-sensitive), and every value must be bitwise identical.
assert out.stride() == ref.stride()
assert torch.equal(out, ref)
@torch.no_grad()
@pytest.mark.parametrize("pads_temporal_only", [False, True])
def test_cached_conv_chunk_loop_bitwise(pads_temporal_only) -> None:
"""The fused conv-input/compact-cache scheme must reproduce the original
clone/cat bookkeeping bitwise across a chunked decode, including the
first-chunk zero fill and the "Rep" marker start used by WanResample."""
torch.cuda.manual_seed(0)
c = 64
if pads_temporal_only:
conv = WanCausalConv3d(c, 2 * c, (3, 1, 1), padding=(1, 0, 0))
else:
conv = WanCausalConv3d(c, c, 3, padding=1)
conv = conv.to(device="cuda", dtype=torch.float32)
conv.weight.data = conv.weight.data.contiguous(memory_format=torch.channels_last_3d)
chunks = [_cl3d((1, c, 1, 10, 14), torch.float32) for _ in range(4)]
def run(force_fallback, start):
cache = [start]
outs = []
if force_fallback:
orig = wanvae.cat_pad_channels_last_3d
wanvae.cat_pad_channels_last_3d = None
try:
for x in chunks:
outs.append(_run_cached_causal_conv(conv, x, cache, 0))
finally:
if force_fallback:
wanvae.cat_pad_channels_last_3d = orig
return outs, cache[0]
for start in (None, "Rep"):
fused_outs, fused_cache = run(False, start)
ref_outs, ref_cache = run(True, start)
for got, want in zip(fused_outs, ref_outs, strict=True):
assert torch.equal(got, want)
got_payload = _cache_payload(fused_cache)
assert got_payload is not None and got_payload.shape[2] == CACHE_T
# Reference cache holds the last CACHE_T unpadded frames.
assert torch.equal(got_payload, ref_cache[:, :, -CACHE_T:])
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_wan_vae_fastpath.py
-79
View File
@@ -1,79 +0,0 @@
"""Wan VAE decoder fast path: fused-kernel numerics and gate dispatch
(the lossless off-path must stay bit-exact)."""
import sys
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from sglang.kernels.ops.diffusion.triton.wan_rmsnorm_silu import wan_rmsnorm_silu
from sglang.multimodal_gen.runtime.models.vaes.wan_vae_cuda_opt import (
FusedWanRMSNormSiLU,
VaeFastPathGate,
)
from sglang.multimodal_gen.runtime.models.vaes.wanvae import WanRMS_norm
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def _cl3d(shape, dtype):
return torch.randn(shape, device="cuda", dtype=dtype).contiguous(
memory_format=torch.channels_last_3d
)
@torch.no_grad()
@pytest.mark.parametrize(
"x_dtype,affine_dtype,atol,rtol",
[
(torch.float32, torch.float32, 1e-5, 1e-5), # FastWan2.2 fp32 decode
(torch.bfloat16, torch.float32, 1.5e-1, 3e-2), # Wan2.1 bf16 autocast
],
)
def test_kernel_numerics(x_dtype, affine_dtype, atol, rtol) -> None:
torch.cuda.manual_seed(0)
x = _cl3d((1, 96, 3, 10, 14), x_dtype)
gamma = torch.randn((96, 1, 1, 1), device="cuda", dtype=affine_dtype)
for bias in (None, torch.randn_like(gamma)):
expected = F.silu(
F.normalize(x, dim=1) * 96**0.5 * gamma + (0 if bias is None else bias)
)
actual = wan_rmsnorm_silu(x, gamma, bias)
assert actual is not None and actual.dtype == expected.dtype
assert actual.stride() == x.stride()
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
@torch.no_grad()
def test_fused_module_gate_dispatch() -> None:
# Gate off must stay bit-exact; gate on must route to the fused kernel.
torch.cuda.manual_seed(0)
norm = WanRMS_norm(96, images=False).to(device="cuda", dtype=torch.bfloat16)
norm.gamma.add_(torch.randn_like(norm.gamma))
gate = VaeFastPathGate()
fused = FusedWanRMSNormSiLU(norm, gate)
# Parameter names must not change (weight transfer matches by name).
assert [n for n, _ in fused.named_parameters()] == ["gamma"]
x = _cl3d((1, 96, 3, 10, 14), torch.bfloat16)
assert torch.equal(fused(x), nn.SiLU()(norm(x)))
gate.enabled = True
expected = wan_rmsnorm_silu(x, norm.gamma, rms_scale=float(norm.scale))
assert torch.equal(fused(x), expected)
@torch.no_grad()
def test_kernel_rejects_empty_input() -> None:
x = torch.empty(1, 96, 0, 2, 2, device="cuda", dtype=torch.bfloat16).to(
memory_format=torch.channels_last_3d
)
gamma = torch.ones(96, 1, 1, 1, device="cuda", dtype=torch.bfloat16)
assert wan_rmsnorm_silu(x, gamma) is None
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))