diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 6faa255f28..ea70f58a21 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -287,6 +287,13 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = ( _CUDA, "Paired in-place Helios transposed Q/K RoPE.", ), + ( + "diffusion.complex_rope", + KernelBackend.TRITON, + "rope.complex_rope_triton:fused_complex_rope", + _CUDA, + "Paired RoPE preserving PyTorch complex64 multiplication rounding.", + ), ( "diffusion.hunyuan_qkv_rope_pack", KernelBackend.TRITON, @@ -294,6 +301,13 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = ( _CUDA, "HunyuanVideo QKV pack + RoPE.", ), + ( + "diffusion.rmsnorm_preserve_reduction", + KernelBackend.TRITON, + "norm.rmsnorm_preserve_reduction:rmsnorm_preserve_reduction", + _CUDA, + "Cast-before-weight RMSNorm preserving the native FP32 mean reduction.", + ), ( "diffusion.silu_mul", KernelBackend.TRITON, @@ -543,6 +557,8 @@ _EXPORTS: dict[str, str] = { "try_fused_bias_mul_add": "sglang.kernels.kda_kernels.norm_scale_shift_jit", "try_fused_bias_scale_residual_norm_scale_shift": "sglang.kernels.kda_kernels.norm_scale_shift_jit", "triton_one_pass_rms_norm": "norm.rmsnorm_onepass_triton", + "can_use_rmsnorm_preserve_reduction": "norm.rmsnorm_preserve_reduction", + "rmsnorm_preserve_reduction": "norm.rmsnorm_preserve_reduction", "can_use_fused_rmsnorm_scale_shift": "norm.rmsnorm_scale_shift_bitexact", "can_use_fused_scale_residual_rmsnorm_scale_shift": "norm.rmsnorm_scale_shift_bitexact", "fused_rmsnorm_scale_shift_bitexact": "norm.rmsnorm_scale_shift_bitexact", @@ -598,6 +614,8 @@ _EXPORTS: dict[str, str] = { "can_use_helios_qk_rope": "rope.helios_qk_rope_jit", "fused_inplace_helios_qk_rope": "rope.helios_qk_rope_jit", "apply_rotary_embedding": "rope.rotary_triton", + "can_use_fused_complex_rope": "rope.complex_rope_triton", + "fused_complex_rope": "rope.complex_rope_triton", # Tensor layout transformations fused with downstream quantization "try_flux2_token_cat_fp8": "sglang.kernels.kda_kernels.flux2_token_cat_fp8_triton", # Activation-function fusions