Fix diffusion fallback guards and validation (#23335)

python/sglang/jit_kernel/tests/diffusion/test_fused_norm_scale_shift.py

@@ -9,6 +9,7 @@ from torch import Tensor
 from sglang.jit_kernel.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
 
@@ -125,6 +126,16 @@ def _make_tensor(index_mode: str, shape: Tuple, dtype: torch.dtype):
     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],