From 0aedc5678b4bf8aec0b77a2be85aa11abe176e5e Mon Sep 17 00:00:00 2001 From: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> Date: Wed, 20 May 2026 05:33:54 +0100 Subject: [PATCH] loader: yield filtered MTP weights lazily to avoid OOM hang on multi-layer EAGLE (#25748) --- python/sglang/srt/model_loader/loader.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 6294908bb..358ac5cf3 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -583,10 +583,12 @@ class DefaultModelLoader(BaseModelLoader): @classmethod def _filter_mtp_weights( cls, weights_iterator, prefix: str, draft_model_idx: int - ) -> Tuple[Tuple[str, torch.Tensor], ...]: - """Filter MTP (Multi-Token Prediction) weights to keep only the - specified draft model layer and remap it to layer 0.""" - filtered_weights = [] + ) -> Generator[Tuple[str, torch.Tensor], None, None]: + """Filter MTP weights to keep only the specified draft model layer + and remap it to layer 0. Yields lazily so the upstream buffered + iterator's sliding window actually bounds CPU memory — eager + materialization caused page-reclaim hangs on large MoE checkpoints + with multi-layer EAGLE.""" for name, tensor in weights_iterator: match = cls._MTP_PATTERN.match(name) if match is not None: @@ -596,8 +598,7 @@ class DefaultModelLoader(BaseModelLoader): new_name = name.replace(match.group(), "model.mtp.layers.0.") else: new_name = name - filtered_weights.append((prefix + new_name, tensor)) - return tuple(filtered_weights) + yield (prefix + new_name, tensor) def _get_all_weights( self,