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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user