345 lines
14 KiB
Python
345 lines
14 KiB
Python
"""Step-by-step Pipeline Test: CuTeDSL runner components vs reference.
|
|
|
|
Loads real layer 0 weights from DeepSeek-V4-Pro-NVFP4.
|
|
Tests each pipeline stage independently:
|
|
1. Token sorting & expert assignment
|
|
2. L1 GEMM (gate+up)
|
|
3. SwiGLU activation (with swiglu_limit clamping)
|
|
4. L2 GEMM (down_proj)
|
|
5. Scatter-add with routing weights
|
|
6. Full runner vs reference
|
|
|
|
Strategy: Comment out stages to isolate bugs, then uncomment one by one.
|
|
"""
|
|
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 — toggle which stages to test
|
|
# ============================================================
|
|
MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
|
|
LAYER_IDX = 0
|
|
NUM_EXPERTS = 48 # local experts per rank (256/8=32, but model uses 48)
|
|
HIDDEN_SIZE = 7168
|
|
INTERMEDIATE_SIZE = 3072 # per routed expert (18432 is shared expert)
|
|
NUM_TOKENS = 8
|
|
TOP_K = 6
|
|
SWIGLU_LIMIT = 10.0
|
|
DEVICE = "cuda"
|
|
|
|
# Which stages to enable (uncomment incrementally to find bugs)
|
|
ENABLE_SORT = True
|
|
ENABLE_L1_GEMM = True
|
|
ENABLE_SWIGLU = True
|
|
ENABLE_L2_GEMM = True
|
|
ENABLE_SCATTER = True
|
|
ENABLE_FULL_RUNNER = True
|
|
|
|
# ============================================================
|
|
# Weight loading (from layertest.py pattern)
|
|
# ============================================================
|
|
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():
|
|
# Match both "layers.X." and "model.layers.X."
|
|
if f"layers.{layer_idx}." in k and "mlp.experts" in k:
|
|
# Normalize: strip "model." prefix if present
|
|
norm_key = k.removeprefix("model.")
|
|
tensors[norm_key] = v
|
|
return tensors
|
|
|
|
|
|
def prepare_nvfp4_weights(nvfp4_tensors, layer_idx, expert_indices, intermediate_size):
|
|
"""Prepare weights via direct view-cast (same as layertest)."""
|
|
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 for reference computation.
|
|
|
|
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)
|
|
|
|
# scale_e4m3 is (N, K_sf) where K_sf = K//16
|
|
K_sf = scale_e4m3.shape[1]
|
|
scale_2d = scale_e4m3.float().repeat_interleave(K // K_sf, dim=1) # (N, K)
|
|
|
|
dequant = bf16_vals * scale_2d * global_scale
|
|
return dequant.to(torch.bfloat16)
|
|
|
|
|
|
# ============================================================
|
|
# Reference pipeline (step by step, BF16)
|
|
# ============================================================
|
|
def reference_moe_bf16(hidden_states, nvfp4_tensors, layer_idx, expert_indices, topk_ids, topk_weights, swiglu_limit):
|
|
"""BF16 reference: dequantize weights, run MoE step by step."""
|
|
num_tokens = hidden_states.shape[0]
|
|
top_k = topk_ids.shape[1]
|
|
output = torch.zeros(num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE)
|
|
|
|
# Store intermediate results for comparison
|
|
intermediates = {}
|
|
|
|
# Sort tokens by expert (for comparing with runner's sorted approach)
|
|
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]
|
|
|
|
intermediates['sorted_ids'] = sorted_ids
|
|
intermediates['sorted_token_ids'] = sorted_token_ids
|
|
intermediates['sorted_weights'] = sorted_weights
|
|
|
|
# Expert offsets
|
|
expert_id_range = torch.arange(len(expert_indices), device=DEVICE)
|
|
tokens_per_expert = torch.zeros(len(expert_indices), dtype=torch.int32, device=DEVICE)
|
|
for i, e in enumerate(expert_indices):
|
|
tokens_per_expert[i] = (sorted_ids == i).sum()
|
|
expert_offsets = torch.zeros(len(expert_indices) + 1, dtype=torch.int32, device=DEVICE)
|
|
expert_offsets[1:] = tokens_per_expert.cumsum(0)
|
|
|
|
intermediates['expert_offsets'] = expert_offsets
|
|
intermediates['tokens_per_expert'] = tokens_per_expert
|
|
|
|
# Gather hidden states for sorted tokens
|
|
slot_hidden = hidden_states[sorted_token_ids]
|
|
intermediates['slot_hidden'] = slot_hidden
|
|
|
|
# Per-expert computation
|
|
l1_out_all = []
|
|
activated_all = []
|
|
l2_out_all = []
|
|
|
|
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] # (T, H)
|
|
|
|
# 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()
|
|
|
|
gate_bf16 = dequantize_nvfp4_weight(gate_w, gate_sf, gate_gs) # (intermediate, hidden)
|
|
up_bf16 = dequantize_nvfp4_weight(up_w, up_sf, up_gs) # (intermediate, hidden)
|
|
|
|
gate = x @ gate_bf16.T # (T, intermediate)
|
|
up = x @ up_bf16.T # (T, intermediate)
|
|
|
|
l1_out = torch.cat([gate, up], dim=1) # (T, 2*intermediate)
|
|
l1_out_all.append((start, end, l1_out))
|
|
|
|
# SwiGLU
|
|
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
|
|
activated_all.append((start, end, activated))
|
|
|
|
# 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_bf16 = dequantize_nvfp4_weight(down_w, down_sf, down_gs) # (hidden, intermediate)
|
|
l2_out = activated @ down_bf16.T # (T, H)
|
|
else:
|
|
l2_out = activated[:, :HIDDEN_SIZE]
|
|
|
|
l2_out_all.append((start, end, l2_out))
|
|
|
|
# Scatter-add
|
|
weighted = l2_out * sorted_weights[start:end].unsqueeze(1).to(l2_out.dtype)
|
|
output.scatter_add_(0, sorted_token_ids[start:end].unsqueeze(1).expand(-1, HIDDEN_SIZE), weighted)
|
|
|
|
intermediates['l1_out_all'] = l1_out_all
|
|
intermediates['activated_all'] = activated_all
|
|
intermediates['l2_out_all'] = l2_out_all
|
|
intermediates['output'] = output
|
|
|
|
return intermediates
|
|
|
|
|
|
# ============================================================
|
|
# Main test
|
|
# ============================================================
|
|
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: sort={ENABLE_SORT} L1={ENABLE_L1_GEMM} swiglu={ENABLE_SWIGLU} L2={ENABLE_L2_GEMM} scatter={ENABLE_SCATTER} full={ENABLE_FULL_RUNNER}")
|
|
|
|
# Load real weights
|
|
print("\n[1/6] Loading checkpoint...")
|
|
nvfp4_tensors = load_layer_tensors(MODEL_PATH, LAYER_IDX)
|
|
print(f" {len(nvfp4_tensors)} tensors loaded")
|
|
|
|
# Figure out expert indices for this rank
|
|
# layer 0 has experts 0-255, we use first NUM_EXPERTS
|
|
expert_indices = list(range(NUM_EXPERTS))
|
|
print(f" Using experts: {expert_indices[:5]}... (first 5 of {NUM_EXPERTS})")
|
|
|
|
print("\n[2/6] Preparing NVFP4 weights (direct view-cast)...")
|
|
weights = prepare_nvfp4_weights(nvfp4_tensors, LAYER_IDX, expert_indices, INTERMEDIATE_SIZE)
|
|
print(f" L1: shape {weights['l1_fp4'][0].shape} dtype {weights['l1_fp4'][0].dtype}")
|
|
print(f" L2: shape {weights['l2_fp4'][0].shape} dtype {weights['l2_fp4'][0].dtype}")
|
|
|
|
# Create input
|
|
hidden_states = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) * 2.0
|
|
|
|
# Realistic top-k
|
|
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
|
|
|
|
# ---- Reference (BF16) ----
|
|
print("\n[3/6] Running BF16 reference pipeline...")
|
|
ref = reference_moe_bf16(hidden_states, nvfp4_tensors, LAYER_IDX, expert_indices, topk_ids, topk_weights, SWIGLU_LIMIT)
|
|
print(f" L1 samples: {len(ref['l1_out_all'])} experts with tokens")
|
|
if ref['l1_out_all']:
|
|
_, _, l1 = ref['l1_out_all'][0]
|
|
print(f" L1 out[0]: amax={l1.amax().item():.4f} mean={l1.mean().item():.4f}")
|
|
if ref['activated_all']:
|
|
_, _, act = ref['activated_all'][0]
|
|
print(f" activated[0]: amax={act.amax().item():.4f} mean={act.mean().item():.4f}")
|
|
if ref['l2_out_all']:
|
|
_, _, l2 = ref['l2_out_all'][0]
|
|
print(f" L2 out[0]: amax={l2.amax().item():.4f} mean={l2.mean().item():.4f}")
|
|
print(f" Final: amax={ref['output'].amax().item():.4f} mean={ref['output'].mean().item():.4f}")
|
|
|
|
# ---- CuTeDSL Runner ----
|
|
print("\n[4/6] Creating 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)
|
|
|
|
print("\n[5/6] Running CuTeDSL runner (with warmup gs)...")
|
|
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()}")
|
|
|
|
# ---- Comparison ----
|
|
print("\n[6/6] Comparing runner vs BF16 reference...")
|
|
ref_out = ref['output']
|
|
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" Cosine: {cos:.6f} MSE: {mse:.6e}")
|
|
|
|
# Per-token
|
|
low_cos = 0
|
|
for i in range(NUM_TOKENS):
|
|
cos_i = F.cosine_similarity(ref_out[i].unsqueeze(0), runner_out[i].unsqueeze(0)).item()
|
|
if cos_i < 0.95:
|
|
low_cos += 1
|
|
if low_cos <= 5:
|
|
print(f" Token {i}: cosine={cos_i:.4f}")
|
|
|
|
if cos >= 0.98:
|
|
print(f"\n✅ PASS: cosine {cos:.6f}")
|
|
elif cos >= 0.90:
|
|
print(f"\n⚠️ MARGINAL: cosine {cos:.6f}")
|
|
else:
|
|
print(f"\n❌ FAIL: cosine {cos:.6f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|