diff --git a/NEXT_OPTIMIZATION.md b/NEXT_OPTIMIZATION.md new file mode 100644 index 00000000..107150d1 --- /dev/null +++ b/NEXT_OPTIMIZATION.md @@ -0,0 +1,344 @@ +pre-remap the weight scale factors (SFB) once, and stop remapping/allocating them inside every GEMM. + +This is lower hanging than rewriting stage_activation in Triton, because weights are static after load. Activation scales SFA still need per-call remap, but SFB does not. + +Current hot path in cutlass_nvfp4_gemm.cu: +``` +cutlass::device_memory::allocation sfa_cutlass(sfa_size); +cutlass::device_memory::allocation sfb_cutlass(sfb_size); + +cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream); +cudaMemsetAsync(sfb_cutlass.get(), 0, sfb_size * sizeof(ElementSF), stream); + +remap_sf_to_cutlass_kernel<<<...>>>(SFA_ptr, sfa_cutlass.get(), ...); +remap_sf_to_cutlass_kernel<<<...>>>(SFB_ptr, sfb_cutlass.get(), ...); +``` + + + +Target shape + +Keep this per GEMM: +``` +SFA row-major activation scales +→ remap dynamically +``` + +Change this: + +``` +SFB weight scales +→ already CUTLASS-remapped +→ pass directly to GEMM +``` + +So the GEMM run becomes: + +``` +// still dynamic +remap SFA + +// no remap +use prepacked_SFB_cutlass directly +``` + +Step 1: add a prepack_sfb CUDA entrypoint + +Add a new exported function next to cutlass_nvfp4_gemm_run. + +Something like: +``` +extern "C" int cutlass_nvfp4_prepack_sfb_run( + const void* SFB_ptr, + void* SFB_cutlass_ptr, + int M, int N, int K, + cudaStream_t stream +) { + using Sm1xxBlkScaledConfig = + typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + + LayoutSFB layout_SFB = + Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB( + cute::make_shape(M, N, K, 1)); + + using ElementSF = + typename Gemm::GemmKernel::CollectiveMainloop::ElementSF; + + int sfb_size = cute::size(layout_SFB); + int K_sf = K / InputSFVectorSize; + + cudaMemsetAsync( + static_cast(SFB_cutlass_ptr), + 0, + sfb_size * sizeof(ElementSF), + stream); + + int block = 256; + remap_sf_to_cutlass_kernel<<<(sfb_size + block - 1) / block, block, 0, stream>>>( + static_cast(SFB_ptr), + static_cast(SFB_cutlass_ptr), + layout_SFB, + N, + K_sf, + true // SFB source is (K_sf, N) + ); + + return 0; +} +``` + +You’ll also want a size query helper: +``` +extern "C" int cutlass_nvfp4_sfb_size( + int M, int N, int K, + int* out_size +) { + using Sm1xxBlkScaledConfig = + typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + + LayoutSFB layout_SFB = + Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB( + cute::make_shape(M, N, K, 1)); + + *out_size = cute::size(layout_SFB); + return 0; +} +``` + +Important: pass the same M, N, K geometry you’ll use for the GEMM at first. Later you can test whether SFB layout is independent of M; I suspect it is effectively N/K-driven, but don’t assume that until you print sizes for several M. + +Step 2: expose it in pytorch_binding.cpp + +Add declarations: + +``` +extern "C" int cutlass_nvfp4_sfb_size( + int M, int N, int K, + int* out_size +); + +extern "C" int cutlass_nvfp4_prepack_sfb_run( + const void* SFB_ptr, + void* SFB_cutlass_ptr, + int M, int N, int K, + cudaStream_t stream +); +``` + +Then add a Python-visible wrapper: +``` +torch::Tensor prepack_sfb( + torch::Tensor SFB, + int64_t M, + int64_t N, + int64_t K +) { + int size = 0; + int rc = cutlass_nvfp4_sfb_size( + static_cast(M), + static_cast(N), + static_cast(K), + &size + ); + TORCH_CHECK(rc == 0, "sfb_size failed"); + + auto out = torch::empty( + {size}, + torch::dtype(SFB.dtype()).device(SFB.device()) + ); + + auto stream = c10::cuda::getCurrentCUDAStream(); + + rc = cutlass_nvfp4_prepack_sfb_run( + SFB.data_ptr(), + out.data_ptr(), + static_cast(M), + static_cast(N), + static_cast(K), + stream.stream() + ); + + TORCH_CHECK(rc == 0, "prepack_sfb failed"); + + return out; +} +``` + +Register it: +``` +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", &cutlass_nvfp4_gemm_forward, + "CUTLASS NVFP4 block-scaled GEMM forward"); + + m.def("prepack_sfb", &prepack_sfb, + "Pre-remap SFB weight scales into CUTLASS layout"); +} +``` + +Step 3: add a GEMM path that accepts prepacked SFB +Add a second C API entrypoint, or a boolean flag, for: +``` +cutlass_nvfp4_gemm_run_prepacked_sfb(...) +``` + +Inside it, delete this: +``` +cutlass::device_memory::allocation sfb_cutlass(sfb_size); +cudaMemsetAsync(sfb_cutlass.get(), 0, ...); +remap_sf_to_cutlass_kernel<<<...>>>(SFB_ptr, sfb_cutlass.get(), ...); +``` + +and use the passed pointer directly: +``` +static_cast(SFB_cutlass_ptr), layout_SFB +``` + +The relevant section becomes: + +``` +cutlass::device_memory::allocation sfa_cutlass(sfa_size); +cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream); + +remap_sf_to_cutlass_kernel<<<(sfa_size + block - 1) / block, block, 0, stream>>>( + static_cast(SFA_ptr), + sfa_cutlass.get(), + layout_SFA, + M, + K_sf, + false); + +typename Gemm::Arguments arguments { + cutlass::gemm::GemmUniversalMode::kGemm, + {M, N, K, 1}, + { + static_cast(A_ptr), stride_A, + static_cast(B_ptr), stride_B, + sfa_cutlass.get(), layout_SFA, + static_cast(SFB_cutlass_ptr), layout_SFB + }, + { + { alpha, beta }, + nullptr, stride_C, + static_cast(D_ptr), stride_D + } +}; +``` + +Step 4: prepack in Python after weight transform + +In weight_transform.py, you currently return: +``` +return (l1_weight_out, l1_sf_out), (l2_weight_out, l2_sf_out) +``` + +Do not do the prepack directly there unless extension import order is clean. Safer first pass: do it lazily in nvfp4_mega_moe_full() once. + +Add helper: + +``` +def _maybe_prepack_weight_sf(weights, weight_sf, N, K, tag): + cache_attr = f"_prepacked_{tag}" + + cached = getattr(_maybe_prepack_weight_sf, cache_attr, None) + if cached is not None: + return cached + + from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm import _C + + E = weight_sf.shape[0] + packed = [] + + # Use M=128 initially to match the MMA tile boundary. + # Later test if SFB size/layout is stable across M. + M_for_layout = 128 + + for e in range(E): + packed.append( + _C.prepack_sfb( + weight_sf[e], + M_for_layout, + N, + K, + ) + ) + + packed = torch.stack(packed, dim=0).contiguous() + setattr(_maybe_prepack_weight_sf, cache_attr, packed) + return packed +``` + +Then before L1/L2 calls: + +``` +l1_N = l1_w.shape[2] +l1_K = l1_w.shape[1] * 2 +l1_sf_prepacked = _maybe_prepack_weight_sf(l1_w, l1_sf, l1_N, l1_K, "l1") + +l2_N = l2_w.shape[2] +l2_K = l2_w.shape[1] * 2 +l2_sf_prepacked = _maybe_prepack_weight_sf(l2_w, l2_sf, l2_N, l2_K, "l2") +``` + +Then pass l1_sf_prepacked / l2_sf_prepacked into the new prepacked-SFB GEMM path. + +Validation test + +Before trusting it, compare old vs new on one expert: + +``` +old = cutlass_nvfp4_blockscaled_gemm( + expert_x, + expert_x_sf, + expert_w, + expert_w_sf, + M_expert, + N, + K, + alpha=alpha, +) + +sfb_pre = _C.prepack_sfb(expert_w_sf, 128, N, K) + +new = cutlass_nvfp4_blockscaled_gemm_prepacked_sfb( + expert_x, + expert_x_sf, + expert_w, + sfb_pre, + M_expert, + N, + K, + alpha=alpha, +) + +print("max diff", (old - new).abs().max()) +print("cos", torch.nn.functional.cosine_similarity( + old.flatten().float(), + new.flatten().float(), + dim=0, +)) +``` + +Expected: +``` +max diff = 0 or tiny BF16-level difference +cos ≈ 1.0 +``` + +If it fails only when M_expert != 128, then LayoutSFB is M-dependent. In that case, cache by an M bucket: +``` +bucket_m = ((M_expert + 127) // 128) * 128 +cache_key = (tag, bucket_m, N, K) +``` + +But test first. If SFB layout size and values are stable across M, keep one prepack per expert per layer. + +Why this is worth doing now + +This removes, per active expert GEMM: +``` +1 cudaMalloc/free-ish allocation for SFB +1 cudaMemsetAsync for SFB padding +1 remap kernel launch for SFB +``` + +Across L1 + L2 and many routed experts, that’s a very clean win without touching routing semantics, quantization, or the actual MMA tile. \ No newline at end of file diff --git a/README.md b/README.md index 8ed2ddb8..1da1713a 100644 --- a/README.md +++ b/README.md @@ -284,7 +284,7 @@ The CUTLASS extension builds inside the container during `pip install` of the nv ## Known Issues -1. **MoE dispatch is slow** — `cutlass_grouped_nvfp4_gemm` uses a Python loop over 48 experts with per-token scatter/gather. Needs a proper grouped GEMM or at least CUDA-side dispatch. +1. ~~**MoE dispatch is slow**~~ — Fixed. Slot-based `index_add_` replaces the Python double loop over tokens×topk. Routing weights applied once at final scatter. Per-expert loop still exists (gather+GEMM) but scatter is vectorized. 2. **stage_activation is Python** — Re-quantization from L1 BF16 output to FP4 for L2 input runs in PyTorch. Should use the Triton staging kernel for speed and consistency with vLLM's built-in staging. diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu index 04b262b4..d2be3bab 100644 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu @@ -213,6 +213,110 @@ int cutlass_nvfp4_gemm_run( return 0; } +///////////////////////////////////////////////////////////////////////////////////////////////// +// SFB prepack: pre-remap weight scale factors once at load time +///////////////////////////////////////////////////////////////////////////////////////////////// + +extern "C" int cutlass_nvfp4_sfb_size( + int M, int N, int K, + int* out_size +) { + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1)); + *out_size = cute::size(layout_SFB); + return 0; +} + +extern "C" int cutlass_nvfp4_prepack_sfb_run( + const void* SFB_ptr, + void* SFB_cutlass_ptr, + int M, int N, int K, + cudaStream_t stream +) { + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1)); + using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF; + + int sfb_size = cute::size(layout_SFB); + int K_sf = K / InputSFVectorSize; + + cudaMemsetAsync(static_cast(SFB_cutlass_ptr), 0, sfb_size * sizeof(ElementSF), stream); + + int block = 256; + remap_sf_to_cutlass_kernel<<<(sfb_size + block - 1) / block, block, 0, stream>>>( + static_cast(SFB_ptr), + static_cast(SFB_cutlass_ptr), + layout_SFB, + N, K_sf, true // SFB source is (K_sf, N) + ); + + return 0; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// +// GEMM with prepacked SFB — skips SFB allocation, memset, and remap +///////////////////////////////////////////////////////////////////////////////////////////////// + +extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb( + const void* A_ptr, const void* SFA_ptr, + const void* B_ptr, const void* SFB_cutlass_ptr, + void* D_ptr, + int M, int N, int K, + float alpha, float beta, + cudaStream_t stream +) { + StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1)); + StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1)); + StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1)); + StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1)); + + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1)); + LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1)); + + using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA; + using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB; + using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF; + + int sfa_size = cute::size(layout_SFA); + int K_sf = K / InputSFVectorSize; + + // Only remap SFA (activation scales) — SFB is prepacked + cutlass::device_memory::allocation sfa_cutlass(sfa_size); + cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream); + + int block = 256; + remap_sf_to_cutlass_kernel<<<(sfa_size + block - 1) / block, block, 0, stream>>>( + static_cast(SFA_ptr), sfa_cutlass.get(), layout_SFA, M, K_sf, false); + + typename Gemm::Arguments arguments { + cutlass::gemm::GemmUniversalMode::kGemm, + {M, N, K, 1}, + { + static_cast(A_ptr), stride_A, + static_cast(B_ptr), stride_B, + sfa_cutlass.get(), layout_SFA, + static_cast(SFB_cutlass_ptr), layout_SFB + }, + { + { alpha, beta }, + nullptr, stride_C, + static_cast(D_ptr), stride_D + } + }; + + Gemm gemm; + CUTLASS_CHECK(gemm.can_implement(arguments)); + + size_t workspace_size = Gemm::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + CUTLASS_CHECK(gemm.initialize(arguments, workspace.get(), stream)); + CUTLASS_CHECK(gemm.run(stream)); + + return 0; +} + } // extern "C" #endif diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py index a9a64ee1..1894b895 100644 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py @@ -21,23 +21,46 @@ def cutlass_nvfp4_blockscaled_gemm( A_packed, # (M, K_half) int8 packed E2M1 SFA, # scale factors for A (float8_e4m3fn) B_packed, # (K_half, N) int8 packed E2M1, column-major for CUTLASS - SFB, # scale factors for B (sf_k, N) float8_e4m3fn, column-major for CUTLASS + SFB, # scale factors for B — either (sf_k, N) float8_e4m3fn row-major, or prepacked CUTLASS layout M, N, K, # Problem dimensions (K in FP4 elements) alpha=1.0, # fp32 scalar applied in epilogue: D = alpha * A @ B + beta * C + sfb_prepacked=False, # True if SFB is already in CUTLASS layout ): - """Single NVFP4 block-scaled GEMM using CUTLASS.""" + """Single NVFP4 block-scaled GEMM using CUTLASS. + + If sfb_prepacked=True, SFB is assumed to be in CUTLASS interleaved layout + (from prepack_sfb) and the SFB remap is skipped. + """ if not _CUTLASS_AVAILABLE: raise RuntimeError("CUTLASS NVFP4 GEMM extension not available") - return _C.forward(A_packed, SFA, B_packed, SFB, M, N, K, alpha) + if sfb_prepacked: + return _C.forward_prepacked_sfb(A_packed, SFA, B_packed, SFB, M, N, K, alpha) + else: + return _C.forward(A_packed, SFA, B_packed, SFB, M, N, K, alpha) + + +def prepack_sfb(SFB, M, N, K): + """Pre-remap SFB weight scales into CUTLASS interleaved layout. + + Call once after weight transform. Returns a tensor that can be passed + to cutlass_nvfp4_blockscaled_gemm with sfb_prepacked=True. + + M is used for layout sizing. Test with different M values to confirm + SFB layout is M-independent; if so, any valid M works (e.g. 128). + """ + if not _CUTLASS_AVAILABLE: + raise RuntimeError("CUTLASS NVFP4 GEMM extension not available") + return _C.prepack_sfb(SFB, M, N, K) def cutlass_grouped_nvfp4_gemm( x_fp4, # (num_tokens, K_half) int8 packed E2M1 x_sf, # (num_tokens, sf_k) float8_e4m3fn block scales weights, # (E_per_rank, K_half, N) int8 packed E2M1, column-major for CUTLASS - weight_sf, # (E_per_rank, sf_k, N) float8_e4m3fn, column-major for CUTLASS + weight_sf, # (E_per_rank, sf_k, N) float8_e4m3fn, column-major — or prepacked (E_per_rank, sfb_size) if sfb_prepacked=True topk_ids, # (num_tokens, NUM_TOPK) int32 — local expert IDs alpha=1.0, # fp32 scalar: D = alpha * A @ B (from stage_activation global scale) + sfb_prepacked=False, # True if weight_sf is already prepacked into CUTLASS layout ): """Per-expert grouped GEMM for MoE dispatch using CUTLASS NVFP4. @@ -56,23 +79,24 @@ def cutlass_grouped_nvfp4_gemm( num_topk = topk_ids.shape[1] # Build slot mapping: which (token, topk) pairs land on local experts? - local_mask = (topk_ids >= 0) & (topk_ids < num_experts) # (num_tokens, num_topk) - slot_token, slot_k = local_mask.nonzero(as_tuple=True) # (num_slots,) - slot_expert = topk_ids[slot_token, slot_k] # (num_slots,) local expert id + local_mask = (topk_ids >= 0) & (topk_ids < num_experts) + slot_token, slot_k = local_mask.nonzero(as_tuple=True) + slot_expert = topk_ids[slot_token, slot_k] num_slots = slot_token.shape[0] if MEGA_MOE_DEBUG: print(f"[cutlass_grouped_gemm] tokens={num_tokens} K={K} N={N} " - f"experts={num_experts} topk={num_topk} slots={num_slots}") + f"experts={num_experts} topk={num_topk} slots={num_slots} " + f"sfb_prepacked={sfb_prepacked}") if num_slots == 0: slot_out = torch.empty(0, N, dtype=torch.bfloat16, device=x_fp4.device) return slot_out, slot_token # Gather activations for all slots - slot_x = x_fp4[slot_token] # (num_slots, K_half) - slot_x_sf = x_sf[slot_token] # (num_slots, sf_k) + slot_x = x_fp4[slot_token] + slot_x_sf = x_sf[slot_token] slot_out = torch.empty(num_slots, N, dtype=torch.bfloat16, device=x_fp4.device) @@ -85,38 +109,26 @@ def cutlass_grouped_nvfp4_gemm( expert_x = slot_x[e_idx] expert_x_sf = slot_x_sf[e_idx] expert_w = weights[e] - expert_w_sf = weight_sf[e] + expert_w_sf = weight_sf[e] # prepacked or raw depending on flag M_expert = e_idx.shape[0] - if e < 3 and M_expert > 0: + if MEGA_MOE_DEBUG and e < 3 and M_expert > 0: print(f"[GEMM-IN] expert={e} M={M_expert} N={N} K={K} " - f"w shape={expert_w.shape} w_sf shape={expert_w_sf.shape} " - f"w absmax={expert_w.view(torch.int8).abs().max().item()} " - f"w_sf range=[{expert_w_sf.to(torch.float32).min().item():.4e}, " - f"{expert_w_sf.to(torch.float32).max().item():.4e}] " - f"w_sf nonzero_frac={(expert_w_sf.view(torch.uint8) != 0).float().mean().item():.4f}") + f"w shape={expert_w.shape} sfb_prepacked={sfb_prepacked}") expert_out = cutlass_nvfp4_blockscaled_gemm( expert_x, expert_x_sf, expert_w, expert_w_sf, M_expert, N, K, alpha=alpha, + sfb_prepacked=sfb_prepacked, ) - torch.cuda.current_stream().synchronize() - - if torch.isnan(expert_out).any() or torch.isinf(expert_out).any(): - raise RuntimeError( - f"expert {e} of {num_experts}: GEMM emitted NaN/Inf. " - f"M={M_expert} N={N} K={K} | " - f"x abs range [{expert_x.view(torch.int8).abs().max().item()}], " - f"x_sf range [{expert_x_sf.to(torch.float32).min().item():.4e}, " - f"{expert_x_sf.to(torch.float32).max().item():.4e}], " - f"w_sf range [{expert_w_sf.to(torch.float32).min().item():.4e}, " - f"{expert_w_sf.to(torch.float32).max().item():.4e}], " - f"x_sf nan_frac={torch.isnan(expert_x_sf.to(torch.float32)).float().mean().item():.4f}, " - f"w_sf nan_frac={torch.isnan(expert_w_sf.to(torch.float32)).float().mean().item():.4f}" - ) + if MEGA_MOE_DEBUG: + if torch.isnan(expert_out).any() or torch.isinf(expert_out).any(): + raise RuntimeError( + f"expert {e} of {num_experts}: GEMM emitted NaN/Inf. " + f"M={M_expert} N={N} K={K}") slot_out[e_idx] = expert_out diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp index a19ec77e..774502df 100644 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp @@ -13,6 +13,27 @@ extern "C" int cutlass_nvfp4_gemm_run( cudaStream_t stream ); +extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb( + const void* A_ptr, const void* SFA_ptr, + const void* B_ptr, const void* SFB_cutlass_ptr, + void* D_ptr, + int M, int N, int K, + float alpha, float beta, + cudaStream_t stream +); + +extern "C" int cutlass_nvfp4_sfb_size( + int M, int N, int K, + int* out_size +); + +extern "C" int cutlass_nvfp4_prepack_sfb_run( + const void* SFB_ptr, + void* SFB_cutlass_ptr, + int M, int N, int K, + cudaStream_t stream +); + torch::Tensor cutlass_nvfp4_gemm_forward( torch::Tensor A_packed, torch::Tensor SFA, @@ -40,6 +61,71 @@ torch::Tensor cutlass_nvfp4_gemm_forward( return D; } +torch::Tensor cutlass_nvfp4_gemm_forward_prepacked_sfb( + torch::Tensor A_packed, + torch::Tensor SFA, + torch::Tensor B_packed, + torch::Tensor SFB_cutlass, + int64_t M, int64_t N, int64_t K, + double alpha = 1.0 +) { + auto D = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A_packed.device())); + + auto stream = c10::cuda::getCurrentCUDAStream(); + cudaStream_t cuda_stream = stream.stream(); + + int rc = cutlass_nvfp4_gemm_run_prepacked_sfb( + A_packed.data_ptr(), SFA.data_ptr(), + B_packed.data_ptr(), SFB_cutlass.data_ptr(), + D.data_ptr(), + static_cast(M), static_cast(N), static_cast(K), + static_cast(alpha), 0.0f, + cuda_stream + ); + + TORCH_CHECK(rc == 0, "CUTLASS NVFP4 GEMM (prepacked SFB) failed with error code ", rc); + + return D; +} + +torch::Tensor prepack_sfb( + torch::Tensor SFB, + int64_t M, + int64_t N, + int64_t K +) { + int size = 0; + int rc = cutlass_nvfp4_sfb_size( + static_cast(M), + static_cast(N), + static_cast(K), + &size + ); + TORCH_CHECK(rc == 0, "sfb_size failed"); + + auto out = torch::empty( + {size}, + torch::dtype(SFB.dtype()).device(SFB.device()) + ); + + auto stream = c10::cuda::getCurrentCUDAStream(); + + rc = cutlass_nvfp4_prepack_sfb_run( + SFB.data_ptr(), + out.data_ptr(), + static_cast(M), + static_cast(N), + static_cast(K), + stream.stream() + ); + + TORCH_CHECK(rc == 0, "prepack_sfb failed"); + + return out; +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("forward", &cutlass_nvfp4_gemm_forward, "CUTLASS NVFP4 block-scaled GEMM forward"); + m.def("forward_prepacked_sfb", &cutlass_nvfp4_gemm_forward_prepacked_sfb, "CUTLASS NVFP4 GEMM forward with prepacked SFB"); + m.def("prepack_sfb", &prepack_sfb, "Pre-remap SFB weight scales into CUTLASS layout"); } diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index 95ff2bda..06217381 100644 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py @@ -89,13 +89,46 @@ MEGA_MOE_DEBUG = int(os.environ.get("MEGA_MOE_DEBUG", "0")) # Main kernel entry points # --------------------------------------------------------------------------- +def _prepack_weight_sf(weight_sf, N, K, tag): + """Lazily prepack SFB weight scales into CUTLASS layout (once per tag). + + Returns a tensor of shape (E, sfb_size) with SFB already in CUTLASS + interleaved layout, skipping the per-call remap+memset+alloc. + """ + cache_attr = f"_prepacked_{tag}" + cached = getattr(_prepack_weight_sf, cache_attr, None) + if cached is not None: + return cached + + from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import prepack_sfb + + E = weight_sf.shape[0] + # M for layout sizing. Test with different M to confirm SFB is M-independent. + # If SFB size changes with M, bucket by M and cache per-bucket. + M_for_layout = 128 + + packed = [] + for e in range(E): + packed.append(prepack_sfb(weight_sf[e], M_for_layout, N, K)) + + packed = torch.stack(packed, dim=0).contiguous() + setattr(_prepack_weight_sf, cache_attr, packed) + + if MEGA_MOE_DEBUG: + print(f"[PREPACK] {tag}: E={E} N={N} K={K} packed_shape={packed.shape} " + f"(was {weight_sf.shape})") + + return packed + + def nvfp4_mega_moe_l1( x_fp4, # (num_tokens, K//2) int8 packed E2M1 x_sf, # (num_tokens, sf_k_groups) float8_e4m3fn l1_weights, # (E_per_rank, K//2, 2*INTER) int8, column-major for CUTLASS - l1_scales, # (E_per_rank, sf_k_groups, 2*INTER) float8_e4m3fn, column-major + l1_scales, # (E_per_rank, sf_k_groups, 2*INTER) float8_e4m3fn, column-major — or prepacked topk_ids, # (num_tokens, NUM_TOPK) int32 — local expert IDs alpha=1.0, # fp32 scalar from stage_activation global scale + sfb_prepacked=False, # True if l1_scales is prepacked CUTLASS layout ): """L1 GEMM: gate_up_proj — slot-based, no routing weights. @@ -110,13 +143,17 @@ def nvfp4_mega_moe_l1( print(f"[nvfp4_moe_l1] tokens={x_fp4.shape[0]} K={K} N={N} native=1") x_sf_fp8 = unpack_ue4m3_u32(x_sf) if x_sf.dtype == torch.uint32 else x_sf - w_sf_fp8 = unpack_ue4m3_u32(l1_scales) if l1_scales.dtype == torch.uint32 else l1_scales + if not sfb_prepacked: + w_sf_fp8 = unpack_ue4m3_u32(l1_scales) if l1_scales.dtype == torch.uint32 else l1_scales + else: + w_sf_fp8 = l1_scales # already prepacked, skip unpack slot_out, slot_token = cutlass_grouped_nvfp4_gemm( x_fp4, x_sf_fp8, l1_weights, w_sf_fp8, topk_ids, alpha=alpha, + sfb_prepacked=sfb_prepacked, ) print(f"[L1-GEMM-OUT] slots={slot_out.shape[0]} N={N} amax={slot_out.abs().max().item():.4e} mean={slot_out.float().mean().item():.4e}") return slot_out, slot_token @@ -126,10 +163,11 @@ def nvfp4_mega_moe_l2( x_fp4, # (num_slots, INTER//2) int8 packed E2M1 x_sf, # (num_slots, sf_k_groups) float8_e4m3fn l2_weights, # (E_per_rank, INTER//2, HIDDEN) int8, column-major for CUTLASS - l2_scales, # (E_per_rank, sf_k_groups, HIDDEN) float8_e4m3fn, column-major + l2_scales, # (E_per_rank, sf_k_groups, HIDDEN) float8_e4m3fn, column-major — or prepacked topk_ids, # (num_tokens, NUM_TOPK) int32 — local expert IDs (for slot mapping) slot_token, # (num_slots,) int64 — token index per slot (from L1) alpha=1.0, # fp32 scalar from stage_activation global scale + sfb_prepacked=False, # True if l2_scales is prepacked CUTLASS layout ): """L2 GEMM: down_proj — slot-based, no routing weights. @@ -144,7 +182,10 @@ def nvfp4_mega_moe_l2( print(f"[nvfp4_moe_l2] slots={x_fp4.shape[0]} K={K} N={N} native=1") x_sf_fp8 = unpack_ue4m3_u32(x_sf) if x_sf.dtype == torch.uint32 else x_sf - w_sf_fp8 = unpack_ue4m3_u32(l2_scales) if l2_scales.dtype == torch.uint32 else l2_scales + if not sfb_prepacked: + w_sf_fp8 = unpack_ue4m3_u32(l2_scales) if l2_scales.dtype == torch.uint32 else l2_scales + else: + w_sf_fp8 = l2_scales # already prepacked # Build local expert IDs per slot (same mapping as L1) num_topk = topk_ids.shape[1] @@ -156,8 +197,9 @@ def nvfp4_mega_moe_l2( slot_out, _ = cutlass_grouped_nvfp4_gemm( x_fp4, x_sf_fp8, l2_weights, w_sf_fp8, - slot_expert_ids, # per-slot expert IDs + slot_expert_ids, alpha=alpha, + sfb_prepacked=sfb_prepacked, ) return slot_out # (num_slots, HIDDEN) bfloat16 @@ -318,11 +360,21 @@ def nvfp4_mega_moe_full( y.zero_() return - # Step 2: L1 GEMM — slot-based, no routing weights + # Prepack SFB weight scales into CUTLASS layout (lazy, once per layer) + l1_N = l1_w.shape[2] + l1_K = l1_w.shape[1] * 2 + l1_sf_prepacked = _prepack_weight_sf(l1_sf, l1_N, l1_K, "l1") + + l2_N = l2_w.shape[2] + l2_K = l2_w.shape[1] * 2 + l2_sf_prepacked = _prepack_weight_sf(l2_sf, l2_N, l2_K, "l2") + + # Step 2: L1 GEMM — slot-based, no routing weights, prepacked SFB l1_slots, _ = nvfp4_mega_moe_l1( - x_fp4, x_sf, l1_w, l1_sf, + x_fp4, x_sf, l1_w, l1_sf_prepacked, topk_ids_local, alpha=l1_global_scale, + sfb_prepacked=True, ) # (num_slots, 2*INTER) bfloat16 if MEGA_MOE_DEBUG: @@ -347,11 +399,12 @@ def nvfp4_mega_moe_full( _l2gs = l2_global_scale if isinstance(l2_global_scale, float) else l2_global_scale.item() print(f"[ALPHA L2] alpha={_l2gs:.4e} l1_sf range [{_l1sf_f32.min().item():.4e}, {_l1sf_f32.max().item():.4e}]") - # Step 5: L2 GEMM — slot-based, no routing weights + # Step 5: L2 GEMM — slot-based, no routing weights, prepacked SFB l2_slots = nvfp4_mega_moe_l2( - l1_fp4, l1_sf_out, l2_w, l2_sf, + l1_fp4, l1_sf_out, l2_w, l2_sf_prepacked, topk_ids_local, slot_token, alpha=l2_global_scale, + sfb_prepacked=True, ) # (num_slots, HIDDEN) bfloat16 if MEGA_MOE_DEBUG: