From 0cdcc4144af80799131e4704205a2302d945e7d8 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 03:13:54 +0000 Subject: [PATCH] refactor: add cutedsl/bridge.py, rewrite layertest to use it bridge.py: clean API for CuTeDSL kernel - quantize_to_nvfp4 / quantize_weight_to_nvfp4 - assemble_scales_2d_side / assemble_scales_3d_side - make_b_k_major (stride conversion) - compute_expert_offsets - run_nvfp4_grouped_gemm (full kernel launch) layertest.py: now uses bridge layer, tests with real DeepSeek-V4 layer 0 weights (7168 hidden, 6144 intermediate). The bridge code will be reused by the vLLM integration layer. --- cutedsl/bridge.py | 274 +++++++++++++++++++++++++++++++++++++++++++++ tests/layertest.py | 264 +++++++++++++++++++------------------------ 2 files changed, 390 insertions(+), 148 deletions(-) create mode 100644 cutedsl/bridge.py diff --git a/cutedsl/bridge.py b/cutedsl/bridge.py new file mode 100644 index 00000000..a061430a --- /dev/null +++ b/cutedsl/bridge.py @@ -0,0 +1,274 @@ +""" +Bridge layer for the CuTeDSL NVFP4 MoE kernel. + +Handles tensor layout conversion from our pipeline's format to what +the ScaledGroupedGemmKernel expects: +- BF16 → NVFP4 quantization (float4_e2m1fn_x2) +- Scale factor assembly (padding + swizzle) +- B tensor K-major stride conversion +- Expert offset computation +""" +import math +import torch +import cutlass +import cutlass.cute as cute +import cutlass.torch as cutlass_torch +import cutlass.utils as utils + +from cutedsl.kernel.moe.torch_scaled_grouped_mm import ( + ScaledGroupedGemmKernel, + pad_and_swizzle_single, + assemble_raw_scales_2d3d_2d_side, + assemble_raw_scales_2d3d_3d_side, + cat_byte_reinterpretable_tensors, + stack_byte_reinterpretable_tensors, +) + +# ── Constants ────────────────────────────────────────────────────────── + +E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] +SF_VEC_SIZE = 16 # NVFP4 block size + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def round_up(a, b): + return ceil_div(a, b) * b + + +# ── Quantization ────────────────────────────────────────────────────── + +def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE): + """Quantize BF16 tensor to NVFP4. + + Args: + x_bf16: (..., D) BF16 tensor + + Returns: + x_fp4: (..., D//2) float4_e2m1fn_x2 — native PyTorch FP4 + x_sf: (..., D//16) float8_e4m3fn — block scales + global_scale: float32 scalar + """ + x_f32 = x_bf16.float() + amax = x_f32.abs().max().clamp(min=1e-8).float() + global_scale = amax / (6.0 * 448.0) + x_norm = x_f32 / global_scale + + last_dim = x_norm.shape[-1] + n_blocks = ceil_div(last_dim, block_size) + + if last_dim % block_size != 0: + pad_size = n_blocks * block_size - last_dim + x_norm = torch.nn.functional.pad(x_norm, (0, pad_size)) + + x_reshaped = x_norm.reshape(*x_norm.shape[:-1], n_blocks, block_size) + block_amax = x_reshaped.abs().amax(dim=-1).clamp(min=1e-8) + block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn) + + # Nearest E2M1 + block_sf_expanded = block_scale.float().unsqueeze(-1) + x_scaled = x_reshaped / block_sf_expanded.clamp(min=1e-8) + + magnitudes = torch.tensor(E2M1_MAGNITUDES, dtype=torch.float32, device=x_bf16.device) + signs = torch.sign(x_scaled) + abs_scaled = x_scaled.abs().unsqueeze(-1) + distances = (abs_scaled - magnitudes).abs() + indices = distances.argmin(dim=-1) + + nibbles = torch.where(signs < 0, indices + 8, indices).to(torch.uint8) + even = nibbles[..., ::2] + odd = nibbles[..., 1::2] + packed = (odd << 4) | even + + packed_shape = list(x_bf16.shape) + packed_shape[-1] = last_dim // 2 + x_fp4 = packed.view(torch.float4_e2m1fn_x2).reshape(packed_shape) + + sf_shape = list(x_bf16.shape[:-1]) + [n_blocks] + block_scale = block_scale.reshape(sf_shape) + + return x_fp4, block_scale, global_scale + + +def quantize_weight_to_nvfp4(w_bf16, block_size=SF_VEC_SIZE): + """Quantize BF16 weight matrix to NVFP4. + + The weight is (K, N) where K is the input dim (packed dimension). + Block scales are computed along K (dim 0). + + Args: + w_bf16: (K, N) BF16 weight matrix + + Returns: + w_fp4: (K//2, N) float4_e2m1fn_x2 — K is the packed dim + w_sf: (K//16, N) float8_e4m3fn — block scales along K + global_scale: float32 scalar + """ + K, N = w_bf16.shape + w_f32 = w_bf16.float() + amax = w_f32.abs().max().clamp(min=1e-8).float() + global_scale = amax / (6.0 * 448.0) + w_norm = w_f32 / global_scale + + k_blocks = ceil_div(K, block_size) + if K % block_size != 0: + w_norm = torch.nn.functional.pad(w_norm, (0, 0, 0, k_blocks * block_size - K)) + + w_reshaped = w_norm.reshape(k_blocks, block_size, N) + w_block_amax = w_reshaped.abs().amax(dim=1).clamp(min=1e-8) + w_sf = (w_block_amax / 6.0).to(torch.float8_e4m3fn) + + w_block_sf = w_sf.float().unsqueeze(1) + w_scaled = w_reshaped / w_block_sf.clamp(min=1e-8) + + magnitudes = torch.tensor(E2M1_MAGNITUDES, dtype=torch.float32, device=w_bf16.device) + signs = torch.sign(w_scaled) + abs_scaled = w_scaled.abs().unsqueeze(-1) + distances = (abs_scaled - magnitudes).abs() + indices = distances.argmin(dim=-1) + nibbles = torch.where(signs < 0, indices + 8, indices).to(torch.uint8) + + even = nibbles[:, ::2, :] + odd = nibbles[:, 1::2, :] + packed = (odd << 4) | even + + w_fp4 = packed.reshape(K // 2, N).view(torch.float4_e2m1fn_x2) + return w_fp4, w_sf, global_scale + + +# ── Scale Factor Assembly ───────────────────────────────────────────── + +def assemble_scales_2d_side(raw_scales): + """Assemble activation scale factors for the 2Dx3D scenario. + + Args: + raw_scales: list of (M_e, K_sf) float8_e4m3fn tensors, one per expert + + Returns: + Assembled and swizzled scale tensor + """ + return assemble_raw_scales_2d3d_2d_side(raw_scales) + + +def assemble_scales_3d_side(raw_scales): + """Assemble weight scale factors for the 2Dx3D scenario. + + Args: + raw_scales: list of (K_sf, N) float8_e4m3fn tensors, one per expert + NOTE: These will be transposed to (N, K_sf) before swizzling, + since the kernel expects N as the non-K dimension. + + Returns: + Assembled and swizzled scale tensor + """ + # Kernel expects (N, K_sf) — transpose before swizzling + transposed = [sf.T.contiguous() for sf in raw_scales] + return assemble_raw_scales_2d3d_3d_side(transposed) + + +# ── Tensor Layout Conversion ────────────────────────────────────────── + +def make_b_k_major(b_tensor): + """Convert B tensor from N-major to K-major layout. + + The kernel expects B with stride (E*K*N, 1, K) — K is contiguous. + torch.stack produces stride (E*K*N, N, 1) — N is contiguous. + + Args: + b_tensor: (experts, K_packed, N_packed) float4_e2m1fn_x2, N-major + + Returns: + Same shape, K-major strides + """ + return b_tensor.permute(0, 2, 1).contiguous().permute(0, 2, 1) + + +def compute_expert_offsets(tokens_per_expert, num_experts, device="cuda"): + """Compute cumulative token offsets for the grouped GEMM. + + Args: + tokens_per_expert: list of int, one per expert + + Returns: + offs: (num_experts,) int32 — cumulative sum + """ + offs = torch.tensor( + [sum(tokens_per_expert[:e+1]) for e in range(num_experts)], + dtype=torch.int32, device=device, + ) + return offs + + +# ── Kernel Launch ───────────────────────────────────────────────────── + +def run_nvfp4_grouped_gemm( + mat_a, # (tokens_sum, K_packed) float4_e2m1fn_x2 + mat_b, # (experts, K_packed, N_packed) float4_e2m1fn_x2, K-major + scale_a, # assembled 2D side (padded + swizzled) + scale_b, # assembled 3D side (padded + swizzled) + expert_offsets, # (experts,) int32 cumulative token offsets + global_scale_a=None, # (experts,) float32 + global_scale_b=None, # (experts,) float32 + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), +): + """Run the CuTeDSL NVFP4 scaled grouped GEMM. + + 2Dx3D: A(tokens, K) x B(experts, K, N) -> C(tokens, N) + """ + num_experts = mat_b.shape[0] + n_dim = mat_b.shape[2] # packed N (in float4 elements) + tokens_sum = mat_a.shape[0] + + out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device) + + kernel = ScaledGroupedGemmKernel( + scenario="2Dx3D", + sf_vec_size=SF_VEC_SIZE, + accumulate_on_output=False, + separate_tensormap_init=True, + consistent_token_padding=False, + mma_tiler_mnk=(*mma_tiler_mn, 256), + cluster_shape_mnk=(*cluster_shape_mn, 1), + ) + + # Convert to CuTe tensors with dynamic layout + def to_cute(t): + ct = cutlass_torch.from_dlpack(t) + return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t)) + + a_c = to_cute(mat_a) + b_c = to_cute(mat_b) + sfa_c = to_cute(scale_a) + sfb_c = to_cute(scale_b) + c_c = to_cute(out) + offs_c = to_cute(expert_offsets) + + workspace_size = kernel.get_workspace_size(num_experts) + workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=mat_a.device) + ws_c = to_cute(workspace) + + gsa_c = to_cute(global_scale_a) if global_scale_a is not None else None + gsb_c = to_cute(global_scale_b) if global_scale_b is not None else None + + import cuda.bindings.driver as cuda + cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1] + max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + compiled = cute.compile( + kernel, a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, + max_active_clusters, stream, + global_scale_a=gsa_c, global_scale_b=gsb_c, + ) + + compiled( + a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, + stream, + global_scale_a=gsa_c, global_scale_b=gsb_c, + ) + torch.cuda.synchronize() + + return out diff --git a/tests/layertest.py b/tests/layertest.py index 672c50ac..1d30b4bc 100644 --- a/tests/layertest.py +++ b/tests/layertest.py @@ -1,12 +1,11 @@ #!/usr/bin/env python3 """ -Layer 0 kernel comparison test: NVFP4 kernel vs BF16 reference. +Layer 0 kernel comparison test: CuTeDSL NVFP4 kernel vs BF16 reference. -No vLLM, no Docker, no tensor parallelism. Just raw weights + our kernel. +No vLLM, no Docker, no tensor parallelism. Just raw weights + CuTeDSL kernel. If cosine < 0.99, the test exits with error. -Usage: - python3 layertest.py +Uses the bridge layer in cutedsl/bridge.py for tensor layout conversion. """ import os @@ -16,6 +15,20 @@ 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" @@ -23,19 +36,18 @@ LAYER_IDX = 0 DEVICE = "cuda" COSINE_THRESHOLD = 0.99 -# E2M1 FP4 lookup table +# 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): - """Find all safetensors shards and return {key: shard_path} mapping.""" 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) @@ -46,7 +58,6 @@ def find_shards(model_dir): with safe_open(sf, framework="pt") as f: for key in f.keys(): key_to_shard[key] = sf - return key_to_shard @@ -54,54 +65,35 @@ 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 -def print_layer_keys(tensors, label, max_keys=20): - """Print sorted tensor keys with shapes and dtypes (first N).""" - print(f"\n {label} — {len(tensors)} tensors") - sorted_keys = sorted(tensors.keys()) - for key in sorted_keys[:max_keys]: - t = tensors[key] - print(f" {key}: dtype={t.dtype} shape={tuple(t.shape)}") - if len(sorted_keys) > max_keys: - print(f" ... ({len(sorted_keys) - max_keys} more)") - - -# ── NVFP4 Dequantization ────────────────────────────────────────────── +# ── 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) @@ -114,17 +106,14 @@ def dequantize_nvfp4_experts(nvfp4_tensors, layer_idx, expert_indices): 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 @@ -136,130 +125,116 @@ def moe_forward_bf16(hidden_states, experts, expert_ids, expert_weights): 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 -# ── NVFP4 Kernel MoE Forward ────────────────────────────────────────── +# ── CuTeDSL NVFP4 Kernel MoE Forward ────────────────────────────────── def moe_forward_nvfp4(hidden_states, nvfp4_tensors, layer_idx, expert_ids, expert_weights): - """Run MoE forward pass using our NVFP4 kernel.""" - from nvfp4_megamoe_kernel import ( - stage_activation, - nvfp4_mega_moe_full, - transform_nvfp4_weights_for_mega_moe, - get_symm_buffer_for_nvfp4_mega_moe, - ) - + """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. - intermediate_half = 3072 - hidden_half = hidden_size // 2 - - l1_weights, l1_scales, l1_global_scales = [], [], [] - l2_weights, l2_scales, l2_global_scales = [], [], [] - + w_fp4_list = [] + w_sf_list = [] + w_gs_list = [] + for e in unique_experts: - gate_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight"].view(torch.int8).to(DEVICE) - gate_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale"].to(DEVICE) - gate_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale_2"].item() - up_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight"].view(torch.int8).to(DEVICE) - up_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE) - up_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale_2"].item() + # 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 - l1_w = torch.cat([gate_w, up_w], dim=0) - l1_sf = torch.cat([gate_sf, up_sf], dim=0) - l1_gs = torch.tensor([gate_gs, up_gs], dtype=torch.float32, device=DEVICE) - l1_weights.append(l1_w) - l1_scales.append(l1_sf) - l1_global_scales.append(l1_gs) - - down_w_key = f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight" - if down_w_key in nvfp4_tensors: - down_w = nvfp4_tensors[down_w_key].view(torch.int8).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() - else: - down_w = torch.zeros(hidden_size, intermediate_half, dtype=torch.int8, device=DEVICE) - down_sf = torch.ones(hidden_size, intermediate_half // 16, dtype=torch.float8_e4m3fn, device=DEVICE) - down_gs = 1.0 - - l2_weights.append(down_w) - l2_scales.append(down_sf) - l2_global_scales.append(torch.tensor([down_gs], dtype=torch.float32, device=DEVICE)) - - l1_w = torch.stack(l1_weights) - l1_sf = torch.stack(l1_scales) - l1_gs = torch.stack(l1_global_scales) - l2_w = torch.stack(l2_weights) - l2_sf = torch.stack(l2_scales) - l2_gs = torch.stack(l2_global_scales) - - (l1_w, l1_sf, l1_global_sf), (l2_w, l2_sf, l2_global_sf) = \ - transform_nvfp4_weights_for_mega_moe( - (l1_w, l1_sf), (l2_w, l2_sf), - l1_weight_scale_2=l1_gs, l2_weight_scale_2=l2_gs, + 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(), ) - - num_slots = num_tokens * top_k - slot_expert = torch.zeros(num_slots, dtype=torch.int32, device=DEVICE) - slot_token = torch.zeros(num_slots, dtype=torch.int64, device=DEVICE) - slot_weight = torch.zeros(num_slots, dtype=torch.float32, device=DEVICE) - - for t in range(num_tokens): - for k in range(top_k): - slot = t * top_k + k - slot_expert[slot] = expert_map[expert_ids[t, k].item()] - slot_token[slot] = t - slot_weight[slot] = expert_weights[t, k].item() - - symm_buffer = get_symm_buffer_for_nvfp4_mega_moe( - group=None, num_experts=num_experts, max_num_tokens=num_tokens, - top_k=top_k, hidden_size=hidden_size, intermediate_size=6144, + 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)] ) - - x_fp4, x_sf, input_global_scale = stage_activation(hidden_states) - symm_buffer.x[:num_tokens].copy_(x_fp4) - symm_buffer.x_sf[:num_tokens].copy_(x_sf) - symm_buffer.input_global_scale = input_global_scale - symm_buffer.topk_idx[:num_tokens].copy_(expert_ids[:, 0:1].expand(-1, top_k)) - symm_buffer.topk_weights[:num_tokens].copy_(expert_weights) - symm_buffer.experts_start_idx = 0 - - y = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) - nvfp4_mega_moe_full( - y, - (l1_w, l1_sf, l1_global_sf), - (l2_w, l2_sf, l2_global_sf), - symm_buffer, + 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 y + + return out # ── Main ─────────────────────────────────────────────────────────────── @@ -270,62 +245,55 @@ def main(): 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) - print_layer_keys(nvfp4_tensors, "NVFP4 checkpoint", max_keys=5) - - # Verify weight_scale dtype - for e in expert_indices[:1]: - for proj in ["gate_proj", "up_proj", "down_proj"]: - key = f"layers.{LAYER_IDX}.mlp.experts.{e}.{proj}.weight_scale" - if key in nvfp4_tensors: - dt = nvfp4_tensors[key].dtype - assert dt == torch.float8_e4m3fn, f"{proj}.weight_scale dtype={dt}, expected float8_e4m3fn" - print(f" {proj}.weight_scale dtype = {dt} ✓") - + 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 reference forward pass ── print("\n Running BF16 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}") - print(f" First token first 8: {[f'{v:.4f}' for v in ref_output[0, :8].tolist()]}") - + del nvfp4_experts_bf16 torch.cuda.empty_cache() - - # ── NVFP4 kernel forward pass ── - print("\n Running NVFP4 kernel...") + + # ── CuTeDSL NVFP4 kernel forward pass ── + print("\n Running CuTeDSL NVFP4 kernel (first run compiles, ~1-2 min)...") kernel_output = moe_forward_nvfp4(hidden_states, nvfp4_tensors, LAYER_IDX, expert_ids, expert_weights) print(f" Kernel: amax={kernel_output.abs().max():.4f} mean={kernel_output.float().mean():.6f}") - print(f" First token first 8: {[f'{v:.4f}' for v in kernel_output[0, :8].tolist()]}") - + # ── 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"\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)