From a26587dd4e0ea6aff9da6fa015e690b926e91dbb Mon Sep 17 00:00:00 2001 From: YC Yen-Ching Tseng Date: Mon, 8 Jun 2026 17:42:39 +0800 Subject: [PATCH] [AMD][diffusion] Add FlyDSL fused normalization kernels for ROCm diffusion models optimization (#22786) --- .../diffusion/flydsl/fused_residual_norm.py | 902 ++++++++++++++++++ .../tests/diffusion/test_flydsl_fused_norm.py | 133 +++ .../runtime/layers/layernorm.py | 72 +- 3 files changed, 1099 insertions(+), 8 deletions(-) create mode 100644 python/sglang/jit_kernel/diffusion/flydsl/fused_residual_norm.py create mode 100644 python/sglang/jit_kernel/tests/diffusion/test_flydsl_fused_norm.py diff --git a/python/sglang/jit_kernel/diffusion/flydsl/fused_residual_norm.py b/python/sglang/jit_kernel/diffusion/flydsl/fused_residual_norm.py new file mode 100644 index 000000000..f7c666d14 --- /dev/null +++ b/python/sglang/jit_kernel/diffusion/flydsl/fused_residual_norm.py @@ -0,0 +1,902 @@ +"""FlyDSL fused normalization kernels for AMD ROCm (gfx950). + +Provides two fused kernels: + - flydsl_fused_residual_norm_scale_shift: + residual_add + gate_mul + RMSNorm/LayerNorm + scale·shift + - flydsl_norm_scale_shift: + RMSNorm/LayerNorm + scale·shift + +Both kernels use register-cache optimization: Phase 2 (scale·shift) +reuses f32 intermediate values from Phase 1 (norm) registers instead +of re-reading from HBM, saving ~20% bandwidth. +""" + +from typing import Optional, Tuple + +import flydsl.compiler as flyc +import flydsl.expr as fx +import torch +from flydsl._mlir import ir +from flydsl._mlir.dialects import arith as arith_ops +from flydsl._mlir.dialects import gpu as _gpu +from flydsl._mlir.dialects import math as math_ops +from flydsl._mlir.dialects import memref as _memref +from flydsl._mlir.dialects import ( + scf, +) +from flydsl._mlir.dialects import vector as _vector +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, const_expr, range_constexpr +from flydsl.expr.arith import ArithValue, CmpIPredicate +from flydsl.expr.typing import Int32, T + +WARP_SIZE = 64 +_VEC = 8 +_NUM_WAVES = 10 +FLYDSL_NORM_MIN_ALIGNED_DIM = WARP_SIZE * _NUM_WAVES * _VEC # 5120 + + +def _v(x): + return buffer_ops._unwrap_value(x) + + +def _build_fused_norm_module(D: int, is_rms: bool, has_gate: bool, has_weight: bool): + VEC = _VEC + NUM_WAVES = _NUM_WAVES + BLOCK = NUM_WAVES * WARP_SIZE + assert ( + D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0 + ), f"FlyDSL fused_residual_norm requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}" + NUM_ITERS = D // (BLOCK * VEC) + + @flyc.kernel(known_block_size=[BLOCK, 1, 1]) + def flydsl_fused_residual_norm_ss_kernel( + y_ptr: fx.Tensor, + res_out_ptr: fx.Tensor, + res_ptr: fx.Tensor, + x_ptr: fx.Tensor, + gate_ptr: fx.Tensor, + weight_ptr: fx.Tensor, + bias_ptr: fx.Tensor, + scale_ptr: fx.Tensor, + shift_ptr: fx.Tensor, + total_rows: Int32, + gate_stride: Int32, + scale_stride: Int32, + shift_stride: Int32, + ): + row = fx.block_idx.x + tid = fx.thread_idx.x + + i32 = T.i32 + f32 = T.f32 + bf16 = T.bf16 + vec_f32_t = ir.VectorType.get([VEC], f32) + vec_bf16_t = ir.VectorType.get([VEC], bf16) + + y_rsrc = buffer_ops.create_buffer_resource(y_ptr, max_size=True) + ro_rsrc = buffer_ops.create_buffer_resource(res_out_ptr, max_size=True) + r_rsrc = buffer_ops.create_buffer_resource(res_ptr, max_size=True) + x_rsrc = buffer_ops.create_buffer_resource(x_ptr, max_size=True) + g_rsrc = buffer_ops.create_buffer_resource(gate_ptr, max_size=True) + w_rsrc = buffer_ops.create_buffer_resource(weight_ptr, max_size=True) + b_rsrc = buffer_ops.create_buffer_resource(bias_ptr, max_size=True) + sc_rsrc = buffer_ops.create_buffer_resource(scale_ptr, max_size=True) + sh_rsrc = buffer_ops.create_buffer_resource(shift_ptr, max_size=True) + + row_i32 = ArithValue(row) + tid_i32 = ArithValue(tid) + D_i32 = arith.constant(D, type=i32) + row_off = row_i32 * D_i32 + gate_row_off = row_i32 * ArithValue(gate_stride) + scale_row_off = row_i32 * ArithValue(scale_stride) + shift_row_off = row_i32 * ArithValue(shift_stride) + + c_zero_f32 = arith.constant(0.0, type=f32) + c_one_f32 = arith.constant(1.0, type=f32) + eps_val = arith.constant(1e-6, type=f32) + D_float = arith.constant(float(D), type=f32) + + # LDS + LDS_SLOTS = NUM_WAVES * 2 + 2 + ws_attr = ir.Attribute.parse("#gpu.address_space") + lds_i8_type = ir.MemRefType.get( + [ir.ShapedType.get_dynamic_size()], T.i8, memory_space=ws_attr + ) + lds_f32_type = ir.MemRefType.get([LDS_SLOTS], f32, memory_space=ws_attr) + lds_i8 = _gpu.DynamicSharedMemoryOp(lds_i8_type).result + byte_zero = arith_ops.ConstantOp( + ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), 0) + ).result + lds = _memref.ViewOp(lds_f32_type, lds_i8, byte_zero, []).result + + lane_id = tid_i32 % arith.constant(WARP_SIZE, type=i32) + wave_id = tid_i32 // arith.constant(WARP_SIZE, type=i32) + + # Phase 1: residual + gate*x, accumulate stats, save f32 in registers + _saved_ro_f32 = [] + partial_sum = _v(c_zero_f32) + partial_sum_sq = _v(c_zero_f32) + + for it in range_constexpr(NUM_ITERS): + col = tid_i32 * arith.constant(VEC, type=i32) + arith.constant( + it * BLOCK * VEC, type=i32 + ) + off = row_off + col + + r_vec = buffer_ops.buffer_load(r_rsrc, off, vec_width=VEC, dtype=bf16) + x_vec = buffer_ops.buffer_load(x_rsrc, off, vec_width=VEC, dtype=bf16) + r_f32 = arith_ops.ExtFOp(vec_f32_t, _v(r_vec)).result + x_f32 = arith_ops.ExtFOp(vec_f32_t, _v(x_vec)).result + + if const_expr(has_gate): + g_off = gate_row_off + col + g_vec = buffer_ops.buffer_load(g_rsrc, g_off, vec_width=VEC, dtype=bf16) + g_f32 = arith_ops.ExtFOp(vec_f32_t, _v(g_vec)).result + gx = arith_ops.MulFOp(g_f32, x_f32).result + ro_f32 = arith_ops.AddFOp(r_f32, gx).result + else: + ro_f32 = arith_ops.AddFOp(r_f32, x_f32).result + + ro_bf16 = arith_ops.TruncFOp(vec_bf16_t, ro_f32).result + buffer_ops.buffer_store(ro_bf16, ro_rsrc, off) + _saved_ro_f32.append(ro_f32) + + if const_expr(not is_rms): + v_sum = _vector.ReductionOp( + f32, _vector.CombiningKind.ADD, ro_f32 + ).result + partial_sum = arith_ops.AddFOp(partial_sum, v_sum).result + ro_sq = arith_ops.MulFOp(ro_f32, ro_f32).result + v_sum_sq = _vector.ReductionOp(f32, _vector.CombiningKind.ADD, ro_sq).result + partial_sum_sq = arith_ops.AddFOp(partial_sum_sq, v_sum_sq).result + + # Intra-wave shuffle + width_c = _v(arith.constant(WARP_SIZE, type=i32)) + w_sum = partial_sum + w_sq = partial_sum_sq + for sh in [32, 16, 8, 4, 2, 1]: + off_sh = _v(arith.constant(sh, type=i32)) + if const_expr(not is_rms): + peer_sum = _gpu.ShuffleOp( + w_sum, off_sh, width_c, mode=_gpu.ShuffleMode.XOR + ).shuffleResult + w_sum = arith_ops.AddFOp(w_sum, peer_sum).result + peer_sq = _gpu.ShuffleOp( + w_sq, off_sh, width_c, mode=_gpu.ShuffleMode.XOR + ).shuffleResult + w_sq = arith_ops.AddFOp(w_sq, peer_sq).result + + # Cross-wave LDS + lane_0 = arith.cmpi(CmpIPredicate.eq, lane_id, arith.constant(0, type=i32)) + wave_idx = arith_ops.IndexCastOp(ir.IndexType.get(), _v(wave_id)).result + + _if_lane0 = scf.IfOp(lane_0) + with ir.InsertionPoint(_if_lane0.then_block): + if const_expr(not is_rms): + _memref.StoreOp(w_sum, lds, [wave_idx]) + sq_slot = arith_ops.AddIOp( + wave_idx, + arith_ops.ConstantOp( + ir.IndexType.get(), + ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES), + ).result, + ).result + _memref.StoreOp(w_sq, lds, [sq_slot]) + scf.YieldOp([]) + _gpu.BarrierOp() + + wave_0 = arith.cmpi(CmpIPredicate.eq, wave_id, arith.constant(0, type=i32)) + active = arith.andi( + wave_0, + arith.cmpi(CmpIPredicate.ult, lane_id, arith.constant(NUM_WAVES, type=i32)), + ) + lane_idx = arith_ops.IndexCastOp(ir.IndexType.get(), _v(lane_id)).result + lane_idx_sq = arith_ops.AddIOp( + lane_idx, + arith_ops.ConstantOp( + ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES) + ).result, + ).result + + if const_expr(is_rms): + _if_active = scf.IfOp(active, [f32], has_else=True) + with ir.InsertionPoint(_if_active.then_block): + sq_val = _memref.LoadOp(lds, [lane_idx_sq]).result + scf.YieldOp([sq_val]) + with ir.InsertionPoint(_if_active.else_block): + scf.YieldOp([_v(c_zero_f32)]) + loaded_sq = _if_active.results[0] + loaded_sum = _v(c_zero_f32) + else: + _if_active = scf.IfOp(active, [f32, f32], has_else=True) + with ir.InsertionPoint(_if_active.then_block): + s_val = _memref.LoadOp(lds, [lane_idx]).result + sq_val = _memref.LoadOp(lds, [lane_idx_sq]).result + scf.YieldOp([s_val, sq_val]) + with ir.InsertionPoint(_if_active.else_block): + scf.YieldOp([_v(c_zero_f32), _v(c_zero_f32)]) + loaded_sum = _if_active.results[0] + loaded_sq = _if_active.results[1] + + final_sum = loaded_sum + final_sq = loaded_sq + for sh in [32, 16, 8, 4, 2, 1]: + off_sh = _v(arith.constant(sh, type=i32)) + if const_expr(not is_rms): + ps = _gpu.ShuffleOp( + final_sum, off_sh, width_c, mode=_gpu.ShuffleMode.XOR + ).shuffleResult + final_sum = arith_ops.AddFOp(final_sum, ps).result + pq = _gpu.ShuffleOp( + final_sq, off_sh, width_c, mode=_gpu.ShuffleMode.XOR + ).shuffleResult + final_sq = arith_ops.AddFOp(final_sq, pq).result + + both_0 = arith.andi(wave_0, lane_0) + final_sum_slot = arith_ops.ConstantOp( + ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES * 2) + ).result + final_sq_slot = arith_ops.ConstantOp( + ir.IndexType.get(), + ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES * 2 + 1), + ).result + _if_both = scf.IfOp(both_0) + with ir.InsertionPoint(_if_both.then_block): + if const_expr(not is_rms): + _memref.StoreOp(final_sum, lds, [final_sum_slot]) + _memref.StoreOp(final_sq, lds, [final_sq_slot]) + scf.YieldOp([]) + _gpu.BarrierOp() + + if const_expr(not is_rms): + total_sum = _memref.LoadOp(lds, [final_sum_slot]).result + else: + total_sum = _v(c_zero_f32) + total_sq = _memref.LoadOp(lds, [final_sq_slot]).result + + # Norm + d_f = _v(D_float) + eps_v = _v(eps_val) + if const_expr(is_rms): + var = arith_ops.DivFOp(total_sq, d_f).result + var_eps = arith_ops.AddFOp(var, eps_v).result + rstd = math_ops.RsqrtOp(var_eps).result + mean = _v(c_zero_f32) + else: + mean = arith_ops.DivFOp(total_sum, d_f).result + mean_sq = arith_ops.MulFOp(mean, mean).result + var = arith_ops.SubFOp( + arith_ops.DivFOp(total_sq, d_f).result, mean_sq + ).result + var_eps = arith_ops.AddFOp(var, eps_v).result + rstd = math_ops.RsqrtOp(var_eps).result + + # Phase 2: normalize using register-cached f32 values (no HBM re-read) + mean_splat = _vector.BroadcastOp(vec_f32_t, mean).result + rstd_splat = _vector.BroadcastOp(vec_f32_t, rstd).result + one_splat = _vector.BroadcastOp(vec_f32_t, _v(c_one_f32)).result + + for it in range_constexpr(NUM_ITERS): + col = tid_i32 * arith.constant(VEC, type=i32) + arith.constant( + it * BLOCK * VEC, type=i32 + ) + off = row_off + col + + ro_f32 = _saved_ro_f32[it] + + if const_expr(is_rms): + x_hat = arith_ops.MulFOp(ro_f32, rstd_splat).result + else: + centered = arith_ops.SubFOp(ro_f32, mean_splat).result + x_hat = arith_ops.MulFOp(centered, rstd_splat).result + + if const_expr(has_weight): + w_vec = buffer_ops.buffer_load(w_rsrc, col, vec_width=VEC, dtype=bf16) + w_f32 = arith_ops.ExtFOp(vec_f32_t, _v(w_vec)).result + x_hat = arith_ops.MulFOp(x_hat, w_f32).result + b_vec = buffer_ops.buffer_load(b_rsrc, col, vec_width=VEC, dtype=bf16) + b_f32 = arith_ops.ExtFOp(vec_f32_t, _v(b_vec)).result + x_hat = arith_ops.AddFOp(x_hat, b_f32).result + + sc_off = scale_row_off + col + sc_vec = buffer_ops.buffer_load(sc_rsrc, sc_off, vec_width=VEC, dtype=bf16) + sc_f32 = arith_ops.ExtFOp(vec_f32_t, _v(sc_vec)).result + sc_p1 = arith_ops.AddFOp(one_splat, sc_f32).result + x_hat = arith_ops.MulFOp(x_hat, sc_p1).result + + sh_off = shift_row_off + col + sh_vec = buffer_ops.buffer_load(sh_rsrc, sh_off, vec_width=VEC, dtype=bf16) + sh_f32 = arith_ops.ExtFOp(vec_f32_t, _v(sh_vec)).result + y_f32 = arith_ops.AddFOp(x_hat, sh_f32).result + + y_bf16 = arith_ops.TruncFOp(vec_bf16_t, y_f32).result + buffer_ops.buffer_store(y_bf16, y_rsrc, off) + + @flyc.jit + def launch_fused_norm( + y: fx.Tensor, + res_out: fx.Tensor, + res: fx.Tensor, + x: fx.Tensor, + gate: fx.Tensor, + weight: fx.Tensor, + bias: fx.Tensor, + scale: fx.Tensor, + shift: fx.Tensor, + total_rows: fx.Int32, + gate_stride: fx.Int32, + scale_stride: fx.Int32, + shift_stride: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + pass + grid_x = arith.index_cast(T.index, total_rows) + launcher = flydsl_fused_residual_norm_ss_kernel( + y, + res_out, + res, + x, + gate, + weight, + bias, + scale, + shift, + total_rows, + gate_stride, + scale_stride, + shift_stride, + ) + LDS_BYTES = (NUM_WAVES * 2 + 2) * 4 + launcher.launch( + grid=(grid_x, 1, 1), block=(BLOCK, 1, 1), smem=LDS_BYTES, stream=stream + ) + + return launch_fused_norm + + +_COMPILE_CACHE = {} + + +def _get_or_compile(D, is_rms, has_gate, has_weight, args): + key = (D, is_rms, has_gate, has_weight) + if key not in _COMPILE_CACHE: + launcher = _build_fused_norm_module(D, is_rms, has_gate, has_weight) + cf = flyc.compile(launcher, *args) + _COMPILE_CACHE[key] = cf + return _COMPILE_CACHE[key] + + +def _to_bf16(t): + """Convert to bf16 only if not already bf16.""" + return t if t.dtype == torch.bfloat16 else t.to(torch.bfloat16) + + +def _prep_slices(t, B, L, C): + """Prepare per-batch tensor slices and kernel row_stride. + + Returns (slices, row_stride) where: + slices[b] = tensor to pass to kernel for batch b + row_stride = 0 (broadcast: all rows share one row) or C (per-row data) + """ + t = _to_bf16(t) + + if t.numel() < C: + row = t.flatten()[0].expand(C).contiguous().unsqueeze(0) + return [row] * B, 0 + + if t.dim() == 1: + return [t.unsqueeze(0).contiguous()] * B, 0 + + if t.dim() == 2: + if t.shape[0] == 1: + return [t.contiguous()] * B, 0 + return [t.contiguous()] * B, C + + if t.dim() == 3: + if t.shape[0] == 1 and t.shape[1] == 1: + return [t.reshape(1, C).contiguous()] * B, 0 + if t.shape[1] == 1: + t_c = t.contiguous() + return [t_c[b] for b in range(B)], 0 + t_exp = t.expand(B, L, C).contiguous() + return [t_exp[b] for b in range(B)], C + + if t.dim() == 4: + nf = t.shape[1] + fs = L // nf + t_exp = t.expand(B, nf, fs, C).reshape(B, L, C).contiguous() + return [t_exp[b] for b in range(B)], C + + t_exp = t.reshape(B, L, C).contiguous() + return [t_exp[b] for b in range(B)], C + + +def _ensure_bf16_contig(t): + """Return bf16-contiguous view, avoiding copies when possible.""" + if t.dtype == torch.bfloat16 and t.is_contiguous(): + return t + return _to_bf16(t).contiguous() + + +@torch.library.custom_op( + "sglang::flydsl_fused_residual_norm_scale_shift", mutates_args=() +) +def flydsl_fused_residual_norm_scale_shift( + residual: torch.Tensor, + x: torch.Tensor, + gate: Optional[torch.Tensor], + weight: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + scale: torch.Tensor, + shift: torch.Tensor, + norm_type: str, + eps: float = 1e-6, +) -> Tuple[torch.Tensor, torch.Tensor]: + B, L, C = x.shape + rows = B * L + bf16 = torch.bfloat16 + + x_2d = _ensure_bf16_contig(x).reshape(rows, C) + res_2d = _ensure_bf16_contig(residual).reshape(rows, C) + y = torch.empty_like(x_2d) + res_out = torch.empty_like(x_2d) + + has_gate = gate is not None + if has_gate: + g_slices, g_stride = _prep_slices(gate, B, L, C) + else: + g_slices, g_stride = [x_2d[:1]] * B, 0 + + has_weight = weight is not None + weight_c = ( + _ensure_bf16_contig(weight) + if has_weight + else torch.empty(C, device=x.device, dtype=bf16) + ) + bias_c = ( + _ensure_bf16_contig(bias) + if bias is not None + else torch.zeros(C, device=x.device, dtype=bf16) + ) + + sc_slices, sc_stride = _prep_slices(scale, B, L, C) + sh_slices, sh_stride = _prep_slices(shift, B, L, C) + + is_rms = norm_type == "rms" + stream = torch.cuda.current_stream() + + dummy_args = ( + y[:L], + res_out[:L], + res_2d[:L], + x_2d[:L], + g_slices[0], + weight_c, + bias_c, + sc_slices[0], + sh_slices[0], + L, + g_stride, + sc_stride, + sh_stride, + stream, + ) + cf = _get_or_compile(C, is_rms, has_gate, has_weight, dummy_args) + + for b in range(B): + s, e = b * L, (b + 1) * L + cf( + y[s:e], + res_out[s:e], + res_2d[s:e], + x_2d[s:e], + g_slices[b], + weight_c, + bias_c, + sc_slices[b], + sh_slices[b], + L, + g_stride, + sc_stride, + sh_stride, + stream, + ) + + return y.view(B, L, C), res_out.view(B, L, C) + + +@flydsl_fused_residual_norm_scale_shift.register_fake +def _fake_flydsl_fused_residual_norm( + residual, + x, + gate, + weight, + bias, + scale, + shift, + norm_type, + eps=1e-6, +): + B, L, C = x.shape + bf16 = torch.bfloat16 + y = torch.empty(B, L, C, device=x.device, dtype=bf16) + res_out = torch.empty(B, L, C, device=x.device, dtype=bf16) + return y, res_out + + +############################################################################### +# _NormScaleShift kernel: norm(x) * (1+scale) + shift (no residual path) +############################################################################### + + +def _build_norm_scale_shift_module(D: int, is_rms: bool, has_weight: bool): + VEC = _VEC + NUM_WAVES = _NUM_WAVES + BLOCK = NUM_WAVES * WARP_SIZE + assert ( + D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0 + ), f"FlyDSL norm_scale_shift requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}" + NUM_ITERS = D // (BLOCK * VEC) + + @flyc.kernel(known_block_size=[BLOCK, 1, 1]) + def flydsl_norm_scale_shift_kernel( + y_ptr: fx.Tensor, + x_ptr: fx.Tensor, + weight_ptr: fx.Tensor, + bias_ptr: fx.Tensor, + scale_ptr: fx.Tensor, + shift_ptr: fx.Tensor, + total_rows: Int32, + scale_stride: Int32, + shift_stride: Int32, + ): + row = fx.block_idx.x + tid = fx.thread_idx.x + + i32 = T.i32 + f32 = T.f32 + bf16 = T.bf16 + vec_f32_t = ir.VectorType.get([VEC], f32) + vec_bf16_t = ir.VectorType.get([VEC], bf16) + + y_rsrc = buffer_ops.create_buffer_resource(y_ptr, max_size=True) + x_rsrc = buffer_ops.create_buffer_resource(x_ptr, max_size=True) + w_rsrc = buffer_ops.create_buffer_resource(weight_ptr, max_size=True) + b_rsrc = buffer_ops.create_buffer_resource(bias_ptr, max_size=True) + sc_rsrc = buffer_ops.create_buffer_resource(scale_ptr, max_size=True) + sh_rsrc = buffer_ops.create_buffer_resource(shift_ptr, max_size=True) + + row_i32 = ArithValue(row) + tid_i32 = ArithValue(tid) + D_i32 = arith.constant(D, type=i32) + row_off = row_i32 * D_i32 + scale_row_off = row_i32 * ArithValue(scale_stride) + shift_row_off = row_i32 * ArithValue(shift_stride) + + c_zero_f32 = arith.constant(0.0, type=f32) + c_one_f32 = arith.constant(1.0, type=f32) + eps_val = arith.constant(1e-6, type=f32) + D_float = arith.constant(float(D), type=f32) + + LDS_SLOTS = NUM_WAVES * 2 + 2 + ws_attr = ir.Attribute.parse("#gpu.address_space") + lds_i8_type = ir.MemRefType.get( + [ir.ShapedType.get_dynamic_size()], T.i8, memory_space=ws_attr + ) + lds_f32_type = ir.MemRefType.get([LDS_SLOTS], f32, memory_space=ws_attr) + lds_i8 = _gpu.DynamicSharedMemoryOp(lds_i8_type).result + byte_zero = arith_ops.ConstantOp( + ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), 0) + ).result + lds = _memref.ViewOp(lds_f32_type, lds_i8, byte_zero, []).result + + lane_id = tid_i32 % arith.constant(WARP_SIZE, type=i32) + wave_id = tid_i32 // arith.constant(WARP_SIZE, type=i32) + + # Phase 1: load x, accumulate stats, save f32 in registers + _saved_x_f32 = [] + partial_sum = _v(c_zero_f32) + partial_sum_sq = _v(c_zero_f32) + + for it in range_constexpr(NUM_ITERS): + col = tid_i32 * arith.constant(VEC, type=i32) + arith.constant( + it * BLOCK * VEC, type=i32 + ) + off = row_off + col + + x_vec = buffer_ops.buffer_load(x_rsrc, off, vec_width=VEC, dtype=bf16) + x_f32 = arith_ops.ExtFOp(vec_f32_t, _v(x_vec)).result + _saved_x_f32.append(x_f32) + + if const_expr(not is_rms): + v_sum = _vector.ReductionOp( + f32, _vector.CombiningKind.ADD, x_f32 + ).result + partial_sum = arith_ops.AddFOp(partial_sum, v_sum).result + x_sq = arith_ops.MulFOp(x_f32, x_f32).result + v_sum_sq = _vector.ReductionOp(f32, _vector.CombiningKind.ADD, x_sq).result + partial_sum_sq = arith_ops.AddFOp(partial_sum_sq, v_sum_sq).result + + # Intra-wave shuffle reduction + width_c = _v(arith.constant(WARP_SIZE, type=i32)) + w_sum = partial_sum + w_sq = partial_sum_sq + for sh in [32, 16, 8, 4, 2, 1]: + off_sh = _v(arith.constant(sh, type=i32)) + if const_expr(not is_rms): + peer_sum = _gpu.ShuffleOp( + w_sum, off_sh, width_c, mode=_gpu.ShuffleMode.XOR + ).shuffleResult + w_sum = arith_ops.AddFOp(w_sum, peer_sum).result + peer_sq = _gpu.ShuffleOp( + w_sq, off_sh, width_c, mode=_gpu.ShuffleMode.XOR + ).shuffleResult + w_sq = arith_ops.AddFOp(w_sq, peer_sq).result + + # Cross-wave LDS reduction + lane_0 = arith.cmpi(CmpIPredicate.eq, lane_id, arith.constant(0, type=i32)) + wave_idx = arith_ops.IndexCastOp(ir.IndexType.get(), _v(wave_id)).result + + _if_lane0 = scf.IfOp(lane_0) + with ir.InsertionPoint(_if_lane0.then_block): + if const_expr(not is_rms): + _memref.StoreOp(w_sum, lds, [wave_idx]) + sq_slot = arith_ops.AddIOp( + wave_idx, + arith_ops.ConstantOp( + ir.IndexType.get(), + ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES), + ).result, + ).result + _memref.StoreOp(w_sq, lds, [sq_slot]) + scf.YieldOp([]) + _gpu.BarrierOp() + + wave_0 = arith.cmpi(CmpIPredicate.eq, wave_id, arith.constant(0, type=i32)) + active = arith.andi( + wave_0, + arith.cmpi(CmpIPredicate.ult, lane_id, arith.constant(NUM_WAVES, type=i32)), + ) + lane_idx = arith_ops.IndexCastOp(ir.IndexType.get(), _v(lane_id)).result + lane_idx_sq = arith_ops.AddIOp( + lane_idx, + arith_ops.ConstantOp( + ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES) + ).result, + ).result + + if const_expr(is_rms): + _if_active = scf.IfOp(active, [f32], has_else=True) + with ir.InsertionPoint(_if_active.then_block): + sq_val = _memref.LoadOp(lds, [lane_idx_sq]).result + scf.YieldOp([sq_val]) + with ir.InsertionPoint(_if_active.else_block): + scf.YieldOp([_v(c_zero_f32)]) + loaded_sq = _if_active.results[0] + loaded_sum = _v(c_zero_f32) + else: + _if_active = scf.IfOp(active, [f32, f32], has_else=True) + with ir.InsertionPoint(_if_active.then_block): + s_val = _memref.LoadOp(lds, [lane_idx]).result + sq_val = _memref.LoadOp(lds, [lane_idx_sq]).result + scf.YieldOp([s_val, sq_val]) + with ir.InsertionPoint(_if_active.else_block): + scf.YieldOp([_v(c_zero_f32), _v(c_zero_f32)]) + loaded_sum = _if_active.results[0] + loaded_sq = _if_active.results[1] + + final_sum = loaded_sum + final_sq = loaded_sq + for sh in [32, 16, 8, 4, 2, 1]: + off_sh = _v(arith.constant(sh, type=i32)) + if const_expr(not is_rms): + ps = _gpu.ShuffleOp( + final_sum, off_sh, width_c, mode=_gpu.ShuffleMode.XOR + ).shuffleResult + final_sum = arith_ops.AddFOp(final_sum, ps).result + pq = _gpu.ShuffleOp( + final_sq, off_sh, width_c, mode=_gpu.ShuffleMode.XOR + ).shuffleResult + final_sq = arith_ops.AddFOp(final_sq, pq).result + + both_0 = arith.andi(wave_0, lane_0) + final_sum_slot = arith_ops.ConstantOp( + ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES * 2) + ).result + final_sq_slot = arith_ops.ConstantOp( + ir.IndexType.get(), + ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES * 2 + 1), + ).result + _if_both = scf.IfOp(both_0) + with ir.InsertionPoint(_if_both.then_block): + if const_expr(not is_rms): + _memref.StoreOp(final_sum, lds, [final_sum_slot]) + _memref.StoreOp(final_sq, lds, [final_sq_slot]) + scf.YieldOp([]) + _gpu.BarrierOp() + + if const_expr(not is_rms): + total_sum = _memref.LoadOp(lds, [final_sum_slot]).result + else: + total_sum = _v(c_zero_f32) + total_sq = _memref.LoadOp(lds, [final_sq_slot]).result + + d_f = _v(D_float) + eps_v = _v(eps_val) + if const_expr(is_rms): + var = arith_ops.DivFOp(total_sq, d_f).result + var_eps = arith_ops.AddFOp(var, eps_v).result + rstd = math_ops.RsqrtOp(var_eps).result + mean = _v(c_zero_f32) + else: + mean = arith_ops.DivFOp(total_sum, d_f).result + mean_sq = arith_ops.MulFOp(mean, mean).result + var = arith_ops.SubFOp( + arith_ops.DivFOp(total_sq, d_f).result, mean_sq + ).result + var_eps = arith_ops.AddFOp(var, eps_v).result + rstd = math_ops.RsqrtOp(var_eps).result + + # Phase 2: normalize from register cache + scale_shift → single output + mean_splat = _vector.BroadcastOp(vec_f32_t, mean).result + rstd_splat = _vector.BroadcastOp(vec_f32_t, rstd).result + one_splat = _vector.BroadcastOp(vec_f32_t, _v(c_one_f32)).result + + for it in range_constexpr(NUM_ITERS): + col = tid_i32 * arith.constant(VEC, type=i32) + arith.constant( + it * BLOCK * VEC, type=i32 + ) + off = row_off + col + + x_f32 = _saved_x_f32[it] + + if const_expr(is_rms): + x_hat = arith_ops.MulFOp(x_f32, rstd_splat).result + else: + centered = arith_ops.SubFOp(x_f32, mean_splat).result + x_hat = arith_ops.MulFOp(centered, rstd_splat).result + + if const_expr(has_weight): + w_vec = buffer_ops.buffer_load(w_rsrc, col, vec_width=VEC, dtype=bf16) + w_f32 = arith_ops.ExtFOp(vec_f32_t, _v(w_vec)).result + x_hat = arith_ops.MulFOp(x_hat, w_f32).result + b_vec = buffer_ops.buffer_load(b_rsrc, col, vec_width=VEC, dtype=bf16) + b_f32 = arith_ops.ExtFOp(vec_f32_t, _v(b_vec)).result + x_hat = arith_ops.AddFOp(x_hat, b_f32).result + + sc_off = scale_row_off + col + sc_vec = buffer_ops.buffer_load(sc_rsrc, sc_off, vec_width=VEC, dtype=bf16) + sc_f32 = arith_ops.ExtFOp(vec_f32_t, _v(sc_vec)).result + sc_p1 = arith_ops.AddFOp(one_splat, sc_f32).result + x_hat = arith_ops.MulFOp(x_hat, sc_p1).result + + sh_off = shift_row_off + col + sh_vec = buffer_ops.buffer_load(sh_rsrc, sh_off, vec_width=VEC, dtype=bf16) + sh_f32 = arith_ops.ExtFOp(vec_f32_t, _v(sh_vec)).result + y_f32 = arith_ops.AddFOp(x_hat, sh_f32).result + + y_bf16 = arith_ops.TruncFOp(vec_bf16_t, y_f32).result + buffer_ops.buffer_store(y_bf16, y_rsrc, off) + + @flyc.jit + def launch_norm_ss( + y: fx.Tensor, + x: fx.Tensor, + weight: fx.Tensor, + bias: fx.Tensor, + scale: fx.Tensor, + shift: fx.Tensor, + total_rows: fx.Int32, + scale_stride: fx.Int32, + shift_stride: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + pass + grid_x = arith.index_cast(T.index, total_rows) + launcher = flydsl_norm_scale_shift_kernel( + y, + x, + weight, + bias, + scale, + shift, + total_rows, + scale_stride, + shift_stride, + ) + LDS_BYTES = (NUM_WAVES * 2 + 2) * 4 + launcher.launch( + grid=(grid_x, 1, 1), block=(BLOCK, 1, 1), smem=LDS_BYTES, stream=stream + ) + + return launch_norm_ss + + +_NSS_COMPILE_CACHE = {} + + +def _get_or_compile_nss(D, is_rms, has_weight, args): + key = ("nss", D, is_rms, has_weight) + if key not in _NSS_COMPILE_CACHE: + launcher = _build_norm_scale_shift_module(D, is_rms, has_weight) + cf = flyc.compile(launcher, *args) + _NSS_COMPILE_CACHE[key] = cf + return _NSS_COMPILE_CACHE[key] + + +@torch.library.custom_op("sglang::flydsl_norm_scale_shift", mutates_args=()) +def flydsl_norm_scale_shift( + x: torch.Tensor, + weight: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + scale: torch.Tensor, + shift: torch.Tensor, + norm_type: str, + eps: float = 1e-6, +) -> torch.Tensor: + B, L, C = x.shape + rows = B * L + bf16 = torch.bfloat16 + + x_2d = _ensure_bf16_contig(x).reshape(rows, C) + y = torch.empty_like(x_2d) + + has_weight = weight is not None + weight_c = ( + _ensure_bf16_contig(weight) + if has_weight + else torch.empty(C, device=x.device, dtype=bf16) + ) + bias_c = ( + _ensure_bf16_contig(bias) + if bias is not None + else torch.zeros(C, device=x.device, dtype=bf16) + ) + + sc_slices, sc_stride = _prep_slices(scale, B, L, C) + sh_slices, sh_stride = _prep_slices(shift, B, L, C) + + is_rms = norm_type == "rms" + stream = torch.cuda.current_stream() + + dummy_args = ( + y[:L], + x_2d[:L], + weight_c, + bias_c, + sc_slices[0], + sh_slices[0], + L, + sc_stride, + sh_stride, + stream, + ) + cf = _get_or_compile_nss(C, is_rms, has_weight, dummy_args) + + for b in range(B): + s, e = b * L, (b + 1) * L + cf( + y[s:e], + x_2d[s:e], + weight_c, + bias_c, + sc_slices[b], + sh_slices[b], + L, + sc_stride, + sh_stride, + stream, + ) + + return y.view(B, L, C) + + +@flydsl_norm_scale_shift.register_fake +def _fake_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps=1e-6): + B, L, C = x.shape + return torch.empty(B, L, C, device=x.device, dtype=torch.bfloat16) diff --git a/python/sglang/jit_kernel/tests/diffusion/test_flydsl_fused_norm.py b/python/sglang/jit_kernel/tests/diffusion/test_flydsl_fused_norm.py new file mode 100644 index 000000000..bfb67cb43 --- /dev/null +++ b/python/sglang/jit_kernel/tests/diffusion/test_flydsl_fused_norm.py @@ -0,0 +1,133 @@ +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, suite="jit-kernel-unit-test-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.jit_kernel.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.jit_kernel.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"])) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index ebd803e59..3b73421d4 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -33,6 +33,7 @@ _is_npu = current_platform.is_npu() _is_musa = current_platform.is_musa() _is_cpu = current_platform.is_cpu() _is_xpu = current_platform.is_xpu() +_use_rocm_flydsl = get_bool_env_var("SGLANG_USE_ROCM_FLYDSL") _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _is_cuda or _is_xpu: @@ -506,10 +507,39 @@ class _ScaleResidualNormScaleShift(CustomOp): self.eps, ) - def forward_hip(self, *args, **kwargs): - # ROCm does not support CUDA/CUTLASS-based fused kernels yet, - # so we fall back to the native PyTorch implementation. - return self.forward_native(*args, **kwargs) + def forward_hip( + self, + residual: torch.Tensor, + x: torch.Tensor, + gate: torch.Tensor | int, + shift: torch.Tensor, + scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not _use_rocm_flydsl: + return self.forward_native(residual, x, gate, shift, scale) + + try: + from sglang.jit_kernel.diffusion.flydsl.fused_residual_norm import ( + FLYDSL_NORM_MIN_ALIGNED_DIM, + flydsl_fused_residual_norm_scale_shift, + ) + except ImportError: + return self.forward_native(residual, x, gate, shift, scale) + + if x.shape[-1] % FLYDSL_NORM_MIN_ALIGNED_DIM != 0: + return self.forward_native(residual, x, gate, shift, scale) + + return flydsl_fused_residual_norm_scale_shift( + residual.contiguous(), + x.contiguous(), + gate.contiguous() if isinstance(gate, torch.Tensor) else None, + _ensure_contiguous(getattr(self.norm, "weight", None)), + _ensure_contiguous(getattr(self.norm, "bias", None)), + scale.contiguous(), + shift.contiguous(), + self.norm_type, + self.eps, + ) def forward_musa(self, *args, **kwargs): # MUSA does not support CUDA/CUTLASS-based fused kernels yet, @@ -648,10 +678,36 @@ class _NormScaleShift(CustomOp): self.eps, ) - def forward_hip(self, *args, **kwargs): - # ROCm does not support CUDA/CUTLASS-based fused kernels yet, - # so we fall back to the native PyTorch implementation. - return self.forward_native(*args, **kwargs) + def forward_hip( + self, + x: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + ) -> torch.Tensor: + if not _use_rocm_flydsl: + return self.forward_native(x, shift, scale) + + try: + from sglang.jit_kernel.diffusion.flydsl.fused_residual_norm import ( + FLYDSL_NORM_MIN_ALIGNED_DIM, + flydsl_norm_scale_shift, + ) + except ImportError: + return self.forward_native(x, shift, scale) + + if x.shape[-1] % FLYDSL_NORM_MIN_ALIGNED_DIM != 0: + return self.forward_native(x, shift, scale) + + result = flydsl_norm_scale_shift( + x.contiguous(), + _ensure_contiguous(getattr(self.norm, "weight", None)), + _ensure_contiguous(getattr(self.norm, "bias", None)), + scale.contiguous(), + shift.contiguous(), + self.norm_type, + self.eps, + ) + return result.to(x.dtype) def forward_musa(self, *args, **kwargs): # MUSA does not support CUDA/CUTLASS-based fused kernels yet,