diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index a440a42f..5e234263 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -142,15 +142,22 @@ class CuTeDSLMoERunner: self.max_num_tokens * self.top_k, device=self.device ) - # Padded hidden/activated: sized for num_experts * max_chunks * 128 (the actual padded size) - max_rows_per_expert = self._max_chunks_per_expert * 128 - padded_max_slots = self.num_experts * max_rows_per_expert - self._padded_hidden_buf = torch.zeros( - padded_max_slots, self.hidden_size, dtype=torch.bfloat16, device=self.device - ) - self._padded_activated_buf = torch.zeros( - padded_max_slots, self.intermediate_size, dtype=torch.bfloat16, device=self.device - ) + # Padded hidden/activated: SHARED across all runners (not per-layer) + # These are only used during run() which is sequential across layers. + # Per-layer allocation would be 72 MB × 60 layers = 4.3 GB → OOM. + device_key = str(self.device) + if not hasattr(CuTeDSLMoERunner, '_shared_padded_bufs'): + CuTeDSLMoERunner._shared_padded_bufs = {} + if device_key not in CuTeDSLMoERunner._shared_padded_bufs: + CuTeDSLMoERunner._shared_padded_bufs[device_key] = { + 'hidden': torch.zeros( + padded_max_slots, self.hidden_size, dtype=torch.bfloat16, device=self.device + ), + 'activated': torch.zeros( + padded_max_slots, self.intermediate_size, dtype=torch.bfloat16, device=self.device + ), + } + self._shared_bufs = CuTeDSLMoERunner._shared_padded_bufs[device_key] # Padded expert offsets buffer: [0, max_rows, 2*max_rows, ...] (fixed) self._padded_expert_offsets_buf = torch.zeros( @@ -389,7 +396,7 @@ class CuTeDSLMoERunner: # -- Gather hidden states into slot order, pad per expert -- slot_hidden = hidden_states[sorted_token_ids] total_padded_slots = self.num_experts * max_rows_per_expert - padded_hidden = self._padded_hidden_buf[:total_padded_slots] + padded_hidden = self._shared_bufs['hidden'][:total_padded_slots] padded_hidden.zero_() row_indices = self._row_indices_buf[:num_slots] expert_assign = torch.searchsorted( @@ -427,7 +434,7 @@ class CuTeDSLMoERunner: activated = torch.nn.functional.silu(gate) * up # === L2: down === - padded_activated = self._padded_activated_buf[:total_padded_slots] + padded_activated = self._shared_bufs['activated'][:total_padded_slots] padded_activated.zero_() padded_activated[padded_dst] = activated diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 2c10d7d4..8070bedb 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -499,14 +499,11 @@ class DeepseekV4MegaMoEExperts(nn.Module): l2_gs.append(down_gs) # Create CuTeDSL runner with directly-cast weights - # Max tokens for buffer pre-allocation: use cudagraph max capture size - # (not scheduler max which can be 8192, causing OOM with 60 layers) - max_cg_size = getattr(self, '_cudagraph_max_capture_size', 512) self._cutedsl_runner = CuTeDSLMoERunner( num_experts=self.num_local_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size, - max_num_tokens=min(self.max_num_tokens, max_cg_size), + max_num_tokens=self.max_num_tokens, top_k=self.top_k, device=l1_fp4[0].device, experts_start_idx=self.experts_start_idx,