some optimizations

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<ElementSF*>(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<const ElementSF*>(SFB_ptr),
+        static_cast<ElementSF*>(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<ElementSF> 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<const ElementSF*>(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<const ArrayElementA*>(A_ptr), stride_A,
+            static_cast<const ArrayElementB*>(B_ptr), stride_B,
+            sfa_cutlass.get(), layout_SFA,
+            static_cast<const ElementSF*>(SFB_cutlass_ptr), layout_SFB
+        },
+        {
+          { alpha, beta },
+          nullptr, stride_C,
+          static_cast<typename Gemm::GemmKernel::CollectiveEpilogue::ElementD*>(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<uint8_t> workspace(workspace_size);
+
+    CUTLASS_CHECK(gemm.initialize(arguments, workspace.get(), stream));
+    CUTLASS_CHECK(gemm.run(stream));
+
+    return 0;
+}
+
 } // extern "C"
 
 #endif

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
 

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<int>(M), static_cast<int>(N), static_cast<int>(K),
+        static_cast<float>(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<int>(M),
+        static_cast<int>(N),
+        static_cast<int>(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<int>(M),
+        static_cast<int>(N),
+        static_cast<int>(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");
 }

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: