diff --git a/python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py b/python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py index a786dba6e..495b1fce0 100644 --- a/python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py +++ b/python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py @@ -201,8 +201,6 @@ class ScaleResidualNormScaleShift: value = tXrX.load() if cutlass.const_expr(isinstance(tGrG, cute.Tensor)): value = tGrG.load() * value - elif cutlass.const_expr(isinstance(tGrG, cutlass.Int32)): - value = tGrG * value # Compute: value = value + if cutlass.const_expr(isinstance(tRrR, cute.Tensor)): value = value + tRrR.load() @@ -276,7 +274,7 @@ def validate_gate(t: Union[torch.Tensor, int], B: int, S: int, D: int): validate_scale_shift(t, B, S, D) -@torch._dynamo.disable # Disable Dynamo tracing +@torch.library.custom_op("sglang::fused_norm_scale_shift", mutates_args=()) def fused_norm_scale_shift( x: torch.Tensor, weight: Optional[torch.Tensor], @@ -285,8 +283,7 @@ def fused_norm_scale_shift( shift: torch.Tensor, norm_type: str, eps: float = 1e-5, - stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT), -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> torch.Tensor: """ Fuse: norm(x) * (1 + scale) + shift where norm is either layernorm or rmsnorm. @@ -301,6 +298,7 @@ def fused_norm_scale_shift( D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per thread and avoid predicated loads (e.g., bounds checks such as `index < D`). """ + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) # Tensor Validation BSD = x.shape validate_x(x, *BSD) @@ -319,14 +317,12 @@ def fused_norm_scale_shift( scale = broadcast_tensor_for_bsfd(scale, *x.shape) # handle various shapes shift = broadcast_tensor_for_bsfd(shift, *x.shape) # handle various shapes # Use scalar placeholders for None tensors as a workaround, since the CuTe DSL - # TVM-FFI backend does not support None parameters. Unless explicitly handled - # (e.g., for gate), scalar values do not result in code generation and have no - # impact on runtime performance. + # TVM-FFI backend does not support None parameters. scalar values do not result + # in code generation and have no impact on runtime performance. weight = 1 if weight is None else weight bias = 0 if bias is None else bias ResOut, Residual, Gate = 0, 0, 1 torch_tensors = [y, ResOut, Residual, x, Gate, weight, bias, scale, shift] - cute_tensor_args = [to_cute_arg(t) for t in torch_tensors] # Compile cache hash_key = ScaleResidualNormScaleShift.make_hash_key(norm_type, *torch_tensors) compiled_fn = _COMPILE_CACHE.get(hash_key) @@ -338,24 +334,31 @@ def fused_norm_scale_shift( ) _COMPILE_CACHE[hash_key] = compiled_fn # Execute - compiled_fn(*cute_tensor_args, eps, stream) + compiled_fn(*torch_tensors, eps, stream) return y else: raise ValueError(f'norm_type must be one of "layer" and "rms"') -@torch._dynamo.disable # Disable Dynamo tracing +@fused_norm_scale_shift.register_fake +def _fused_norm_scale_shift_fake(x, weight, bias, scale, shift, norm_type, eps): + y = x.new_empty(x.shape) + return y + + +@torch.library.custom_op( + "sglang::fused_scale_residual_norm_scale_shift", mutates_args=() +) def fused_scale_residual_norm_scale_shift( residual: torch.Tensor, x: torch.Tensor, - gate: Union[Optional[torch.Tensor], int], + gate: Optional[torch.Tensor], # Union[Optional[torch.Tensor], int] indeed weight: Optional[torch.Tensor], bias: Optional[torch.Tensor], scale: torch.Tensor, shift: torch.Tensor, norm_type: str, eps: float = 1e-5, - stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT), ) -> Tuple[torch.Tensor, torch.Tensor]: """ Fuse: norm(residual + gate * x) * (1 + scale) + shift @@ -363,7 +366,7 @@ def fused_scale_residual_norm_scale_shift( Expects: - residual, x: [B, S, D] - - gate: None, 1, [1], [D], [1/B, D], [1/B, 1/S, D] or [B, F, 1, D] + - gate: None, [1], [D], [1/B, D], [1/B, 1/S, D] or [B, F, 1, D] - weight/bias: None, [D] - scale/shift: [1], [D], [1/B, D], [1/B, 1/S, D] or [B, F, 1, D] - norm_type: str, "layer" or "rms" @@ -381,8 +384,10 @@ def fused_scale_residual_norm_scale_shift( validate_weight_bias(bias, *BSD) validate_scale_shift(scale, *BSD) validate_scale_shift(shift, *BSD) - if norm_type == "layer" or norm_type == "rms": + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + # if norm_type == "layer" or norm_type == "rms": D = x.shape[-1] if D % 256 != 0 or D > 8192: raise ValueError( @@ -394,14 +399,12 @@ def fused_scale_residual_norm_scale_shift( scale = broadcast_tensor_for_bsfd(scale, *x.shape) # handle various shapes shift = broadcast_tensor_for_bsfd(shift, *x.shape) # handle various shapes # Use scalar placeholders for None tensors as a workaround, since the CuTe DSL - # TVM-FFI backend does not support None parameters. Unless explicitly handled - # (e.g., for gate), scalar values do not result in code generation and have no - # impact on runtime performance. + # TVM-FFI backend does not support None parameters. scalar values do not result + # in code generation and have no impact on runtime performance. gate = 1 if gate is None else gate weight = 1 if weight is None else weight bias = 0 if bias is None else bias torch_tensors = [y, resi_out, residual, x, gate, weight, bias, scale, shift] - cute_tensor_args = [to_cute_arg(t) for t in torch_tensors] # Compile cache hash_key = ScaleResidualNormScaleShift.make_hash_key(norm_type, *torch_tensors) compiled_fn = _COMPILE_CACHE.get(hash_key) @@ -413,7 +416,16 @@ def fused_scale_residual_norm_scale_shift( ) _COMPILE_CACHE[hash_key] = compiled_fn # Execute - compiled_fn(*cute_tensor_args, eps, stream) + compiled_fn(*torch_tensors, eps, stream) return y, resi_out else: raise ValueError(f'norm_type must be one of "layer" and "rms"') + + +@fused_scale_residual_norm_scale_shift.register_fake +def _fused_scale_residual_norm_scale_shift_fake( + residual, x, gate, weight, bias, scale, shift, norm_type, eps +): + y = x.new_empty(x.shape) + residual_out = x.new_empty(x.shape) + return y, residual_out diff --git a/python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py b/python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py index 73a815f5e..443edf3cc 100644 --- a/python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py +++ b/python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py @@ -115,8 +115,6 @@ def fused_scale_residual_norm_scale_shift_ref( def _make_tensor(index_mode: str, shape: Tuple, dtype: torch.dtype): - if index_mode == "int1": - return 1 if index_mode == "NAT": return None return torch.randn(*SHAPE_MAP[index_mode](*shape), device=DEVICE, dtype=dtype) @@ -229,7 +227,7 @@ class TestFusedScaleResidualNormScaleShift: scale_mode=index_mode, shift_mode=index_mode, norm_type=norm_type ) - @pytest.mark.parametrize("index_mode", INDEX_MODES + ["int1"]) + @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) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 6998bffdd..f855795f2 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -319,6 +319,11 @@ class _ScaleResidualNormScaleShift(CustomOp): fused_scale_residual_norm_scale_shift, ) + if isinstance(gate, int) and gate != 1: + raise ValueError( + f"Only gate value of 1 is supported for int type, but got {gate}" + ) + return fused_scale_residual_norm_scale_shift( residual.contiguous(), x.contiguous(),