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
This commit is contained in:
2026-05-20 03:16:34 +00:00
parent b84f2f7bf9
commit 9f0c1b8c5d

View File

@@ -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