diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index aea180a5..34f5ec74 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -45,7 +45,11 @@ class CuTeDSLMoERunner: self._l2_gsb = None def _ensure_stacked(self): - """Lazily stack weight tensors into the format the kernel expects.""" + """Lazily stack weight tensors into the format the kernel expects. + + After stacking, the per-expert lists are freed to avoid holding + two copies of ~175GB of weight data in GPU memory. + """ if self._l1_mat_b is not None: return self._l1_mat_b = make_b_k_major(torch.stack(self.l1_fp4)) @@ -54,6 +58,13 @@ class CuTeDSLMoERunner: self._l2_scale_b = assemble_scales_3d_side(self.l2_sf) self._l1_gsb = torch.tensor(self.l1_gs, dtype=torch.float32, device=self.device) self._l2_gsb = torch.tensor(self.l2_gs, dtype=torch.float32, device=self.device) + # Free per-expert lists — stacked tensors are the only copy now + self.l1_fp4 = None + self.l1_sf = None + self.l1_gs = None + self.l2_fp4 = None + self.l2_sf = None + self.l2_gs = None def prepare_weights_direct(self, l1_fp4, l1_sf, l1_gs, l2_fp4, l2_sf, l2_gs): """Set weights directly from checkpoint (no dequant→requant)."""