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

fix: rewrite layertest cleanly, test full MoE pipeline

Browse Source
This commit is contained in:
biondizzle
2026-05-16 03:23:33 +00:00
parent 09ff5c5b98
commit 6139cd6ff5
1 changed files with 16 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files

145
tests/layertest.py
View File

@@ -1,13 +1,10 @@
#!/usr/bin/env python3
"""
Layer 0 kernel comparison test: CuTeDSL NVFP4 kernel vs BF16 reference.
Layer 0 full MoE pipeline test: CuTeDSL NVFP4 vs BF16 reference.
No vLLM, no Docker, no tensor parallelism. Just raw weights + CuTeDSL kernel.
If cosine < 0.99, the test exits with error.
Uses the bridge layer in cutedsl/bridge.py for tensor layout conversion.
Tests the complete pipeline: L1→SiLU→L2→scatter
If cosine < 0.99, exits with error.
"""
import os
import sys
import json
@@ -15,42 +12,25 @@ import glob
import torch
from safetensors import safe_open
# Add repo root to path
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, REPO_ROOT)
from cutedsl.bridge import (
quantize_to_nvfp4,
quantize_weight_to_nvfp4,
assemble_scales_2d_side,
assemble_scales_3d_side,
make_b_k_major,
compute_expert_offsets,
run_nvfp4_grouped_gemm,
)
from cutedsl.moe_pipeline import (
stage_activation,
prepare_nvfp4_moe_weights,
run_nvfp4_moe,
)
# ── Constants ──────────────────────────────────────────────────────────
NVFP4_MODEL_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
LAYER_IDX = 0
DEVICE = "cuda"
COSINE_THRESHOLD = 0.99
# E2M1 FP4 lookup table (for BF16 dequant reference)
E2M1_LUT = torch.tensor([
0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
-0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
], dtype=torch.float32)
# ── Checkpoint loading ─────────────────────────────────────────────────
def find_shards(model_dir):
index_path = os.path.join(model_dir, "model.safetensors.index.json")
key_to_shard = {}
@@ -68,7 +48,6 @@ def find_shards(model_dir):
def load_layer_tensors(model_dir, layer_idx):
"""Load all tensors for a specific layer. Keys normalized (no 'model.' prefix)."""
key_to_shard = find_shards(model_dir)
layer_prefix = f"layers.{layer_idx}."
shard_to_keys = {}
@@ -85,10 +64,7 @@ def load_layer_tensors(model_dir, layer_idx):
return tensors
# ── NVFP4 Dequantization (BF16 reference) ─────────────────────────────
def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale):
"""Dequantize NVFP4 (E2M1 + E4M3 + global) to BF16."""
device = packed_uint8.device
lut = E2M1_LUT.to(device)
lower = lut[(packed_uint8 & 0x0F).long()]
@@ -104,7 +80,6 @@ def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale):
def dequantize_nvfp4_experts(nvfp4_tensors, layer_idx, expert_indices):
"""Dequantize expert weights from NVFP4 checkpoint → BF16."""
experts = {}
for e in expert_indices:
expert = {}
@@ -124,10 +99,7 @@ def dequantize_nvfp4_experts(nvfp4_tensors, layer_idx, expert_indices):
return experts
# ── BF16 MoE Forward ───────────────────────────────────────────────────
def moe_forward_bf16(hidden_states, experts, expert_ids, expert_weights):
"""Run MoE forward pass in pure BF16 (torch.matmul)."""
num_tokens, hidden_size = hidden_states.shape
top_k = expert_ids.shape[1]
output = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE)
@@ -149,81 +121,6 @@ def moe_forward_bf16(hidden_states, experts, expert_ids, expert_weights):
return output
# ── CuTeDSL NVFP4 Kernel MoE Forward ──────────────────────────────────
def moe_forward_nvfp4_l1_only(slot_hidden, nvfp4_tensors, layer_idx, expert_indices, tokens_per_expert):
"""Run L1 (gate+up) GEMM using CuTeDSL.
slot_hidden is already laid out slot-major: [expert0_tokens | expert1_tokens | ...]
"""
num_slots, hidden_size = slot_hidden.shape
num_experts = len(expert_indices)
# Quantize activation
x_fp4, x_sf, x_igs = quantize_to_nvfp4(slot_hidden)
# Load and quantize weights
w_fp4_list = []
w_sf_list = []
w_gs_list = []
for e in expert_indices:
gate_w_key = f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight"
gate_sf_key = f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale"
gate_gs_key = f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale_2"
up_w_key = f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight"
up_sf_key = f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale"
up_gs_key = f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale_2"
gate_w_bf16 = dequantize_nvfp4_weight(
nvfp4_tensors[gate_w_key].to(DEVICE),
nvfp4_tensors[gate_sf_key].to(DEVICE),
nvfp4_tensors[gate_gs_key].item(),
)
up_w_bf16 = dequantize_nvfp4_weight(
nvfp4_tensors[up_w_key].to(DEVICE),
nvfp4_tensors[up_sf_key].to(DEVICE),
nvfp4_tensors[up_gs_key].item(),
)
# Fuse gate + up, transpose to (K=hidden, N=6144)
fused = torch.cat([gate_w_bf16, up_w_bf16], dim=0) # (6144, 7168)
l1_w_bf16 = fused.T # (7168, 6144)
l1_w_fp4, l1_w_sf, l1_w_gs = quantize_weight_to_nvfp4(l1_w_bf16)
w_fp4_list.append(l1_w_fp4)
w_sf_list.append(l1_w_sf)
w_gs_list.append(l1_w_gs)
# Stack and convert to K-major
mat_b = make_b_k_major(torch.stack(w_fp4_list))
# Assemble scale factors
# scale_a: per-expert activation scales, split by expert offsets
x_sf_parts = []
offset = 0
for tpe in tokens_per_expert:
x_sf_parts.append(x_sf[offset:offset+tpe])
offset += tpe
scale_a = assemble_scales_2d_side(x_sf_parts)
scale_b = assemble_scales_3d_side(w_sf_list)
# Expert offsets
expert_offsets = compute_expert_offsets(tokens_per_expert, num_experts)
# Global scales
global_scale_a = torch.tensor([x_igs] * num_experts, dtype=torch.float32, device=DEVICE)
global_scale_b = torch.tensor(w_gs_list, dtype=torch.float32, device=DEVICE)
# Run kernel
out = run_nvfp4_grouped_gemm(
mat_a=x_fp4, mat_b=mat_b,
scale_a=scale_a, scale_b=scale_b,
expert_offsets=expert_offsets,
global_scale_a=global_scale_a, global_scale_b=global_scale_b,
)
return out
def main():
torch.manual_seed(42)
expert_indices = [0, 1, 2]
@@ -231,70 +128,60 @@ def main():
num_tokens = 4
hidden_size = 7168
# ── Load NVFP4 checkpoint ──
print("=" * 70)
print(" Loading NVFP4 checkpoint layer 0")
print("=" * 70)
nvfp4_tensors = load_layer_tensors(NVFP4_MODEL_DIR, LAYER_IDX)
expert_keys = [k for k in sorted(nvfp4_tensors.keys()) if 'experts.0.' in k]
print(f" {len(nvfp4_tensors)} tensors loaded")
for key in expert_keys[:3]:
t = nvfp4_tensors[key]
print(f" {key}: dtype={t.dtype} shape={tuple(t.shape)}")
# ── Prepare NVFP4 weights ──
print("
Preparing NVFP4 weights (dequant → re-quant)...")
# Prepare weights
print("\n Preparing NVFP4 weights...")
weights = prepare_nvfp4_moe_weights(nvfp4_tensors, LAYER_IDX, expert_indices)
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}")
# ── Dequantize → BF16 reference ──
print("
Dequantizing NVFP4 → BF16 reference...")
# Dequantize for BF16 reference
print("\n Dequantizing NVFP4 -> BF16 reference...")
nvfp4_experts_bf16 = dequantize_nvfp4_experts(nvfp4_tensors, LAYER_IDX, expert_indices)
# ── Create test input ──
# Test input
hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) * 2.0
expert_ids = torch.tensor([[0, 1]] * num_tokens, dtype=torch.int32, device=DEVICE)
expert_weights = torch.tensor([[0.6, 0.4]] * num_tokens, dtype=torch.float32, device=DEVICE)
# ── BF16 full MoE reference ──
print("
Running BF16 MoE reference...")
# BF16 reference
print("\n Running BF16 MoE reference...")
ref_output = moe_forward_bf16(hidden_states, nvfp4_experts_bf16, expert_ids, expert_weights)
print(f" BF16 ref: amax={ref_output.abs().max():.4f} mean={ref_output.float().mean():.6f}")
del nvfp4_experts_bf16
torch.cuda.empty_cache()
# ── CuTeDSL NVFP4 full MoE pipeline ──
print("
Running CuTeDSL NVFP4 MoE pipeline (first run compiles, ~1-2 min)...")
# CuTeDSL NVFP4 pipeline
print("\n Running CuTeDSL NVFP4 MoE pipeline (first run compiles)...")
kernel_output = run_nvfp4_moe(
hidden_states, expert_ids, expert_weights,
weights, expert_indices,
)
print(f" Kernel: amax={kernel_output.abs().max():.4f} mean={kernel_output.float().mean():.6f}")
# ── Compare ──
# Compare
cosine = torch.nn.functional.cosine_similarity(
kernel_output.flatten().unsqueeze(0).float(),
ref_output.flatten().unsqueeze(0).float(),
).item()
mse = (kernel_output.float() - ref_output.float()).pow(2).mean().item()
print(f"
{'=' * 70}")
print(f"\n{'=' * 70}")
print(f" RESULT: cosine={cosine:.6f} MSE={mse:.6e}")
print(f"{'=' * 70}")
if cosine < COSINE_THRESHOLD:
print(f" ❌ FAIL: cosine {cosine:.6f} < {COSINE_THRESHOLD}")
print(f" FAIL: cosine {cosine:.6f} < {COSINE_THRESHOLD}")
sys.exit(1)
else:
print(f" ✅ PASS: cosine {cosine:.6f} >= {COSINE_THRESHOLD}")
print(f" PASS: cosine {cosine:.6f} >= {COSINE_THRESHOLD}")
if __name__ == "__main__":