biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Shared expert: dedicated CuTeDSL runner with proper scale assembly

Browse Source 
- 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
This commit is contained in:
biondizzle
2026-05-18 20:08:34 +00:00
parent b3451c74f8
commit c1aa4af123
2 changed files with 248 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files

266
cutedsl/shared_expert_pipeline.py
View File

@@ -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

79
tests/test_shared_expert.py
View File

@@ -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: