diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 2b5b67c0..de93657b 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -76,6 +76,10 @@ class CuTeDSLMoERunner: self._per_expert_scale_bufs_l2 = None self._padded_x_sf_buf_l1 = None self._padded_x_sf_buf_l2 = None + self._l1_gsa_buf = None + self._l2_gsa_buf = None + self._output_buf = None + self._row_indices_buf = None self._buffers_allocated = False def _fill_token_indices(self): @@ -114,6 +118,20 @@ class CuTeDSLMoERunner: self.num_experts * 128, padded_cols_l2, dtype=torch.float16, device=self.device ).to(torch.float8_e4m3fn) + # Pre-allocated global_scale_a buffers (filled via .fill_(), no torch.full during capture) + self._l1_gsa_buf = torch.zeros(self.num_experts, dtype=torch.float32, device=self.device) + self._l2_gsa_buf = torch.zeros(self.num_experts, dtype=torch.float32, device=self.device) + + # Pre-allocated output buffer + self._output_buf = torch.zeros( + self.max_num_tokens, self.hidden_size, dtype=torch.bfloat16, device=self.device + ) + + # Row indices for scale assembly (max_num_tokens * top_k slots) + self._row_indices_buf = torch.arange( + self.max_num_tokens * self.top_k, device=self.device + ) + self._buffers_allocated = True def _ensure_stacked(self): @@ -204,7 +222,7 @@ class CuTeDSLMoERunner: # Tokens beyond 128 per expert wrap to zero rows (harmless — zero scale # means zero contribution in GEMM). All indexing is fixed-shape. total_rows = x_sf.shape[0] - row_indices = torch.arange(total_rows, device=x_sf.device) + row_indices = self._row_indices_buf[:total_rows] expert_assign = torch.searchsorted( expert_offsets[1:], row_indices, right=True ).clamp(max=num_experts - 1) @@ -360,10 +378,7 @@ class CuTeDSLMoERunner: x_sf, expert_offsets[:self.num_experts + 1], self._padded_x_sf_buf_l1, self._per_expert_scale_bufs_l1 ) - l1_gsa = torch.full( - (self.num_experts,), self._l1_activation_global_scale, - dtype=torch.float32, device=device - ) + l1_gsa = self._l1_gsa_buf.fill_(self._l1_activation_global_scale) l1_out = run_nvfp4_grouped_gemm( mat_a=x_fp4, mat_b=self._l1_mat_b, @@ -386,10 +401,7 @@ class CuTeDSLMoERunner: l2_x_sf, expert_offsets[:self.num_experts + 1], self._padded_x_sf_buf_l2, self._per_expert_scale_bufs_l2 ) - l2_gsa = torch.full( - (self.num_experts,), self._l2_activation_global_scale, - dtype=torch.float32, device=device - ) + l2_gsa = self._l2_gsa_buf.fill_(self._l2_activation_global_scale) l2_out = run_nvfp4_grouped_gemm( mat_a=l2_x_fp4, mat_b=self._l2_mat_b, @@ -399,7 +411,8 @@ class CuTeDSLMoERunner: ) # === Scatter -> final output === - y = torch.zeros(num_tokens, self.hidden_size, dtype=torch.bfloat16, device=device) + y = self._output_buf[:num_tokens] + y.zero_() weighted_out = l2_out * sorted_weights.unsqueeze(1).to(l2_out.dtype) y.scatter_add_( 0,