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

Simplify pipeline test: BF16 ref + bridge ref + full runner

Browse Source
This commit is contained in:
biondizzle
2026-05-17 21:20:41 +00:00
parent 6eade5e7f8
commit a100bd11c1
1 changed files with 165 additions and 420 deletions
Download Patch File Download Diff File Expand all files Collapse all files

585
tests/test_pipeline_real_weights.py
View File

@@ -1,23 +1,19 @@
"""Step-by-step Pipeline Test: Isolate each component of CuTeDSL runner.
"""Pipeline Test: Step-by-step using CuTeDSL bridge + BF16 reference.
Loads real layer 0 weights from DeepSeek-V4-Pro-NVFP4.
Tests each pipeline stage independently against a BF16 reference:
Stage 1: Token sort + expert assignment
Stage 2: L1 GEMM (gate+up)
Stage 3: SwiGLU activation (with swiglu_limit clamping)
Stage 4: L2 GEMM (down_proj)
Stage 5: Scatter-add with routing weights
Stage 6: Full runner end-to-end
Tests each stage of the NVFP4 MoE pipeline:
1. Token sort + expert assignment
2. L1 GEMM (gate+up)
3. SwiGLU activation
4. L2 GEMM (down_proj)
5. Scatter-add
6. Full runner end-to-end
Strategy: Stages are tested incrementally. Enable STAGE_START to begin at that stage.
All stages from STAGE_START through STAGE_END are tested.
Set STAGE_START=1 STAGE_END=1 to test only stage 1, etc.
Incrementally enable stages with STAGE_START/STAGE_END.
"""
import torch
import torch.nn.functional as F
import sys
import os
import math
import glob
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
@@ -29,19 +25,16 @@ MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
LAYER_IDX = 0
NUM_EXPERTS = 48
HIDDEN_SIZE = 7168
INTERMEDIATE_SIZE = 3072 # per routed expert
INTERMEDIATE_SIZE = 3072
NUM_TOKENS = 8
TOP_K = 6
SWIGLU_LIMIT = 10.0
DEVICE = "cuda"
# Stage control (1-6)
STAGE_START = 1
STAGE_END = 2
STAGE_END = 6
# ============================================================
# Weight loading
# ============================================================
def load_layer_tensors(model_dir, layer_idx):
tensors = {}
for sf in glob.glob(os.path.join(model_dir, "*.safetensors")):
@@ -54,10 +47,92 @@ def load_layer_tensors(model_dir, layer_idx):
return tensors
def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale):
"""Dequantize NVFP4 to BF16. Input: (N, K_packed), scale: (N, K_sf)."""
lut = torch.tensor([
0., 0.5, 1., 1.5, 2., 3., 4., 6.,
-0., -0.5, -1., -1.5, -2., -3., -4., -6.
], dtype=torch.float32, device=packed_uint8.device)
lower = lut[(packed_uint8 & 0x0F).long()]
upper = lut[((packed_uint8 >> 4) & 0x0F).long()]
N = packed_uint8.shape[0]
K = packed_uint8.shape[1] * 2
bf16_vals = torch.stack([lower, upper], dim=-1).reshape(N, K)
K_sf = scale_e4m3.shape[1]
scale_2d = scale_e4m3.float().repeat_interleave(K // K_sf, dim=1)
return (bf16_vals * scale_2d * global_scale).to(torch.bfloat16)
def bf16_moe_reference(hidden_states, nvfp4_tensors, layer_idx, expert_indices, topk_ids, topk_weights, swiglu_limit):
"""Full BF16 reference MoE. Returns output + per-expert intermediates."""
num_tokens = hidden_states.shape[0]
top_k = topk_ids.shape[1]
output = torch.zeros(num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE)
# Per-expert intermediates (keyed by local expert index)
expert_data = {}
for i, e in enumerate(expert_indices):
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()
gate_bf16 = dequantize_nvfp4_weight(gate_w, gate_sf, gate_gs)
up_bf16 = dequantize_nvfp4_weight(up_w, up_sf, up_gs)
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_bf16 = dequantize_nvfp4_weight(down_w, down_sf, down_gs)
else:
down_bf16 = None
# Collect tokens for this expert
mask = (topk_ids == i) # (num_tokens, top_k)
token_rows, k_rows = torch.where(mask)
if token_rows.numel() == 0:
continue
x = hidden_states[token_rows] # (T, H)
gate = x @ gate_bf16.T
up = x @ up_bf16.T
l1_out = torch.cat([gate, up], dim=1)
gate_silu = F.silu(gate)
if swiglu_limit is not None:
gate_silu = gate_silu.clamp(max=swiglu_limit)
up = up.clamp(min=-swiglu_limit, max=swiglu_limit)
activated = gate_silu * up
if down_bf16 is not None:
l2_out = activated @ down_bf16.T
else:
l2_out = activated[:, :HIDDEN_SIZE]
# Scatter
weights = topk_weights[token_rows, k_rows]
weighted = l2_out * weights.unsqueeze(1).to(l2_out.dtype)
output.scatter_add_(0, token_rows.unsqueeze(1).expand(-1, HIDDEN_SIZE), weighted)
expert_data[i] = {
'tokens': token_rows,
'x': x,
'gate': gate, 'up': up, 'l1_out': l1_out,
'activated': activated,
'l2_out': l2_out,
}
return output, expert_data
def prepare_nvfp4_weights(nvfp4_tensors, layer_idx, expert_indices, intermediate_size):
l1_fp4, l1_sf, l1_gs = [], [], []
l2_fp4, l2_sf, l2_gs = [], [], []
for e in expert_indices:
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)
@@ -65,436 +140,106 @@ def prepare_nvfp4_weights(nvfp4_tensors, layer_idx, expert_indices, intermediate
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()
fused_w = torch.cat([gate_w, up_w], dim=0)
fused_w_fp4 = fused_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()
fused_w = torch.cat([gate_w, up_w], dim=0).view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()
fused_sf = torch.cat([gate_sf, up_sf], dim=0).permute(1, 0).contiguous()
l1_max_gs = max(gate_gs, up_gs)
if gate_gs != up_gs:
fused_sf_f32 = fused_sf.float()
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)
sf32 = fused_sf.float()
sf32[:, :intermediate_size] *= (gate_gs / l1_max_gs)
sf32[:, intermediate_size:] *= (up_gs / l1_max_gs)
fused_sf = sf32.to(torch.float8_e4m3fn)
l1_fp4.append(fused_w); l1_sf.append(fused_sf); l1_gs.append(l1_max_gs)
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()
l2_fp4.append(down_w_fp4)
l2_sf.append(down_sf.permute(1, 0).contiguous())
l2_gs.append(down_gs)
dw = nvfp4_tensors[down_key].to(DEVICE)
dsf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE)
dgs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale_2"].item()
l2_fp4.append(dw.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous())
l2_sf.append(dsf.permute(1, 0).contiguous()); l2_gs.append(dgs)
else:
l2_fp4.append(torch.zeros(intermediate_size // 2, HIDDEN_SIZE, dtype=torch.float4_e2m1fn_x2, device=DEVICE))
l2_sf.append(torch.ones(intermediate_size // 16, HIDDEN_SIZE, dtype=torch.float8_e4m3fn, device=DEVICE))
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,
}
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 dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale):
"""Dequantize NVFP4 weight to BF16.
packed_uint8: (N, K_packed) where K_packed = K//2
scale_e4m3: (N, K_sf) where K_sf = K//16
Returns: (N, K) BF16
"""
lut = torch.tensor([
0., 0.5, 1., 1.5, 2., 3., 4., 6.,
-0., -0.5, -1., -1.5, -2., -3., -4., -6.
], dtype=torch.float32, device=packed_uint8.device)
lower = lut[(packed_uint8 & 0x0F).long()]
upper = lut[((packed_uint8 >> 4) & 0x0F).long()]
N = packed_uint8.shape[0]
K = packed_uint8.shape[1] * 2
bf16_vals = torch.stack([lower, upper], dim=-1).reshape(N, K)
K_sf = scale_e4m3.shape[1]
scale_2d = scale_e4m3.float().repeat_interleave(K // K_sf, dim=1)
dequant = bf16_vals * scale_2d * global_scale
return dequant.to(torch.bfloat16)
# ============================================================
# Stage tests
# ============================================================
def test_stage1_sort(hidden_states, topk_ids, topk_weights, expert_indices):
"""Stage 1: Token sorting & expert assignment."""
print("\n--- Stage 1: Token Sort & Expert Assignment ---")
num_tokens = hidden_states.shape[0]
top_k = topk_ids.shape[1]
num_experts = len(expert_indices)
# Reference: simple sort by expert ID
flat_ids = topk_ids.reshape(-1)
flat_weights = topk_weights.reshape(-1)
sort_idx = flat_ids.argsort(stable=True)
sorted_ids = flat_ids[sort_idx]
sorted_weights = flat_weights[sort_idx]
token_indices = torch.arange(num_tokens, device=DEVICE).unsqueeze(1).expand(-1, top_k).reshape(-1)
sorted_token_ids = token_indices[sort_idx]
expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=DEVICE)
for i in range(num_experts):
expert_offsets[i + 1] = (sorted_ids == i).sum()
expert_offsets[1:] = expert_offsets[1:].cumsum(0)
tokens_per_expert = expert_offsets[1:] - expert_offsets[:-1]
print(f" Tokens per expert: min={tokens_per_expert.min().item()} max={tokens_per_expert.max().item()} total={tokens_per_expert.sum().item()}")
print(f" Expert offsets: {expert_offsets.tolist()}")
print(f" Sorted token IDs (first 20): {sorted_token_ids[:20].tolist()}")
return {
'sorted_ids': sorted_ids,
'sorted_token_ids': sorted_token_ids,
'sorted_weights': sorted_weights,
'expert_offsets': expert_offsets,
'tokens_per_expert': tokens_per_expert,
'slot_hidden': hidden_states[sorted_token_ids],
}
def test_stage2_l1_gemm(slot_hidden, expert_offsets, nvfp4_tensors, layer_idx, expert_indices, weights):
"""Stage 2: L1 GEMM (gate+up) using CuTeDSL bridge directly."""
print("\n--- Stage 2: L1 GEMM (gate+up) ---")
from cutedsl.bridge import (
quantize_to_nvfp4, run_nvfp4_grouped_gemm,
assemble_scales_3d_side, make_b_k_major,
)
num_experts = len(expert_indices)
# Stack weights for GEMM
l1_mat_b = torch.stack(weights['l1_fp4'])
l1_scale_b = torch.stack(weights['l1_sf'])
l1_gsb = torch.tensor(weights['l1_gs'], dtype=torch.float32, device=DEVICE)
# Make B-K major
l1_mat_b = make_b_k_major(l1_mat_b)
l1_scale_b = assemble_scales_3d_side(l1_scale_b)
# Quantize activation with dynamic gs
x_fp4, x_sf, l1_gs_dyn = quantize_to_nvfp4(slot_hidden)
l1_gsa = torch.full((num_experts,), l1_gs_dyn, dtype=torch.float32, device=DEVICE)
l1_scale_a = assemble_scales_3d_side(x_sf, expert_offsets[:num_experts+1], num_experts)
# Run GEMM
l1_out = run_nvfp4_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[1:],
global_scale_a=l1_gsa, global_scale_b=l1_gsb,
)
print(f" L1 gs (dynamic): {l1_gs_dyn:.6f}")
print(f" L1 out: shape={l1_out.shape} amax={l1_out.amax().item():.4f} mean={l1_out.mean().item():.4f}")
print(f" L1 out NaN: {torch.isnan(l1_out).any().item()} Inf: {torch.isinf(l1_out).any().item()}")
# BF16 reference for first expert
ref_l1_parts = []
for i, e in enumerate(expert_indices):
start = expert_offsets[i].item()
end = expert_offsets[i + 1].item()
if start == end:
continue
x = slot_hidden[start:end]
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()
gate_bf16 = dequantize_nvfp4_weight(gate_w, gate_sf, gate_gs)
up_bf16 = dequantize_nvfp4_weight(up_w, up_sf, up_gs)
gate_ref = x @ gate_bf16.T
up_ref = x @ up_bf16.T
ref_l1_parts.append((start, end, gate_ref, up_ref))
# Compare L1 output for first expert that has tokens
if ref_l1_parts:
start, end, gate_ref, up_ref = ref_l1_parts[0]
l1_gate = l1_out[start:end, :INTERMEDIATE_SIZE]
l1_up = l1_out[start:end, INTERMEDIATE_SIZE:]
cos_gate = F.cosine_similarity(gate_ref.flatten().unsqueeze(0), l1_gate.flatten().unsqueeze(0)).item()
cos_up = F.cosine_similarity(up_ref.flatten().unsqueeze(0), l1_up.flatten().unsqueeze(0)).item()
print(f" L1 vs BF16 (expert {expert_indices[0]}, {end-start} tokens):")
print(f" gate: cosine={cos_gate:.6f} ref_amax={gate_ref.amax().item():.4f} run_amax={l1_gate.amax().item():.4f}")
print(f" up: cosine={cos_up:.6f} ref_amax={up_ref.amax().item():.4f} run_amax={l1_up.amax().item():.4f}")
return l1_out, ref_l1_parts, l1_gs_dyn
def test_stage3_swiglu(l1_out, ref_l1_parts, swiglu_limit):
"""Stage 3: SwiGLU activation with clamping."""
print("\n--- Stage 3: SwiGLU Activation ---")
# Runner path
gate = l1_out[:, :INTERMEDIATE_SIZE]
up = l1_out[:, INTERMEDIATE_SIZE:]
gate_silu = F.silu(gate)
if swiglu_limit is not None:
gate_silu = gate_silu.clamp(max=swiglu_limit)
up = up.clamp(min=-swiglu_limit, max=swiglu_limit)
activated = gate_silu * up
print(f" activated: shape={activated.shape} amax={activated.amax().item():.4f} mean={activated.mean().item():.4f}")
print(f" gate_silu amax: {gate_silu.amax().item():.4f} up amax: {up.amax().item():.4f}")
# BF16 reference
if ref_l1_parts:
start, end, gate_ref, up_ref = ref_l1_parts[0]
gate_silu_ref = F.silu(gate_ref)
if swiglu_limit is not None:
gate_silu_ref = gate_silu_ref.clamp(max=swiglu_limit)
up_ref = up_ref.clamp(min=-swiglu_limit, max=swiglu_limit)
activated_ref = gate_silu_ref * up_ref
act_runner = activated[start:end]
cos = F.cosine_similarity(activated_ref.flatten().unsqueeze(0), act_runner.flatten().unsqueeze(0)).item()
print(f" vs BF16 (expert 0): cosine={cos:.6f}")
return activated
def test_stage4_l2_gemm(activated, expert_offsets, nvfp4_tensors, layer_idx, expert_indices, weights):
"""Stage 4: L2 GEMM (down_proj)."""
print("\n--- Stage 4: L2 GEMM (down_proj) ---")
from cutedsl.bridge import (
quantize_to_nvfp4, run_nvfp4_grouped_gemm,
assemble_scales_3d_side, make_b_k_major,
)
num_experts = len(expert_indices)
l2_mat_b = torch.stack(weights['l2_fp4'])
l2_scale_b = torch.stack(weights['l2_sf'])
l2_gsb = torch.tensor(weights['l2_gs'], dtype=torch.float32, device=DEVICE)
l2_mat_b = make_b_k_major(l2_mat_b)
l2_scale_b = assemble_scales_3d_side(l2_scale_b)
l2_x_fp4, l2_x_sf, l2_gs_dyn = quantize_to_nvfp4(activated)
l2_gsa = torch.full((num_experts,), l2_gs_dyn, dtype=torch.float32, device=DEVICE)
l2_scale_a = assemble_scales_3d_side(l2_x_sf, expert_offsets[:num_experts+1], num_experts)
l2_out = run_nvfp4_grouped_gemm(
mat_a=l2_x_fp4, mat_b=l2_mat_b,
scale_a=l2_scale_a, scale_b=l2_scale_b,
expert_offsets=expert_offsets[1:],
global_scale_a=l2_gsa, global_scale_b=l2_gsb,
)
print(f" L2 gs (dynamic): {l2_gs_dyn:.6f}")
print(f" L2 out: shape={l2_out.shape} amax={l2_out.amax().item():.4f} mean={l2_out.mean().item():.4f}")
# BF16 reference for first expert
if expert_offsets[1] > 0:
e = expert_indices[0]
start = 0
end = expert_offsets[1].item()
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_bf16 = dequantize_nvfp4_weight(down_w, down_sf, down_gs)
# Need activated reference — use the one we computed in stage 3
gate = l2_out[:end] # Will compare against runner's L2
ref_l2 = activated[start:end] @ down_bf16.T
cos = F.cosine_similarity(ref_l2.flatten().unsqueeze(0), gate.flatten().unsqueeze(0)).item()
print(f" vs BF16 (expert 0): cosine={cos:.6f} ref_amax={ref_l2.amax().item():.4f} run_amax={gate.amax().item():.4f}")
return l2_out, l2_gs_dyn
def test_stage5_scatter(l2_out, expert_offsets, sorted_token_ids, sorted_weights, num_tokens):
"""Stage 5: Scatter-add with routing weights."""
print("\n--- Stage 5: Scatter-add ---")
output = torch.zeros(num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE)
weighted_out = l2_out * sorted_weights.unsqueeze(1).to(l2_out.dtype)
output.scatter_add_(0, sorted_token_ids.unsqueeze(1).expand(-1, HIDDEN_SIZE), weighted_out)
print(f" Output: shape={output.shape} amax={output.amax().item():.4f} mean={output.mean().item():.4f}")
return output
def test_stage6_full_runner(hidden_states, topk_weights, topk_ids, weights, expert_indices):
"""Stage 6: Full CuTeDSL runner end-to-end."""
print("\n--- Stage 6: Full CuTeDSL Runner ---")
from vllm.nvfp4_cutedsl import CuTeDSLMoERunner
runner = CuTeDSLMoERunner(
num_experts=NUM_EXPERTS, hidden_size=HIDDEN_SIZE,
intermediate_size=INTERMEDIATE_SIZE, max_num_tokens=NUM_TOKENS,
top_k=TOP_K, device=DEVICE,
)
runner.l1_fp4 = weights['l1_fp4']
runner.l1_sf = weights['l1_sf']
runner.l1_gs = weights['l1_gs']
runner.l2_fp4 = weights['l2_fp4']
runner.l2_sf = weights['l2_sf']
runner.l2_gs = weights['l2_gs']
runner.set_swiglu_limit(SWIGLU_LIMIT)
with torch.no_grad():
runner.compute_activation_global_scales(hidden_states, topk_weights, topk_ids)
print(f" Warmup gs: L1={runner._l1_activation_global_scale:.6f} L2={runner._l2_activation_global_scale:.6f}")
runner_out = runner.run(hidden_states, topk_weights, topk_ids)
print(f" Runner: amax={runner_out.amax().item():.4f} mean={runner_out.mean().item():.4f}")
print(f" NaN: {torch.isnan(runner_out).any().item()} Inf: {torch.isinf(runner_out).any().item()}")
return runner_out
# ============================================================
# Main
# ============================================================
def main():
torch.cuda.set_device(0)
torch.manual_seed(42)
print(f"=== Step-by-Step Pipeline Test ===")
print(f" Experts: {NUM_EXPERTS}, H={HIDDEN_SIZE}, I={INTERMEDIATE_SIZE}")
print(f" Tokens: {NUM_TOKENS}, top_k={TOP_K}, swiglu_limit={SWIGLU_LIMIT}")
print(f" Stages: {STAGE_START}-{STAGE_END}")
print(f"=== Pipeline Test (stages {STAGE_START}-{STAGE_END}) ===")
print(f" {NUM_EXPERTS} experts, H={HIDDEN_SIZE}, I={INTERMEDIATE_SIZE}, T={NUM_TOKENS}, top_k={TOP_K}")
# Load weights
print("\nLoading checkpoint...")
nvfp4_tensors = load_layer_tensors(MODEL_PATH, LAYER_IDX)
print(f" {len(nvfp4_tensors)} tensors loaded")
expert_indices = list(range(NUM_EXPERTS))
print("Preparing NVFP4 weights...")
weights = prepare_nvfp4_weights(nvfp4_tensors, LAYER_IDX, expert_indices, INTERMEDIATE_SIZE)
print(f" L1: shape {weights['l1_fp4'][0].shape}")
print(f" L2: shape {weights['l2_fp4'][0].shape}")
# Create input
hidden_states = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) * 2.0
topk_ids = torch.zeros(NUM_TOKENS, TOP_K, dtype=torch.int64, device=DEVICE)
for i in range(NUM_TOKENS):
experts_perm = torch.randperm(NUM_EXPERTS)[:TOP_K]
topk_ids[i] = experts_perm
topk_ids[i] = torch.randperm(NUM_EXPERTS)[:TOP_K]
topk_weights = torch.ones(NUM_TOKENS, TOP_K, dtype=torch.float32, device=DEVICE) / TOP_K
# Run stages
sort_data = None
l1_out = None
activated = None
l2_out = None
pipeline_out = None
# BF16 reference
print("\n--- BF16 Reference ---")
ref_out, ref_expert = bf16_moe_reference(hidden_states, nvfp4_tensors, LAYER_IDX, expert_indices, topk_ids, topk_weights, SWIGLU_LIMIT)
print(f" Output: amax={ref_out.amax().item():.4f} mean={ref_out.mean().item():.4f}")
for i in list(ref_expert.keys())[:3]:
d = ref_expert[i]
print(f" Expert {i}: {d['tokens'].numel()} tokens, l1_amax={d['l1_out'].amax().item():.4f} act_amax={d['activated'].amax().item():.4f} l2_amax={d['l2_out'].amax().item():.4f}")
# Full CuTeDSL runner
if STAGE_END >= 6:
print("\n--- Stage 6: Full CuTeDSL Runner ---")
from vllm.nvfp4_cutedsl import CuTeDSLMoERunner
runner = CuTeDSLMoERunner(
num_experts=NUM_EXPERTS, hidden_size=HIDDEN_SIZE,
intermediate_size=INTERMEDIATE_SIZE, max_num_tokens=NUM_TOKENS,
top_k=TOP_K, device=DEVICE,
)
runner.l1_fp4 = weights['l1_fp4']; runner.l1_sf = weights['l1_sf']; runner.l1_gs = weights['l1_gs']
runner.l2_fp4 = weights['l2_fp4']; runner.l2_sf = weights['l2_sf']; runner.l2_gs = weights['l2_gs']
runner.set_swiglu_limit(SWIGLU_LIMIT)
with torch.no_grad():
runner.compute_activation_global_scales(hidden_states, topk_weights, topk_ids)
print(f" Warmup gs: L1={runner._l1_activation_global_scale:.6f} L2={runner._l2_activation_global_scale:.6f}")
runner_out = runner.run(hidden_states, topk_weights, topk_ids)
print(f" Runner: amax={runner_out.amax().item():.4f} mean={runner_out.mean().item():.4f}")
print(f" NaN: {torch.isnan(runner_out).any().item()}")
if STAGE_START <= 1:
sort_data = test_stage1_sort(hidden_states, topk_ids, topk_weights, expert_indices)
if STAGE_START <= 2 and STAGE_END >= 2:
if sort_data is None:
print("\n[Stage 2] Skipped (need stage 1 data)")
else:
l1_out, ref_l1_parts, l1_gs_dyn = test_stage2_l1_gemm(
sort_data['slot_hidden'], sort_data['expert_offsets'],
nvfp4_tensors, LAYER_IDX, expert_indices, weights
)
if STAGE_START <= 3 and STAGE_END >= 3:
if l1_out is None:
print("\n[Stage 3] Skipped (need stage 2 data)")
else:
activated = test_stage3_swiglu(l1_out, ref_l1_parts, SWIGLU_LIMIT)
if STAGE_START <= 4 and STAGE_END >= 4:
if activated is None or sort_data is None:
print("\n[Stage 4] Skipped (need stages 2-3 data)")
else:
l2_out, l2_gs_dyn = test_stage4_l2_gemm(
activated, sort_data['expert_offsets'],
nvfp4_tensors, LAYER_IDX, expert_indices, weights
)
if STAGE_START <= 5 and STAGE_END >= 5:
if l2_out is None or sort_data is None:
print("\n[Stage 5] Skipped (need stages 1-4 data)")
else:
pipeline_out = test_stage5_scatter(
l2_out, sort_data['expert_offsets'],
sort_data['sorted_token_ids'], sort_data['sorted_weights'], NUM_TOKENS
)
if STAGE_START <= 6 and STAGE_END >= 6:
runner_out = test_stage6_full_runner(hidden_states, topk_weights, topk_ids, weights, expert_indices)
# Compare against pipeline reference
if pipeline_out is not None:
cos = F.cosine_similarity(pipeline_out.flatten().unsqueeze(0), runner_out.flatten().unsqueeze(0)).item()
print(f"\n Pipeline vs Runner: cosine={cos:.6f}")
# Also compare against full BF16 reference
print("\n Full BF16 reference for comparison...")
ref_out = torch.zeros(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE)
for i, e in enumerate(expert_indices):
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()
gate_bf16 = dequantize_nvfp4_weight(gate_w, gate_sf, gate_gs)
up_bf16 = dequantize_nvfp4_weight(up_w, up_sf, up_gs)
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_bf16 = dequantize_nvfp4_weight(down_w, down_sf, down_gs)
else:
down_bf16 = None
for t in range(NUM_TOKENS):
for k in range(TOP_K):
eid = topk_ids[t, k].item()
if eid != i:
continue
w = topk_weights[t, k].item()
x = hidden_states[t]
gate = x @ gate_bf16.T
up = x @ up_bf16.T
gate_silu = F.silu(gate).clamp(max=SWIGLU_LIMIT) if SWIGLU_LIMIT else F.silu(gate)
if SWIGLU_LIMIT:
up = up.clamp(min=-SWIGLU_LIMIT, max=SWIGLU_LIMIT)
act = gate_silu * up
if down_bf16 is not None:
y = act @ down_bf16.T
else:
y = act[:HIDDEN_SIZE]
ref_out[t] += w * y
cos = F.cosine_similarity(ref_out.flatten().unsqueeze(0), runner_out.flatten().unsqueeze(0)).item()
mse = (ref_out - runner_out).pow(2).mean().item()
print(f"\n Runner vs BF16: cosine={cos:.6f} MSE={mse:.6e}")
if cos >= 0.98:
print(f" ✅ PASS")
elif cos >= 0.90:
print(f" ⚠️ MARGINAL")
else:
print(f" ❌ FAIL")
print(f" vs BF16: cosine={cos:.6f}")
for t in range(NUM_TOKENS):
ct = F.cosine_similarity(ref_out[t].unsqueeze(0), runner_out[t].unsqueeze(0)).item()
if ct < 0.9:
print(f" Token {t}: cosine={ct:.4f} ref_max={ref_out[t].amax().item():.4f} run_max={runner_out[t].amax().item():.4f}")
# layertest-style bridge reference (should match runner if runner is correct)
if STAGE_END >= 1 and STAGE_START <= 1:
print("\n--- Bridge Reference (run_nvfp4_moe) ---")
from cutedsl.bridge import run_nvfp4_moe
# layertest uses 3 experts — let's use same subset for quick test
small_experts = list(range(min(3, NUM_EXPERTS)))
small_weights = prepare_nvfp4_weights(nvfp4_tensors, LAYER_IDX, small_experts, INTERMEDIATE_SIZE)
small_topk = torch.zeros(NUM_TOKENS, 2, dtype=torch.int32, device=DEVICE)
for i in range(NUM_TOKENS):
small_topk[i] = torch.tensor([0, 1], dtype=torch.int32)
small_tw = torch.tensor([[0.6, 0.4]] * NUM_TOKENS, dtype=torch.float32, device=DEVICE)
bridge_out = run_nvfp4_moe(hidden_states, small_topk, small_tw, small_weights, small_experts)
# BF16 ref for same subset
ref3, _ = bf16_moe_reference(hidden_states, nvfp4_tensors, LAYER_IDX, small_experts, small_topk, small_tw, SWIGLU_LIMIT)
cos3 = F.cosine_similarity(ref3.flatten().unsqueeze(0), bridge_out.flatten().unsqueeze(0)).item()
print(f" Bridge (3 experts) vs BF16: cosine={cos3:.6f}")
if cos3 >= 0.98:
print(" ✅ Bridge reference works correctly")
if __name__ == "__main__":