From 0c878b3a9e3d330ba2163e523de0109ed2764f65 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 19:54:04 +0000 Subject: [PATCH] temp: restore old layertest+bridge for cosine comparison --- cutedsl/bridge.py | 209 +++++----------------------------------- cutedsl/moe_pipeline.py | 19 ++-- tests/layertest.py | 76 +-------------- 3 files changed, 36 insertions(+), 268 deletions(-) diff --git a/cutedsl/bridge.py b/cutedsl/bridge.py index bd9969e5..a061430a 100644 --- a/cutedsl/bridge.py +++ b/cutedsl/bridge.py @@ -7,31 +7,8 @@ the ScaledGroupedGemmKernel expects: - Scale factor assembly (padding + swizzle) - B tensor K-major stride conversion - Expert offset computation - -CUDA-graph-compatible: no .item() calls, no torch.cuda.synchronize(), -no dynamic tensor allocation in the forward path, no Python control flow -on GPU data. """ import math -import threading - -# Cached LUT for E2M1 quantization (created once per device, cudagraph-safe) -_NVFP4_STEP_LUT_CACHE = {} -_NVFP4_STEP_LUT_LOCK = threading.Lock() - - -def _get_step_to_idx_lut(device): - """Get or create the E2M1 step-to-index LUT for the given device. - - Cached per device to avoid CPU→CUDA copies during cudagraph capture. - """ - with _NVFP4_STEP_LUT_LOCK: - if device not in _NVFP4_STEP_LUT_CACHE: - _NVFP4_STEP_LUT_CACHE[device] = torch.as_tensor( - [0, 1, 2, 3, 4, 4, 5, 5, 6, 6, 6, 7, 7], - dtype=torch.int8, device=device, - ) - return _NVFP4_STEP_LUT_CACHE[device] import torch import cutlass import cutlass.cute as cute @@ -65,14 +42,7 @@ def round_up(a, b): def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE): """Quantize BF16 tensor to NVFP4. - - NOTE: This function is NOT cudagraph-safe because it uses .max() - which forces a CPU-GPU sync. It should only be called during - weight preparation (offline), NOT during the forward pass. - For activation quantization during forward, use - quantize_activation_nvfp4() instead (cudagraph-safe, fixed global scale). - Args: x_bf16: (..., D) BF16 tensor @@ -97,15 +67,15 @@ def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_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 — memory-efficient clamp approach + # 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().clamp(max=6.0) - - half_steps = (abs_scaled * 2.0).round().clamp(0, 12).to(torch.int8) - step_to_idx = _get_step_to_idx_lut(x_bf16.device) - indices = step_to_idx[half_steps.long()] + 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] @@ -122,62 +92,6 @@ def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE): return x_fp4, block_scale, global_scale -def quantize_activation_nvfp4(x_bf16, global_scale, block_size=SF_VEC_SIZE): - """Quantize BF16 activation tensor to NVFP4 (cudagraph-safe). - - Unlike quantize_to_nvfp4(), this takes a pre-computed global_scale - instead of computing it via .max() (which forces CPU-GPU sync). - The global_scale should be computed once during warmup and cached. - - All operations are pure GPU with no CPU-GPU syncs. - - Args: - x_bf16: (..., D) BF16 tensor - global_scale: float32 scalar (pre-computed, NOT from .max()) - block_size: NVFP4 block size - - Returns: - x_fp4: (..., D//2) float4_e2m1fn_x2 - x_sf: (..., D//16) float8_e4m3fn - """ - x_f32 = x_bf16.float() - 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) - - block_sf_expanded = block_scale.float().unsqueeze(-1) - x_scaled = x_reshaped / block_sf_expanded.clamp(min=1e-8) - signs = torch.sign(x_scaled) - abs_scaled = x_scaled.abs().clamp(max=6.0) - - half_steps = (abs_scaled * 2.0).round().clamp(0, 12).to(torch.int8) - step_to_idx = _get_step_to_idx_lut(x_bf16.device) - indices = step_to_idx[half_steps.long()] - - 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 - - def quantize_weight_to_nvfp4(w_bf16, block_size=SF_VEC_SIZE): """Quantize BF16 weight matrix to NVFP4. @@ -287,86 +201,6 @@ def compute_expert_offsets(tokens_per_expert, num_experts, device="cuda"): return offs -# ── Compiled Kernel Cache ───────────────────────────────────────────── - -_compiled_kernel_cache = {} - - -def _get_compiled_kernel(num_experts, device, mma_tiler_mn, cluster_shape_mn): - """Get or compile the CuTeDSL grouped GEMM kernel (cached by shape config). - - The kernel compilation is deterministic for a given (num_experts, device, tiler, cluster) - config, so we cache it to avoid recompiling on every forward call. - """ - cache_key = (num_experts, str(device), mma_tiler_mn, cluster_shape_mn) - if cache_key in _compiled_kernel_cache: - return _compiled_kernel_cache[cache_key] - - 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), - ) - - 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) - - # We need dummy tensors to compile against — use minimal sizes - # The compiled kernel works with dynamic shapes via mark_layout_dynamic - K_packed = 256 # minimal - N_packed = 256 - tokens = 1 - dummy_a = torch.zeros(tokens, K_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2) - dummy_b = torch.zeros(num_experts, K_packed, N_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2) - dummy_sfa = torch.zeros(1, 1, dtype=torch.float16, device=device).to(torch.float8_e4m3fn) - dummy_sfb = torch.zeros(1, 1, dtype=torch.float16, device=device).to(torch.float8_e4m3fn) - dummy_c = torch.zeros(tokens, N_packed, dtype=torch.bfloat16, device=device) - dummy_offs = torch.zeros(num_experts, dtype=torch.int32, device=device) - ws_size = kernel.get_workspace_size(num_experts) - dummy_ws = torch.full((ws_size,), 255, dtype=torch.uint8, device=device) - dummy_gsa = torch.ones(num_experts, dtype=torch.float32, device=device) - dummy_gsb = torch.ones(num_experts, dtype=torch.float32, device=device) - - def to_cute(t): - ct = cutlass_torch.from_dlpack(t) - return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t)) - - compiled = cute.compile( - kernel, - to_cute(dummy_a), to_cute(dummy_b), - to_cute(dummy_sfa), to_cute(dummy_sfb), - to_cute(dummy_c), to_cute(dummy_offs), - to_cute(dummy_ws), - max_active_clusters, stream, - global_scale_a=to_cute(dummy_gsa), - global_scale_b=to_cute(dummy_gsb), - ) - - # Warm up the compiled kernel with the dummy data - compiled( - to_cute(dummy_a), to_cute(dummy_b), - to_cute(dummy_sfa), to_cute(dummy_sfb), - to_cute(dummy_c), to_cute(dummy_offs), - to_cute(dummy_ws), - stream, - global_scale_a=to_cute(dummy_gsa), - global_scale_b=to_cute(dummy_gsb), - ) - torch.cuda.synchronize() - - # Free dummies - del dummy_a, dummy_b, dummy_sfa, dummy_sfb, dummy_c, dummy_offs, dummy_ws, dummy_gsa, dummy_gsb - - _compiled_kernel_cache[cache_key] = (compiled, kernel, max_active_clusters) - return compiled, kernel, max_active_clusters - - # ── Kernel Launch ───────────────────────────────────────────────────── def run_nvfp4_grouped_gemm( @@ -383,19 +217,21 @@ def run_nvfp4_grouped_gemm( """Run the CuTeDSL NVFP4 scaled grouped GEMM. 2Dx3D: A(tokens, K) x B(experts, K, N) -> C(tokens, N) - - CUDA-graph-compatible: uses cached compiled kernel, no synchronize(), - no cute.compile() in the forward path. """ num_experts = mat_b.shape[0] - n_dim = mat_b.shape[2] # N dimension (logical, not packed — float4_e2m1fn_x2 packs along K, not N) + n_dim = mat_b.shape[2] # packed N (in float4 elements) tokens_sum = mat_a.shape[0] - device = mat_a.device - out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=device) + out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device) - compiled, kernel, max_active_clusters = _get_compiled_kernel( - num_experts, device, mma_tiler_mn, cluster_shape_mn + 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 @@ -411,21 +247,28 @@ def run_nvfp4_grouped_gemm( offs_c = to_cute(expert_offsets) workspace_size = kernel.get_workspace_size(num_experts) - workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=device) + 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, ) - # NOTE: No torch.cuda.synchronize() here — cudagraph capture forbids it. - # The caller is responsible for any needed synchronization. + torch.cuda.synchronize() return out diff --git a/cutedsl/moe_pipeline.py b/cutedsl/moe_pipeline.py index 899a4708..52bbf6a8 100644 --- a/cutedsl/moe_pipeline.py +++ b/cutedsl/moe_pipeline.py @@ -187,9 +187,6 @@ def run_nvfp4_moe( # Global scales: alpha = igs * weight_gs for each expert l1_global_scale_a = torch.tensor([x_igs] * num_experts, dtype=torch.float32, device=device) l1_global_scale_b = torch.tensor(weights['l1_gs'], dtype=torch.float32, device=device) - print(f" L1 global_scale_a: {l1_global_scale_a.tolist()}", flush=True) - print(f" L1 global_scale_b: {l1_global_scale_b.tolist()}", flush=True) - print(f" alpha (a*b): {(l1_global_scale_a * l1_global_scale_b).tolist()}", flush=True) # Run L1 GEMM l1_out = run_nvfp4_grouped_gemm( @@ -197,20 +194,16 @@ def run_nvfp4_moe( scale_a=l1_scale_a, scale_b=l1_scale_b, expert_offsets=expert_offsets, global_scale_a=l1_global_scale_a, global_scale_b=l1_global_scale_b, - ) # (num_slots, 2*intermediate) BF16 - print(f" L1 GEMM output: shape={l1_out.shape}, amax={l1_out.abs().amax().item():.4f}", flush=True) + ) # (num_slots, intermediate) BF16 # ════════════════════════════════════════════════════════════════ # SiLU(gate) * up (BF16 — nonlinear requires BF16) # ════════════════════════════════════════════════════════════════ - # L1 output is (tokens, 2*intermediate) — gate and up fused - intermediate_size = l1_out.shape[1] // 2 - gate = l1_out[:, :intermediate_size] - up = l1_out[:, intermediate_size:] - print(f" gate: shape={gate.shape}, amax={gate.abs().amax().item():.4f}", flush=True) - print(f" up: shape={up.shape}, amax={up.abs().amax().item():.4f}", flush=True) - activated = torch.nn.functional.silu(gate) * up # (num_slots, intermediate) BF16 - print(f" After SiLU(gate)*up: shape={activated.shape}, amax={activated.abs().amax().item():.4f}", flush=True) + intermediate = l1_out.shape[1] + half = intermediate // 2 # 3072 + gate = l1_out[:, :half] + up = l1_out[:, half:] + activated = torch.nn.functional.silu(gate) * up # (num_slots, half) BF16 # ════════════════════════════════════════════════════════════════ # L2: down projection (NVFP4 × NVFP4 → BF16) diff --git a/tests/layertest.py b/tests/layertest.py index 83e645d8..23e108b9 100644 --- a/tests/layertest.py +++ b/tests/layertest.py @@ -16,6 +16,7 @@ REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, REPO_ROOT) from cutedsl.moe_pipeline import ( + prepare_nvfp4_moe_weights, run_nvfp4_moe, ) @@ -120,75 +121,6 @@ def moe_forward_bf16(hidden_states, experts, expert_ids, expert_weights): return output -def prepare_nvfp4_weights_direct(nvfp4_tensors, layer_idx, expert_indices, intermediate_size): - """Prepare weights via direct view-cast (no BF16 round-trip). - - Checkpoint uint8 → float4_e2m1fn_x2 (byte-preserving). - Block scales float8_e4m3fn → used directly. - Global scales float32 → used directly. - - For L1 (gate+up fused): normalize dual global scales to max, fold ratio - into block scales via float32 (one multiply + float8 round-trip on ratio only). - """ - l1_fp4, l1_sf, l1_gs = [], [], [] - l2_fp4, l2_sf, l2_gs = [], [], [] - - for e in expert_indices: - # 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() - - # Fuse gate+up along N, transpose to K-major - fused_w = torch.cat([gate_w, up_w], dim=0) # (2*intermediate, hidden//2) uint8 - fused_w_fp4 = fused_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous() - # (hidden//2, 2*intermediate) — K=hidden packed, N=2*intermediate - - # Fuse block scales: checkpoint is (N, K_sf), bridge expects (K_sf, N) - fused_sf = torch.cat([gate_sf, up_sf], dim=0) # (2*intermediate, hidden//16) = (N, K_sf) - fused_sf = fused_sf.permute(1, 0).contiguous() # → (K_sf, N) - - # Normalize dual global scales - l1_max_gs = max(gate_gs, up_gs) - if gate_gs != up_gs: - fused_sf_f32 = fused_sf.float() - # Gate is first intermediate cols, up is second (after transpose) - 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) - - # 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_w_fp4 = down_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous() - # (intermediate//2, hidden) — K=intermediate packed, N=hidden - - l2_fp4.append(down_w_fp4) - l2_sf.append(down_sf.permute(1, 0).contiguous()) # (N, K_sf) → (K_sf, N) - l2_gs.append(down_gs) - else: - # Expert 211 has no down_proj - l2_fp4.append(torch.zeros(3072 // 2, 7168, dtype=torch.float4_e2m1fn_x2, device=DEVICE)) - l2_sf.append(torch.ones(3072 // 16, 7168, dtype=torch.float8_e4m3fn, device=DEVICE)) # (K_sf, N) - 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 main(): torch.manual_seed(42) expert_indices = [0, 1, 2] @@ -203,9 +135,9 @@ def main(): nvfp4_tensors = load_layer_tensors(NVFP4_MODEL_DIR, LAYER_IDX) print(f" {len(nvfp4_tensors)} tensors loaded") - # Prepare weights — DIRECT PATH (no BF16 round-trip) - print("\n Preparing NVFP4 weights (direct view-cast)...") - weights = prepare_nvfp4_weights_direct(nvfp4_tensors, LAYER_IDX, expert_indices, 3072) + # Prepare weights + print("\n Preparing NVFP4 weights...") + 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}")