biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

restore: new bridge/moe_pipeline/layertest

Browse Source
This commit is contained in:
biondizzle
2026-05-16 19:55:19 +00:00
parent 0c878b3a9e
commit cc75a55bd9
3 changed files with 268 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

209
cutedsl/bridge.py
View File

@@ -7,8 +7,31 @@ the ScaledGroupedGemmKernel expects:
- Scale factor assembly (padding + swizzle)
- B tensor K-major stride conversion
- Expert offset computation
CUDA-graph-compatible: no .item() calls, no torch.cuda.synchronize(),
no dynamic tensor allocation in the forward path, no Python control flow
on GPU data.
"""
import math
import threading
# Cached LUT for E2M1 quantization (created once per device, cudagraph-safe)
_NVFP4_STEP_LUT_CACHE = {}
_NVFP4_STEP_LUT_LOCK = threading.Lock()
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.
"""
with _NVFP4_STEP_LUT_LOCK:
if device not in _NVFP4_STEP_LUT_CACHE:
_NVFP4_STEP_LUT_CACHE[device] = torch.as_tensor(
[0, 1, 2, 3, 4, 4, 5, 5, 6, 6, 6, 7, 7],
dtype=torch.int8, device=device,
)
return _NVFP4_STEP_LUT_CACHE[device]
import torch
import cutlass
import cutlass.cute as cute
@@ -42,7 +65,14 @@ def round_up(a, b):
def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE):
"""Quantize BF16 tensor to NVFP4.
NOTE: This function is NOT cudagraph-safe because it uses .max()
which forces a CPU-GPU sync. It should only be called during
weight preparation (offline), NOT during the forward pass.
For activation quantization during forward, use
quantize_activation_nvfp4() instead (cudagraph-safe, fixed global scale).
Args:
x_bf16: (..., D) BF16 tensor
@@ -67,15 +97,15 @@ def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE):
block_amax = x_reshaped.abs().amax(dim=-1).clamp(min=1e-8)
block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn)
# Nearest E2M1
# Nearest E2M1 — memory-efficient clamp approach
block_sf_expanded = block_scale.float().unsqueeze(-1)
x_scaled = x_reshaped / block_sf_expanded.clamp(min=1e-8)
magnitudes = torch.tensor(E2M1_MAGNITUDES, dtype=torch.float32, device=x_bf16.device)
signs = torch.sign(x_scaled)
abs_scaled = x_scaled.abs().unsqueeze(-1)
distances = (abs_scaled - magnitudes).abs()
indices = distances.argmin(dim=-1)
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]
@@ -92,6 +122,62 @@ def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE):
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).
The global_scale should be computed once during warmup and cached.
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).clamp(min=1e-8)
block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn)
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.
@@ -201,6 +287,86 @@ def compute_expert_offsets(tokens_per_expert, num_experts, device="cuda"):
return offs
# ── Compiled Kernel Cache ─────────────────────────────────────────────
_compiled_kernel_cache = {}
def _get_compiled_kernel(num_experts, device, mma_tiler_mn, cluster_shape_mn):
"""Get or compile the CuTeDSL grouped GEMM kernel (cached by shape config).
The kernel compilation is deterministic for a given (num_experts, device, tiler, cluster)
config, so we cache it to avoid recompiling on every forward call.
"""
cache_key = (num_experts, str(device), mma_tiler_mn, cluster_shape_mn)
if cache_key in _compiled_kernel_cache:
return _compiled_kernel_cache[cache_key]
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),
)
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)
# We need dummy tensors to compile against — use minimal sizes
# The compiled kernel works with dynamic shapes via mark_layout_dynamic
K_packed = 256 # minimal
N_packed = 256
tokens = 1
dummy_a = torch.zeros(tokens, K_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2)
dummy_b = torch.zeros(num_experts, K_packed, N_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2)
dummy_sfa = torch.zeros(1, 1, dtype=torch.float16, device=device).to(torch.float8_e4m3fn)
dummy_sfb = torch.zeros(1, 1, dtype=torch.float16, device=device).to(torch.float8_e4m3fn)
dummy_c = torch.zeros(tokens, N_packed, dtype=torch.bfloat16, device=device)
dummy_offs = torch.zeros(num_experts, dtype=torch.int32, device=device)
ws_size = kernel.get_workspace_size(num_experts)
dummy_ws = torch.full((ws_size,), 255, dtype=torch.uint8, device=device)
dummy_gsa = torch.ones(num_experts, dtype=torch.float32, device=device)
dummy_gsb = torch.ones(num_experts, dtype=torch.float32, device=device)
def to_cute(t):
ct = cutlass_torch.from_dlpack(t)
return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
compiled = cute.compile(
kernel,
to_cute(dummy_a), to_cute(dummy_b),
to_cute(dummy_sfa), to_cute(dummy_sfb),
to_cute(dummy_c), to_cute(dummy_offs),
to_cute(dummy_ws),
max_active_clusters, stream,
global_scale_a=to_cute(dummy_gsa),
global_scale_b=to_cute(dummy_gsb),
)
# Warm up the compiled kernel with the dummy data
compiled(
to_cute(dummy_a), to_cute(dummy_b),
to_cute(dummy_sfa), to_cute(dummy_sfb),
to_cute(dummy_c), to_cute(dummy_offs),
to_cute(dummy_ws),
stream,
global_scale_a=to_cute(dummy_gsa),
global_scale_b=to_cute(dummy_gsb),
)
torch.cuda.synchronize()
# Free dummies
del dummy_a, dummy_b, dummy_sfa, dummy_sfb, dummy_c, dummy_offs, dummy_ws, dummy_gsa, dummy_gsb
_compiled_kernel_cache[cache_key] = (compiled, kernel, max_active_clusters)
return compiled, kernel, max_active_clusters
# ── Kernel Launch ─────────────────────────────────────────────────────
def run_nvfp4_grouped_gemm(
@@ -217,21 +383,19 @@ def run_nvfp4_grouped_gemm(
"""Run the CuTeDSL NVFP4 scaled grouped GEMM.
2Dx3D: A(tokens, K) x B(experts, K, N) -> C(tokens, N)
CUDA-graph-compatible: uses cached compiled kernel, no synchronize(),
no cute.compile() in the forward path.
"""
num_experts = mat_b.shape[0]
n_dim = mat_b.shape[2] # packed N (in float4 elements)
n_dim = mat_b.shape[2] # N dimension (logical, not packed — float4_e2m1fn_x2 packs along K, not N)
tokens_sum = mat_a.shape[0]
device = mat_a.device
out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device)
out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, 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),
compiled, kernel, max_active_clusters = _get_compiled_kernel(
num_experts, device, mma_tiler_mn, cluster_shape_mn
)
# Convert to CuTe tensors with dynamic layout
@@ -247,28 +411,21 @@ def run_nvfp4_grouped_gemm(
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)
workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=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,
)
torch.cuda.synchronize()
# NOTE: No torch.cuda.synchronize() here — cudagraph capture forbids it.
# The caller is responsible for any needed synchronization.
return out

19
cutedsl/moe_pipeline.py
View File

@@ -187,6 +187,9 @@ def run_nvfp4_moe(
# Global scales: alpha = igs * weight_gs for each expert
l1_global_scale_a = torch.tensor([x_igs] * num_experts, dtype=torch.float32, device=device)
l1_global_scale_b = torch.tensor(weights['l1_gs'], dtype=torch.float32, device=device)
print(f" L1 global_scale_a: {l1_global_scale_a.tolist()}", flush=True)
print(f" L1 global_scale_b: {l1_global_scale_b.tolist()}", flush=True)
print(f" alpha (a*b): {(l1_global_scale_a * l1_global_scale_b).tolist()}", flush=True)
# Run L1 GEMM
l1_out = run_nvfp4_grouped_gemm(
@@ -194,16 +197,20 @@ def run_nvfp4_moe(
scale_a=l1_scale_a, scale_b=l1_scale_b,
expert_offsets=expert_offsets,
global_scale_a=l1_global_scale_a, global_scale_b=l1_global_scale_b,
) # (num_slots, intermediate) BF16
) # (num_slots, 2*intermediate) BF16
print(f" L1 GEMM output: shape={l1_out.shape}, amax={l1_out.abs().amax().item():.4f}", flush=True)
# ════════════════════════════════════════════════════════════════
# SiLU(gate) * up (BF16 — nonlinear requires BF16)
# ════════════════════════════════════════════════════════════════
intermediate = l1_out.shape[1]
half = intermediate // 2 # 3072
gate = l1_out[:, :half]
up = l1_out[:, half:]
activated = torch.nn.functional.silu(gate) * up # (num_slots, half) BF16
# L1 output is (tokens, 2*intermediate) — gate and up fused
intermediate_size = l1_out.shape[1] // 2
gate = l1_out[:, :intermediate_size]
up = l1_out[:, intermediate_size:]
print(f" gate: shape={gate.shape}, amax={gate.abs().amax().item():.4f}", flush=True)
print(f" up: shape={up.shape}, amax={up.abs().amax().item():.4f}", flush=True)
activated = torch.nn.functional.silu(gate) * up # (num_slots, intermediate) BF16
print(f" After SiLU(gate)*up: shape={activated.shape}, amax={activated.abs().amax().item():.4f}", flush=True)
# ════════════════════════════════════════════════════════════════
# L2: down projection (NVFP4 × NVFP4 → BF16)

76
tests/layertest.py
View File

@@ -16,7 +16,6 @@ REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, REPO_ROOT)
from cutedsl.moe_pipeline import (
prepare_nvfp4_moe_weights,
run_nvfp4_moe,
)
@@ -121,6 +120,75 @@ def moe_forward_bf16(hidden_states, experts, expert_ids, expert_weights):
return output
def prepare_nvfp4_weights_direct(nvfp4_tensors, layer_idx, expert_indices, intermediate_size):
"""Prepare weights via direct view-cast (no BF16 round-trip).
Checkpoint uint8 → float4_e2m1fn_x2 (byte-preserving).
Block scales float8_e4m3fn → used directly.
Global scales float32 → used directly.
For L1 (gate+up fused): normalize dual global scales to max, fold ratio
into block scales via float32 (one multiply + float8 round-trip on ratio only).
"""
l1_fp4, l1_sf, l1_gs = [], [], []
l2_fp4, l2_sf, l2_gs = [], [], []
for e in expert_indices:
# L1: gate + up
gate_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight"].to(DEVICE)
up_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight"].to(DEVICE)
gate_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale"].to(DEVICE)
up_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE)
gate_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale_2"].item()
up_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale_2"].item()
# Fuse gate+up along N, transpose to K-major
fused_w = torch.cat([gate_w, up_w], dim=0) # (2*intermediate, hidden//2) uint8
fused_w_fp4 = fused_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()
# (hidden//2, 2*intermediate) — K=hidden packed, N=2*intermediate
# Fuse block scales: checkpoint is (N, K_sf), bridge expects (K_sf, N)
fused_sf = torch.cat([gate_sf, up_sf], dim=0) # (2*intermediate, hidden//16) = (N, K_sf)
fused_sf = fused_sf.permute(1, 0).contiguous() # → (K_sf, N)
# Normalize dual global scales
l1_max_gs = max(gate_gs, up_gs)
if gate_gs != up_gs:
fused_sf_f32 = fused_sf.float()
# Gate is first intermediate cols, up is second (after transpose)
fused_sf_f32[:, :intermediate_size] *= (gate_gs / l1_max_gs)
fused_sf_f32[:, intermediate_size:] *= (up_gs / l1_max_gs)
fused_sf = fused_sf_f32.to(torch.float8_e4m3fn)
l1_fp4.append(fused_w_fp4)
l1_sf.append(fused_sf)
l1_gs.append(l1_max_gs)
# L2: down
down_key = f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight"
if down_key in nvfp4_tensors:
down_w = nvfp4_tensors[down_key].to(DEVICE)
down_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE)
down_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale_2"].item()
down_w_fp4 = down_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()
# (intermediate//2, hidden) — K=intermediate packed, N=hidden
l2_fp4.append(down_w_fp4)
l2_sf.append(down_sf.permute(1, 0).contiguous()) # (N, K_sf) → (K_sf, N)
l2_gs.append(down_gs)
else:
# Expert 211 has no down_proj
l2_fp4.append(torch.zeros(3072 // 2, 7168, dtype=torch.float4_e2m1fn_x2, device=DEVICE))
l2_sf.append(torch.ones(3072 // 16, 7168, dtype=torch.float8_e4m3fn, device=DEVICE)) # (K_sf, N)
l2_gs.append(1.0)
return {
'l1_fp4': l1_fp4, 'l1_sf': l1_sf, 'l1_gs': l1_gs,
'l2_fp4': l2_fp4, 'l2_sf': l2_sf, 'l2_gs': l2_gs,
}
def main():
torch.manual_seed(42)
expert_indices = [0, 1, 2]
@@ -135,9 +203,9 @@ def main():
nvfp4_tensors = load_layer_tensors(NVFP4_MODEL_DIR, LAYER_IDX)
print(f" {len(nvfp4_tensors)} tensors loaded")
# Prepare weights
print("\n Preparing NVFP4 weights...")
weights = prepare_nvfp4_moe_weights(nvfp4_tensors, LAYER_IDX, expert_indices)
# Prepare weights — DIRECT PATH (no BF16 round-trip)
print("\n Preparing NVFP4 weights (direct view-cast)...")
weights = prepare_nvfp4_weights_direct(nvfp4_tensors, LAYER_IDX, expert_indices, 3072)
print(f" L1: {len(weights['l1_fp4'])} experts, shape {weights['l1_fp4'][0].shape}")
print(f" L2: {len(weights['l2_fp4'])} experts, shape {weights['l2_fp4'][0].shape}")