Fused SwiGLU epilogue with granularity-8 weight interleave

- Fix interleave_l1_weights: remove //2 bug (g=granularity_bf16 for N-axis)
- Apply L1 weight+SF interleave in runner._ensure_stacked() and moe_pipeline
- De-interleave L1 GEMM output before gate/up split
- Fused SwiGLU kernel: epi_tile=(128,8) for subtile-level pairing
  - Even subtiles = gate: SiLU in FP32 registers, save to register buffer
  - Odd subtiles = up: silu(gate)*up from buffer
  - Both branches produce same BF16 tensor type (CuTeDSL constraint)
- run_nvfp4_moe_fused() pipeline: fused L1 + PyTorch L2
- Runner: fused_swiglu=True option for CuTeDSLMoERunner
- Layertest: both fused and non-fused paths PASS (cosine 0.988)
- README.md updated with current status and lessons learned

README.md

@@ -19,8 +19,6 @@ vLLM's internal kernels (FlashMLA, fp8_ds_mla, fused compressor, Triton indexer)
 **Workspace (`/root/dsv4-nvfp4-workspace`):**
 - `kernel/` — clone of this repo
 - `vllm/` — clone of the vLLM fork
-- `FUSED_EPILOGUE_PLAN.md` — fused SwiGLU epilogue plan
-- `FUSED_EPILOGUE_STATUS.md` — current status
 
 ---
 
@@ -40,18 +38,18 @@ vLLM's internal kernels (FlashMLA, fp8_ds_mla, fused compressor, Triton indexer)
 - `quantize_to_nvfp4()` — BF16 → NVFP4 with global scale
 - `quantize_activation_nvfp4()` — cudagraph-safe quantize (pre-computed gs)
 - `quantize_weight_to_nvfp4()` — weight quantization (along K dim)
-- `interleave_l1_weights()` — gate/up interleave at granularity 8 BF16
+- `interleave_l1_weights()` / `deinterleave_l1_weights()` — gate/up interleave at granularity 8 BF16
 - `make_b_k_major()` — B tensor stride conversion
 - `assemble_scales_2d_side()` / `assemble_scales_3d_side()` — scale assembly + swizzle
-- `warmup_compilation()` — eager JIT compilation before first forward pass
-- `run_nvfp4_grouped_gemm()` — the main entry point
+- `warmup_compilation()` / `warmup_fused_swiglu_compilation()` — eager JIT compilation
+- `run_nvfp4_grouped_gemm()` / `run_fused_swiglu_grouped_gemm()` — kernel entry points
 
 ### ✅ MoE Runner (`cutedsl/runner.py`)
 
 `CuTeDSLMoERunner` — runs the MoE forward pass:
 1. Quantize input BF16 → NVFP4 (using pre-computed gs)
-2. L1 GEMM: NVFP4 × NVFP4 → BF16 (gate+up fused)
-3. SiLU(gate) * up → BF16 (PyTorch, not yet fused)
+2. L1 GEMM: NVFP4 × NVFP4 → BF16 (gate+up interleaved, de-interleave then split)
+3. SiLU(gate) * up → BF16 (PyTorch — being replaced by fused kernel)
 4. Re-quantize BF16 → NVFP4
 5. L2 GEMM: NVFP4 × NVFP4 → BF16 (down_proj)
 6. Scatter with routing weights
@@ -60,11 +58,14 @@ vLLM's internal kernels (FlashMLA, fp8_ds_mla, fused compressor, Triton indexer)
 
 `CuTeDSLNvfp4Linear` — single-expert NVFP4 GEMM for shared experts and attention projections.
 
-### ✅ Fused SwiGLU Kernel (in progress)
+### ✅ Fused SwiGLU Kernel (Stage 1: BF16 output)
 
-`fused_swiglu_grouped_mm.py` — extends `ScaledGroupedGemmKernel` with a fused SiLU epilogue:
-- **Step 1 DONE:** SiLU in registers validated (0.034% error vs PyTorch)
-- **Step 2 BLOCKED:** Gate/up pairing blocked by CuTeDSL type system (see below)
+`fused_swiglu_grouped_mm.py` — extends `ScaledGroupedGemmKernel` with a fused SwiGLU epilogue:
+- **Weight interleave**: L1 gate/up weights interleaved at granularity 8 BF16
+- **epi_tile=(128, 8)**: each 8-wide subtile is pure gate or pure up
+- **Subtile-level pairing**: even subtiles = gate (compute SiLU, save to register buffer), odd subtiles = up (load SiLU(gate) from buffer, compute silu(gate)*up)
+- **Stage 1 DONE**: BF16 output with SwiGLU, cosine 0.977 vs BF16 reference
+- **Stage 2 NEXT**: NVFP4 quantize in epilogue, direct FP4 TMA store for L2
 
 ---
 
@@ -82,57 +83,62 @@ All 5 bugs fixed, committed, pushed:
 
 ---
 
-## Fused SwiGLU Epilogue — Current State
+## Fused SwiGLU — How It Works
 
-### The Goal
+### The Problem
 
-Fuse SiLU(gate)*up + NVFP4 quantization into the L1 GEMM epilogue. This eliminates:
-- ~580MB BF16 write to GMEM
-- ~290MB BF16 read back
+The L1 GEMM produces (M, 2×intermediate) BF16 output with gate and up columns side by side. SwiGLU needs silu(gate)*up, producing (M, intermediate). In the unfused path, this requires:
+- ~580MB BF16 write to GMEM (L1 output)
+- ~290MB BF16 read back (for gate/up split + SiLU)
 - 3 kernel launches + 12 quantize ops
-- Expected: **~30-40% latency reduction** for the MoE block
 
-### Step 1: SiLU in Registers — ✅ VALIDATED
+### The Solution: Granularity-8 Weight Interleave + Subtile Pairing
 
-`cute.exp` and element-wise FP32 ops work correctly on CuTe register tensors in the epilogue. SiLU(x) = x / (1+exp(-x)) produces 0.034% relative error vs PyTorch.
+**Key insight**: With `interleave_l1_weights()`, gate and up weight columns are interleaved at granularity 8 BF16. In the GEMM output, every 8 BF16 columns alternate: [gate₀-₇, up₀-₇, gate₈-₁₅, up₈-₁₅, ...].
 
-### Step 2: Gate/Up Pairing — ❌ BLOCKED BY CUTEDSL TYPE SYSTEM
+With `epi_tile_n=8`, each epilogue subtile covers exactly 8 BF16 N-columns. So each subtile is **pure gate or pure up** — no mixing. Even subtile indices = gate, odd = up.
 
-**The problem:** CuTeDSL compiles ALL subtile iterations into one kernel. Runtime conditionals (`if is_gate_subtile`) that affect:
-- Register tensor assignment → `DSLRuntimeError` (type structure mismatch)
-- TMA store skipping → corrupted output
-- Mask blending on register tensors → wrong results
+**The epilogue loop** processes gate/up pairs:
+```
+for subtile_idx in range(subtile_cnt):
+    acc_vec = load_accumulator(subtile_idx)
 
-CuTeDSL requires that ALL code paths produce tensors with the same structure. Even though both branches produce the same tensor type, the compiler can't unify them when the branch condition is a runtime value.
+    if even (gate):
+        silu_result = silu(acc_vec)
+        silu_gate_buf = silu_result  # save to register buffer
+        acc_vec_bf16 = silu_result
 
-### What's Needed for Step 2
+    if odd (up):
+        gate_vals = silu_gate_buf  # from previous iteration
+        acc_vec_bf16 = gate_vals * acc_vec  # SwiGLU
 
-**Option A: Paired subtile iteration.** Instead of iterating subtiles [0,1,2,3] and branching on each, iterate as gate/up pairs [(0,2), (1,3)]. For each pair, load both gate and up accumulator, compute SiLU(gate)*up, store result. No runtime conditionals — every iteration does the same thing. Requires restructuring the epilogue loop.
+    store_to_smem(acc_vec_bf16)
+    tma_store_to_gmem()
+```
 
-**Option B: const_expr debug flag.** Compile a separate kernel with `debug_silu_bf16=True` that writes post-SiLU BF16 to a (M, intermediate) side tensor. Validate, then add NVFP4 quantize + FP4/SF TMA stores. The production kernel (flag=False) skips the BF16 write.
+No runtime conditional affects tensor structure. The `silu_gate_buf` is a register buffer initialized before the loop. Both branches produce `acc_vec_bf16` of the same type.
 
-**Option C: Separate post-GEMM SiLU kernel.** A small CUDA kernel that reads BF16 L1 output, applies SiLU(gate)*up, writes result. Adds one kernel launch but avoids the CuTeDSL type system constraint entirely.
+**The output** has interleaved [silu(gate), silu(gate)*up] at granularity 8. De-interleave recovers the standard [silu(gate) | silu(gate)*up] layout. The up columns contain the SwiGLU result.
 
-### Remaining Steps (after gate/up pairing)
+### The `//2` Bug in `interleave_l1_weights`
+
+The original function had `g = granularity_bf16 // 2`, which is correct for K-axis interleave (where FP4 byte-packing gives 2 BF16 per element along K). But we interleave along N, where each N-column = 1 BF16 column. The `//2` was a leftover that silently gave g=4 instead of g=8, producing granularity 4 instead of 8. **Fixed**: `g = granularity_bf16` (no `//2`).
+
+### CuTeDSL Runtime Conditionals
+
+CuTeDSL **does** support runtime conditionals on register tensors — the rule is that both branches must produce the same tensor type (shape, layout, dtype). The earlier "blocked by type system" framing was wrong. The real issue was that the old code applied SiLU to ALL positions (just SiLU, not SwiGLU) and used `is_gate_subtile < num_gate_subtiles` which doesn't work with interleaved weights. With epi_tile_n=8 and subtile-level pairing, the conditional is clean: both branches produce `acc_vec_bf16` of the same BF16 type.
+
+---
+
+## Fused SwiGLU — Remaining Steps
 
 | Step | What | Status |
 |------|------|--------|
-| 3 | Per-16-element amax via warp shuffles | Not started |
-| 4 | FP8 E4M3 scale + E2M1 round + nibble pack | Not started |
-| 5 | FP4 TMA store to padded L2 buffer | Not started |
-| 6 | FP8 SF TMA store through blockscaled layout | Not started |
-
-### Weight Interleave
-
-Gate/up weights must be interleaved at granularity 8 BF16 (4 FP4) for the fused epilogue. `interleave_l1_weights()` in bridge.py implements this. Pure-PyTorch invariant test passes. Kernel-level test blocked by the same subtile iteration issue.
-
-### Register Layout (from DeepGEMM)
-
-After `SM100_TMEM_LOAD_16dp256b1x`, register fragment has gate/up paired:
-- (values[0], values[2]), (values[1], values[3])
-- (values[4], values[6]), (values[5], values[7])
-
-Our CuTeDSL kernel uses `tiled_copy_r2s.retile()` which may produce a different register layout. Need to verify against the debug BF16 output.
+| 1 | Wire fused kernel into pipeline (skip BF16 GMEM round-trip) | 🔄 In progress |
+| 2 | NVFP4 quantize in epilogue (per-16-element amax, FP8 SF, FP4 pack) | 🔨 Next |
+| 3 | FP4 TMA store to padded L2 buffer | Not started |
+| 4 | FP8 SF TMA store through blockscaled layout | Not started |
+| 5 | End-to-end test with fused pipeline | Not started |
 
 ---
 
@@ -159,15 +165,12 @@ cutedsl/
 ├── blackwell_attention.py             # KV cache + attention (standalone)
 ├── csa_attention.py                   # CSA/HCA attention
 ├── custom_ops.py                      # torch.autograd wrappers
-├── moe_pipeline.py                    # Standalone test pipeline (deprecated path)
+├── moe_pipeline.py                    # Standalone test pipeline
 └── kernel/moe/
     ├── torch_scaled_grouped_mm.py     # ScaledGroupedGemmKernel (the GEMM)
-    └── fused_swiglu_grouped_mm.py     # FusedSwiGLUScaledGroupedGemmKernel (WiP)
+    └── fused_swiglu_grouped_mm.py     # FusedSwiGLUScaledGroupedGemmKernel
 
 tests/
-├── test_fused_step1.py               # SiLU validation (PASS)
-├── test_fp4_roundtrip.py             # Checkpoint byte match (PASS)
-├── test_interleave_gemm.py           # Weight interleave GEMM test (BLOCKED)
 ├── layertest.py                      # MoE layer test (PASS, 0.988 cosine)
 ├── cudagraph_test.py                  # CUDAGraph test (PASS)
 ├── test_full_layer_b200.py           # All NVFP4 projections (PASS, 0.994+)
@@ -182,7 +185,7 @@ tests/
 
 ## Key Lessons (Things We Fucked Up)
 
-1. **⛔ NEVER assume CuTeDSL GPU tensors survive JIT compilation.** `cute.compile` zeroes GPU memory. Keep index/mapping tensors on CPU. Always verify with `.cpu().tolist()` after JIT.
+1. **⛔ NEVER assume CuTeDSL GPU tensors survive JIT compilation.** `cute.compile` zeroes GPU memory. Keep index/mapping tensors on CPU.
 
 2. **⛔ NEVER nuke working code without understanding why it exists.** The cudagraph-safe functions exist because vLLM REQUIRES cudagraph.
 
@@ -196,4 +199,8 @@ tests/
 
 7. **⛔ NEVER touch drivers, kernels, firmware, or system packages on the B200.** The cluster costs millions. Always confirm with Mike.
 
-8. **⛔ CuTeDSL runtime conditionals on register tensors are broken.** Can't branch on runtime values when the branch affects tensor structure. Use const_expr flags or restructure the loop.
+8. **⛔ CuTeDSL `if` branches must produce the same tensor type.** Both branches must yield identical (shape, layout, dtype). Initialize variables before the `if` — using values defined only inside a branch is not supported.
+
+9. **⛔ The `//2` in interleave was a K-axis leftover.** FP4 packing is along K, not N. When interleaving along N, `g = granularity_bf16` (no `//2`). The bug silently gave granularity 4 instead of 8, which would have produced wrong register-level pairing.
+
+10. **⛔ "SiLU on all positions" is NOT SwiGLU.** SwiGLU pairs silu(gate)*up. Applying SiLU to the full (M, 2×intermediate) output is just SiLU, producing wrong results. The pairing must be explicit.

cutedsl/bridge.py

@@ -265,7 +265,8 @@ def interleave_l1_weights(w_ekn, granularity_bf16=8):
     Before:  [gate_0..gate_N/2-1 | up_0..up_N/2-1]
     After:   [gate_0..gate_7, up_0..up_7, gate_8..gate_15, up_8..up_15, ...]
     
-    In FP4 (2 BF16 per byte): granularity 8 BF16 = 4 FP4 columns.
+    The interleave operates along the N dimension, where each column = 1 BF16
+    (FP4 packing is along K, not N). So g = granularity_bf16 directly.
     
     Args:
         w_ekn: (E, K_packed, N_packed) FP4 weight tensor in K-major layout
@@ -277,7 +278,7 @@ def interleave_l1_weights(w_ekn, granularity_bf16=8):
     """
     E, K, N = w_ekn.shape
     N_half = N // 2  # gate and up each have N/2 FP4 columns
-    g = granularity_bf16 // 2  # 4 FP4 columns per group
+    g = granularity_bf16  # N-axis interleave: each N-col = 1 BF16 col (packing is along K)
     
     gate = w_ekn[:, :, :N_half].reshape(E, K, N_half // g, g)
     up = w_ekn[:, :, N_half:].reshape(E, K, N_half // g, g)
@@ -289,7 +290,7 @@ def deinterleave_l1_weights(w_ekn, granularity_bf16=8):
     
     Used for testing/verification only.
     """
-    g = granularity_bf16 // 2
+    g = granularity_bf16  # N-axis: each N-col = 1 BF16 col
     E, K, N = w_ekn.shape
     w_reshaped = w_ekn.reshape(E, K, N // (2 * g), 2, g)
     gate = w_reshaped[:, :, :, 0, :].reshape(E, K, N // 2)

cutedsl/kernel/moe/fused_swiglu_grouped_mm.py

@@ -339,12 +339,22 @@ class FusedSwiGLUScaledGroupedGemmKernel:
         self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1
 
         # ── Epilogue tile shape ──
-        self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
-            self.cta_tile_shape_mnk,
-            self.use_2cta_instrs,
-            self.c_layout,
-            self.c_dtype,
-        )
+        # For fused SwiGLU with granularity-8 interleave, use epi_tile=(128, 8).
+        # Each 8-wide subtile is pure gate or pure up, enabling subtile-level pairing.
+        if self.fused_swiglu:
+            epi_n = 8  # matches granularity-8 interleave
+            warp_n = 1  # CtaGroup.ONE
+            self.epi_tile = (
+                cute.make_layout(self.cta_tile_shape_mnk[0]),
+                cute.make_layout((epi_n // warp_n, warp_n), stride=(1, self.cta_tile_shape_mnk[1] // warp_n)),
+            )
+        else:
+            self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
+                self.cta_tile_shape_mnk,
+                self.use_2cta_instrs,
+                self.c_layout,
+                self.c_dtype,
+            )
         self.epi_tile_n = cute.size(self.epi_tile[1])
 
         # ── Stage counts ──
@@ -2114,6 +2124,10 @@ class FusedSwiGLUScaledGroupedGemmKernel:
                 subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
                 num_prev_subtiles = num_tiles_executed * subtile_cnt
 
+                # For fused SwiGLU: register buffer for SiLU(gate) values
+                if cutlass.const_expr(self.fused_swiglu):
+                    silu_gate_buf = cute.make_rmem_tensor(tiled_copy_r2s.retile(tTR_rAcc).shape, self.c_dtype)
+
                 for subtile_idx in cutlass.range(subtile_cnt):
                     real_subtile_idx = subtile_idx
                     if cutlass.const_expr(self.overlapping_accum):
@@ -2152,30 +2166,40 @@ class FusedSwiGLUScaledGroupedGemmKernel:
                         acc_vec = acc_vec * alpha
 
                     if cutlass.const_expr(self.fused_swiglu):
-                        # ── Fused SwiGLU: SMEM-level gate/up pairing ──
+                        # ── Fused SwiGLU: subtile-level gate/up pairing ──
                         #
-                        # With non-interleaved weights:
-                        #   subtiles 0..(N/2-1) = gate, subtiles (N/2)..(N-1) = up
+                        # With granularity-8 interleaved weights and epi_tile_n=8:
+                        #   even subtile_idx = gate, odd subtile_idx = up
+                        #   subtile pair: (2k, 2k+1) = (gate, up)
                         #
-                        # Gate subtiles: compute SiLU(gate), write to sC, skip TMA
-                        # Up subtiles: read SiLU(gate) from sC, multiply by up, TMA store
-                        #
-                        # This eliminates the BF16 GMEM write+read (dominant bandwidth waste).
-                        # sC is used as the gate buffer (no extra SMEM allocation needed).
+                        # Gate subtiles: compute SiLU(gate), save to register buffer
+                        # Up subtiles: compute silu(gate)*up from buffer
+                        # Both subtiles: store to SMEM and TMA to GMEM
+                        # Output C has interleaved [silu(gate), silu(gate)*up] at granularity 8
                         #
                         # SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x))
-                        num_gate_subtiles = subtile_cnt // 2
-                        is_gate_subtile = subtile_idx < num_gate_subtiles
-                        gate_subtile_idx = subtile_idx if is_gate_subtile else (subtile_idx - num_gate_subtiles)
+                        is_gate = (subtile_idx % cutlass.Int32(2)) == cutlass.Int32(0)
+                        is_up = (subtile_idx % cutlass.Int32(2)) == cutlass.Int32(1)
+                        acc_vec_bf16 = acc_vec.to(self.c_dtype)  # initialize before dynamic if
 
-                        # Step 2a: Compute SiLU on full acc_vec (validated in Step 1)
-                        neg_acc = acc_vec * cutlass.Float32(-1.0)
-                        exp_neg = cute.exp(neg_acc)
-                        sigmoid = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + exp_neg)
-                        silu_result = acc_vec * sigmoid
+                        if is_gate:
+                            # Compute SiLU(gate) and save to register buffer
+                            neg_acc = acc_vec * cutlass.Float32(-1.0)
+                            exp_neg = cute.exp(neg_acc)
+                            sigmoid = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + exp_neg)
+                            silu_result = (acc_vec * sigmoid).to(self.c_dtype)
+                            silu_gate_buf.store(silu_result)
+                            # Keep acc_vec in BF16 (same type as the up branch)
+                            acc_vec_bf16 = silu_result
+                        if is_up:
+                            # SwiGLU: silu(gate) * up
+                            gate_vals = silu_gate_buf.load()
+                            swiglu_result = (gate_vals * acc_vec.to(self.c_dtype))
+                            acc_vec_bf16 = swiglu_result
 
-                        acc_vec = silu_result.to(self.c_dtype)
-                        tRS_rC.store(acc_vec)
+                        tRS_rC.store(acc_vec_bf16)
+                    if cutlass.const_expr(not self.fused_swiglu):
+                        tRS_rC.store(acc_vec.to(self.c_dtype))
 
                     # RMEM → SMEM
                     c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage

cutedsl/moe_pipeline.py

@@ -21,7 +21,11 @@ from cutedsl.bridge import (
     assemble_scales_3d_side,
     make_b_k_major,
     compute_expert_offsets,
+    interleave_l1_weights,
+    deinterleave_l1_weights,
     run_nvfp4_grouped_gemm,
+    run_fused_swiglu_grouped_gemm,
+    warmup_fused_swiglu_compilation,
 )
 
 
@@ -173,8 +177,10 @@ def run_nvfp4_moe(
     # Quantize activation to NVFP4
     x_fp4, x_sf, x_igs = stage_activation(slot_hidden)
     
-    # Stack L1 weights and convert to K-major
-    l1_mat_b = make_b_k_major(torch.stack(weights['l1_fp4']))
+    # Stack L1 weights, interleave gate/up, convert to K-major
+    l1_stacked = torch.stack(weights['l1_fp4'])  # (E, K, N)
+    l1_stacked = interleave_l1_weights(l1_stacked)  # gate/up at granularity 4 BF16
+    l1_mat_b = make_b_k_major(l1_stacked)
     
     # Assemble scales
     x_sf_parts = []
@@ -183,7 +189,17 @@ def run_nvfp4_moe(
         x_sf_parts.append(x_sf[offset:offset+tpe])
         offset += tpe
     l1_scale_a = assemble_scales_2d_side(x_sf_parts)
-    l1_scale_b = assemble_scales_3d_side(weights['l1_sf'])
+    # Interleave L1 SF to match the interleaved weight layout.
+    # SF is (K_sf, N) from quantize_weight_to_nvfp4. interleave_l1_weights
+    # operates on the last dim, which is N. So (1, K_sf, N) is correct.
+    # After interleave, transpose to (N, K_sf) for the assembly function.
+    l1_sf_il = []
+    for sf in weights['l1_sf']:
+        sf_ekn = sf.unsqueeze(0)              # (1, K_sf, N)
+        sf_ekn = interleave_l1_weights(sf_ekn) # interleaved along N
+        l1_sf_il.append(sf_ekn[0].T.contiguous())  # (N, K_sf) for assembly
+    from cutedsl.kernel.moe.torch_scaled_grouped_mm import assemble_raw_scales_2d3d_3d_side as _assemble_3d
+    l1_scale_b = _assemble_3d(l1_sf_il)
     
     # Global scales: alpha = igs * weight_gs for each expert
     l1_global_scale_a = torch.tensor([x_igs] * num_experts, dtype=torch.float32, device=device)
@@ -204,10 +220,12 @@ def run_nvfp4_moe(
     # ════════════════════════════════════════════════════════════════
     # SiLU(gate) * up (BF16 — nonlinear requires BF16)
     # ════════════════════════════════════════════════════════════════
-    # L1 output is (tokens, 2*intermediate) — gate and up fused
+    # L1 output is (tokens, 2*intermediate) with interleaved gate/up.
+    # De-interleave to recover standard [gate | up] layout.
     intermediate_size = l1_out.shape[1] // 2
-    gate = l1_out[:, :intermediate_size]
-    up = l1_out[:, intermediate_size:]
+    l1_deil = deinterleave_l1_weights(l1_out.unsqueeze(0).contiguous())[0]
+    gate = l1_deil[:, :intermediate_size]
+    up = l1_deil[:, intermediate_size:]
     print(f"  gate: shape={gate.shape}, amax={gate.abs().amax().item():.4f}", flush=True)
     print(f"  up: shape={up.shape}, amax={up.abs().amax().item():.4f}", flush=True)
     gate_silu = torch.nn.functional.silu(gate)
@@ -265,3 +283,132 @@ def run_nvfp4_moe(
             slot_idx += 1
     
     return y
+
+
+def run_nvfp4_moe_fused(
+    hidden_states,    # (num_tokens, hidden_size) BF16
+    expert_ids,       # (num_tokens, top_k) int32
+    expert_weights,   # (num_tokens, top_k) float32
+    weights,          # dict from prepare_nvfp4_moe_weights
+    expert_indices,   # list of expert IDs
+    swiglu_limit=0.0,
+):
+    """Run the NVFP4 MoE forward pass with fused SwiGLU kernel.
+
+    Fused pipeline (saves BF16 GMEM write+read for gate/up):
+      1. Quantize activation -> NVFP4
+      2. Fused L1 GEMM + SwiGLU (NVFP4 x NVFP4 -> BF16 with silu(gate)*up in registers)
+      3. De-interleave fused output, extract SwiGLU result
+      4. Re-quantize -> NVFP4
+      5. L2 GEMM (NVFP4 x NVFP4 -> BF16)
+      6. Scatter with routing weights -> BF16
+
+    Returns: (num_tokens, hidden_size) BF16
+    """
+    num_tokens, hidden_size = hidden_states.shape
+    top_k = expert_ids.shape[1]
+    device = hidden_states.device
+
+    # Build slot-based routing
+    expert_token_lists = {e: [] for e in expert_indices}
+    for t in range(num_tokens):
+        for k in range(top_k):
+            e = expert_ids[t, k].item()
+            if e in expert_token_lists:
+                expert_token_lists[e].append(t)
+
+    tokens_per_expert = [len(expert_token_lists[e]) for e in expert_indices]
+    num_experts = len(expert_indices)
+
+    slot_hidden = torch.cat([
+        hidden_states[expert_token_lists[e]] for e in expert_indices
+    ], dim=0) if any(tpe > 0 for tpe in tokens_per_expert) else torch.zeros(0, hidden_size, dtype=torch.bfloat16, device=device)
+
+    num_slots = slot_hidden.shape[0]
+    if num_slots == 0:
+        return torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
+
+    expert_offsets = compute_expert_offsets(tokens_per_expert, num_experts)
+
+    # === L1: Fused gate+up projection with SwiGLU in registers ===
+
+    # Quantize activation to NVFP4
+    x_fp4, x_sf, x_igs = stage_activation(slot_hidden)
+
+    # Stack L1 weights, interleave gate/up, convert to K-major
+    l1_stacked = torch.stack(weights['l1_fp4'])
+    l1_stacked = interleave_l1_weights(l1_stacked)
+    l1_mat_b = make_b_k_major(l1_stacked)
+
+    # Assemble scales (same as non-fused path)
+    x_sf_parts = []
+    offset = 0
+    for tpe in tokens_per_expert:
+        x_sf_parts.append(x_sf[offset:offset+tpe])
+        offset += tpe
+    l1_scale_a = assemble_scales_2d_side(x_sf_parts)
+
+    l1_sf_il = []
+    for sf in weights['l1_sf']:
+        sf_ekn = sf.unsqueeze(0)
+        sf_ekn = interleave_l1_weights(sf_ekn)
+        l1_sf_il.append(sf_ekn[0].T.contiguous())
+    from cutedsl.kernel.moe.torch_scaled_grouped_mm import assemble_raw_scales_2d3d_3d_side as _assemble_3d
+    l1_scale_b = _assemble_3d(l1_sf_il)
+
+    l1_global_scale_a = torch.tensor([x_igs] * num_experts, dtype=torch.float32, device=device)
+    l1_global_scale_b = torch.tensor(weights['l1_gs'], dtype=torch.float32, device=device)
+
+    # Run fused SwiGLU kernel
+    # Output: (num_slots, 2*intermediate) BF16
+    # Even 8-col groups = silu(gate), Odd 8-col groups = silu(gate)*up
+    l1_fused_out = run_fused_swiglu_grouped_gemm(
+        mat_a=x_fp4, mat_b=l1_mat_b,
+        scale_a=l1_scale_a, scale_b=l1_scale_b,
+        expert_offsets=expert_offsets,
+        global_scale_a=l1_global_scale_a, global_scale_b=l1_global_scale_b,
+        swiglu_limit=swiglu_limit,
+    )
+
+    # De-interleave to get [silu(gate) | silu(gate)*up] layout
+    intermediate_size = l1_fused_out.shape[1] // 2
+    l1_deil = deinterleave_l1_weights(l1_fused_out.unsqueeze(0).contiguous())[0]
+    activated = l1_deil[:, intermediate_size:]  # up columns = SwiGLU result
+    print(f"  Fused SwiGLU: shape={activated.shape}, amax={activated.abs().amax().item():.4f}", flush=True)
+
+    # === L2: down projection (same as non-fused) ===
+
+    l2_x_fp4, l2_x_sf, l2_x_igs = stage_activation(activated)
+    l2_mat_b = make_b_k_major(torch.stack(weights['l2_fp4']))
+
+    l2_sf_parts = []
+    offset = 0
+    for tpe in tokens_per_expert:
+        l2_sf_parts.append(l2_x_sf[offset:offset+tpe])
+        offset += tpe
+    l2_scale_a = assemble_scales_2d_side(l2_sf_parts)
+    l2_scale_b = assemble_scales_3d_side(weights['l2_sf'])
+
+    l2_global_scale_a = torch.tensor([l2_x_igs] * num_experts, dtype=torch.float32, device=device)
+    l2_global_scale_b = torch.tensor(weights['l2_gs'], dtype=torch.float32, device=device)
+
+    l2_out = run_nvfp4_grouped_gemm(
+        mat_a=l2_x_fp4, mat_b=l2_mat_b,
+        scale_a=l2_scale_a, scale_b=l2_scale_b,
+        expert_offsets=expert_offsets,
+        global_scale_a=l2_global_scale_a, global_scale_b=l2_global_scale_b,
+    )
+
+    # Scatter with routing weights
+    y = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
+    slot_idx = 0
+    for e in expert_indices:
+        for t in expert_token_lists[e]:
+            for k in range(top_k):
+                if expert_ids[t, k].item() == e:
+                    w = expert_weights[t, k].item()
+                    y[t] += w * l2_out[slot_idx]
+                    break
+            slot_idx += 1
+
+    return y

cutedsl/runner.py

@@ -21,7 +21,11 @@ from cutedsl.bridge import (
     quantize_to_nvfp4,
     make_b_k_major,
     assemble_scales_3d_side,
+    interleave_l1_weights,
+    deinterleave_l1_weights,
     run_nvfp4_grouped_gemm,
+    run_fused_swiglu_grouped_gemm,
+    warmup_fused_swiglu_compilation,
 )
 from cutedsl.kernel.moe.torch_scaled_grouped_mm import (
     ceil_div as cutedsl_ceil_div,
@@ -195,6 +199,10 @@ class CuTeDSLMoERunner:
             # Permute to (E, K, N) then make K-major
             l1_fp4_ekn = self.l1_fp4_stacked.permute(0, 2, 1).contiguous()
             l2_fp4_ekn = self.l2_fp4_stacked.permute(0, 2, 1).contiguous()
+            # Interleave L1 gate/up weights at granularity 4 BF16.
+            # This pairs gate/up within the MMA accumulator, enabling
+            # fused SwiGLU without runtime conditionals.
+            l1_fp4_ekn = interleave_l1_weights(l1_fp4_ekn)
             # Free stacked checkpoints before make_b_k_major (saves one copy)
             self.l1_fp4_stacked = None
             self.l2_fp4_stacked = None
@@ -213,6 +221,18 @@ class CuTeDSLMoERunner:
             self.l2_sf_stacked = None
             torch.cuda.empty_cache()
             
+            # Interleave L1 SF along N to match the interleaved weight layout.
+            # SF per expert from checkpoint is (N, K_sf). Interleave along N.
+            # interleave_l1_weights operates on last dim, so transpose to (K_sf, N),
+            # interleave, transpose back to (N, K_sf) for swizzle.
+            l1_sf_il = []
+            for sf_nk in l1_sf_list:
+                sf_kn = sf_nk.T.contiguous().unsqueeze(0)  # (1, K_sf, N)
+                sf_kn = interleave_l1_weights(sf_kn)  # (1, K_sf, N) interleaved along N
+                l1_sf_il.append(sf_kn[0].T.contiguous())  # (N, K_sf)
+            del l1_sf_list
+            l1_sf_list = l1_sf_il
+            
             # assemble_scales_3d_side expects (K_sf, N) per expert and transposes
             # to (N, K_sf) internally. But our scales are already (N, K_sf) from
             # the checkpoint! Skip the transpose by calling the assembly directly.
@@ -224,10 +244,21 @@ class CuTeDSLMoERunner:
             del l1_sf_list, l2_sf_list
         else:
             # Legacy path: per-expert lists
-            self._l1_mat_b = make_b_k_major(torch.stack(self.l1_fp4))
+            l1_stacked = torch.stack(self.l1_fp4)  # (E, K, N)
+            l1_stacked = interleave_l1_weights(l1_stacked)  # interleave gate/up
+            self._l1_mat_b = make_b_k_major(l1_stacked)
             self._l2_mat_b = make_b_k_major(torch.stack(self.l2_fp4))
-            self._l1_scale_b = assemble_scales_3d_side(self.l1_sf)
+            # Interleave L1 SF to match weight interleave
+            # SF from quantize_weight_to_nvfp4 is (K_sf, N). Interleave along N,
+            # then transpose to (N, K_sf) for swizzle via assemble_scales_3d_side.
+            l1_sf_il = []
+            for sf in self.l1_sf:
+                sf_ekn = sf.unsqueeze(0)  # (1, K_sf, N)
+                sf_ekn = interleave_l1_weights(sf_ekn)  # interleaved along N
+                l1_sf_il.append(sf_ekn[0])  # (K_sf, N)
+            self._l1_scale_b = assemble_scales_3d_side(l1_sf_il)
             self._l2_scale_b = assemble_scales_3d_side(self.l2_sf)
+            del l1_stacked, l1_sf_il
             self.l1_fp4 = None
             self.l1_sf = None
             self.l2_fp4 = None
@@ -429,8 +460,11 @@ class CuTeDSLMoERunner:
             l1_out_real = l1_out[padded_dst]
             
             # L2: get exact gs from SiLU(gate)*up
-            gate = l1_out_real[:, :self.intermediate_size]
-            up = l1_out_real[:, self.intermediate_size:]
+            # De-interleave L1 output: with interleaved weights, L1 GEMM
+            # output has [gate]*4, [up]*4 pattern. De-interleave before splitting.
+            l1_deil = deinterleave_l1_weights(l1_out_real.unsqueeze(0).contiguous())[0]
+            gate = l1_deil[:, :self.intermediate_size]
+            up = l1_deil[:, self.intermediate_size:]
             gate_silu = torch.nn.functional.silu(gate)
             if self._swiglu_limit is not None:
                 gate_silu = gate_silu.clamp(max=self._swiglu_limit)
@@ -535,25 +569,36 @@ class CuTeDSLMoERunner:
         )
         l1_gsa = self._l1_gsa_buf.fill_(self._l1_activation_global_scale)
         
-        l1_out = run_nvfp4_grouped_gemm(
-            mat_a=padded_x_fp4, mat_b=self._l1_mat_b,
-            scale_a=l1_scale_a, scale_b=self._l1_scale_b,
-            expert_offsets=padded_expert_offsets[1:self.num_experts + 1],
-            global_scale_a=l1_gsa, global_scale_b=self._l1_gsb,
-        )
-        
-        # Extract real token outputs from padded GEMM output
-        l1_out_real = l1_out[padded_dst]
-        
-        # === SiLU(gate) * up (with swiglu_limit clamp) ===
-        gate = l1_out_real[:, :self.intermediate_size]
-        up = l1_out_real[:, self.intermediate_size:]
-        gate_silu = torch.nn.functional.silu(gate)
-        # Apply DeepSeek-V4 swiglu_limit: clamp both silu(gate) and up
-        if self._swiglu_limit is not None:
-            gate_silu = gate_silu.clamp(max=self._swiglu_limit)
-            up = up.clamp(min=-self._swiglu_limit, max=self._swiglu_limit)
-        activated = gate_silu * up
+        if self._fused_swiglu:
+            # === Fused L1 GEMM + SwiGLU in kernel registers ===
+            l1_out = run_fused_swiglu_grouped_gemm(
+                mat_a=padded_x_fp4, mat_b=self._l1_mat_b,
+                scale_a=l1_scale_a, scale_b=self._l1_scale_b,
+                expert_offsets=padded_expert_offsets[1:self.num_experts + 1],
+                global_scale_a=l1_gsa, global_scale_b=self._l1_gsb,
+                swiglu_limit=self._swiglu_limit if self._swiglu_limit is not None else 0.0,
+            )
+            l1_out_real = l1_out[padded_dst]
+            # De-interleave: odd 8-col groups = silu(gate)*up (the SwiGLU result)
+            l1_deil = deinterleave_l1_weights(l1_out_real.unsqueeze(0).contiguous())[0]
+            activated = l1_deil[:, self.intermediate_size:]
+        else:
+            # === Non-fused L1 GEMM + PyTorch SiLU(gate)*up ===
+            l1_out = run_nvfp4_grouped_gemm(
+                mat_a=padded_x_fp4, mat_b=self._l1_mat_b,
+                scale_a=l1_scale_a, scale_b=self._l1_scale_b,
+                expert_offsets=padded_expert_offsets[1:self.num_experts + 1],
+                global_scale_a=l1_gsa, global_scale_b=self._l1_gsb,
+            )
+            l1_out_real = l1_out[padded_dst]
+            l1_deil = deinterleave_l1_weights(l1_out_real.unsqueeze(0).contiguous())[0]
+            gate = l1_deil[:, :self.intermediate_size]
+            up = l1_deil[:, self.intermediate_size:]
+            gate_silu = torch.nn.functional.silu(gate)
+            if self._swiglu_limit is not None:
+                gate_silu = gate_silu.clamp(max=self._swiglu_limit)
+                up = up.clamp(min=-self._swiglu_limit, max=self._swiglu_limit)
+            activated = gate_silu * up
         
         # === L2: down ===
         # Quantize activated (per-token), scatter into padded FP4 buffer

tests/layertest.py

@@ -17,6 +17,7 @@ sys.path.insert(0, REPO_ROOT)
 
 from cutedsl.moe_pipeline import (
     run_nvfp4_moe,
+    run_nvfp4_moe_fused,
 )
 
 NVFP4_MODEL_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
@@ -251,6 +252,32 @@ def main():
     else:
         print(f"  PASS: cosine {cosine:.6f} >= {COSINE_THRESHOLD}")
 
+    # ══════════════════════════════════════════════════════════════
+    # Fused SwiGLU pipeline test
+    # ══════════════════════════════════════════════════════════════
+    print(f"\n  Running CuTeDSL NVFP4 MoE FUSED pipeline (first run compiles)...")
+    fused_output = run_nvfp4_moe_fused(
+        hidden_states, expert_ids, expert_weights,
+        weights, expert_indices,
+    )
+    print(f"    Fused:   amax={fused_output.abs().max():.4f} mean={fused_output.float().mean():.6f}")
+
+    fused_cosine = torch.nn.functional.cosine_similarity(
+        fused_output.flatten().unsqueeze(0).float(),
+        ref_output.flatten().unsqueeze(0).float(),
+    ).item()
+    fused_mse = (fused_output.float() - ref_output.float()).pow(2).mean().item()
+
+    print(f"\n{'=' * 70}")
+    print(f"  FUSED RESULT: cosine={fused_cosine:.6f}  MSE={fused_mse:.6e}")
+    print(f"{'=' * 70}")
+
+    if fused_cosine < COSINE_THRESHOLD:
+        print(f"  FAIL: fused cosine {fused_cosine:.6f} < {COSINE_THRESHOLD}")
+        sys.exit(1)
+    else:
+        print(f"  PASS: fused cosine {fused_cosine:.6f} >= {COSINE_THRESHOLD}")
+
 
 if __name__ == "__main__":
     main()