From c1aa4af123d83082c3c03b316657f7dd1bf4985f Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 18 May 2026 20:08:34 +0000 Subject: [PATCH] Shared expert: dedicated CuTeDSL runner with proper scale assembly - CuTeDSLSharedExpertRunner: num_groups=1 GEMM, no scatter/routing - _assemble_scales_single_group: pad to 128 rows + Blackwell swizzle - All buffers pre-allocated for cudagraph compatibility - Updated test to use dedicated runner instead of MoE runner hack --- cutedsl/shared_expert_pipeline.py | 266 +++++++++++++++++++++++------- tests/test_shared_expert.py | 79 +++++---- 2 files changed, 248 insertions(+), 97 deletions(-) diff --git a/cutedsl/shared_expert_pipeline.py b/cutedsl/shared_expert_pipeline.py index 390ce45b..6c098178 100644 --- a/cutedsl/shared_expert_pipeline.py +++ b/cutedsl/shared_expert_pipeline.py @@ -9,30 +9,41 @@ Pipeline: 3. SiLU(gate) * up → BF16 4. Re-quantize: BF16 → NVFP4 (using warmup gs) 5. L2 GEMM: NVFP4_act × NVFP4_weight(down) → BF16 + +Unlike MoE, there's no routing, no scatter, no expert offsets. +All tokens go through the same expert (the shared expert). +Scale assembly is just: quantize activation → pad to 128-row alignment → Blackwell swizzle. + +CUDA-graph-compatible: all buffers pre-allocated, no CPU-GPU syncs, +no dynamic shapes. Padding rows are zeros that contribute nothing to GEMM output. """ import torch -import os -import sys - -sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')) from cutedsl.bridge import ( quantize_activation_nvfp4, - assemble_scales_3d_side, + quantize_to_nvfp4, make_b_k_major, + assemble_scales_3d_side, run_nvfp4_grouped_gemm, ) +from cutedsl.kernel.moe.torch_scaled_grouped_mm import ( + ceil_div as cutedsl_ceil_div, + pad_and_swizzle_single, +) class CuTeDSLSharedExpertRunner: - """NVFP4 shared expert runner using CuTeDSL GEMM (num_groups=1).""" - + """NVFP4 shared expert runner using CuTeDSL GEMM (num_groups=1). + + CUDA-graph-compatible: all buffers pre-allocated, no CPU-GPU syncs. + """ + def __init__( self, hidden_size: int, intermediate_size: int, - max_num_tokens: int, + max_num_tokens: int = 8192, device: str = "cuda", swiglu_limit: float = 10.0, ): @@ -41,7 +52,7 @@ class CuTeDSLSharedExpertRunner: self.max_num_tokens = max_num_tokens self.device = device self.swiglu_limit = swiglu_limit - + # Weights (set after construction, then call finalize_weights) self.l1_fp4 = None self.l1_sf = None @@ -49,7 +60,7 @@ class CuTeDSLSharedExpertRunner: self.l2_fp4 = None self.l2_sf = None self.l2_gs = None - + # Processed weights (set by finalize_weights) self._l1_mat_b = None self._l2_mat_b = None @@ -57,27 +68,39 @@ class CuTeDSLSharedExpertRunner: self._l2_scale_b = None self._l1_gsb = None self._l2_gsb = None - + # Activation global scales (set by compute_activation_global_scales) - self._l1_activation_global_scale = 1.0 / 2688 - self._l2_activation_global_scale = 1.0 / 2688 - + self._l1_activation_global_scale = 1.0 / (6.0 * 448.0) + self._l2_activation_global_scale = 1.0 / (6.0 * 448.0) + + # Pre-allocated cudagraph buffers (set in _allocate_buffers) + self._padded_x_fp4_buf_l1 = None + self._padded_x_sf_buf_l1 = None + self._padded_x_fp4_buf_l2 = None + self._padded_x_sf_buf_l2 = None + self._l1_gsa_buf = None + self._l2_gsa_buf = None + self._expert_offsets_buf = None + self._buffers_allocated = False + + import os print(f"[CLAWMINE] SharedExpert init: hidden={hidden_size} intermediate={intermediate_size} " - f"max_tokens={max_num_tokens} pid={os.getpid()}") - + f"max_tokens={max_num_tokens} pid={os.getpid()}", flush=True) + def set_swiglu_limit(self, limit: float): self.swiglu_limit = limit - + def finalize_weights(self): """Process weights for CuTeDSL GEMM. Must be called after setting l1/l2 weights.""" # Stack weights and convert to K-major - self._l1_mat_b = make_b_k_major(torch.stack(self.l1_fp4)) + # l1_fp4/l2_fp4 are lists with 1 element (the shared expert) + self._l1_mat_b = make_b_k_major(torch.stack(self.l1_fp4)) # (1, K_packed, N_packed) self._l2_mat_b = make_b_k_major(torch.stack(self.l2_fp4)) - self._l1_scale_b = assemble_scales_3d_side(self.l1_sf) + self._l1_scale_b = assemble_scales_3d_side(self.l1_sf) # (1, N, K_sf_padded) self._l2_scale_b = assemble_scales_3d_side(self.l2_sf) self._l1_gsb = torch.tensor(self.l1_gs, dtype=torch.float32, device=self.device) self._l2_gsb = torch.tensor(self.l2_gs, dtype=torch.float32, device=self.device) - + # Free raw weights self.l1_fp4 = None self.l1_sf = None @@ -85,65 +108,186 @@ class CuTeDSLSharedExpertRunner: self.l2_fp4 = None self.l2_sf = None self.l2_gs = None - - def compute_activation_global_scales(self, hidden_states): - """Compute activation global scales from a warmup forward.""" + + def _allocate_buffers(self): + """Pre-allocate all buffers at max size for cudagraph compatibility.""" + max_rows = cutedsl_ceil_div(self.max_num_tokens, 128) * 128 # pad to 128 + + # L1: hidden_size packed, L2: intermediate_size packed + self._padded_x_fp4_buf_l1 = torch.zeros( + max_rows, self.hidden_size // 2, dtype=torch.uint8, device=self.device + ).view(torch.float4_e2m1fn_x2) + self._padded_x_fp4_buf_l2 = torch.zeros( + max_rows, self.intermediate_size // 2, dtype=torch.uint8, device=self.device + ).view(torch.float4_e2m1fn_x2) + + # Padded scale buffers (need same padded dimensions as pad_and_swizzle_single produces) + K_sf_l1 = cutedsl_ceil_div(self.hidden_size, 16) + padded_cols_l1 = cutedsl_ceil_div(K_sf_l1, 4) * 4 + K_sf_l2 = cutedsl_ceil_div(self.intermediate_size, 16) + padded_cols_l2 = cutedsl_ceil_div(K_sf_l2, 4) * 4 + self._padded_x_sf_buf_l1 = torch.zeros( + max_rows, padded_cols_l1, dtype=torch.float16, device=self.device + ).to(torch.float8_e4m3fn) + self._padded_x_sf_buf_l2 = torch.zeros( + max_rows, padded_cols_l2, dtype=torch.float16, device=self.device + ).to(torch.float8_e4m3fn) + + # Global scale buffers + self._l1_gsa_buf = torch.zeros(1, dtype=torch.float32, device=self.device) + self._l2_gsa_buf = torch.zeros(1, dtype=torch.float32, device=self.device) + + # Expert offsets for num_groups=1: just [num_tokens_padded] + # The GEMM expects expert_offsets as (num_experts,) cumulative offsets + # For 1 expert: offsets = [num_tokens] (just one element) + self._expert_offsets_buf = torch.zeros(1, dtype=torch.int32, device=self.device) + + self._buffers_allocated = True + + def _ensure_initialized(self): + """Lazily initialize stacked weights and buffers.""" + if self._l1_mat_b is None: + self.finalize_weights() + if not self._buffers_allocated: + self._allocate_buffers() + + def _assemble_scales_single_group(self, x_sf, padded_x_sf_buf): + """Assemble 2D-side activation scales for num_groups=1. + + For a single group, scale assembly is just: + 1. Zero the padded buffer + 2. Copy x_sf into the top rows + 3. Apply pad_and_swizzle_single (pads to 128 rows + Blackwell swizzle) + + Returns the swizzled 1D scale tensor. + """ + padded_x_sf = padded_x_sf_buf + padded_x_sf.zero_() + num_rows, num_cols = x_sf.shape + padded_x_sf[:num_rows, :num_cols] = x_sf + return pad_and_swizzle_single(padded_x_sf) + + def compute_activation_global_scales(self, hidden_states_sample): + """Compute activation global scales from a warmup forward pass. + + Called BEFORE cudagraph capture. Uses quantize_to_nvfp4 to get + the exact global_scale from the data, then runs L1 to compute + L2 gs from actual SiLU(gate)*up output. + """ + self._ensure_initialized() + with torch.no_grad(): - act_amax = hidden_states.amax().item() - self._l1_activation_global_scale = 1.0 / (act_amax * 1.5) if act_amax > 0 else 1.0 / 2688 - - # Run L1 to get intermediate - l1_out = self._run_l1(hidden_states) + # L1: exact gs from quantize_to_nvfp4 + _, _, l1_gs = quantize_to_nvfp4(hidden_states_sample) + self._l1_activation_global_scale = l1_gs + + # Run L1 GEMM to get intermediate for L2 gs + num_tokens = hidden_states_sample.shape[0] + l1_out = self._run_l1(hidden_states_sample) if l1_out is not None and not torch.isnan(l1_out).any(): gate = l1_out[:, :self.intermediate_size] up = l1_out[:, self.intermediate_size:] - gate_silu = torch.nn.functional.silu(gate).clamp(max=self.swiglu_limit) - up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) - intermediate = gate_silu * up - int_amax = intermediate.amax().item() - self._l2_activation_global_scale = 1.0 / (int_amax * 1.5) if int_amax > 0 else 1.0 / 2688 - + gate_silu = torch.nn.functional.silu(gate) + if self.swiglu_limit is not None: + gate_silu = gate_silu.clamp(max=self.swiglu_limit) + up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) + activated = gate_silu * up + _, _, l2_gs = quantize_to_nvfp4(activated) + self._l2_activation_global_scale = l2_gs + + print(f" SharedExpert warmup: L1 gs={self._l1_activation_global_scale:.10f} " + f"L2 gs={self._l2_activation_global_scale:.10f}", flush=True) + def _run_l1(self, hidden_states: torch.Tensor) -> torch.Tensor: """L1 GEMM: activation × gate_up_weight → BF16.""" - gs = self._l1_activation_global_scale - + num_tokens = hidden_states.shape[0] + padded_rows = cutedsl_ceil_div(num_tokens, 128) * 128 + # Quantize activation - x_fp4, x_sf = quantize_activation_nvfp4(hidden_states, gs) - - # A-side scales: combine activation block scales with weight block scales - # For num_groups=1, we need to match dimensions - # x_sf shape: [num_tokens, hidden_size // 16] - # l1_scale_b shape: [1, num_scale_rows, hidden_size] - # Need to pad x_sf to match l1_scale_b dimensions - + x_fp4, x_sf = quantize_activation_nvfp4( + hidden_states, self._l1_activation_global_scale + ) + + # Scatter x_fp4 into padded buffer + padded_x_fp4 = self._padded_x_fp4_buf_l1 + padded_x_fp4.view(torch.uint8).zero_() + padded_x_fp4.view(torch.uint8)[:num_tokens] = x_fp4.view(torch.uint8) + + # Assemble A-side scales + scale_a = self._assemble_scales_single_group(x_sf, self._padded_x_sf_buf_l1) + + # Expert offsets: [padded_rows] for 1 group + expert_offsets = self._expert_offsets_buf + expert_offsets.fill_(padded_rows) + + # Global scales + gsa = self._l1_gsa_buf.fill_(self._l1_activation_global_scale) + # Run GEMM out = run_nvfp4_grouped_gemm( - x_fp4, self._l1_mat_b[0], x_sf, self._l1_scale_b[0], - num_groups=1, + mat_a=padded_x_fp4, + mat_b=self._l1_mat_b, + scale_a=scale_a, + scale_b=self._l1_scale_b, + expert_offsets=expert_offsets, + global_scale_a=gsa, + global_scale_b=self._l1_gsb, ) - return out - + + # Extract real token outputs + return out[:num_tokens] + def _run_l2(self, intermediate: torch.Tensor) -> torch.Tensor: """L2 GEMM: intermediate × down_weight → BF16.""" - gs = self._l2_activation_global_scale - - x_fp4, x_sf = quantize_activation_nvfp4(intermediate, gs) - - out = run_nvfp4_grouped_gemm( - x_fp4, self._l2_mat_b[0], x_sf, self._l2_scale_b[0], - num_groups=1, + num_tokens = intermediate.shape[0] + padded_rows = cutedsl_ceil_div(num_tokens, 128) * 128 + + # Quantize activation + x_fp4, x_sf = quantize_activation_nvfp4( + intermediate, self._l2_activation_global_scale ) - return out - + + # Scatter into padded buffer + padded_x_fp4 = self._padded_x_fp4_buf_l2 + padded_x_fp4.view(torch.uint8).zero_() + padded_x_fp4.view(torch.uint8)[:num_tokens] = x_fp4.view(torch.uint8) + + # Assemble A-side scales + scale_a = self._assemble_scales_single_group(x_sf, self._padded_x_sf_buf_l2) + + # Expert offsets + expert_offsets = self._expert_offsets_buf + expert_offsets.fill_(padded_rows) + + # Global scales + gsa = self._l2_gsa_buf.fill_(self._l2_activation_global_scale) + + # Run GEMM + out = run_nvfp4_grouped_gemm( + mat_a=padded_x_fp4, + mat_b=self._l2_mat_b, + scale_a=scale_a, + scale_b=self._l2_scale_b, + expert_offsets=expert_offsets, + global_scale_a=gsa, + global_scale_b=self._l2_gsb, + ) + + return out[:num_tokens] + def run(self, hidden_states: torch.Tensor) -> torch.Tensor: """Full shared expert forward: L1 → SiLU → L2 → output.""" + self._ensure_initialized() + l1_out = self._run_l1(hidden_states) - + gate = l1_out[:, :self.intermediate_size] up = l1_out[:, self.intermediate_size:] - gate_silu = torch.nn.functional.silu(gate).clamp(max=self.swiglu_limit) - up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) + gate_silu = torch.nn.functional.silu(gate) + if self.swiglu_limit is not None: + gate_silu = gate_silu.clamp(max=self.swiglu_limit) + up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) intermediate = gate_silu * up - + output = self._run_l2(intermediate) return output diff --git a/tests/test_shared_expert.py b/tests/test_shared_expert.py index cb05ecd3..f26a0eff 100644 --- a/tests/test_shared_expert.py +++ b/tests/test_shared_expert.py @@ -1,7 +1,7 @@ -"""Standalone test: Shared expert using CuTeDSL MoE runner with 1 expert. +"""Standalone test: Shared expert using CuTeDSL dedicated runner. -The shared expert is just "1 expert, no routing, top_k=1". -We reuse the existing CuTeDSLMoERunner with num_experts=1. +Tests the CuTeDSLSharedExpertRunner for the shared expert path. +Compares against BF16 dequantized reference. Usage: python3 test_shared_expert.py """ @@ -16,7 +16,7 @@ DEVICE = "cuda:0" LAYER_IDX = 0 HIDDEN_SIZE = 7168 HC_MULT = 4 -HC_DIM = HC_MULT * HIDDEN_SIZE +HC_DIM = HC_MULT * HIDDEN_SIZE # 28672 INTERMEDIATE_SIZE = 3072 SWIGLU_LIMIT = 10.0 NUM_TOKENS = 4 @@ -37,6 +37,7 @@ def load_tensor(key, wm, model_dir): def dequant_nvfp4(packed_uint8, scale_e4m3, global_scale): + """Dequantize NVFP4 weight to BF16 for reference.""" device = packed_uint8.device lut = E2M1_LUT.to(device) lower = lut[(packed_uint8 & 0x0F).long()] @@ -56,17 +57,17 @@ def main(): torch.manual_seed(42) sys.path.insert(0, "/root/nvfp4-megamoe-kernel") - from vllm.nvfp4_cutedsl import CuTeDSLMoERunner + from cutedsl.shared_expert_pipeline import CuTeDSLSharedExpertRunner with open(os.path.join(MODEL_PATH, "model.safetensors.index.json")) as f: wm = json.load(f)["weight_map"] P = lambda key: load_tensor(key, wm, MODEL_PATH).to(DEVICE) - print("=== Shared Expert Test (CuTeDSL MoE runner, 1 expert) ===\n") + print("=== Shared Expert Test (CuTeDSL SharedExpertRunner) ===\n") # Load shared expert weights prefix = f"model.layers.{LAYER_IDX}.mlp.shared_experts" - + gate_w = P(f"{prefix}.gate_proj.weight") gate_sf = P(f"{prefix}.gate_proj.weight_scale") gate_gs = P(f"{prefix}.gate_proj.weight_scale_2").item() @@ -82,51 +83,56 @@ def main(): print(f"down_proj: shape={down_w.shape} gs={down_gs:.8f}") # Stack gate + up into gate_up_proj (same format as MoE L1) + # gate/up weights are (intermediate, hidden) uint8 packed gate_up_w = torch.cat([gate_w, up_w], dim=0) gate_up_sf = torch.cat([gate_sf, up_sf], dim=0) mgs = max(gate_gs, up_gs) if gate_gs != up_gs: sf32 = gate_up_sf.float() - sf32[:, :INTERMEDIATE_SIZE] *= (gate_gs / mgs) - sf32[:, INTERMEDIATE_SIZE:] *= (up_gs / mgs) + sf32[:INTERMEDIATE_SIZE] *= (gate_gs / mgs) + sf32[INTERMEDIATE_SIZE:] *= (up_gs / mgs) gate_up_sf = sf32.to(torch.float8_e4m3fn) - # Convert to CuTeDSL format - l1_fp4 = gate_up_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous() - l1_sf = gate_up_sf.permute(1, 0).contiguous() - l2_fp4 = down_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous() - l2_sf = down_sf.permute(1, 0).contiguous() + # Convert to CuTeDSL format: + # Checkpoint weights are (out_features, in_features) uint8 packed + # We need float4_e2m1fn_x2 with (out_features, in_features // 2) after view + # Then permute to (in_features // 2, out_features) for K-major (K=in_features) + l1_fp4 = [gate_up_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()] + l1_sf = [gate_up_sf.permute(1, 0).contiguous()] + l2_fp4 = [down_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()] + l2_sf = [down_sf.permute(1, 0).contiguous()] - # Create MoE runner with 1 expert - runner = CuTeDSLMoERunner( - num_experts=1, hidden_size=HC_DIM, - intermediate_size=INTERMEDIATE_SIZE, max_num_tokens=8192, - top_k=1, device=DEVICE, + # Create runner + runner = CuTeDSLSharedExpertRunner( + hidden_size=HC_DIM, + intermediate_size=INTERMEDIATE_SIZE, + max_num_tokens=8192, + device=DEVICE, + swiglu_limit=SWIGLU_LIMIT, ) - runner.l1_fp4 = [l1_fp4] - runner.l1_sf = [l1_sf] + runner.l1_fp4 = l1_fp4 + runner.l1_sf = l1_sf runner.l1_gs = [mgs] - runner.l2_fp4 = [l2_fp4] - runner.l2_sf = [l2_sf] + runner.l2_fp4 = l2_fp4 + runner.l2_sf = l2_sf runner.l2_gs = [down_gs] - runner.set_swiglu_limit(SWIGLU_LIMIT) + runner.finalize_weights() - # Warmup + # Warmup to compute activation global scales dummy = torch.randn(NUM_TOKENS, HC_DIM, dtype=torch.bfloat16, device=DEVICE) * 2.0 - dummy_topk_ids = torch.zeros(NUM_TOKENS, 1, dtype=torch.int64, device=DEVICE) - dummy_topk_weights = torch.ones(NUM_TOKENS, 1, dtype=torch.float32, device=DEVICE) - runner.compute_activation_global_scales(dummy, dummy_topk_weights, dummy_topk_ids) - print(f"Warmup gs: L1={runner._l1_activation_global_scale:.6f} L2={runner._l2_activation_global_scale:.6f}") + runner._ensure_initialized() + runner.compute_activation_global_scales(dummy) + print(f"Warmup gs: L1={runner._l1_activation_global_scale:.6f} " + f"L2={runner._l2_activation_global_scale:.6f}") # Run CuTeDSL print("\n--- CuTeDSL Forward ---") hidden = torch.randn(NUM_TOKENS, HC_DIM, dtype=torch.bfloat16, device=DEVICE) * 2.0 - topk_ids = torch.zeros(NUM_TOKENS, 1, dtype=torch.int64, device=DEVICE) - topk_weights = torch.ones(NUM_TOKENS, 1, dtype=torch.float32, device=DEVICE) - + with torch.no_grad(): - output = runner.run(hidden, topk_weights, topk_ids) - print(f"CuTeDSL output: amax={output.amax():.4f} NaN={torch.isnan(output).any()}") + output = runner.run(hidden) + print(f"CuTeDSL output: shape={output.shape} amax={output.amax():.4f} " + f"NaN={torch.isnan(output).any()}") # BF16 reference print("\n--- BF16 Reference ---") @@ -142,10 +148,11 @@ def main(): intermediate = gate_silu * up ref_output = intermediate @ down_bf16.T - print(f"BF16 ref: amax={ref_output.amax():.4f}") + print(f"BF16 ref: shape={ref_output.shape} amax={ref_output.amax():.4f}") # Compare - cos = F.cosine_similarity(ref_output.flatten().unsqueeze(0), output.flatten().unsqueeze(0)).item() + cos = F.cosine_similarity(ref_output.flatten().unsqueeze(0), + output.flatten().unsqueeze(0)).item() mse = (ref_output - output).pow(2).mean().item() print(f"\n=== RESULT: cosine={cos:.6f} MSE={mse:.6e} ===") if cos >= 0.98: