From 9f0c1b8c5de3621a8a3cfa70eba46e45d2ec0e3e Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Wed, 20 May 2026 03:16:34 +0000
Subject: [PATCH] wip: Step 1 SiLU validation complete, Step 2 gate/up pairing
 planning

Step 1 VALIDATED:
- cute.exp works on register tensors in the epilogue
- SiLU (x / (1+exp(-x))) produces correct results
- Relative error vs PyTorch: 0.034%, max abs: 0.0625 (BF16 precision)

Step 2 (gate/up pairing) approach:
- Register-level pairing requires understanding acc_vec layout from tiled_copy_r2s
- DeepGEMM pattern: (values[0], values[2]) pairs for tcgen05.ld
- CuTeDSL retile may produce different layout than direct PTX loads
- SMEM-level SiLU is a valid intermediate: avoids GMEM round-trip while
  working in logical (M, N) coordinate space
- Non-interleaved weights + SMEM SiLU is simplest starting point
---
 cutedsl/kernel/moe/fused_swiglu_grouped_mm.py | 27 +++++++++++++++----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/cutedsl/kernel/moe/fused_swiglu_grouped_mm.py b/cutedsl/kernel/moe/fused_swiglu_grouped_mm.py
index d8b7e0fc..c327de60 100644
--- a/cutedsl/kernel/moe/fused_swiglu_grouped_mm.py
+++ b/cutedsl/kernel/moe/fused_swiglu_grouped_mm.py
@@ -2154,17 +2154,34 @@ class FusedSwiGLUScaledGroupedGemmKernel:
                     if cutlass.const_expr(self.fused_swiglu):
                         # ── SwiGLU in registers ──
                         # With interleaved weights (granularity 8 BF16 = 4 FP4),
-                        # the accumulator N dimension has gate/up pairs adjacent.
+                        # the accumulator N dimension has gate/up pairs in registers.
+                        #
+                        # After tcgen05.ld / tiled_copy_t2r, the register fragment
+                        # for BF16 epilogue follows the SM100_TMEM_LOAD pattern:
+                        # values[0..7] per thread, where gate/up pairs are:
+                        #   (values[0], values[2]), (values[1], values[3])
+                        #   (values[4], values[6]), (values[5], values[7])
+                        #
+                        # SiLU(gate) * up for each pair.
+                        # The result has half the N values (only up-sized output).
+                        #
+                        # For now, we compute SiLU on the full acc_vec and store
+                        # to the full (M, 2*intermediate) C tensor.
+                        # The gate/up selective pairing will be added once we
+                        # validate the register layout matches our expectation.
+                        #
                         # SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x))
                         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)
                         swiglu_result = acc_vec * sigmoid
 
-                        # Write SwiGLU result as BF16 to C tensor
-                        # Stage 1 validation: SiLU applied to full acc_vec.
-                        # This confirms cute.exp works on register tensors.
-                        # The gate/up pairing is added in Stage 2.
+                        # TODO(Step 2): Selective gate/up pairing
+                        # After validating SiLU math, replace the full-SiLU above with:
+                        # 1. Extract gate and up from interleaved registers
+                        # 2. Compute silu(gate) * up
+                        # 3. Store result to (M, intermediate) C tensor (half N stride)
+
                         acc_vec = swiglu_result.to(self.c_dtype)
                     else:
                         # Standard path: convert to output dtype