bug fix

src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py

@@ -71,32 +71,36 @@ def cutlass_grouped_nvfp4_gemm(
         slot_out:    (num_slots, N) bfloat16 — per-slot GEMM results
         slot_token:  (num_slots,) int64 — token index for each slot
     """
-    num_tokens = x_fp4.shape[0]
+    num_slots = x_fp4.shape[0]
     K_half = x_fp4.shape[1]
     K = K_half * 2
     N = weights.shape[2]
     num_experts = weights.shape[0]
-    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)
-    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"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
+        slot_token_out = torch.empty(0, dtype=torch.int64, device=x_fp4.device)
+        return slot_out, slot_token_out
 
-    # Gather activations for all slots
-    slot_x = x_fp4[slot_token]
-    slot_x_sf = x_sf[slot_token]
+    # topk_ids is either:
+    #   2D (num_tokens, num_topk) from L1 — build slot mapping
+    #   1D (num_slots,) from L2 — already per-slot expert IDs
+    if topk_ids.dim() == 2:
+        num_tokens = topk_ids.shape[0]
+        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]
+    else:
+        # 1D per-slot expert IDs — slot_token is just arange
+        slot_expert = topk_ids
+        slot_token = torch.arange(num_slots, device=x_fp4.device)
+
+    if MEGA_MOE_DEBUG:
+        print(f"[cutlass_grouped_gemm] slots={num_slots} K={K} N={N} "
+              f"experts={num_experts} sfb_prepacked={sfb_prepacked}")
+
+    slot_x = x_fp4
+    slot_x_sf = x_sf
 
     slot_out = torch.empty(num_slots, N, dtype=torch.bfloat16, device=x_fp4.device)
 
@@ -109,7 +113,7 @@ 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]  # prepacked or raw depending on flag
+        expert_w_sf = weight_sf[e]
         M_expert = e_idx.shape[0]
 
         if MEGA_MOE_DEBUG and e < 3 and M_expert > 0:

src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py

@@ -187,8 +187,7 @@ def nvfp4_mega_moe_l2(
     else:
         w_sf_fp8 = l2_scales  # already prepacked
 
-    # Build local expert IDs per slot (same mapping as L1)
-    num_topk = topk_ids.shape[1]
+    # Build 1D per-slot expert IDs (same slot ordering as L1)
     num_experts = l2_weights.shape[0]
     local_mask = (topk_ids >= 0) & (topk_ids < num_experts)
     _, slot_k = local_mask.nonzero(as_tuple=True)
@@ -197,7 +196,7 @@ def nvfp4_mega_moe_l2(
     slot_out, _ = cutlass_grouped_nvfp4_gemm(
         x_fp4, x_sf_fp8,
         l2_weights, w_sf_fp8,
-        slot_expert_ids,
+        slot_expert_ids,  # 1D per-slot expert IDs — GEMM handles directly
         alpha=alpha,
         sfb_prepacked=sfb_prepacked,
     )