[AMD][Diffusion] Migrate FlyDSL fused norm kernels to the v0.3.0 stable API (#36349)

Co-authored-by: Phil Li <haicli@amd.com>

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.expr import const_expr, range_constexpr
-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
 
 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):
+# 128-bit vector copies over bf16 elements.
-    return buffer_ops._unwrap_value(x)
+_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,
+        total_rows: fx.Int32,
-        gate_stride: Int32,
+        gate_stride: fx.Int32,
-        scale_stride: Int32,
+        scale_stride: fx.Int32,
-        shift_stride: 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
+        n_float = float(D)
-        f32 = T.f32
+        c_zero = fx.Float32(0.0)
-        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)
+        lds = fx.SharedAllocator().allocate(SharedStorage).peek()
-        ro_rsrc = buffer_ops.create_buffer_resource(res_out_ptr, max_size=True)
+        s_sum = lds.s_sum.view(fx.make_layout(NUM_WAVES, 1))
-        r_rsrc = buffer_ops.create_buffer_resource(res_ptr, max_size=True)
+        s_sq = lds.s_sq.view(fx.make_layout(NUM_WAVES, 1))
-        x_rsrc = buffer_ops.create_buffer_resource(x_ptr, max_size=True)
+        s_final = lds.s_final.view(fx.make_layout(2, 1))
-        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)
+        copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), _ELEM_BITS)
-        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)
+        y_div = _row_div(y_ptr, row)
-        c_one_f32 = arith.constant(1.0, type=f32)
+        ro_div = _row_div(res_out_ptr, row)
-        eps_val = arith.constant(1e-6, type=f32)
+        r_div = _row_div(res_ptr, row)
-        D_float = arith.constant(float(D), type=f32)
+        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
+        def wave_reduce_add(val):
-        LDS_SLOTS = NUM_WAVES * 2 + 2
+            w = val
-        ws_attr = ir.Attribute.parse("#gpu.address_space<workgroup>")
+            for i in range_constexpr(len(_SHUFFLE_OFFSETS)):
-        lds_i8_type = ir.MemRefType.get(
+                w = w + fx.gpu.shuffle_xor(w, _SHUFFLE_OFFSETS[i], WARP_SIZE)
-            [ir.ShapedType.get_dynamic_size()], T.i8, memory_space=ws_attr
+            return w
-        )
-        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)
+        # Phase 1: residual + gate*x, accumulate stats, keep f32 in registers.
-        wave_id = tid_i32 // arith.constant(WARP_SIZE, type=i32)
+        saved_ro = []
-
+        partial_sum = c_zero
-        # Phase 1: residual + gate*x, accumulate stats, save f32 in registers
+        partial_sq = c_zero
-        _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(
+            idx = tid + it * BLOCK
-                it * BLOCK * VEC, type=i32
-            )
-            off = row_off + col
 
-            r_vec = buffer_ops.buffer_load(r_rsrc, off, vec_width=VEC, dtype=bf16)
+            r_f32 = _load_vec(copy_atom, r_div, idx).to(fx.Float32)
-            x_vec = buffer_ops.buffer_load(x_rsrc, off, vec_width=VEC, dtype=bf16)
+            x_f32 = _load_vec(copy_atom, x_div, idx).to(fx.Float32)
-            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_f32 = _load_vec(copy_atom, g_div, idx).to(fx.Float32)
-                g_vec = buffer_ops.buffer_load(g_rsrc, g_off, vec_width=VEC, dtype=bf16)
+                ro_f32 = r_f32 + g_f32 * x_f32
-                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_f32 = r_f32 + x_f32
 
-            ro_bf16 = arith_ops.TruncFOp(vec_bf16_t, ro_f32).result
+            _store_vec(copy_atom, ro_div, idx, ro_f32.to(fx.BFloat16))
-            buffer_ops.buffer_store(ro_bf16, ro_rsrc, off)
+            saved_ro.append(ro_f32)
-            _saved_ro_f32.append(ro_f32)
 
             if const_expr(not is_rms):
-                v_sum = _vector.ReductionOp(
+                partial_sum = partial_sum + ro_f32.reduce(fx.ReductionOp.ADD)
-                    f32, _vector.CombiningKind.ADD, ro_f32
+            partial_sq = partial_sq + (ro_f32 * ro_f32).reduce(fx.ReductionOp.ADD)
-                ).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
+        # Stage 1: 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 = wave_reduce_add(partial_sum)
-                    w_sum, off_sh, width_c, mode=_gpu.ShuffleMode.XOR
+        w_sq = wave_reduce_add(partial_sq)
-                ).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
+        if lane_id == 0:
-        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])
+                fx.memref_store(w_sum, s_sum, wave_id)
-            sq_slot = arith_ops.AddIOp(
+            fx.memref_store(w_sq, s_sq, wave_id)
-                wave_idx,
+        fx.gpu.barrier()
-                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))
+        # Stage 2: wave 0 folds the per-wave partials and publishes the results.
-        active = arith.andi(
+        if wave_id == 0:
-            wave_0,
+            in_range = lane_id < NUM_WAVES
-            arith.cmpi(CmpIPredicate.ult, lane_id, arith.constant(NUM_WAVES, type=i32)),
+            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)
                 )
-        lane_idx = arith_ops.IndexCastOp(ir.IndexType.get(), _v(lane_id)).result
+            v_sq = wave_reduce_add(
-        lane_idx_sq = arith_ops.AddIOp(
+                in_range.select(fx.memref_load(s_sq, lane_safe), c_zero)
-            lane_idx,
+            )
-            arith_ops.ConstantOp(
+            if lane_id == 0:
-                ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES)
+                if const_expr(not is_rms):
-            ).result,
+                    fx.memref_store(v_sum, s_final, 0)
-        ).result
+                fx.memref_store(v_sq, s_final, 1)
+        fx.gpu.barrier()
+
+        total_sq = fx.memref_load(s_final, 1)
 
         if const_expr(is_rms):
-            _if_active = scf.IfOp(active, [f32], has_else=True)
+            rstd = fx.math.rsqrt(total_sq / n_float + _EPS)
-            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)
+            mean = fx.memref_load(s_final, 0) / n_float
-            with ir.InsertionPoint(_if_active.then_block):
+            var = total_sq / n_float - mean * mean
-                s_val = _memref.LoadOp(lds, [lane_idx]).result
+            rstd = fx.math.rsqrt(var + _EPS)
-                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
 
+        # Phase 2: normalize from the register cache (no HBM re-read).
         for it in range_constexpr(NUM_ITERS):
-            col = tid_i32 * arith.constant(VEC, type=i32) + arith.constant(
+            idx = tid + it * BLOCK
-                it * BLOCK * VEC, type=i32
+            ro_f32 = saved_ro[it]
-            )
-            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
+                x_hat = ro_f32 * rstd
             else:
-                centered = arith_ops.SubFOp(ro_f32, mean_splat).result
+                x_hat = (ro_f32 - mean) * rstd
-                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)
+                x_hat = x_hat * _load_vec(copy_atom, w_div, idx).to(fx.Float32)
-                w_f32 = arith_ops.ExtFOp(vec_f32_t, _v(w_vec)).result
+                x_hat = x_hat + _load_vec(copy_atom, b_div, idx).to(fx.Float32)
-                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_f32 = _load_vec(copy_atom, sc_div, idx).to(fx.Float32)
-            sc_vec = buffer_ops.buffer_load(sc_rsrc, sc_off, vec_width=VEC, dtype=bf16)
+            x_hat = x_hat * (sc_f32 + 1.0)
-            sc_f32 = arith_ops.ExtFOp(vec_f32_t, _v(sc_vec)).result
+            y_f32 = x_hat + _load_vec(copy_atom, sh_div, idx).to(fx.Float32)
-            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
+            _store_vec(copy_atom, y_div, idx, y_f32.to(fx.BFloat16))
-            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(
@@ -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=(total_rows, 1, 1), block=(BLOCK, 1, 1), stream=stream)
-        launcher.launch(
-            grid=(grid_x, 1, 1), block=(BLOCK, 1, 1), smem=LDS_BYTES, 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,
+        total_rows: fx.Int32,
-        scale_stride: Int32,
+        scale_stride: fx.Int32,
-        shift_stride: 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
+        n_float = float(D)
-        f32 = T.f32
+        c_zero = fx.Float32(0.0)
-        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)
+        lds = fx.SharedAllocator().allocate(SharedStorage).peek()
-        x_rsrc = buffer_ops.create_buffer_resource(x_ptr, max_size=True)
+        s_sum = lds.s_sum.view(fx.make_layout(NUM_WAVES, 1))
-        w_rsrc = buffer_ops.create_buffer_resource(weight_ptr, max_size=True)
+        s_sq = lds.s_sq.view(fx.make_layout(NUM_WAVES, 1))
-        b_rsrc = buffer_ops.create_buffer_resource(bias_ptr, max_size=True)
+        s_final = lds.s_final.view(fx.make_layout(2, 1))
-        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)
+        copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), _ELEM_BITS)
-        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)
+        y_div = _row_div(y_ptr, row)
-        c_one_f32 = arith.constant(1.0, type=f32)
+        x_div = _row_div(x_ptr, row)
-        eps_val = arith.constant(1e-6, type=f32)
+        if const_expr(has_weight):
-        D_float = arith.constant(float(D), type=f32)
+            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
+        def wave_reduce_add(val):
-        ws_attr = ir.Attribute.parse("#gpu.address_space<workgroup>")
+            w = val
-        lds_i8_type = ir.MemRefType.get(
+            for i in range_constexpr(len(_SHUFFLE_OFFSETS)):
-            [ir.ShapedType.get_dynamic_size()], T.i8, memory_space=ws_attr
+                w = w + fx.gpu.shuffle_xor(w, _SHUFFLE_OFFSETS[i], WARP_SIZE)
-        )
+            return w
-        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)
+        # Phase 1: load x, accumulate stats, keep f32 in registers.
-        wave_id = tid_i32 // arith.constant(WARP_SIZE, type=i32)
+        saved_x = []
-
+        partial_sum = c_zero
-        # Phase 1: load x, accumulate stats, save f32 in registers
+        partial_sq = c_zero
-        _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(
+            idx = tid + it * BLOCK
-                it * BLOCK * VEC, type=i32
+            x_f32 = _load_vec(copy_atom, x_div, idx).to(fx.Float32)
+            saved_x.append(x_f32)
+
+            if const_expr(not is_rms):
+                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):
+            w_sum = wave_reduce_add(partial_sum)
+        w_sq = wave_reduce_add(partial_sq)
+
+        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()
+
+        # 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)
                 )
-            off = row_off + col
+            v_sq = wave_reduce_add(
-
+                in_range.select(fx.memref_load(s_sq, lane_safe), c_zero)
-            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
+            if lane_id == 0:
-        lane_idx_sq = arith_ops.AddIOp(
+                if const_expr(not is_rms):
-            lane_idx,
+                    fx.memref_store(v_sum, s_final, 0)
-            arith_ops.ConstantOp(
+                fx.memref_store(v_sq, s_final, 1)
-                ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES)
+        fx.gpu.barrier()
-            ).result,
+
-        ).result
+        total_sq = fx.memref_load(s_final, 1)
 
         if const_expr(is_rms):
-            _if_active = scf.IfOp(active, [f32], has_else=True)
+            rstd = fx.math.rsqrt(total_sq / n_float + _EPS)
-            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)
+            mean = fx.memref_load(s_final, 0) / n_float
-            with ir.InsertionPoint(_if_active.then_block):
+            var = total_sq / n_float - mean * mean
-                s_val = _memref.LoadOp(lds, [lane_idx]).result
+            rstd = fx.math.rsqrt(var + _EPS)
-                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
 
+        # Phase 2: normalize from the register cache + scale_shift.
         for it in range_constexpr(NUM_ITERS):
-            col = tid_i32 * arith.constant(VEC, type=i32) + arith.constant(
+            idx = tid + it * BLOCK
-                it * BLOCK * VEC, type=i32
+            x_f32 = saved_x[it]
-            )
-            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
+                x_hat = x_f32 * rstd
             else:
-                centered = arith_ops.SubFOp(x_f32, mean_splat).result
+                x_hat = (x_f32 - mean) * rstd
-                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)
+                x_hat = x_hat * _load_vec(copy_atom, w_div, idx).to(fx.Float32)
-                w_f32 = arith_ops.ExtFOp(vec_f32_t, _v(w_vec)).result
+                x_hat = x_hat + _load_vec(copy_atom, b_div, idx).to(fx.Float32)
-                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_f32 = _load_vec(copy_atom, sc_div, idx).to(fx.Float32)
-            sc_vec = buffer_ops.buffer_load(sc_rsrc, sc_off, vec_width=VEC, dtype=bf16)
+            x_hat = x_hat * (sc_f32 + 1.0)
-            sc_f32 = arith_ops.ExtFOp(vec_f32_t, _v(sc_vec)).result
+            y_f32 = x_hat + _load_vec(copy_atom, sh_div, idx).to(fx.Float32)
-            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
+            _store_vec(copy_atom, y_div, idx, y_f32.to(fx.BFloat16))
-            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(
@@ -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=(total_rows, 1, 1), block=(BLOCK, 1, 1), stream=stream)
-        launcher.launch(
-            grid=(grid_x, 1, 1), block=(BLOCK, 1, 1), smem=LDS_BYTES, stream=stream
-        )
 
     return launch_norm_ss
 

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):
+def _flydsl_ops():
-    if residual is not None:
+    """Resolve the FlyDSL exports, skipping when the installed FlyDSL is too old.
-        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
 
+    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])
+    The kernel module raises ImportError when the stable FlyDSL surface it needs
-@pytest.mark.parametrize(
+    is absent, which is also how ``layernorm.py`` detects that it must fall back
-    "norm_type,B,L",
+    to the native path.  A runner image predating that surface should skip here
-    [("rms", 1, 16), ("rms", 2, 16), ("layer", 2, 16), ("rms", 1, 90000)],
+    rather than report a kernel regression.
-)
+    """
-def test_flydsl_norm_scale_shift(with_residual, norm_type, B, L):
     _require_rocm()
-    # Imported inside the test: the FlyDSL compiler only exists on ROCm, and
+    try:
-    # 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,
         )
+    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)
+        _mk((FLYDSL_D,), torch.float32)
-        if norm_type == "layer"
+        if (has_weight and norm_type == "layer")
         else None
     )
-    scale = torch.randn(B, 1, FLYDSL_D, device=DEVICE, dtype=torch.bfloat16)
+    scale = _mk(ss_shape)
-    shift = torch.randn(B, 1, FLYDSL_D, device=DEVICE, dtype=torch.bfloat16)
+    shift = _mk(ss_shape)
 
-    if with_residual:
+    y, res_out = fused_op(
-        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(
+    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, res_ref, atol=5e-2, rtol=5e-2)
+    torch.testing.assert_close(res_out, 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
-        )
     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"]))