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.
This commit is contained in:
274
cutedsl/bridge.py
Normal file
274
cutedsl/bridge.py
Normal file
@@ -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
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user