biondizzle
5385de3142
L1 produces (tokens, 6144) gate+up, not (tokens, 7168) hidden. Compare against BF16 L1 reference only.
317 lines
13 KiB
Python
317 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Layer 0 kernel comparison test: CuTeDSL NVFP4 kernel 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.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
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,
|
|
)
|
|
|
|
# ── 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 = {}
|
|
if os.path.exists(index_path):
|
|
with open(index_path) as f:
|
|
index = json.load(f)
|
|
for key, shard in index["weight_map"].items():
|
|
key_to_shard[key] = os.path.join(model_dir, shard)
|
|
else:
|
|
for sf in glob.glob(os.path.join(model_dir, "*.safetensors")):
|
|
with safe_open(sf, framework="pt") as f:
|
|
for key in f.keys():
|
|
key_to_shard[key] = sf
|
|
return key_to_shard
|
|
|
|
|
|
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 = {}
|
|
for key, shard in key_to_shard.items():
|
|
norm_key = key.removeprefix("model.")
|
|
if not norm_key.startswith(layer_prefix):
|
|
continue
|
|
shard_to_keys.setdefault(shard, []).append((key, norm_key))
|
|
tensors = {}
|
|
for shard, keys in shard_to_keys.items():
|
|
with safe_open(shard, framework="pt") as f:
|
|
for orig_key, norm_key in keys:
|
|
tensors[norm_key] = f.get_tensor(orig_key)
|
|
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()]
|
|
upper = lut[((packed_uint8 >> 4) & 0x0F).long()]
|
|
out_features = packed_uint8.shape[0]
|
|
in_features = packed_uint8.shape[1] * 2
|
|
unpacked = torch.empty(out_features, in_features, dtype=torch.float32, device=device)
|
|
unpacked[:, 0::2] = lower
|
|
unpacked[:, 1::2] = upper
|
|
block_scale = scale_e4m3.float()
|
|
block_expanded = block_scale.repeat_interleave(16, dim=1)[:, :in_features]
|
|
return (unpacked * block_expanded * global_scale).to(torch.bfloat16)
|
|
|
|
|
|
def dequantize_nvfp4_experts(nvfp4_tensors, layer_idx, expert_indices):
|
|
"""Dequantize expert weights from NVFP4 checkpoint → BF16."""
|
|
experts = {}
|
|
for e in expert_indices:
|
|
expert = {}
|
|
for proj in ["gate_proj", "up_proj", "down_proj"]:
|
|
weight_key = f"layers.{layer_idx}.mlp.experts.{e}.{proj}.weight"
|
|
scale_key = f"layers.{layer_idx}.mlp.experts.{e}.{proj}.weight_scale"
|
|
gs_key = f"layers.{layer_idx}.mlp.experts.{e}.{proj}.weight_scale_2"
|
|
if weight_key not in nvfp4_tensors:
|
|
if proj == "down_proj" and e == 211:
|
|
continue
|
|
raise KeyError(f"Missing {weight_key}")
|
|
weight = nvfp4_tensors[weight_key].to(DEVICE)
|
|
scale = nvfp4_tensors[scale_key].to(DEVICE)
|
|
global_scale = nvfp4_tensors[gs_key].item()
|
|
expert[proj] = dequantize_nvfp4_weight(weight, scale, global_scale)
|
|
experts[e] = expert
|
|
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)
|
|
for t in range(num_tokens):
|
|
for k in range(top_k):
|
|
e = expert_ids[t, k].item()
|
|
w = expert_weights[t, k].item()
|
|
if e not in experts:
|
|
continue
|
|
x = hidden_states[t]
|
|
gate = x @ experts[e]["gate_proj"].T
|
|
up = x @ experts[e]["up_proj"].T
|
|
activated = torch.nn.functional.silu(gate) * up
|
|
if "down_proj" in experts[e]:
|
|
y = activated @ experts[e]["down_proj"].T
|
|
else:
|
|
y = activated[:hidden_size]
|
|
output[t] += w * y
|
|
return output
|
|
|
|
|
|
# ── CuTeDSL NVFP4 Kernel MoE Forward ──────────────────────────────────
|
|
|
|
def moe_forward_nvfp4_l1_only(hidden_states, nvfp4_tensors, layer_idx, expert_ids, expert_weights):
|
|
"""Run MoE forward pass using the CuTeDSL NVFP4 kernel via bridge."""
|
|
num_tokens, hidden_size = hidden_states.shape
|
|
top_k = expert_ids.shape[1]
|
|
|
|
# Map expert IDs to local indices
|
|
unique_experts = sorted(set(expert_ids.flatten().tolist()))
|
|
num_experts = len(unique_experts)
|
|
expert_map = {e: i for i, e in enumerate(unique_experts)}
|
|
|
|
# ── Step 1: Quantize activation ──
|
|
x_fp4, x_sf, x_igs = quantize_to_nvfp4(hidden_states)
|
|
|
|
# ── Step 2: Load and quantize weights from checkpoint ──
|
|
# Checkpoint weight is (N, K//2) uint8, scale is (N, K//16) float8_e4m3fn
|
|
# We need to dequantize to BF16 first, then re-quantize with our pipeline
|
|
# (the checkpoint format is the same NVFP4, but we need to use our quantizer
|
|
# for the bridge to produce correct tensor layouts)
|
|
#
|
|
# Actually, we can load the checkpoint weights directly as float4_e2m1fn_x2
|
|
# and the scales as float8_e4m3fn. Just need to reshape.
|
|
|
|
w_fp4_list = []
|
|
w_sf_list = []
|
|
w_gs_list = []
|
|
|
|
for e in unique_experts:
|
|
# L1: gate + up fused → (2*3072, 3584) packed
|
|
# For now, dequantize checkpoint to BF16 then re-quantize
|
|
# This ensures the FP4 values match our quantization convention
|
|
|
|
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: (6144, 7168) → quantize as (K=7168, N=6144)
|
|
# Kernel expects B: (experts, K, N) with K=hidden, N=intermediate
|
|
fused_l1 = torch.cat([gate_w_bf16, up_w_bf16], dim=0) # (6144, 7168)
|
|
# B is (K, N) where K=hidden=7168, N=6144
|
|
l1_w_bf16 = fused_l1.T # (7168, 6144) — K=7168 is dim 0
|
|
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 weights and convert to K-major
|
|
mat_b = torch.stack(w_fp4_list) # (experts, K//2, N) N-major
|
|
mat_b = make_b_k_major(mat_b) # (experts, K//2, N) K-major
|
|
|
|
# Assemble scale factors
|
|
scale_a = assemble_scales_2d_side(
|
|
[x_sf[e*top_k:(e+1)*top_k] for e in range(num_experts)]
|
|
)
|
|
scale_b = assemble_scales_3d_side(w_sf_list)
|
|
|
|
# Expert offsets
|
|
tokens_per_expert = [top_k] * num_experts # simplified: each expert gets top_k tokens
|
|
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 the 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
|
|
|
|
|
|
# ── Main ───────────────────────────────────────────────────────────────
|
|
|
|
def main():
|
|
torch.manual_seed(42)
|
|
expert_indices = [0, 1, 2]
|
|
top_k = 2
|
|
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 and LAYER_IDX == 0]
|
|
print(f" {len(nvfp4_tensors)} tensors loaded")
|
|
for key in expert_keys[:5]:
|
|
t = nvfp4_tensors[key]
|
|
print(f" {key}: dtype={t.dtype} shape={tuple(t.shape)}")
|
|
|
|
# ── Dequantize → BF16 reference ──
|
|
print("\n Dequantizing NVFP4 → BF16...")
|
|
nvfp4_experts_bf16 = dequantize_nvfp4_experts(nvfp4_tensors, LAYER_IDX, expert_indices)
|
|
for e in expert_indices[:2]:
|
|
for proj, w in nvfp4_experts_bf16[e].items():
|
|
print(f" Expert {e} {proj}: shape={tuple(w.shape)} amax={w.abs().max():.4f}")
|
|
|
|
# ── Create 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 L1 reference (gate+up only) ──
|
|
print("\n Running BF16 L1 reference...")
|
|
ref_l1 = torch.zeros(num_tokens, 6144, dtype=torch.bfloat16, device=DEVICE)
|
|
for t in range(num_tokens):
|
|
for k in range(top_k):
|
|
e = expert_ids[t, k].item()
|
|
w = expert_weights[t, k].item()
|
|
if e not in nvfp4_experts_bf16:
|
|
continue
|
|
x = hidden_states[t]
|
|
gate = x @ nvfp4_experts_bf16[e]["gate_proj"].T # (3072,)
|
|
up = x @ nvfp4_experts_bf16[e]["up_proj"].T # (3072,)
|
|
ref_l1[t] += w * torch.cat([gate, up])
|
|
|
|
print(f" BF16 L1 ref: amax={ref_l1.abs().max():.4f} mean={ref_l1.float().mean():.6f}")
|
|
|
|
del nvfp4_experts_bf16
|
|
torch.cuda.empty_cache()
|
|
|
|
# ── CuTeDSL NVFP4 L1 kernel ──
|
|
print("\n Running CuTeDSL NVFP4 L1 kernel (first run compiles, ~1-2 min)...")
|
|
kernel_l1 = moe_forward_nvfp4_l1_only(hidden_states, nvfp4_tensors, LAYER_IDX, expert_ids, expert_weights)
|
|
print(f" Kernel L1: amax={kernel_l1.abs().max():.4f} mean={kernel_l1.float().mean():.6f}")
|
|
|
|
# ── Compare ──
|
|
cosine = torch.nn.functional.cosine_similarity(
|
|
kernel_l1.flatten().unsqueeze(0).float(),
|
|
ref_l1.flatten().unsqueeze(0).float(),
|
|
).item()
|
|
mse = (kernel_l1.float() - ref_l1.float()).pow(2).mean().item()
|
|
|
|
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}")
|
|
sys.exit(1)
|
|
else:
|
|
print(f" ✅ PASS: cosine {cosine:.6f} >= {COSINE_THRESHOLD}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|