817 lines
31 KiB
Python
817 lines
31 KiB
Python
"""
|
|
Bridge layer for the CuTeDSL NVFP4 MoE kernel.
|
|
|
|
Handles tensor layout conversion from our pipeline's format to what
|
|
the ScaledGroupedGemmKernel expects:
|
|
- BF16 → NVFP4 quantization (float4_e2m1fn_x2)
|
|
- Scale factor assembly (padding + swizzle)
|
|
- B tensor K-major stride conversion
|
|
- Expert offset computation
|
|
"""
|
|
import math
|
|
import torch
|
|
import cutlass
|
|
import cutlass.cute as cute
|
|
import cutlass.torch as cutlass_torch
|
|
import cutlass.utils as utils
|
|
|
|
from cutedsl.kernel.moe.torch_scaled_grouped_mm import (
|
|
ScaledGroupedGemmKernel,
|
|
pad_and_swizzle_single,
|
|
assemble_raw_scales_2d3d_2d_side,
|
|
assemble_raw_scales_2d3d_3d_side,
|
|
cat_byte_reinterpretable_tensors,
|
|
stack_byte_reinterpretable_tensors,
|
|
)
|
|
|
|
# ── Constants ──────────────────────────────────────────────────────────
|
|
|
|
E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
|
|
|
|
# Cache compiled kernels + pre-allocated workspace by cache_key
|
|
# Each entry: {'compiled': callable, 'workspace': Tensor, 'workspace_size': int}
|
|
#
|
|
# Key design decisions (Bug #1 fix):
|
|
# - cute.compile does NOT corrupt GPU memory (verified 2026-05-20 on B200).
|
|
# The original _needs_token_refill hack was a misdiagnosis. The real bug
|
|
# was elsewhere (likely OOB write or weight loading).
|
|
# - Workspace is pre-allocated per cache entry during warmup_compilation()
|
|
# and reused on subsequent calls. No torch.full() in the hot path.
|
|
# - CuTe tensor wrappers (from_dlpack + mark_layout_dynamic) are cheap
|
|
# metadata wrappers. We re-create them per call from real tensors.
|
|
# Caching them would hold stale references to tensors that get freed.
|
|
_compiled_kernel_cache = {}
|
|
|
|
# Cached LUT for E2M1 quantization (created once per device, cudagraph-safe)
|
|
_NVFP4_STEP_LUT_CACHE = {}
|
|
def _get_step_to_idx_lut(device):
|
|
"""Get or create the E2M1 step-to-index LUT for the given device.
|
|
|
|
Cached per device to avoid CPU->CUDA copies during cudagraph capture.
|
|
Must be pre-populated during warmup (before torch.compile/cudagraph capture)
|
|
so the lock is never entered on the compiled path.
|
|
"""
|
|
# Fast path: already cached — no lock needed (torch.compile-safe)
|
|
if device in _NVFP4_STEP_LUT_CACHE:
|
|
return _NVFP4_STEP_LUT_CACHE[device]
|
|
# Slow path: first call, create the LUT
|
|
lut = torch.as_tensor(
|
|
[0, 1, 2, 3, 4, 4, 5, 5, 6, 6, 6, 7, 7],
|
|
dtype=torch.int8, device=device,
|
|
)
|
|
_NVFP4_STEP_LUT_CACHE[device] = lut
|
|
return lut
|
|
SF_VEC_SIZE = 16 # NVFP4 block size
|
|
|
|
|
|
def ceil_div(a, b):
|
|
return (a + b - 1) // b
|
|
|
|
|
|
def round_up(a, b):
|
|
return ceil_div(a, b) * b
|
|
|
|
|
|
# ── Quantization ──────────────────────────────────────────────────────
|
|
|
|
def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE):
|
|
"""Quantize BF16 tensor to NVFP4.
|
|
|
|
Args:
|
|
x_bf16: (..., D) BF16 tensor
|
|
|
|
Returns:
|
|
x_fp4: (..., D//2) float4_e2m1fn_x2 — native PyTorch FP4
|
|
x_sf: (..., D//16) float8_e4m3fn — block scales
|
|
global_scale: float32 scalar
|
|
"""
|
|
x_f32 = x_bf16.float()
|
|
amax = x_f32.abs().max().clamp(min=1e-8).float()
|
|
global_scale = amax / (6.0 * 448.0)
|
|
x_norm = x_f32 / global_scale
|
|
|
|
last_dim = x_norm.shape[-1]
|
|
n_blocks = ceil_div(last_dim, block_size)
|
|
|
|
if last_dim % block_size != 0:
|
|
pad_size = n_blocks * block_size - last_dim
|
|
x_norm = torch.nn.functional.pad(x_norm, (0, pad_size))
|
|
|
|
x_reshaped = x_norm.reshape(*x_norm.shape[:-1], n_blocks, block_size)
|
|
block_amax = x_reshaped.abs().amax(dim=-1)
|
|
# Detect zero blocks and underflow blocks (amax > 0 but too small for FP8).
|
|
# Smallest positive FP8 e4m3fn is 2^-9 ≈ 1.95e-3. If amax/6 < this,
|
|
# the block scale underflows to 0, and dividing x by the clamped 1e-8
|
|
# inflates values into nonzero FP4 buckets — producing wrong results.
|
|
zero_block = block_amax < (6.0 * 2.0 ** -9) # < ~0.0117
|
|
# Zero out x for zero/underflow blocks before division.
|
|
# This ensures x_scaled = 0 → FP4 nibbles = 0.
|
|
x_reshaped = torch.where(zero_block.unsqueeze(-1),
|
|
torch.zeros_like(x_reshaped), x_reshaped)
|
|
block_amax = block_amax.clamp(min=1e-8)
|
|
block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn)
|
|
# Force zero/underflow blocks: FP8 scale = 0 (exact zero).
|
|
block_scale = torch.where(zero_block, torch.zeros_like(block_scale), block_scale)
|
|
|
|
# Nearest E2M1
|
|
block_sf_expanded = block_scale.float().unsqueeze(-1)
|
|
x_scaled = x_reshaped / block_sf_expanded.clamp(min=1e-8)
|
|
|
|
signs = torch.sign(x_scaled)
|
|
abs_scaled = x_scaled.abs().clamp(max=6.0)
|
|
|
|
half_steps = (abs_scaled * 2.0).round().clamp(0, 12).to(torch.int8)
|
|
step_to_idx = _get_step_to_idx_lut(x_bf16.device)
|
|
indices = step_to_idx[half_steps.long()]
|
|
|
|
nibbles = torch.where(signs < 0, indices + 8, indices).to(torch.uint8)
|
|
even = nibbles[..., ::2]
|
|
odd = nibbles[..., 1::2]
|
|
packed = (odd << 4) | even
|
|
|
|
packed_shape = list(x_bf16.shape)
|
|
packed_shape[-1] = last_dim // 2
|
|
x_fp4 = packed.view(torch.float4_e2m1fn_x2).reshape(packed_shape)
|
|
|
|
sf_shape = list(x_bf16.shape[:-1]) + [n_blocks]
|
|
block_scale = block_scale.reshape(sf_shape)
|
|
|
|
return x_fp4, block_scale, global_scale
|
|
|
|
|
|
def quantize_activation_nvfp4(x_bf16, global_scale, block_size=SF_VEC_SIZE):
|
|
"""Quantize BF16 activation tensor to NVFP4 (cudagraph-safe).
|
|
|
|
Unlike quantize_to_nvfp4(), this takes a pre-computed global_scale
|
|
instead of computing it via .max() (which forces CPU-GPU sync).
|
|
All operations are pure GPU with no CPU-GPU syncs.
|
|
|
|
Args:
|
|
x_bf16: (..., D) BF16 tensor
|
|
global_scale: float32 scalar (pre-computed, NOT from .max())
|
|
block_size: NVFP4 block size
|
|
|
|
Returns:
|
|
x_fp4: (..., D//2) float4_e2m1fn_x2
|
|
x_sf: (..., D//16) float8_e4m3fn
|
|
"""
|
|
x_f32 = x_bf16.float()
|
|
x_norm = x_f32 / global_scale
|
|
|
|
last_dim = x_norm.shape[-1]
|
|
n_blocks = ceil_div(last_dim, block_size)
|
|
|
|
if last_dim % block_size != 0:
|
|
pad_size = n_blocks * block_size - last_dim
|
|
x_norm = torch.nn.functional.pad(x_norm, (0, pad_size))
|
|
|
|
x_reshaped = x_norm.reshape(*x_norm.shape[:-1], n_blocks, block_size)
|
|
block_amax = x_reshaped.abs().amax(dim=-1)
|
|
# Detect zero blocks and underflow blocks (same threshold as quantize_to_nvfp4).
|
|
zero_block = block_amax < (6.0 * 2.0 ** -9)
|
|
x_reshaped = torch.where(zero_block.unsqueeze(-1),
|
|
torch.zeros_like(x_reshaped), x_reshaped)
|
|
block_amax = block_amax.clamp(min=1e-8)
|
|
block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn)
|
|
block_scale = torch.where(zero_block, torch.zeros_like(block_scale), block_scale)
|
|
|
|
block_sf_expanded = block_scale.float().unsqueeze(-1)
|
|
x_scaled = x_reshaped / block_sf_expanded.clamp(min=1e-8)
|
|
signs = torch.sign(x_scaled)
|
|
abs_scaled = x_scaled.abs().clamp(max=6.0)
|
|
|
|
half_steps = (abs_scaled * 2.0).round().clamp(0, 12).to(torch.int8)
|
|
step_to_idx = _get_step_to_idx_lut(x_bf16.device)
|
|
indices = step_to_idx[half_steps.long()]
|
|
|
|
nibbles = torch.where(signs < 0, indices + 8, indices).to(torch.uint8)
|
|
even = nibbles[..., ::2]
|
|
odd = nibbles[..., 1::2]
|
|
packed = (odd << 4) | even
|
|
|
|
packed_shape = list(x_bf16.shape)
|
|
packed_shape[-1] = last_dim // 2
|
|
x_fp4 = packed.view(torch.float4_e2m1fn_x2).reshape(packed_shape)
|
|
|
|
sf_shape = list(x_bf16.shape[:-1]) + [n_blocks]
|
|
block_scale = block_scale.reshape(sf_shape)
|
|
|
|
return x_fp4, block_scale
|
|
|
|
|
|
def quantize_weight_to_nvfp4(w_bf16, block_size=SF_VEC_SIZE):
|
|
"""Quantize BF16 weight matrix to NVFP4.
|
|
|
|
The weight is (K, N) where K is the input dim (packed dimension).
|
|
Block scales are computed along K (dim 0).
|
|
|
|
Args:
|
|
w_bf16: (K, N) BF16 weight matrix
|
|
|
|
Returns:
|
|
w_fp4: (K//2, N) float4_e2m1fn_x2 — K is the packed dim
|
|
w_sf: (K//16, N) float8_e4m3fn — block scales along K
|
|
global_scale: float32 scalar
|
|
"""
|
|
K, N = w_bf16.shape
|
|
w_f32 = w_bf16.float()
|
|
amax = w_f32.abs().max().clamp(min=1e-8).float()
|
|
global_scale = amax / (6.0 * 448.0)
|
|
w_norm = w_f32 / global_scale
|
|
|
|
k_blocks = ceil_div(K, block_size)
|
|
if K % block_size != 0:
|
|
w_norm = torch.nn.functional.pad(w_norm, (0, 0, 0, k_blocks * block_size - K))
|
|
|
|
w_reshaped = w_norm.reshape(k_blocks, block_size, N)
|
|
w_block_amax = w_reshaped.abs().amax(dim=1)
|
|
# Detect zero blocks and underflow blocks (same threshold).
|
|
zero_block = w_block_amax < (6.0 * 2.0 ** -9)
|
|
w_reshaped = torch.where(zero_block.unsqueeze(1),
|
|
torch.zeros_like(w_reshaped), w_reshaped)
|
|
w_block_amax = w_block_amax.clamp(min=1e-8)
|
|
w_sf = (w_block_amax / 6.0).to(torch.float8_e4m3fn)
|
|
w_sf = torch.where(zero_block, torch.zeros_like(w_sf), w_sf)
|
|
|
|
w_block_sf = w_sf.float().unsqueeze(1)
|
|
w_scaled = w_reshaped / w_block_sf.clamp(min=1e-8)
|
|
|
|
signs = torch.sign(w_scaled)
|
|
abs_scaled = w_scaled.abs().clamp(max=6.0)
|
|
|
|
half_steps = (abs_scaled * 2.0).round().clamp(0, 12).to(torch.int8)
|
|
step_to_idx = _get_step_to_idx_lut(w_bf16.device)
|
|
indices = step_to_idx[half_steps.long()]
|
|
nibbles = torch.where(signs < 0, indices + 8, indices).to(torch.uint8)
|
|
|
|
even = nibbles[:, ::2, :]
|
|
odd = nibbles[:, 1::2, :]
|
|
packed = (odd << 4) | even
|
|
|
|
w_fp4 = packed.reshape(K // 2, N).view(torch.float4_e2m1fn_x2)
|
|
return w_fp4, w_sf, global_scale
|
|
|
|
|
|
# ── Scale Factor Assembly ─────────────────────────────────────────────
|
|
|
|
def interleave_l1_weights(w_ekn, granularity_bf16=8):
|
|
"""Interleave gate/up weights at granularity 8 in BF16 (4 in FP4).
|
|
|
|
The fused SwiGLU epilogue requires gate/up pairs to be adjacent in the
|
|
MMA accumulator. With interleaved weights, the MMA tile produces
|
|
gate[i*8..i*8+7] and up[i*8..i*8+7] next to each other in registers,
|
|
enabling a single-register SwiGLU without SMEM round-trips.
|
|
|
|
Before: [gate_0..gate_N/2-1 | up_0..up_N/2-1]
|
|
After: [gate_0..gate_7, up_0..up_7, gate_8..gate_15, up_8..up_15, ...]
|
|
|
|
The interleave operates along the N dimension, where each column = 1 BF16
|
|
(FP4 packing is along K, not N). So g = granularity_bf16 directly.
|
|
|
|
Args:
|
|
w_ekn: (E, K_packed, N_packed) FP4 weight tensor in K-major layout
|
|
N_packed = 2*intermediate/2 = intermediate (gate+up fused)
|
|
granularity_bf16: interleave group size in BF16 elements (default 8)
|
|
|
|
Returns:
|
|
(E, K_packed, N_packed) FP4 weight tensor with interleaved gate/up
|
|
"""
|
|
E, K, N = w_ekn.shape
|
|
N_half = N // 2 # gate and up each have N/2 FP4 columns
|
|
g = granularity_bf16 # N-axis interleave: each N-col = 1 BF16 col (packing is along K)
|
|
|
|
gate = w_ekn[:, :, :N_half].reshape(E, K, N_half // g, g)
|
|
up = w_ekn[:, :, N_half:].reshape(E, K, N_half // g, g)
|
|
return torch.stack([gate, up], dim=3).reshape(E, K, N)
|
|
|
|
|
|
def deinterleave_l1_weights(w_ekn, granularity_bf16=8):
|
|
"""De-interleave gate/up weights (inverse of interleave_l1_weights).
|
|
|
|
Used for testing/verification only.
|
|
"""
|
|
g = granularity_bf16 # N-axis: each N-col = 1 BF16 col
|
|
E, K, N = w_ekn.shape
|
|
w_reshaped = w_ekn.reshape(E, K, N // (2 * g), 2, g)
|
|
gate = w_reshaped[:, :, :, 0, :].reshape(E, K, N // 2)
|
|
up = w_reshaped[:, :, :, 1, :].reshape(E, K, N // 2)
|
|
return torch.cat([gate, up], dim=2)
|
|
|
|
|
|
def assemble_scales_2d_side(raw_scales):
|
|
"""Assemble activation scale factors for the 2Dx3D scenario.
|
|
|
|
Args:
|
|
raw_scales: list of (M_e, K_sf) float8_e4m3fn tensors, one per expert
|
|
|
|
Returns:
|
|
Assembled and swizzled scale tensor
|
|
"""
|
|
return assemble_raw_scales_2d3d_2d_side(raw_scales)
|
|
|
|
|
|
def assemble_scales_3d_side(raw_scales):
|
|
"""Assemble weight scale factors for the 2Dx3D scenario.
|
|
|
|
Args:
|
|
raw_scales: list of (K_sf, N) float8_e4m3fn tensors, one per expert
|
|
NOTE: These will be transposed to (N, K_sf) before swizzling,
|
|
since the kernel expects N as the non-K dimension.
|
|
|
|
Returns:
|
|
Assembled and swizzled scale tensor
|
|
"""
|
|
# Kernel expects (N, K_sf) — transpose before swizzling
|
|
transposed = [sf.T.contiguous() for sf in raw_scales]
|
|
return assemble_raw_scales_2d3d_3d_side(transposed)
|
|
|
|
|
|
# ── Tensor Layout Conversion ──────────────────────────────────────────
|
|
|
|
def make_b_k_major(b_tensor):
|
|
"""Convert B tensor from N-major to K-major layout.
|
|
|
|
The kernel expects B with stride (E*K*N, 1, K) — K is contiguous.
|
|
torch.stack produces stride (E*K*N, N, 1) — N is contiguous.
|
|
|
|
Args:
|
|
b_tensor: (experts, K_packed, N_packed) float4_e2m1fn_x2, N-major
|
|
|
|
Returns:
|
|
Same shape, K-major strides
|
|
"""
|
|
return b_tensor.permute(0, 2, 1).contiguous().permute(0, 2, 1)
|
|
|
|
|
|
def compute_expert_offsets(tokens_per_expert, num_experts, device="cuda"):
|
|
"""Compute cumulative token offsets for the grouped GEMM.
|
|
|
|
Args:
|
|
tokens_per_expert: list of int, one per expert
|
|
|
|
Returns:
|
|
offs: (num_experts,) int32 — cumulative sum
|
|
"""
|
|
offs = torch.tensor(
|
|
[sum(tokens_per_expert[:e+1]) for e in range(num_experts)],
|
|
dtype=torch.int32, device=device,
|
|
)
|
|
return offs
|
|
|
|
|
|
# ── Kernel Launch ─────────────────────────────────────────────────────
|
|
|
|
|
|
def warmup_compilation(num_experts, K_packed, N_packed, device,
|
|
mma_tiler_mn=(128, 128), cluster_shape_mn=(1, 1)):
|
|
"""Eagerly JIT-compile the GEMM kernel for a specific shape.
|
|
|
|
Call this BEFORE model weights are loaded to ensure cute.compile
|
|
runs exactly once per shape. The compiled kernel is cached and
|
|
reused by run_nvfp4_grouped_gemm on the forward path.
|
|
|
|
Uses random non-zero data. Zero-filled FP4/FP8 tensors cause
|
|
cudaErrorIllegalInstruction because the GEMM arithmetic hits
|
|
invalid values (division by zero in scale dequantization, NaN
|
|
propagation). Random data produces valid intermediate results
|
|
that exercise the kernel's full arithmetic path.
|
|
|
|
The warmup tensors are freed immediately after compilation.
|
|
Memory cost is minimal (~50MB for typical DeepSeek-V4 shapes).
|
|
|
|
Args:
|
|
num_experts: number of experts (local, after expert parallelism)
|
|
K_packed: K dimension in float4 elements (i.e. K_original // 2)
|
|
N_packed: N dimension in float4 elements (i.e. N_original // 2)
|
|
device: 'cuda:X' or 'cuda'
|
|
mma_tiler_mn: GEMM tiling (default (128,128))
|
|
cluster_shape_mn: cluster shape (default (1,1))
|
|
"""
|
|
cache_key = (num_experts, str(device), mma_tiler_mn, cluster_shape_mn,
|
|
K_packed, N_packed)
|
|
|
|
if cache_key in _compiled_kernel_cache:
|
|
return # Already compiled
|
|
|
|
# Generate VALID FP4 data by quantizing random BF16 through our pipeline.
|
|
# Random uint8 bytes as FP4 bit patterns produce NaN/Inf when the GEMM
|
|
# dequantizes them (FP4 value * FP8 scale * FP32 global scale), which
|
|
# causes cudaErrorIllegalInstruction in the Blackwell MMA hardware.
|
|
# The ONLY safe approach: generate random BF16, quantize through our
|
|
# quantize_to_nvfp4, producing mathematically consistent FP4 + FP8 scales.
|
|
_warmup_a_bf16 = torch.randn(128, K_packed * 2, dtype=torch.bfloat16, device=device) * 0.1
|
|
mat_a, scale_a, _ = quantize_to_nvfp4(_warmup_a_bf16)
|
|
del _warmup_a_bf16
|
|
_warmup_b_bf16 = torch.randn(1, K_packed * 2, N_packed * 2, dtype=torch.bfloat16, device=device) * 0.1 # 1 expert: kernel compiles same regardless of count
|
|
mat_b, scale_b, _ = quantize_to_nvfp4(_warmup_b_bf16)
|
|
del _warmup_b_bf16
|
|
out = torch.zeros(128, N_packed, dtype=torch.bfloat16, device=device)
|
|
expert_offsets = torch.full((1,), 128, dtype=torch.int32, device=device)
|
|
global_scale_a = torch.ones(1, dtype=torch.float32, device=device)
|
|
global_scale_b = torch.ones(1, dtype=torch.float32, device=device)
|
|
|
|
kernel = ScaledGroupedGemmKernel(
|
|
scenario="2Dx3D",
|
|
sf_vec_size=SF_VEC_SIZE,
|
|
accumulate_on_output=False,
|
|
separate_tensormap_init=True,
|
|
consistent_token_padding=False,
|
|
mma_tiler_mnk=(*mma_tiler_mn, 256),
|
|
cluster_shape_mnk=(*cluster_shape_mn, 1),
|
|
)
|
|
|
|
def to_cute(t):
|
|
ct = cutlass_torch.from_dlpack(t)
|
|
return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
|
|
|
|
a_c = to_cute(mat_a)
|
|
b_c = to_cute(mat_b)
|
|
sfa_c = to_cute(scale_a)
|
|
sfb_c = to_cute(scale_b)
|
|
c_c = to_cute(out)
|
|
offs_c = to_cute(expert_offsets)
|
|
|
|
workspace_size = kernel.get_workspace_size(num_experts)
|
|
workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=device)
|
|
ws_c = to_cute(workspace)
|
|
|
|
gsa_c = to_cute(global_scale_a)
|
|
gsb_c = to_cute(global_scale_b)
|
|
|
|
import cuda.bindings.driver as cuda
|
|
cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1]
|
|
max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
|
|
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
|
|
compiled = cute.compile(
|
|
kernel, a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,
|
|
max_active_clusters, stream,
|
|
global_scale_a=gsa_c, global_scale_b=gsb_c,
|
|
)
|
|
|
|
# Warmup run (required by CuTeDSL)
|
|
compiled(
|
|
a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, stream,
|
|
global_scale_a=gsa_c, global_scale_b=gsb_c,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
# Cache compiled kernel + pre-allocated workspace
|
|
# (NOT CuTe wrappers — they hold refs to dummy tensors)
|
|
_compiled_kernel_cache[cache_key] = {
|
|
'compiled': compiled,
|
|
'workspace': workspace, # pre-allocated, reused
|
|
'workspace_size': workspace_size,
|
|
}
|
|
|
|
# Free dummy data tensors (workspace is kept in cache)
|
|
del mat_a, mat_b, scale_a, scale_b, out, expert_offsets
|
|
del global_scale_a, global_scale_b
|
|
del a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, gsa_c, gsb_c
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
|
|
def run_nvfp4_grouped_gemm(
|
|
mat_a, # (tokens_sum, K_packed) float4_e2m1fn_x2
|
|
mat_b, # (experts, K_packed, N_packed) float4_e2m1fn_x2, K-major
|
|
scale_a, # assembled 2D side (padded + swizzled)
|
|
scale_b, # assembled 3D side (padded + swizzled)
|
|
expert_offsets, # (experts,) int32 cumulative token offsets
|
|
global_scale_a=None, # (experts,) float32
|
|
global_scale_b=None, # (experts,) float32
|
|
mma_tiler_mn=(128, 128),
|
|
cluster_shape_mn=(1, 1),
|
|
):
|
|
"""Run the CuTeDSL NVFP4 scaled grouped GEMM.
|
|
|
|
2Dx3D: A(tokens, K) x B(experts, K, N) -> C(tokens, N)
|
|
|
|
Kernel is compiled once (either via warmup_compilation() during init
|
|
or lazily on first call), then cached. Workspace is pre-allocated
|
|
during warmup and reused — no torch.full() in the hot path.
|
|
No torch.cuda.synchronize() or .item() in the forward path.
|
|
"""
|
|
num_experts = mat_b.shape[0]
|
|
n_dim = mat_b.shape[2]
|
|
tokens_sum = mat_a.shape[0]
|
|
|
|
out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device)
|
|
|
|
cache_key = (num_experts, str(mat_a.device), mma_tiler_mn, cluster_shape_mn,
|
|
mat_a.shape[1], mat_b.shape[2])
|
|
|
|
if cache_key not in _compiled_kernel_cache:
|
|
# Lazy compilation — safety net if warmup_compilation wasn't called.
|
|
# This should NOT happen in production (warmup is called during init).
|
|
kernel = ScaledGroupedGemmKernel(
|
|
scenario="2Dx3D",
|
|
sf_vec_size=SF_VEC_SIZE,
|
|
accumulate_on_output=False,
|
|
separate_tensormap_init=True,
|
|
consistent_token_padding=False,
|
|
mma_tiler_mnk=(*mma_tiler_mn, 256),
|
|
cluster_shape_mnk=(*cluster_shape_mn, 1),
|
|
)
|
|
|
|
def to_cute(t):
|
|
ct = cutlass_torch.from_dlpack(t)
|
|
return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
|
|
|
|
a_c = to_cute(mat_a)
|
|
b_c = to_cute(mat_b)
|
|
sfa_c = to_cute(scale_a)
|
|
sfb_c = to_cute(scale_b)
|
|
c_c = to_cute(out)
|
|
offs_c = to_cute(expert_offsets)
|
|
|
|
workspace_size = kernel.get_workspace_size(num_experts)
|
|
workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=mat_a.device)
|
|
ws_c = to_cute(workspace)
|
|
|
|
gsa_c = to_cute(global_scale_a) if global_scale_a is not None else None
|
|
gsb_c = to_cute(global_scale_b) if global_scale_b is not None else None
|
|
|
|
import cuda.bindings.driver as cuda
|
|
cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1]
|
|
max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
|
|
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
|
|
compiled = cute.compile(
|
|
kernel, a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,
|
|
max_active_clusters, stream,
|
|
global_scale_a=gsa_c, global_scale_b=gsb_c,
|
|
)
|
|
compiled(
|
|
a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, stream,
|
|
global_scale_a=gsa_c, global_scale_b=gsb_c,
|
|
)
|
|
|
|
_compiled_kernel_cache[cache_key] = {
|
|
'compiled': compiled,
|
|
'workspace': workspace,
|
|
'workspace_size': workspace_size,
|
|
}
|
|
|
|
# --- Invoke cached kernel ---
|
|
entry = _compiled_kernel_cache[cache_key]
|
|
compiled = entry['compiled']
|
|
workspace = entry['workspace']
|
|
|
|
# Re-create CuTe wrappers from current tensors each call.
|
|
# This is cheap (metadata only, no GPU work) and avoids stale
|
|
# references to tensors from previous calls that may have been freed.
|
|
def to_cute(t):
|
|
ct = cutlass_torch.from_dlpack(t)
|
|
return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
|
|
|
|
a_c = to_cute(mat_a)
|
|
b_c = to_cute(mat_b)
|
|
sfa_c = to_cute(scale_a)
|
|
sfb_c = to_cute(scale_b)
|
|
c_c = to_cute(out)
|
|
offs_c = to_cute(expert_offsets)
|
|
ws_c = to_cute(workspace)
|
|
|
|
gsa_c = to_cute(global_scale_a) if global_scale_a is not None else None
|
|
gsb_c = to_cute(global_scale_b) if global_scale_b is not None else None
|
|
|
|
import cuda.bindings.driver as cuda
|
|
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
|
|
compiled(
|
|
a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, stream,
|
|
global_scale_a=gsa_c, global_scale_b=gsb_c,
|
|
)
|
|
|
|
return out
|
|
|
|
|
|
# ── Fused SwiGLU GEMM (Stage 1: SiLU in registers, BF16 output) ──────
|
|
|
|
# Cache for fused kernel (separate from standard GEMM cache)
|
|
_fused_kernel_cache = {}
|
|
|
|
|
|
def warmup_fused_swiglu_compilation(num_experts, K_packed, N_packed, device,
|
|
swiglu_limit=0.0,
|
|
mma_tiler_mn=(128, 128),
|
|
cluster_shape_mn=(1, 1)):
|
|
"""Eagerly JIT-compile the fused SwiGLU GEMM kernel.
|
|
|
|
Must be called during model initialization. See warmup_compilation()
|
|
for the standard GEMM equivalent.
|
|
"""
|
|
from cutedsl.kernel.moe.fused_swiglu_grouped_mm import FusedSwiGLUScaledGroupedGemmKernel
|
|
|
|
cache_key = ('fused', num_experts, str(device), mma_tiler_mn, cluster_shape_mn,
|
|
K_packed, N_packed, swiglu_limit)
|
|
|
|
if cache_key in _fused_kernel_cache:
|
|
return
|
|
|
|
# Generate VALID FP4 data by quantizing random BF16 (same as warmup_compilation)
|
|
_warmup_a_bf16 = torch.randn(128, K_packed * 2, dtype=torch.bfloat16, device=device) * 0.1
|
|
mat_a, scale_a, _ = quantize_to_nvfp4(_warmup_a_bf16)
|
|
del _warmup_a_bf16
|
|
_warmup_b_bf16 = torch.randn(1, K_packed * 2, N_packed * 2, dtype=torch.bfloat16, device=device) * 0.1 # 1 expert: kernel compiles same regardless of count
|
|
mat_b, scale_b, _ = quantize_to_nvfp4(_warmup_b_bf16)
|
|
del _warmup_b_bf16
|
|
# BF16 output (Stage 1: we still write BF16)
|
|
# The fused kernel writes intermediate (N/2) since gate+up → silu result
|
|
out = torch.zeros(128, N_packed, dtype=torch.bfloat16, device=device)
|
|
expert_offsets = torch.full((1,), 128, dtype=torch.int32, device=device)
|
|
global_scale_a = torch.ones(1, dtype=torch.float32, device=device)
|
|
global_scale_b = torch.ones(1, dtype=torch.float32, device=device)
|
|
|
|
kernel = FusedSwiGLUScaledGroupedGemmKernel(
|
|
scenario="2Dx3D",
|
|
sf_vec_size=SF_VEC_SIZE,
|
|
accumulate_on_output=False,
|
|
separate_tensormap_init=True,
|
|
consistent_token_padding=False,
|
|
mma_tiler_mnk=(*mma_tiler_mn, 256),
|
|
cluster_shape_mnk=(*cluster_shape_mn, 1),
|
|
fused_swiglu=True,
|
|
swiglu_limit=swiglu_limit,
|
|
)
|
|
|
|
def to_cute(t):
|
|
ct = cutlass_torch.from_dlpack(t)
|
|
return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
|
|
|
|
a_c = to_cute(mat_a)
|
|
b_c = to_cute(mat_b)
|
|
sfa_c = to_cute(scale_a)
|
|
sfb_c = to_cute(scale_b)
|
|
c_c = to_cute(out)
|
|
offs_c = to_cute(expert_offsets)
|
|
workspace_size = kernel.get_workspace_size(num_experts)
|
|
workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=device)
|
|
ws_c = to_cute(workspace)
|
|
gsa_c = to_cute(global_scale_a)
|
|
gsb_c = to_cute(global_scale_b)
|
|
|
|
import cuda.bindings.driver as cuda
|
|
cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1]
|
|
max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
|
|
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
|
|
compiled = cute.compile(
|
|
kernel, a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,
|
|
max_active_clusters, stream,
|
|
global_scale_a=gsa_c, global_scale_b=gsb_c,
|
|
)
|
|
compiled(
|
|
a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, stream,
|
|
global_scale_a=gsa_c, global_scale_b=gsb_c,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
_fused_kernel_cache[cache_key] = {
|
|
'compiled': compiled,
|
|
'workspace': workspace,
|
|
'workspace_size': workspace_size,
|
|
}
|
|
|
|
del mat_a, mat_b, scale_a, scale_b, out, expert_offsets
|
|
del global_scale_a, global_scale_b
|
|
del a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, gsa_c, gsb_c
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def run_fused_swiglu_grouped_gemm(
|
|
mat_a, # (tokens_sum, K_packed) float4_e2m1fn_x2
|
|
mat_b, # (experts, K_packed, N_packed) float4_e2m1fn_x2, K-major
|
|
scale_a, # assembled 2D side (padded + swizzled)
|
|
scale_b, # assembled 3D side (padded + swizzled)
|
|
expert_offsets, # (experts,) int32 cumulative token offsets
|
|
global_scale_a=None, # (experts,) float32
|
|
global_scale_b=None, # (experts,) float32
|
|
swiglu_limit=0.0,
|
|
mma_tiler_mn=(128, 128),
|
|
cluster_shape_mn=(1, 1),
|
|
):
|
|
"""Run the fused SwiGLU NVFP4 scaled grouped GEMM.
|
|
|
|
Stage 1: SiLU is applied to the full accumulator in registers,
|
|
then written as BF16 to C. Gate/up pairing is not yet implemented.
|
|
"""
|
|
from cutedsl.kernel.moe.fused_swiglu_grouped_mm import FusedSwiGLUScaledGroupedGemmKernel
|
|
|
|
num_experts = mat_b.shape[0]
|
|
n_dim = mat_b.shape[2]
|
|
tokens_sum = mat_a.shape[0]
|
|
|
|
out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device)
|
|
|
|
cache_key = ('fused', num_experts, str(mat_a.device), mma_tiler_mn, cluster_shape_mn,
|
|
mat_a.shape[1], mat_b.shape[2], swiglu_limit)
|
|
|
|
if cache_key not in _fused_kernel_cache:
|
|
# Lazy compilation
|
|
kernel = FusedSwiGLUScaledGroupedGemmKernel(
|
|
scenario="2Dx3D",
|
|
sf_vec_size=SF_VEC_SIZE,
|
|
accumulate_on_output=False,
|
|
separate_tensormap_init=True,
|
|
consistent_token_padding=False,
|
|
mma_tiler_mnk=(*mma_tiler_mn, 256),
|
|
cluster_shape_mnk=(*cluster_shape_mn, 1),
|
|
fused_swiglu=True,
|
|
swiglu_limit=swiglu_limit,
|
|
)
|
|
|
|
def to_cute(t):
|
|
ct = cutlass_torch.from_dlpack(t)
|
|
return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
|
|
|
|
a_c = to_cute(mat_a)
|
|
b_c = to_cute(mat_b)
|
|
sfa_c = to_cute(scale_a)
|
|
sfb_c = to_cute(scale_b)
|
|
c_c = to_cute(out)
|
|
offs_c = to_cute(expert_offsets)
|
|
workspace_size = kernel.get_workspace_size(num_experts)
|
|
workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=mat_a.device)
|
|
ws_c = to_cute(workspace)
|
|
gsa_c = to_cute(global_scale_a) if global_scale_a is not None else None
|
|
gsb_c = to_cute(global_scale_b) if global_scale_b is not None else None
|
|
|
|
import cuda.bindings.driver as cuda
|
|
cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1]
|
|
max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
|
|
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
|
|
compiled = cute.compile(
|
|
kernel, a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c,
|
|
max_active_clusters, stream,
|
|
global_scale_a=gsa_c, global_scale_b=gsb_c,
|
|
)
|
|
compiled(
|
|
a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, stream,
|
|
global_scale_a=gsa_c, global_scale_b=gsb_c,
|
|
)
|
|
|
|
_fused_kernel_cache[cache_key] = {
|
|
'compiled': compiled,
|
|
'workspace': workspace,
|
|
'workspace_size': workspace_size,
|
|
}
|
|
|
|
entry = _fused_kernel_cache[cache_key]
|
|
compiled = entry['compiled']
|
|
workspace = entry['workspace']
|
|
|
|
def to_cute(t):
|
|
ct = cutlass_torch.from_dlpack(t)
|
|
return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
|
|
|
|
a_c = to_cute(mat_a)
|
|
b_c = to_cute(mat_b)
|
|
sfa_c = to_cute(scale_a)
|
|
sfb_c = to_cute(scale_b)
|
|
c_c = to_cute(out)
|
|
offs_c = to_cute(expert_offsets)
|
|
ws_c = to_cute(workspace)
|
|
gsa_c = to_cute(global_scale_a) if global_scale_a is not None else None
|
|
gsb_c = to_cute(global_scale_b) if global_scale_b is not None else None
|
|
|
|
import cuda.bindings.driver as cuda
|
|
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
|
|
compiled(
|
|
a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, stream,
|
|
global_scale_a=gsa_c, global_scale_b=gsb_c,
|
|
)
|
|
|
|
return out
|
|
|
|
|
|
|
|
def deinterleave_quantize_nvfp4_cuda(fused_bf16, intermediate, global_scale, granularity=8):
|
|
"""De-interleave + quantize fused SwiGLU output using a custom CUDA kernel.
|
|
|
|
Single kernel launch, no Python loop. 4x faster than the Python path.
|
|
|
|
Args:
|
|
fused_bf16: (M, 2*intermediate) BF16 — fused L1 output with interleaved gate/up
|
|
intermediate: intermediate dimension (e.g., 3072)
|
|
global_scale: pre-computed global scale for quantization
|
|
granularity: interleave granularity in BF16 columns (default 8)
|
|
|
|
Returns:
|
|
x_fp4: (M, intermediate//2) float4_e2m1fn_x2 — quantized SwiGLU
|
|
x_sf: (M, intermediate//16) float8_e4m3fn — block scales
|
|
"""
|
|
from torch.utils.cpp_extension import load
|
|
import os
|
|
kernel_dir = os.path.join(os.path.dirname(__file__), "kernels")
|
|
mod = load(
|
|
name="deinterleave_quantize_nvfp4",
|
|
sources=[os.path.join(kernel_dir, "deinterleave_quantize.cu")],
|
|
extra_cuda_cflags=["-gencode=arch=compute_100a,code=sm_100a"],
|
|
verbose=False,
|
|
)
|
|
return mod.deinterleave_quantize_nvfp4(fused_bf16, intermediate, granularity, global_scale)
|