fix: allocate token_indices before CuTeDSL JIT compilation

CuTeDSL's cute.compile appears to corrupt GPU memory state,
causing torch.arange to produce zero-filled tensors when allocated
after the JIT compilation. Moving token_indices allocation before
the weight stacking operations fixes the corruption.

vllm/nvfp4_cutedsl.py

@@ -76,22 +76,10 @@ class CuTeDSLMoERunner:
         self._buffers_allocated = False
 
     def _allocate_buffers(self):
-        """Pre-allocate all buffers at max size for cudagraph compatibility."""
-        max_slots = self.max_num_tokens * self.top_k
+        """Pre-allocate scale buffers at max size for cudagraph compatibility."""
         K_sf = cutedsl_ceil_div(self.hidden_size, 16)
         padded_cols = cutedsl_ceil_div(K_sf, 4) * 4
         
-        # Slot -> token mapping: [0,0,...,0, 1,1,...,1, ...] (top_k repeats)
-        self._token_indices = torch.arange(
-            self.max_num_tokens, device=self.device, dtype=torch.int32
-        ).unsqueeze(1).expand(-1, self.top_k).contiguous().view(-1)
-        
-        self._expert_id_range = torch.arange(self.num_experts, device=self.device)
-        
-        self._expert_offsets_buf = torch.zeros(
-            self.num_experts + 1, dtype=torch.int32, device=self.device
-        )
-        
         # Per-expert scale buffers: separate L1/L2 since K_sf differs
         K_sf_l1 = cutedsl_ceil_div(self.hidden_size, 16)
         padded_cols_l1 = cutedsl_ceil_div(K_sf_l1, 4) * 4
@@ -120,6 +108,20 @@ class CuTeDSLMoERunner:
     def _ensure_stacked(self):
         if self._l1_mat_b is not None:
             return
+        
+        # Allocate token indices FIRST, before any CUDA JIT compilation
+        # (CuTeDSL's cute.compile can corrupt GPU memory state during JIT)
+        if self._token_indices is None:
+            self._token_indices = torch.arange(
+                self.max_num_tokens, device=self.device, dtype=torch.int32
+            ).unsqueeze(1).expand(-1, self.top_k).contiguous().view(-1)
+        if self._expert_id_range is None:
+            self._expert_id_range = torch.arange(self.num_experts, device=self.device)
+        if self._expert_offsets_buf is None:
+            self._expert_offsets_buf = torch.zeros(
+                self.num_experts + 1, dtype=torch.int32, device=self.device
+            )
+        
         self._l1_mat_b = make_b_k_major(torch.stack(self.l1_fp4))
         self._l2_mat_b = make_b_k_major(torch.stack(self.l2_fp4))
         self._l1_scale_b = assemble_scales_3d_side(self.l1_sf)