feat: full NVFP4 MoE pipeline (L1→SiLU→L2→scatter)
cutedsl/moe_pipeline.py: complete pipeline - stage_activation: BF16 → NVFP4 (keeps data in FP4) - L1 GEMM: NVFP4 × NVFP4 → BF16 (gate+up) - SiLU(gate) * up: BF16 (only nonlinear, can't avoid) - Re-quantize: BF16 → NVFP4 (back to native) - L2 GEMM: NVFP4 × NVFP4 → BF16 (down_proj) - Scatter with routing weights → BF16 output layertest.py: now tests the FULL MoE pipeline against BF16 reference. NVFP4-native: both GEMMs use float4_e2m1fn_x2 for A and B, float8_e4m3fn for block scales, float32 for global scales. BF16 only for SiLU activation and final scatter.
This commit is contained in:
@@ -29,6 +29,12 @@ from cutedsl.bridge import (
|
||||
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"
|
||||
@@ -231,76 +237,56 @@ def main():
|
||||
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]
|
||||
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[:5]:
|
||||
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)...")
|
||||
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("\n Dequantizing NVFP4 → BF16...")
|
||||
print("
|
||||
Dequantizing NVFP4 → BF16 reference...")
|
||||
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)
|
||||
|
||||
# ── Build slot-based layout for grouped GEMM ──
|
||||
# The kernel expects activation laid out as [expert_0_tokens | expert_1_tokens | ...]
|
||||
# Each token can appear in multiple experts (top-k routing)
|
||||
num_slots = num_tokens * top_k
|
||||
slot_expert = expert_ids.flatten() # (num_slots,)
|
||||
|
||||
# Build per-expert token lists
|
||||
expert_token_lists = {e: [] for e in expert_indices}
|
||||
for t in range(num_tokens):
|
||||
for k in range(top_k):
|
||||
e = expert_ids[t, k].item()
|
||||
expert_token_lists[e].append(t)
|
||||
|
||||
tokens_per_expert = [len(expert_token_lists[e]) for e in expert_indices]
|
||||
|
||||
# Build slot-major activation: concat tokens for each expert
|
||||
slot_hidden = torch.cat([
|
||||
hidden_states[expert_token_lists[e]] for e in expert_indices
|
||||
], dim=0) # (num_slots, hidden_size)
|
||||
|
||||
expert_offsets = compute_expert_offsets(tokens_per_expert, len(expert_indices))
|
||||
|
||||
# ── BF16 L1 reference (slot-major, matching kernel output) ──
|
||||
print("\n Running BF16 L1 reference...")
|
||||
ref_l1_parts = []
|
||||
for e in expert_indices:
|
||||
for t in expert_token_lists[e]:
|
||||
gate = hidden_states[t] @ nvfp4_experts_bf16[e]["gate_proj"].T
|
||||
up = hidden_states[t] @ nvfp4_experts_bf16[e]["up_proj"].T
|
||||
ref_l1_parts.append(torch.cat([gate, up]))
|
||||
ref_l1 = torch.cat(ref_l1_parts, dim=0) # (num_slots, 6144)
|
||||
print(f" BF16 L1 ref: amax={ref_l1.abs().max():.4f} mean={ref_l1.float().mean():.6f}")
|
||||
# ── BF16 full MoE reference ──
|
||||
print("
|
||||
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 L1 kernel ──
|
||||
print("\n Running CuTeDSL NVFP4 L1 kernel (first run compiles, ~1-2 min)...")
|
||||
kernel_l1 = moe_forward_nvfp4_l1_only(slot_hidden, nvfp4_tensors, LAYER_IDX, expert_indices, tokens_per_expert)
|
||||
print(f" Kernel L1: amax={kernel_l1.abs().max():.4f} mean={kernel_l1.float().mean():.6f}")
|
||||
# ── CuTeDSL NVFP4 full MoE pipeline ──
|
||||
print("
|
||||
Running CuTeDSL NVFP4 MoE pipeline (first run compiles, ~1-2 min)...")
|
||||
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 ──
|
||||
ref_flat = ref_l1.flatten()
|
||||
kernel_flat = kernel_l1.flatten()
|
||||
|
||||
cosine = torch.nn.functional.cosine_similarity(
|
||||
kernel_flat.unsqueeze(0).float(),
|
||||
ref_flat.unsqueeze(0).float(),
|
||||
kernel_output.flatten().unsqueeze(0).float(),
|
||||
ref_output.flatten().unsqueeze(0).float(),
|
||||
).item()
|
||||
mse = (kernel_flat.float() - ref_flat.float()).pow(2).mean().item()
|
||||
mse = (kernel_output.float() - ref_output.float()).pow(2).mean().item()
|
||||
|
||||
print(f"\n{'=' * 70}")
|
||||
print(f"
|
||||
{'=' * 70}")
|
||||
print(f" RESULT: cosine={cosine:.6f} MSE={mse:.6e}")
|
||||
print(f"{'=' * 70}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user