From cdd813cf7eb8bd7ff88ad529a79078e36274b835 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 08:54:52 +0000 Subject: [PATCH] fix: free per-expert weight lists after stacking in CuTeDSL runner _ensure_stacked() creates stacked copies of all weights but never freed the per-expert lists. For 256 experts on a 175GB model, this doubles weight memory to ~350GB, causing OOM. Now the per-expert lists (l1_fp4, l1_sf, l1_gs, l2_fp4, l2_sf, l2_gs) are set to None after stacking, keeping only the single stacked copy. --- vllm/nvfp4_cutedsl.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) 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)."""