""" 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] # Cache compiled kernels + pre-allocated workspace by cache_key # Each entry: {'compiled': callable, 'workspace': Tensor, 'workspace_size': int} # # Key design decisions (Bug #1 fix): # - cute.compile does NOT corrupt GPU memory (verified 2026-05-20 on B200). # The original _needs_token_refill hack was a misdiagnosis. The real bug # was elsewhere (likely OOB write or weight loading). # - Workspace is pre-allocated per cache entry during warmup_compilation() # and reused on subsequent calls. No torch.full() in the hot path. # - CuTe tensor wrappers (from_dlpack + mark_layout_dynamic) are cheap # metadata wrappers. We re-create them per call from real tensors. # Caching them would hold stale references to tensors that get freed. _compiled_kernel_cache = {} # Cached LUT for E2M1 quantization (created once per device, cudagraph-safe) _NVFP4_STEP_LUT_CACHE = {} 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. Must be pre-populated during warmup (before torch.compile/cudagraph capture) so the lock is never entered on the compiled path. """ # Fast path: already cached — no lock needed (torch.compile-safe) if device in _NVFP4_STEP_LUT_CACHE: return _NVFP4_STEP_LUT_CACHE[device] # Slow path: first call, create the LUT lut = torch.as_tensor( [0, 1, 2, 3, 4, 4, 5, 5, 6, 6, 6, 7, 7], dtype=torch.int8, device=device, ) _NVFP4_STEP_LUT_CACHE[device] = lut return lut 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) # Detect zero blocks and underflow blocks (amax > 0 but too small for FP8). # Smallest positive FP8 e4m3fn is 2^-9 ≈ 1.95e-3. If amax/6 < this, # the block scale underflows to 0, and dividing x by the clamped 1e-8 # inflates values into nonzero FP4 buckets — producing wrong results. zero_block = block_amax < (6.0 * 2.0 ** -9) # < ~0.0117 # Zero out x for zero/underflow blocks before division. # This ensures x_scaled = 0 → FP4 nibbles = 0. x_reshaped = torch.where(zero_block.unsqueeze(-1), torch.zeros_like(x_reshaped), x_reshaped) block_amax = block_amax.clamp(min=1e-8) block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn) # Force zero/underflow blocks: FP8 scale = 0 (exact zero). block_scale = torch.where(zero_block, torch.zeros_like(block_scale), block_scale) # Nearest E2M1 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, 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). 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) # Detect zero blocks and underflow blocks (same threshold as quantize_to_nvfp4). zero_block = block_amax < (6.0 * 2.0 ** -9) x_reshaped = torch.where(zero_block.unsqueeze(-1), torch.zeros_like(x_reshaped), x_reshaped) block_amax = block_amax.clamp(min=1e-8) block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn) block_scale = torch.where(zero_block, torch.zeros_like(block_scale), block_scale) 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. 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) # Detect zero blocks and underflow blocks (same threshold). zero_block = w_block_amax < (6.0 * 2.0 ** -9) w_reshaped = torch.where(zero_block.unsqueeze(1), torch.zeros_like(w_reshaped), w_reshaped) w_block_amax = w_block_amax.clamp(min=1e-8) w_sf = (w_block_amax / 6.0).to(torch.float8_e4m3fn) w_sf = torch.where(zero_block, torch.zeros_like(w_sf), w_sf) w_block_sf = w_sf.float().unsqueeze(1) w_scaled = w_reshaped / w_block_sf.clamp(min=1e-8) signs = torch.sign(w_scaled) abs_scaled = w_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(w_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 w_fp4 = packed.reshape(K // 2, N).view(torch.float4_e2m1fn_x2) return w_fp4, w_sf, global_scale # ── Scale Factor Assembly ───────────────────────────────────────────── def interleave_l1_weights(w_ekn, granularity_bf16=8): """Interleave gate/up weights at granularity 8 in BF16 (4 in FP4). The fused SwiGLU epilogue requires gate/up pairs to be adjacent in the MMA accumulator. With interleaved weights, the MMA tile produces gate[i*8..i*8+7] and up[i*8..i*8+7] next to each other in registers, enabling a single-register SwiGLU without SMEM round-trips. Before: [gate_0..gate_N/2-1 | up_0..up_N/2-1] After: [gate_0..gate_7, up_0..up_7, gate_8..gate_15, up_8..up_15, ...] The interleave operates along the N dimension, where each column = 1 BF16 (FP4 packing is along K, not N). So g = granularity_bf16 directly. Args: w_ekn: (E, K_packed, N_packed) FP4 weight tensor in K-major layout N_packed = 2*intermediate/2 = intermediate (gate+up fused) granularity_bf16: interleave group size in BF16 elements (default 8) Returns: (E, K_packed, N_packed) FP4 weight tensor with interleaved gate/up """ E, K, N = w_ekn.shape N_half = N // 2 # gate and up each have N/2 FP4 columns g = granularity_bf16 # N-axis interleave: each N-col = 1 BF16 col (packing is along K) gate = w_ekn[:, :, :N_half].reshape(E, K, N_half // g, g) up = w_ekn[:, :, N_half:].reshape(E, K, N_half // g, g) return torch.stack([gate, up], dim=3).reshape(E, K, N) def deinterleave_l1_weights(w_ekn, granularity_bf16=8): """De-interleave gate/up weights (inverse of interleave_l1_weights). Used for testing/verification only. """ g = granularity_bf16 # N-axis: each N-col = 1 BF16 col E, K, N = w_ekn.shape w_reshaped = w_ekn.reshape(E, K, N // (2 * g), 2, g) gate = w_reshaped[:, :, :, 0, :].reshape(E, K, N // 2) up = w_reshaped[:, :, :, 1, :].reshape(E, K, N // 2) return torch.cat([gate, up], dim=2) 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 warmup_compilation(num_experts, K_packed, N_packed, device, mma_tiler_mn=(128, 128), cluster_shape_mn=(1, 1)): """Eagerly JIT-compile the GEMM kernel for a specific shape. Call this BEFORE model weights are loaded to ensure cute.compile runs exactly once per shape. The compiled kernel is cached and reused by run_nvfp4_grouped_gemm on the forward path. Uses random non-zero data. Zero-filled FP4/FP8 tensors cause cudaErrorIllegalInstruction because the GEMM arithmetic hits invalid values (division by zero in scale dequantization, NaN propagation). Random data produces valid intermediate results that exercise the kernel's full arithmetic path. The warmup tensors are freed immediately after compilation. Memory cost is minimal (~50MB for typical DeepSeek-V4 shapes). Args: num_experts: number of experts (local, after expert parallelism) K_packed: K dimension in float4 elements (i.e. K_original // 2) N_packed: N dimension in float4 elements (i.e. N_original // 2) device: 'cuda:X' or 'cuda' mma_tiler_mn: GEMM tiling (default (128,128)) cluster_shape_mn: cluster shape (default (1,1)) """ cache_key = (num_experts, str(device), mma_tiler_mn, cluster_shape_mn, K_packed, N_packed) if cache_key in _compiled_kernel_cache: return # Already compiled # Generate VALID FP4 data by quantizing random BF16 through our pipeline. # Random uint8 bytes as FP4 bit patterns produce NaN/Inf when the GEMM # dequantizes them (FP4 value * FP8 scale * FP32 global scale), which # causes cudaErrorIllegalInstruction in the Blackwell MMA hardware. # The ONLY safe approach: generate random BF16, quantize through our # quantize_to_nvfp4, producing mathematically consistent FP4 + FP8 scales. _warmup_a_bf16 = torch.randn(128, K_packed * 2, dtype=torch.bfloat16, device=device) * 0.1 mat_a, scale_a, _ = quantize_to_nvfp4(_warmup_a_bf16) del _warmup_a_bf16 _warmup_b_bf16 = torch.randn(1, K_packed * 2, N_packed * 2, dtype=torch.bfloat16, device=device) * 0.1 # 1 expert: kernel compiles same regardless of count mat_b, scale_b, _ = quantize_to_nvfp4(_warmup_b_bf16) del _warmup_b_bf16 out = torch.zeros(128, N_packed, dtype=torch.bfloat16, device=device) expert_offsets = torch.full((1,), 128, dtype=torch.int32, device=device) global_scale_a = torch.ones(1, dtype=torch.float32, device=device) global_scale_b = torch.ones(1, dtype=torch.float32, device=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), ) 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=device) ws_c = to_cute(workspace) gsa_c = to_cute(global_scale_a) gsb_c = to_cute(global_scale_b) 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, ) # Warmup run (required by CuTeDSL) 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() # Cache compiled kernel + pre-allocated workspace # (NOT CuTe wrappers — they hold refs to dummy tensors) _compiled_kernel_cache[cache_key] = { 'compiled': compiled, 'workspace': workspace, # pre-allocated, reused 'workspace_size': workspace_size, } # Free dummy data tensors (workspace is kept in cache) del mat_a, mat_b, scale_a, scale_b, out, expert_offsets del global_scale_a, global_scale_b del a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, gsa_c, gsb_c torch.cuda.empty_cache() 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) Kernel is compiled once (either via warmup_compilation() during init or lazily on first call), then cached. Workspace is pre-allocated during warmup and reused — no torch.full() in the hot path. No torch.cuda.synchronize() or .item() in the forward path. """ num_experts = mat_b.shape[0] n_dim = mat_b.shape[2] tokens_sum = mat_a.shape[0] out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device) cache_key = (num_experts, str(mat_a.device), mma_tiler_mn, cluster_shape_mn, mat_a.shape[1], mat_b.shape[2]) if cache_key not in _compiled_kernel_cache: # Lazy compilation — safety net if warmup_compilation wasn't called. # This should NOT happen in production (warmup is called during init). 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), ) 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, ) _compiled_kernel_cache[cache_key] = { 'compiled': compiled, 'workspace': workspace, 'workspace_size': workspace_size, } # --- Invoke cached kernel --- entry = _compiled_kernel_cache[cache_key] compiled = entry['compiled'] workspace = entry['workspace'] # Re-create CuTe wrappers from current tensors each call. # This is cheap (metadata only, no GPU work) and avoids stale # references to tensors from previous calls that may have been freed. 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) 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 stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) 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, ) return out # ── Fused SwiGLU GEMM (Stage 1: SiLU in registers, BF16 output) ────── # Cache for fused kernel (separate from standard GEMM cache) _fused_kernel_cache = {} def warmup_fused_swiglu_compilation(num_experts, K_packed, N_packed, device, swiglu_limit=0.0, mma_tiler_mn=(128, 128), cluster_shape_mn=(1, 1)): """Eagerly JIT-compile the fused SwiGLU GEMM kernel. Must be called during model initialization. See warmup_compilation() for the standard GEMM equivalent. """ from cutedsl.kernel.moe.fused_swiglu_grouped_mm import FusedSwiGLUScaledGroupedGemmKernel cache_key = ('fused', num_experts, str(device), mma_tiler_mn, cluster_shape_mn, K_packed, N_packed, swiglu_limit) if cache_key in _fused_kernel_cache: return # Generate VALID FP4 data by quantizing random BF16 (same as warmup_compilation) _warmup_a_bf16 = torch.randn(128, K_packed * 2, dtype=torch.bfloat16, device=device) * 0.1 mat_a, scale_a, _ = quantize_to_nvfp4(_warmup_a_bf16) del _warmup_a_bf16 _warmup_b_bf16 = torch.randn(1, K_packed * 2, N_packed * 2, dtype=torch.bfloat16, device=device) * 0.1 # 1 expert: kernel compiles same regardless of count mat_b, scale_b, _ = quantize_to_nvfp4(_warmup_b_bf16) del _warmup_b_bf16 # BF16 output (Stage 1: we still write BF16) # The fused kernel writes intermediate (N/2) since gate+up → silu result out = torch.zeros(128, N_packed, dtype=torch.bfloat16, device=device) expert_offsets = torch.full((1,), 128, dtype=torch.int32, device=device) global_scale_a = torch.ones(1, dtype=torch.float32, device=device) global_scale_b = torch.ones(1, dtype=torch.float32, device=device) kernel = FusedSwiGLUScaledGroupedGemmKernel( 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), fused_swiglu=True, swiglu_limit=swiglu_limit, ) 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=device) ws_c = to_cute(workspace) gsa_c = to_cute(global_scale_a) gsb_c = to_cute(global_scale_b) 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() _fused_kernel_cache[cache_key] = { 'compiled': compiled, 'workspace': workspace, 'workspace_size': workspace_size, } del mat_a, mat_b, scale_a, scale_b, out, expert_offsets del global_scale_a, global_scale_b del a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, gsa_c, gsb_c torch.cuda.empty_cache() def run_fused_swiglu_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 swiglu_limit=0.0, mma_tiler_mn=(128, 128), cluster_shape_mn=(1, 1), ): """Run the fused SwiGLU NVFP4 scaled grouped GEMM. Stage 1: SiLU is applied to the full accumulator in registers, then written as BF16 to C. Gate/up pairing is not yet implemented. """ from cutedsl.kernel.moe.fused_swiglu_grouped_mm import FusedSwiGLUScaledGroupedGemmKernel num_experts = mat_b.shape[0] n_dim = mat_b.shape[2] tokens_sum = mat_a.shape[0] out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device) cache_key = ('fused', num_experts, str(mat_a.device), mma_tiler_mn, cluster_shape_mn, mat_a.shape[1], mat_b.shape[2], swiglu_limit) if cache_key not in _fused_kernel_cache: # Lazy compilation kernel = FusedSwiGLUScaledGroupedGemmKernel( 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), fused_swiglu=True, swiglu_limit=swiglu_limit, ) 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, ) _fused_kernel_cache[cache_key] = { 'compiled': compiled, 'workspace': workspace, 'workspace_size': workspace_size, } entry = _fused_kernel_cache[cache_key] compiled = entry['compiled'] workspace = entry['workspace'] 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) 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 stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) 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, ) return out def deinterleave_quantize_nvfp4_cuda(fused_bf16, intermediate, global_scale, granularity=8): """De-interleave + quantize fused SwiGLU output using a custom CUDA kernel. Single kernel launch, no Python loop. 4x faster than the Python path. Args: fused_bf16: (M, 2*intermediate) BF16 — fused L1 output with interleaved gate/up intermediate: intermediate dimension (e.g., 3072) global_scale: pre-computed global scale for quantization granularity: interleave granularity in BF16 columns (default 8) Returns: x_fp4: (M, intermediate//2) float4_e2m1fn_x2 — quantized SwiGLU x_sf: (M, intermediate//16) float8_e4m3fn — block scales """ from torch.utils.cpp_extension import load import os kernel_dir = os.path.join(os.path.dirname(__file__), "kernels") mod = load( name="deinterleave_quantize_nvfp4", sources=[os.path.join(kernel_dir, "deinterleave_quantize.cu")], extra_cuda_cflags=["-gencode=arch=compute_100a,code=sm_100a"], verbose=False, ) return mod.deinterleave_quantize_nvfp4(fused_bf16, intermediate, granularity, global_scale)