115 lines
4.5 KiB
Python
115 lines
4.5 KiB
Python
"""Empirical register layout analysis for the fused SwiGLU epilogue.
|
|
|
|
With epi_tile=(128, 8), each subtile covers 8 BF16 N-columns and 128 M-rows.
|
|
The TMA store writes the BF16 output to GMEM in a deterministic order.
|
|
|
|
By running the fused kernel with gate=1.0/up=2.0 interleaved weights,
|
|
the output at odd 8-column groups should be silu(gate)*up ≈ 2*silu(1.0),
|
|
and even groups should be silu(gate) ≈ silu(1.0).
|
|
|
|
This script analyzes the GMEM output to understand:
|
|
1. Which 8-column groups are gate vs up (verify interleaving)
|
|
2. The BF16 values at specific (M, N) positions
|
|
3. Whether the subtile pairing is correct
|
|
"""
|
|
import sys, os
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from cutedsl.bridge import (
|
|
quantize_weight_to_nvfp4, quantize_activation_nvfp4,
|
|
make_b_k_major, interleave_l1_weights, deinterleave_l1_weights,
|
|
compute_expert_offsets, run_fused_swiglu_grouped_gemm, quantize_to_nvfp4,
|
|
assemble_scales_2d_side, assemble_scales_3d_side,
|
|
)
|
|
from cutedsl.kernel.moe.torch_scaled_grouped_mm import assemble_raw_scales_2d3d_3d_side
|
|
|
|
torch.manual_seed(42)
|
|
device = "cuda"
|
|
hidden = 7168
|
|
intermediate = 3072
|
|
K_packed = hidden // 2
|
|
|
|
# gate=1.0, up=2.0 — clear signal for interleaving
|
|
gate_w = torch.ones(hidden, intermediate, dtype=torch.bfloat16, device=device)
|
|
up_w = torch.ones(hidden, intermediate, dtype=torch.bfloat16, device=device) * 2.0
|
|
l1_w = torch.cat([gate_w, up_w], dim=1)
|
|
l1_fp4, l1_sf, l1_gs = quantize_weight_to_nvfp4(l1_w)
|
|
|
|
# Interleave
|
|
l1_ekn = interleave_l1_weights(l1_fp4.unsqueeze(0))
|
|
l1_mat_b = make_b_k_major(l1_ekn)
|
|
|
|
# SF interleave
|
|
l1_sf_ekn = l1_sf.unsqueeze(0)
|
|
l1_sf_il = interleave_l1_weights(l1_sf_ekn)
|
|
l1_sf_il_list = [l1_sf_il[0].T.contiguous()]
|
|
l1_scale_b = assemble_raw_scales_2d3d_3d_side(l1_sf_il_list)
|
|
l1_gsb = torch.tensor([l1_gs], dtype=torch.float32, device=device)
|
|
|
|
# Input: 128 tokens, all 0.1
|
|
n_tokens = 128
|
|
hidden_states = torch.ones(n_tokens, hidden, dtype=torch.bfloat16, device=device) * 0.1
|
|
gs_a = 1.0 / 2688.0
|
|
x_fp4, x_sf = quantize_activation_nvfp4(hidden_states, gs_a)
|
|
expert_offsets = torch.tensor([128, 128, 128], dtype=torch.int32, device=device)
|
|
l1_gsa = torch.tensor([gs_a] * 3, dtype=torch.float32, device=device)
|
|
|
|
from cutedsl.kernel.moe.torch_scaled_grouped_mm import ceil_div, pad_and_swizzle_single
|
|
K_sf = ceil_div(K_packed, 8)
|
|
x_sf_parts = [x_sf]
|
|
l1_scale_a = assemble_scales_2d_side(x_sf_parts)
|
|
|
|
# Run fused kernel
|
|
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_gsa, global_scale_b=l1_gsb,
|
|
)
|
|
|
|
print(f"Fused output shape: {fused_out.shape}")
|
|
out0 = fused_out[0].float() # First token, as float32
|
|
|
|
# BF16 reference: silu(1.0) = 0.7311, silu(1.0)*2.0 = 1.4621
|
|
import math
|
|
silu_one = 1.0 / (1.0 + math.exp(-1.0)) # sigmoid(1.0) = 0.7311
|
|
silu_two = 2.0 / (1.0 + math.exp(-2.0)) # sigmoid(2.0) = 1.7616
|
|
|
|
# Compute expected values
|
|
# With input 0.1 and gate=1.0, the gate GEMM output ≈ 7168 * 0.1 * 1.0 * gs_a * gs_b
|
|
# Let's check empirically
|
|
gate_vals = []
|
|
up_vals = []
|
|
for i in range(0, 64, 8):
|
|
chunk = out0[i:i+8].tolist()
|
|
is_gate = all(abs(v - out0[0].item()) < 1.0 for v in chunk)
|
|
label = "gate(silu)" if is_gate else "up(swiglu)"
|
|
if is_gate:
|
|
gate_vals.append(out0[i].item())
|
|
else:
|
|
up_vals.append(out0[i].item())
|
|
print(f" Cols {i:3d}-{i+7:3d}: {[round(v,2) for v in chunk]} → {label}")
|
|
|
|
if gate_vals and up_vals:
|
|
g = gate_vals[0]
|
|
u = up_vals[0]
|
|
print(f"\nGate ≈ {g:.4f}, Up ≈ {u:.4f}")
|
|
print(f"Ratio up/gate ≈ {u/g:.4f}")
|
|
print(f"Expected silu(1.0) ≈ {silu_one:.4f}, silu(1.0)*2.0 ≈ {2*silu_one:.4f}")
|
|
print(f"Actual gate/expected_gate ≈ {g / (7168 * 0.1 * 1.0 * gs_a * l1_gs):.4f}")
|
|
|
|
# Now check the de-interleaved SwiGLU output
|
|
l1_deil = deinterleave_l1_weights(fused_out.unsqueeze(0).contiguous())[0]
|
|
swiglu_result = l1_deil[:, intermediate:]
|
|
silu_gate = l1_deil[:, :intermediate]
|
|
|
|
print(f"\nDe-interleaved silu(gate) amax: {silu_gate.abs().amax():.4f}")
|
|
print(f"De-interleaved SwiGLU amax: {swiglu_result.abs().amax():.4f}")
|
|
print(f"SwiGLU/silu(gate) ratio: {(swiglu_result[0,0] / silu_gate[0,0]):.4f}")
|
|
|
|
# Verify: quantize the SwiGLU result and check it matches the Python quantize path
|
|
x_fp4, x_sf, gs = quantize_to_nvfp4(swiglu_result)
|
|
print(f"\nQuantized SwiGLU: FP4 shape={x_fp4.shape}, SF shape={x_sf.shape}, gs={gs:.8f}")
|
|
print(f"FP4 amax (uint8): {x_fp4.view(torch.uint8).amax()}")
|