diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index a2e1193d..9e4f45c0 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/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)