diff --git a/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py b/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py index ddf5ffc28..39691a992 100644 --- a/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py +++ b/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py @@ -9,6 +9,14 @@ Provides two fused kernels: 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. + +Written against the FlyDSL v0.3.0 stable public API only; see +docs/api_stability.md in the FlyDSL tree for the classification rules. +Raw MLIR dialect builders, compiler-internal contexts, private expr +submodules, and anything under the FlyDSL source-only kernel tree are all +outside the shipped wheel or the stability contract -- do not reintroduce +them here. CI greps this file for those names, so avoid spelling them even +in comments. """ from typing import Optional, Tuple @@ -16,35 +24,132 @@ 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, const_expr, range_constexpr -from flydsl.expr.arith import ArithValue, CmpIPredicate -from flydsl.expr.typing import Int32, T - -try: - from flydsl.expr import buffer_ops -except ImportError: - # flydsl 0.3.0 removed flydsl.expr.buffer_ops and pushed the buffer-resource - # layer down to its consumers; AITER carries the same module, same API. - from aiter.ops.flydsl.kernels import buffer_ops +from flydsl.expr import const_expr, range_constexpr WARP_SIZE = 64 _VEC = 8 _NUM_WAVES = 10 FLYDSL_NORM_MIN_ALIGNED_DIM = WARP_SIZE * _NUM_WAVES * _VEC # 5120 +# Kernel-side epsilon. The public `eps` argument is intentionally ignored, as it +# was before the stable-API migration; changing that is a separate correctness fix. +_EPS = 1e-6 -def _v(x): - return buffer_ops._unwrap_value(x) +# 128-bit vector copies over bf16 elements. +_ELEM_BITS = 16 + +# Reduction order is part of the numerical contract; keep the descending sequence. +_SHUFFLE_OFFSETS = (32, 16, 8, 4, 2, 1) + + +def _require_stable_api() -> None: + """Fail as ImportError when the installed FlyDSL predates the v0.3.0 surface. + + Callers in multimodal_gen/runtime/layers/layernorm.py guard this module with + `except ImportError` and fall back to the native path, so a missing symbol + must surface as ImportError rather than AttributeError. + """ + required = { + "flydsl.expr": ( + "Tensor", + "Stream", + "Int32", + "Float32", + "BFloat16", + "ReductionOp", + "SharedAllocator", + "struct", + "Array", + "make_layout", + "logical_divide", + "slice", + "make_copy_atom", + "copy_atom_call", + "make_rmem_tensor", + "memref_load_vec", + "memref_store_vec", + "memref_load", + "memref_store", + ), + "flydsl.expr.gpu": ("barrier", "shuffle_xor", "block_idx", "thread_idx"), + "flydsl.expr.math": ("rsqrt",), + "flydsl.expr.rocdl": ("make_buffer_tensor", "BufferCopy128b"), + "flydsl.compiler": ("kernel", "jit", "compile"), + } + roots = { + "flydsl.expr": fx, + "flydsl.expr.gpu": fx.gpu, + "flydsl.expr.math": fx.math, + "flydsl.expr.rocdl": fx.rocdl, + "flydsl.compiler": flyc, + } + missing = [ + f"{mod}.{name}" + for mod, names in required.items() + for name in names + if not hasattr(roots[mod], name) + ] + if missing: + raise ImportError( + "FlyDSL is too old for sglang's fused norm kernels; missing stable " + f"v0.3.0 APIs: {', '.join(missing)}" + ) + + +_require_stable_api() + + +def _make_reduction_storage(slots: int): + """LDS layout for the two-stage block reduction. + + s_sum/s_sq hold one partial per wave; s_final holds the two broadcast values + so the final write never aliases a slot that is still being read. + """ + + @fx.struct + class SharedStorage: + s_sum: fx.Array[fx.Float32, slots, 16] + s_sq: fx.Array[fx.Float32, slots, 16] + s_final: fx.Array[fx.Float32, 2, 16] + + return SharedStorage + + +def _load_vec(copy_atom, div_tensor, idx): + r = fx.make_rmem_tensor(_VEC, fx.BFloat16) + fx.copy_atom_call(copy_atom, fx.slice(div_tensor, (None, idx)), r) + return fx.memref_load_vec(r) + + +def _store_vec(copy_atom, div_tensor, idx, val): + r = fx.make_rmem_tensor(_VEC, fx.BFloat16) + fx.memref_store_vec(val, r) + fx.copy_atom_call(copy_atom, r, fx.slice(div_tensor, (None, idx))) + + +def _row_div(tensor, row): + """Vector-partitioned view of one row of a rank-2 tensor.""" + return fx.logical_divide( + fx.slice(fx.rocdl.make_buffer_tensor(tensor), (row, None)), + fx.make_layout(_VEC, 1), + ) + + +def _flat_div(tensor): + """Vector-partitioned view of a rank-1 tensor.""" + return fx.logical_divide( + fx.rocdl.make_buffer_tensor(tensor), fx.make_layout(_VEC, 1) + ) + + +def _bcast_row(row, stride): + """Row index honoring a runtime broadcast stride. + + A broadcast operand arrives as a dense (1, C) tensor and the host signals it + with stride == 0, so selecting row 0 reproduces the previous `row * stride` + addressing exactly. + """ + return (stride != 0).select(row, 0) def _build_fused_norm_module(D: int, is_rms: bool, has_gate: bool, has_weight: bool): @@ -55,6 +160,7 @@ def _build_fused_norm_module(D: int, is_rms: bool, has_gate: bool, has_weight: b f"FlyDSL fused_residual_norm requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}" ) NUM_ITERS = D // (BLOCK * VEC) + SharedStorage = _make_reduction_storage(NUM_WAVES) @flyc.kernel(known_block_size=[BLOCK, 1, 1]) def flydsl_fused_residual_norm_ss_kernel( @@ -67,258 +173,124 @@ def _build_fused_norm_module(D: int, is_rms: bool, has_gate: bool, has_weight: b bias_ptr: fx.Tensor, scale_ptr: fx.Tensor, shift_ptr: fx.Tensor, - total_rows: Int32, - gate_stride: Int32, - scale_stride: Int32, - shift_stride: Int32, + total_rows: fx.Int32, + gate_stride: fx.Int32, + scale_stride: fx.Int32, + shift_stride: fx.Int32, ): row = fx.block_idx.x tid = fx.thread_idx.x + lane_id = tid % WARP_SIZE + wave_id = tid // WARP_SIZE - 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) + n_float = float(D) + c_zero = fx.Float32(0.0) - 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) + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + s_sum = lds.s_sum.view(fx.make_layout(NUM_WAVES, 1)) + s_sq = lds.s_sq.view(fx.make_layout(NUM_WAVES, 1)) + s_final = lds.s_final.view(fx.make_layout(2, 1)) - 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) + copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), _ELEM_BITS) - 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) + y_div = _row_div(y_ptr, row) + ro_div = _row_div(res_out_ptr, row) + r_div = _row_div(res_ptr, row) + x_div = _row_div(x_ptr, row) + if const_expr(has_gate): + g_div = _row_div(gate_ptr, _bcast_row(row, gate_stride)) + if const_expr(has_weight): + w_div = _flat_div(weight_ptr) + b_div = _flat_div(bias_ptr) + sc_div = _row_div(scale_ptr, _bcast_row(row, scale_stride)) + sh_div = _row_div(shift_ptr, _bcast_row(row, shift_stride)) - # 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 + def wave_reduce_add(val): + w = val + for i in range_constexpr(len(_SHUFFLE_OFFSETS)): + w = w + fx.gpu.shuffle_xor(w, _SHUFFLE_OFFSETS[i], WARP_SIZE) + return w - 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) + # Phase 1: residual + gate*x, accumulate stats, keep f32 in registers. + saved_ro = [] + partial_sum = c_zero + partial_sq = c_zero 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 + idx = tid + it * BLOCK - 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 + r_f32 = _load_vec(copy_atom, r_div, idx).to(fx.Float32) + x_f32 = _load_vec(copy_atom, x_div, idx).to(fx.Float32) 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 + g_f32 = _load_vec(copy_atom, g_div, idx).to(fx.Float32) + ro_f32 = r_f32 + g_f32 * x_f32 else: - ro_f32 = arith_ops.AddFOp(r_f32, x_f32).result + ro_f32 = r_f32 + x_f32 - 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) + _store_vec(copy_atom, ro_div, idx, ro_f32.to(fx.BFloat16)) + saved_ro.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() + partial_sum = partial_sum + ro_f32.reduce(fx.ReductionOp.ADD) + partial_sq = partial_sq + (ro_f32 * ro_f32).reduce(fx.ReductionOp.ADD) + # Stage 1: intra-wave shuffle reduction. 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 + w_sum = wave_reduce_add(partial_sum) + w_sq = wave_reduce_add(partial_sq) - # 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 + if lane_id == 0: + if const_expr(not is_rms): + fx.memref_store(w_sum, s_sum, wave_id) + fx.memref_store(w_sq, s_sq, wave_id) + fx.gpu.barrier() - # 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 + # Stage 2: wave 0 folds the per-wave partials and publishes the results. + if wave_id == 0: + in_range = lane_id < NUM_WAVES + lane_safe = in_range.select(lane_id, 0) + if const_expr(not is_rms): + v_sum = wave_reduce_add( + in_range.select(fx.memref_load(s_sum, lane_safe), c_zero) + ) + v_sq = wave_reduce_add( + in_range.select(fx.memref_load(s_sq, lane_safe), c_zero) ) - off = row_off + col + if lane_id == 0: + if const_expr(not is_rms): + fx.memref_store(v_sum, s_final, 0) + fx.memref_store(v_sq, s_final, 1) + fx.gpu.barrier() - ro_f32 = _saved_ro_f32[it] + total_sq = fx.memref_load(s_final, 1) + + if const_expr(is_rms): + rstd = fx.math.rsqrt(total_sq / n_float + _EPS) + else: + mean = fx.memref_load(s_final, 0) / n_float + var = total_sq / n_float - mean * mean + rstd = fx.math.rsqrt(var + _EPS) + + # Phase 2: normalize from the register cache (no HBM re-read). + for it in range_constexpr(NUM_ITERS): + idx = tid + it * BLOCK + ro_f32 = saved_ro[it] if const_expr(is_rms): - x_hat = arith_ops.MulFOp(ro_f32, rstd_splat).result + x_hat = ro_f32 * rstd else: - centered = arith_ops.SubFOp(ro_f32, mean_splat).result - x_hat = arith_ops.MulFOp(centered, rstd_splat).result + x_hat = (ro_f32 - mean) * rstd 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 + x_hat = x_hat * _load_vec(copy_atom, w_div, idx).to(fx.Float32) + x_hat = x_hat + _load_vec(copy_atom, b_div, idx).to(fx.Float32) - 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 + sc_f32 = _load_vec(copy_atom, sc_div, idx).to(fx.Float32) + x_hat = x_hat * (sc_f32 + 1.0) + y_f32 = x_hat + _load_vec(copy_atom, sh_div, idx).to(fx.Float32) - 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) + _store_vec(copy_atom, y_div, idx, y_f32.to(fx.BFloat16)) @flyc.jit def launch_fused_norm( @@ -337,10 +309,6 @@ def _build_fused_norm_module(D: int, is_rms: bool, has_gate: bool, has_weight: b 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, @@ -356,10 +324,7 @@ def _build_fused_norm_module(D: int, is_rms: bool, has_gate: bool, has_weight: b 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 - ) + launcher.launch(grid=(total_rows, 1, 1), block=(BLOCK, 1, 1), stream=stream) return launch_fused_norm @@ -547,6 +512,7 @@ def _build_norm_scale_shift_module(D: int, is_rms: bool, has_weight: bool): f"FlyDSL norm_scale_shift requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}" ) NUM_ITERS = D // (BLOCK * VEC) + SharedStorage = _make_reduction_storage(NUM_WAVES) @flyc.kernel(known_block_size=[BLOCK, 1, 1]) def flydsl_norm_scale_shift_kernel( @@ -556,237 +522,109 @@ def _build_norm_scale_shift_module(D: int, is_rms: bool, has_weight: bool): bias_ptr: fx.Tensor, scale_ptr: fx.Tensor, shift_ptr: fx.Tensor, - total_rows: Int32, - scale_stride: Int32, - shift_stride: Int32, + total_rows: fx.Int32, + scale_stride: fx.Int32, + shift_stride: fx.Int32, ): row = fx.block_idx.x tid = fx.thread_idx.x + lane_id = tid % WARP_SIZE + wave_id = tid // WARP_SIZE - 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) + n_float = float(D) + c_zero = fx.Float32(0.0) - 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) + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + s_sum = lds.s_sum.view(fx.make_layout(NUM_WAVES, 1)) + s_sq = lds.s_sq.view(fx.make_layout(NUM_WAVES, 1)) + s_final = lds.s_final.view(fx.make_layout(2, 1)) - 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) + copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), _ELEM_BITS) - 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) + y_div = _row_div(y_ptr, row) + x_div = _row_div(x_ptr, row) + if const_expr(has_weight): + w_div = _flat_div(weight_ptr) + b_div = _flat_div(bias_ptr) + sc_div = _row_div(scale_ptr, _bcast_row(row, scale_stride)) + sh_div = _row_div(shift_ptr, _bcast_row(row, shift_stride)) - 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 + def wave_reduce_add(val): + w = val + for i in range_constexpr(len(_SHUFFLE_OFFSETS)): + w = w + fx.gpu.shuffle_xor(w, _SHUFFLE_OFFSETS[i], WARP_SIZE) + return w - 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) + # Phase 1: load x, accumulate stats, keep f32 in registers. + saved_x = [] + partial_sum = c_zero + partial_sq = c_zero 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) + idx = tid + it * BLOCK + x_f32 = _load_vec(copy_atom, x_div, idx).to(fx.Float32) + saved_x.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() + partial_sum = partial_sum + x_f32.reduce(fx.ReductionOp.ADD) + partial_sq = partial_sq + (x_f32 * x_f32).reduce(fx.ReductionOp.ADD) + # Stage 1: intra-wave shuffle reduction. 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 + w_sum = wave_reduce_add(partial_sum) + w_sq = wave_reduce_add(partial_sq) - 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 + if lane_id == 0: + if const_expr(not is_rms): + fx.memref_store(w_sum, s_sum, wave_id) + fx.memref_store(w_sq, s_sq, wave_id) + fx.gpu.barrier() - # 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 + # Stage 2: wave 0 folds the per-wave partials and publishes the results. + if wave_id == 0: + in_range = lane_id < NUM_WAVES + lane_safe = in_range.select(lane_id, 0) + if const_expr(not is_rms): + v_sum = wave_reduce_add( + in_range.select(fx.memref_load(s_sum, lane_safe), c_zero) + ) + v_sq = wave_reduce_add( + in_range.select(fx.memref_load(s_sq, lane_safe), c_zero) ) - off = row_off + col + if lane_id == 0: + if const_expr(not is_rms): + fx.memref_store(v_sum, s_final, 0) + fx.memref_store(v_sq, s_final, 1) + fx.gpu.barrier() - x_f32 = _saved_x_f32[it] + total_sq = fx.memref_load(s_final, 1) + + if const_expr(is_rms): + rstd = fx.math.rsqrt(total_sq / n_float + _EPS) + else: + mean = fx.memref_load(s_final, 0) / n_float + var = total_sq / n_float - mean * mean + rstd = fx.math.rsqrt(var + _EPS) + + # Phase 2: normalize from the register cache + scale_shift. + for it in range_constexpr(NUM_ITERS): + idx = tid + it * BLOCK + x_f32 = saved_x[it] if const_expr(is_rms): - x_hat = arith_ops.MulFOp(x_f32, rstd_splat).result + x_hat = x_f32 * rstd else: - centered = arith_ops.SubFOp(x_f32, mean_splat).result - x_hat = arith_ops.MulFOp(centered, rstd_splat).result + x_hat = (x_f32 - mean) * rstd 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 + x_hat = x_hat * _load_vec(copy_atom, w_div, idx).to(fx.Float32) + x_hat = x_hat + _load_vec(copy_atom, b_div, idx).to(fx.Float32) - 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 + sc_f32 = _load_vec(copy_atom, sc_div, idx).to(fx.Float32) + x_hat = x_hat * (sc_f32 + 1.0) + y_f32 = x_hat + _load_vec(copy_atom, sh_div, idx).to(fx.Float32) - 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) + _store_vec(copy_atom, y_div, idx, y_f32.to(fx.BFloat16)) @flyc.jit def launch_norm_ss( @@ -801,10 +639,6 @@ def _build_norm_scale_shift_module(D: int, is_rms: bool, has_weight: bool): 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, @@ -816,10 +650,7 @@ def _build_norm_scale_shift_module(D: int, is_rms: bool, has_weight: bool): 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 - ) + launcher.launch(grid=(total_rows, 1, 1), block=(BLOCK, 1, 1), stream=stream) return launch_norm_ss diff --git a/test/registered/kernels/ops/diffusion/test_norm_flydsl.py b/test/registered/kernels/ops/diffusion/test_norm_flydsl.py index e399477e5..8ff2a4c53 100644 --- a/test/registered/kernels/ops/diffusion/test_norm_flydsl.py +++ b/test/registered/kernels/ops/diffusion/test_norm_flydsl.py @@ -9,20 +9,22 @@ Oracle: an fp32 reference chain, with a tolerance -- the kernel keeps fp32 statistics but reorders the reduction. """ +import os +import subprocess 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") +register_amd_ci(est_time=60, stage="jit-kernel-unit", runner_config="amd") pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU required") DEVICE = "cuda" +FLYDSL_MODULE = "sglang.kernels.ops.diffusion.norm.fused_residual_norm_flydsl" FLYDSL_D = 5120 FLYDSL_EPS = 1e-6 @@ -32,67 +34,259 @@ def _require_rocm(): 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 +def _flydsl_ops(): + """Resolve the FlyDSL exports, skipping when the installed FlyDSL is too old. + Resolved inside each test: the FlyDSL compiler only exists on ROCm, and the + facade imports the submodule the moment an export is named -- a module-level + import would fail collection of this whole file on CUDA. -@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): + The kernel module raises ImportError when the stable FlyDSL surface it needs + is absent, which is also how ``layernorm.py`` detects that it must fall back + to the native path. A runner image predating that surface should skip here + rather than report a kernel regression. + """ _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, - ) + try: + from sglang.kernels.ops.diffusion import ( + flydsl_fused_residual_norm_scale_shift, + flydsl_norm_scale_shift, + ) + except ImportError as exc: + pytest.skip(f"FlyDSL unavailable or too old for the fused norm kernels: {exc}") + return flydsl_fused_residual_norm_scale_shift, flydsl_norm_scale_shift + +def _mk(shape, dtype=torch.bfloat16): + return torch.randn(*shape, device=DEVICE, dtype=dtype) + + +def _ref_rms_norm(x_f32, eps): + var = x_f32.pow(2).mean(-1, keepdim=True) + return x_f32 * torch.rsqrt(var + eps) + + +def _apply_affine(normed, weight, bias, norm_type): + """Reproduce the kernel's affine stage. + + The kernel applies weight and bias together, and skips both when weight is + absent -- a standalone bias is not reachable through the public op. + """ + if weight is None: + return normed + out = normed * weight.float() + if norm_type == "layer" and bias is not None: + out = out + bias.float() + return out + + +def _ref_norm(x_bf16, weight, bias, norm_type, eps): + base = x_bf16.float() + if norm_type == "layer": + mean = base.mean(-1, keepdim=True) + var = base.var(-1, keepdim=True, unbiased=False) + normed = (base - mean) * torch.rsqrt(var + eps) + else: + normed = _ref_rms_norm(base, eps) + return _apply_affine(normed, weight, bias, norm_type) + + +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) + normed = _ref_norm(ref_res_bf16, weight, bias, norm_type, eps) + 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): + normed = _ref_norm(x, weight, bias, norm_type, eps) + return (normed * (1.0 + scale.float()) + shift.float()).to(torch.bfloat16) + + +@pytest.fixture(autouse=True) +def _seed(): 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) + + +# norm_type, B, L, has_gate, has_weight, scale/shift layout +FUSED_CASES = [ + ("rms", 1, 16, True, True, "bcast"), + ("rms", 2, 16, True, True, "bcast"), + ("layer", 2, 16, True, True, "bcast"), + ("rms", 1, 90000, True, True, "bcast"), + # gate-free path: a distinct compiled specialization + ("rms", 2, 16, False, True, "bcast"), + ("layer", 2, 16, False, True, "bcast"), + # no-affine path: the production default + ("rms", 2, 16, True, False, "bcast"), + ("layer", 2, 16, True, False, "bcast"), + # per-row scale/shift exercises the non-zero row stride + ("rms", 2, 16, True, True, "perrow"), + ("layer", 2, 16, True, True, "perrow"), +] + + +@pytest.mark.parametrize("norm_type,B,L,has_gate,has_weight,ss", FUSED_CASES) +def test_flydsl_fused_residual_norm_scale_shift( + norm_type, B, L, has_gate, has_weight, ss +): + fused_op, _ = _flydsl_ops() + + ss_shape = (B, 1, FLYDSL_D) if ss == "bcast" else (B, L, FLYDSL_D) + residual = _mk((B, L, FLYDSL_D)) + x = _mk((B, L, FLYDSL_D)) + gate = _mk((B, 1, FLYDSL_D)) if has_gate else None + weight = _mk((FLYDSL_D,), torch.float32) if has_weight else None bias = ( - torch.randn(FLYDSL_D, device=DEVICE, dtype=torch.float32) - if norm_type == "layer" + _mk((FLYDSL_D,), torch.float32) + if (has_weight and 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) + scale = _mk(ss_shape) + shift = _mk(ss_shape) - 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 - ) + y, res_out = fused_op( + residual, x, gate, weight, bias, scale, shift, norm_type, FLYDSL_EPS + ) + y_ref, res_ref = _ref_fused_residual_norm_ss( + residual, x, gate, weight, bias, scale, shift, norm_type, FLYDSL_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, True, "bcast"), + ("layer", 2, 16, True, "bcast"), + ("rms", 1, 90000, True, "bcast"), + ("layer", 1, 90000, True, "bcast"), + # no-affine path + ("rms", 2, 16, False, "bcast"), + ("layer", 2, 16, False, "bcast"), + # per-row scale/shift + ("rms", 2, 16, True, "perrow"), + ("layer", 2, 16, True, "perrow"), +] + + +@pytest.mark.parametrize("norm_type,B,L,has_weight,ss", NSS_CASES) +def test_flydsl_norm_scale_shift(norm_type, B, L, has_weight, ss): + _, nss_op = _flydsl_ops() + + ss_shape = (B, 1, FLYDSL_D) if ss == "bcast" else (B, L, FLYDSL_D) + x = _mk((B, L, FLYDSL_D)) + weight = _mk((FLYDSL_D,), torch.float32) if has_weight else None + bias = ( + _mk((FLYDSL_D,), torch.float32) + if (has_weight and norm_type == "layer") + else None + ) + scale = _mk(ss_shape) + shift = _mk(ss_shape) + + y = nss_op(x, weight, bias, scale, shift, norm_type, FLYDSL_EPS) + y_ref = _ref_norm_ss(x, weight, bias, scale, shift, norm_type, FLYDSL_EPS) + torch.testing.assert_close(y, y_ref, atol=1.0, rtol=5e-2) + + +def test_flydsl_fused_frame_gate_4d_scale_shift(): + """4D (B, NF, 1, D) scale/shift is expanded to per-row by _prep_slices.""" + fused_op, _ = _flydsl_ops() + + B, L, NF = 2, 16, 4 + residual, x = _mk((B, L, FLYDSL_D)), _mk((B, L, FLYDSL_D)) + gate = _mk((B, 1, FLYDSL_D)) + weight = _mk((FLYDSL_D,), torch.float32) + scale4, shift4 = _mk((B, NF, 1, FLYDSL_D)), _mk((B, NF, 1, FLYDSL_D)) + + y, res_out = fused_op( + residual, x, gate, weight, None, scale4, shift4, "rms", FLYDSL_EPS + ) + + def expand(t): + return t.expand(B, NF, L // NF, FLYDSL_D).reshape(B, L, FLYDSL_D) + + y_ref, res_ref = _ref_fused_residual_norm_ss( + residual, + x, + gate, + weight, + None, + expand(scale4), + expand(shift4), + "rms", + FLYDSL_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) + + +def test_flydsl_multi_iteration_dim(): + """D=10240 drives NUM_ITERS=2, exercising more than one register-cached tile.""" + _, nss_op = _flydsl_ops() + + big_d, B, L = 10240, 1, 16 + x = _mk((B, L, big_d)) + weight = _mk((big_d,), torch.float32) + scale, shift = _mk((B, 1, big_d)), _mk((B, 1, big_d)) + + y = nss_op(x, weight, None, scale, shift, "rms", FLYDSL_EPS) + y_ref = _ref_norm_ss(x, weight, None, scale, shift, "rms", FLYDSL_EPS) + torch.testing.assert_close(y, y_ref, atol=1.0, rtol=5e-2) + + +def test_flydsl_compile_cache_reuse_across_row_counts_and_layouts(): + """Guard for the compile-cache/shape-specialization hazard. + + The cache key is (D, is_rms, has_gate, has_weight) and deliberately excludes + the row count and whether scale/shift are broadcast. One compiled kernel must + therefore stay correct when both of those change between calls. + """ + _, nss_op = _flydsl_ops() + + weight = _mk((FLYDSL_D,), torch.float32) + for L, ss in ((16, "bcast"), (90000, "bcast"), (16, "perrow"), (64, "bcast")): + ss_shape = (1, 1, FLYDSL_D) if ss == "bcast" else (1, L, FLYDSL_D) + x = _mk((1, L, FLYDSL_D)) + scale, shift = _mk(ss_shape), _mk(ss_shape) + y = nss_op(x, weight, None, scale, shift, "rms", FLYDSL_EPS) + y_ref = _ref_norm_ss(x, weight, None, scale, shift, "rms", FLYDSL_EPS) + torch.testing.assert_close( + y, y_ref, atol=1.0, rtol=5e-2, msg=f"L={L} layout={ss}" + ) + + +def test_flydsl_imports_without_flydsl_source_tree(): + """The module must resolve against the installed FlyDSL wheel alone. + + A FlyDSL source checkout on PYTHONPATH makes its `kernels` package + importable and would hide an accidental source-tree dependency, so re-import + in a subprocess with those entries stripped from the path. + """ + _flydsl_ops() + + clean = [ + d + for d in sys.path + if d + and not os.path.isfile(os.path.join(d, "kernels", "common", "buffer_ops.py")) + ] + env = dict(os.environ, PYTHONPATH=os.pathsep.join(clean)) + code = ( + "import importlib, sys;" + f"m = importlib.import_module('{FLYDSL_MODULE}');" + "assert 'kernels' not in sys.modules, 'leaked FlyDSL source-tree kernels package';" + "print('OK', m.FLYDSL_NORM_MIN_ALIGNED_DIM)" + ) + r = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, env=env + ) + assert r.returncode == 0, f"stdout={r.stdout}\nstderr={r.stderr}" + assert f"OK {FLYDSL_D}" in r.stdout, r.stdout + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"]))