502 lines
22 KiB
Python
502 lines
22 KiB
Python
"""Step-by-step Pipeline Test: Isolate each component of CuTeDSL runner.
|
|
|
|
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
|
|
|
|
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.
|
|
"""
|
|
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__)), '..'))
|
|
|
|
# ============================================================
|
|
# CONFIG
|
|
# ============================================================
|
|
MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
|
|
LAYER_IDX = 0
|
|
NUM_EXPERTS = 48
|
|
HIDDEN_SIZE = 7168
|
|
INTERMEDIATE_SIZE = 3072 # per routed expert
|
|
NUM_TOKENS = 8
|
|
TOP_K = 6
|
|
SWIGLU_LIMIT = 10.0
|
|
DEVICE = "cuda"
|
|
|
|
# Stage control (1-6)
|
|
STAGE_START = 1
|
|
STAGE_END = 2
|
|
|
|
# ============================================================
|
|
# Weight loading
|
|
# ============================================================
|
|
def load_layer_tensors(model_dir, layer_idx):
|
|
tensors = {}
|
|
for sf in glob.glob(os.path.join(model_dir, "*.safetensors")):
|
|
from safetensors.torch import load_file
|
|
data = load_file(sf)
|
|
for k, v in data.items():
|
|
if f"layers.{layer_idx}." in k and "mlp.experts" in k:
|
|
norm_key = k.removeprefix("model.")
|
|
tensors[norm_key] = v
|
|
return tensors
|
|
|
|
|
|
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)
|
|
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()
|
|
|
|
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_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)
|
|
|
|
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)
|
|
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,
|
|
}
|
|
|
|
|
|
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}")
|
|
|
|
# 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_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
|
|
|
|
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")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|