clean up

src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py

@@ -283,11 +283,8 @@ def nvfp4_mega_moe_full(
     # vLLM's symm_buffer stores global IDs (0..383) but our weight tensors
     # are indexed by local ID (0..47). Each rank handles a contiguous chunk:
     # rank r gets experts [r*E_per_rank, (r+1)*E_per_rank).
-    # We derive the start index from the first global ID that maps to local 0.
     num_experts_per_rank = l1_w.shape[0]
-    # Find experts_start_idx: the minimum global ID that this rank handles.
-    # All topk_ids in the buffer should fall within this rank's range.
-    experts_start_idx = (topk_ids.min().item() // num_experts_per_rank) * num_experts_per_rank
+    experts_start_idx = symm_buffer.experts_start_idx
     topk_ids_local = topk_ids - experts_start_idx
 
     if MEGA_MOE_DEBUG:

src/nvfp4_megamoe_kernel/symm_buffer.py

@@ -31,6 +31,7 @@ class SymmBuffer:
         self.top_k = top_k
         self.hidden_size = hidden_size
         self.intermediate_size = intermediate_size
+        self.experts_start_idx = 0  # set by caller before kernel invocation
         
         device = torch.cuda.current_device()