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