[BugFix] Fix static_forward_context cleanup and MTP draft loading (vllm-project#11491)

### What this PR does / why we need it?
This PR fixes two issues in `ModelNetLoaderElastic` when loading models
via D2D netloader, especially in MTP / speculative decoding scenarios.
**1. Clear `static_forward_context` before fallback**
When elastic loading fails and the loader falls back to
`DefaultModelLoader`, `initialize_model` may have already registered
attention layers into `compilation_config.static_forward_context`.
Without clearing these registrations, `revert_to_default()` can hit
duplicate-registration errors or leave stale layer state.
Changes:
- Add `_get_static_forward_context()` and
`_clear_static_forward_context()` helpers.
- Clear both the passed `vllm_config` and the global
`get_current_vllm_config()` (when available), with deduplication by
object id.
- Invoke cleanup in the fallback path, right before
`revert_to_default()`.
**2. Improve MTP draft model netloader efficiency and synchronization**
In MTP scenarios, target and draft models are loaded sequentially via
netloader. Two problems were observed:
- **DRAM cache inefficiency for draft model**: When `INT8_CACHE=dram`,
the draft model's `ElasticServer` also cached int8 weights to CPU DRAM.
Draft models are typically much smaller; keeping int8 cache in HBM
avoids unnecessary CPU↔NPU transfers during weight serving and improves
startup/serving efficiency.
- **Race across ranks**: Draft model loading could start before all
ranks finished target model netloader setup. A
`torch.distributed.barrier()` is added after target model loading (only
when `speculative_config` is set) so all ranks synchronize before draft
loading proceeds.
Changes:
- Draft `ElasticServer` uses `int8_cache="hbm"` when user config is not
`"no"` (overrides `"dram"` for draft only).
- Add `_sync_target_netloader_before_draft()` barrier after target model
load completes.
- Draft model path does not trigger the barrier.

fixes vllm-project#11553 

### Does this PR introduce _any_ user-facing change?
No API or CLI changes. 

### How was this patch tested?
Pre-commit passed
Unit tests passed
E2E verified: Elastic server start time has been significantly reduced
under netloader + MTP with 'int8_cache=dram' serving on 4 * Ascend 800I
A3 scenario.
**fix1**
**before:**
<img width="1589" height="804" alt="image"
src="https://github.com/user-attachments/assets/22c71938-d8ed-4e12-aa2d-6451e71fbad7"
/>
**after:**
<img width="1612" height="570" alt="image"
src="https://github.com/user-attachments/assets/07315209-a28b-4649-bdcc-ed92d2f905ca"
/>
curl coordinator:
<img width="1035" height="586" alt="image"
src="https://github.com/user-attachments/assets/ad5975c4-b3d7-494b-b539-8824b71dc686"
/>

**fix2**
**before:**
<img width="1017" height="503" alt="image"
src="https://github.com/user-attachments/assets/155f74f4-a254-4a73-aff7-a203c91c44e2"
/>
**after:**
<img width="1003" height="494" alt="image"
src="https://github.com/user-attachments/assets/569efaf8-7562-4826-91f7-d8148047bfab"
/>
curl coordinator: 
<img width="1032" height="635" alt="image"
src="https://github.com/user-attachments/assets/7a97c969-d5ad-48ef-a9a5-5be484c9a951"
/>

- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@1f486d9

---------

Signed-off-by: yilunh <hanyilun1@huawei.com>
Co-authored-by: yilunh998 <yilunh998@users.noreply.github.com>
Co-authored-by: yilunh <hanyilun1@huawei.com>

Author: yilunh998

tests/ut/model_loader/netloader/test_netloader.py

@@ -39,6 +39,7 @@ class DummyVllmConfig:
     parallel_config = DummyParallelConfig()
     additional_config = None
     quant_config = None
+    speculative_config: object | None = None
 
 
 class DummyModelConfig:
@@ -133,6 +134,30 @@ def test_init_with_invalid_config(monkeypatch):
     assert loader.output_prefix is None
 
 
+def test_clear_static_forward_context_clears_current_vllm_config(monkeypatch):
+    class DummyCompilationConfig:
+        def __init__(self):
+            self.static_forward_context = {}
+
+    class ConfigWithCompilation:
+        def __init__(self):
+            self.compilation_config = DummyCompilationConfig()
+
+    passed_config = ConfigWithCompilation()
+    current_config = ConfigWithCompilation()
+    passed_config.compilation_config.static_forward_context["passed.layer"] = object()
+    current_config.compilation_config.static_forward_context["current.layer"] = object()
+    monkeypatch.setattr(
+        "vllm_ascend.model_loader.netloader.netloader.get_current_vllm_config",
+        lambda: current_config,
+    )
+
+    ModelNetLoaderElastic._clear_static_forward_context(passed_config)
+
+    assert passed_config.compilation_config.static_forward_context == {}
+    assert current_config.compilation_config.static_forward_context == {}
+
+
 @patch("vllm_ascend.model_loader.netloader.netloader.logger")
 def test_load_model_elastic_success(mock_logger, monkeypatch, tmp_path):
     monkeypatch.setattr("torch.distributed.get_rank", lambda: 0)
@@ -219,6 +244,100 @@ def __exit__(self, a, b, c):
     return dummy_model
 
 
+@patch("vllm_ascend.model_loader.netloader.netloader.logger")
+def test_target_model_waits_for_all_netloader_ranks_before_draft(mock_logger, monkeypatch):
+    dummy_model = _patch_loader_common(monkeypatch)
+    monkeypatch.setattr("vllm_ascend.model_loader.netloader.netloader.elastic_load", lambda **kwargs: dummy_model)
+
+    class DummyElasticServer:
+        def __init__(*a, **k):
+            pass
+
+        def start(self):
+            pass
+
+    barrier_calls = []
+    monkeypatch.setattr("vllm_ascend.model_loader.netloader.netloader.ElasticServer", DummyElasticServer)
+    monkeypatch.setattr("torch.distributed.is_available", lambda: True)
+    monkeypatch.setattr("torch.distributed.is_initialized", lambda: True)
+    monkeypatch.setattr("torch.distributed.barrier", lambda: barrier_calls.append("barrier"))
+
+    extra = {
+        "SOURCE": [{"device_id": 0, "sources": ["127.0.0.1:5000"]}],
+        "MODEL": "dummy-model",
+        "LISTEN_PORT": 5555,
+        "INT8_CACHE": "dram",
+    }
+    loader = make_loader_with_config(extra)
+    vllm_config = DummyVllmConfig()
+    vllm_config.speculative_config = object()
+    loader.load_model(vllm_config, DummyModelConfig())
+
+    assert barrier_calls == ["barrier"]
+
+
+@patch("vllm_ascend.model_loader.netloader.netloader.logger")
+def test_failed_target_model_participates_in_barrier_before_error(mock_logger, monkeypatch):
+    _patch_loader_common(monkeypatch)
+    monkeypatch.setattr("vllm_ascend.model_loader.netloader.netloader.elastic_load", lambda **kwargs: None)
+    monkeypatch.setattr(
+        "vllm_ascend.model_loader.netloader.netloader.revert_to_default",
+        lambda *args, **kwargs: (None, False),
+    )
+
+    barrier_calls = []
+    monkeypatch.setattr("torch.distributed.is_available", lambda: True)
+    monkeypatch.setattr("torch.distributed.is_initialized", lambda: True)
+    monkeypatch.setattr("torch.distributed.barrier", lambda: barrier_calls.append("barrier"))
+
+    extra = {
+        "SOURCE": [{"device_id": 0, "sources": ["127.0.0.1:5000"]}],
+        "MODEL": "dummy-model",
+        "LISTEN_PORT": 5555,
+        "INT8_CACHE": "dram",
+    }
+    loader = make_loader_with_config(extra)
+    vllm_config = DummyVllmConfig()
+    vllm_config.speculative_config = object()
+
+    with pytest.raises(RuntimeError, match="NetLoader elastic loads model fails"):
+        loader.load_model(vllm_config, DummyModelConfig())
+
+    assert barrier_calls == ["barrier"]
+
+
+@patch("vllm_ascend.model_loader.netloader.netloader.logger")
+def test_draft_model_does_not_wait_for_target_netloader_barrier(mock_logger, monkeypatch):
+    dummy_model = _patch_loader_common(monkeypatch)
+    monkeypatch.setattr("vllm_ascend.model_loader.netloader.netloader.elastic_load", lambda **kwargs: dummy_model)
+
+    class DummyElasticServer:
+        def __init__(*a, **k):
+            pass
+
+        def start(self):
+            pass
+
+    barrier_calls = []
+    monkeypatch.setattr("vllm_ascend.model_loader.netloader.netloader.ElasticServer", DummyElasticServer)
+    monkeypatch.setattr("torch.distributed.is_available", lambda: True)
+    monkeypatch.setattr("torch.distributed.is_initialized", lambda: True)
+    monkeypatch.setattr("torch.distributed.barrier", lambda: barrier_calls.append("barrier"))
+
+    extra = {
+        "SOURCE": [{"device_id": 0, "sources": ["127.0.0.1:5000"]}],
+        "MODEL": "draft-model",
+        "LISTEN_PORT": 5555,
+        "INT8_CACHE": "dram",
+    }
+    loader = make_loader_with_config(extra)
+    vllm_config = DummyVllmConfig()
+    vllm_config.speculative_config = object()
+    loader.load_model(vllm_config, DummyDraftModelConfig())
+
+    assert barrier_calls == []
+
+
 @patch("vllm_ascend.model_loader.netloader.netloader.logger")
 def test_load_draft_model_elastic_success(mock_logger, monkeypatch, tmp_path):
     dummy_model = _patch_loader_common(monkeypatch)
@@ -229,6 +348,7 @@ def test_load_draft_model_elastic_success(mock_logger, monkeypatch, tmp_path):
     class DummyElasticServer:
         def __init__(self, *args, **kwargs):
             elastic_server_instances.append(self)
+            self.int8_cache = args[7]
             self.group_name = kwargs.get("group_name")
 
         def start(self):
@@ -248,10 +368,41 @@ def start(self):
 
     assert isinstance(result, nn.Module)
     assert loader._draft_elastic_server is elastic_server_instances[0]
+    assert loader._draft_elastic_server.int8_cache == "no"
     assert loader._draft_elastic_server.group_name == "netloader_draft"
     assert not (tmp_path / "output_0.txt").exists()
 
 
+@patch("vllm_ascend.model_loader.netloader.netloader.logger")
+def test_load_draft_model_uses_hbm_when_int8_cache_is_dram(mock_logger, monkeypatch):
+    dummy_model = _patch_loader_common(monkeypatch)
+    monkeypatch.setattr("vllm_ascend.model_loader.netloader.netloader.elastic_load", lambda **kwargs: dummy_model)
+
+    captured = {}
+
+    class DummyElasticServer:
+        def __init__(self, *args, **kwargs):
+            captured["int8_cache"] = args[7]
+            captured["group_name"] = kwargs.get("group_name")
+
+        def start(self):
+            pass
+
+    monkeypatch.setattr("vllm_ascend.model_loader.netloader.netloader.ElasticServer", DummyElasticServer)
+
+    extra = {
+        "SOURCE": [{"device_id": 0, "sources": ["127.0.0.1:5000"]}],
+        "MODEL": "draft-model",
+        "LISTEN_PORT": 5555,
+        "INT8_CACHE": "dram",
+    }
+    loader = make_loader_with_config(extra)
+    loader.load_model(DummyVllmConfig(), DummyDraftModelConfig())
+
+    assert captured["int8_cache"] == "hbm"
+    assert captured["group_name"] == "netloader_draft"
+
+
 @patch("vllm_ascend.model_loader.netloader.netloader.logger")
 def test_load_draft_model_port_offset_and_group_name(mock_logger, monkeypatch, tmp_path):
     dummy_model = _patch_loader_common(monkeypatch)

vllm_ascend/model_loader/netloader/netloader.py

@@ -35,6 +35,12 @@
 
 DRAFT_PORT_OFFSET = 10000
 
+try:
+    # Older vLLM versions may not expose the current-config accessor.
+    from vllm.config import get_current_vllm_config
+except ImportError:
+    get_current_vllm_config = None
+
 
 @register_model_loader("netloader")
 class ModelNetLoaderElastic(BaseModelLoader):
@@ -136,6 +142,61 @@ def _is_draft_model(model_config: ModelConfig) -> bool:
         """Check whether the model_config corresponds to a draft model for speculative decoding."""
         return getattr(model_config, "runner_type", None) == "draft"
 
+    @staticmethod
+    def _sync_target_netloader_before_draft(vllm_config: VllmConfig) -> None:
+        if getattr(vllm_config, "speculative_config", None) is None:
+            return
+        if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+            return
+
+        logger.info("Waiting for all target netloader ranks before loading draft model")
+        barrier_start = time.perf_counter()
+        torch.distributed.barrier()
+        logger.info(
+            "Target netloader barrier before draft model time: %s",
+            time.perf_counter() - barrier_start,
+        )
+
+    @staticmethod
+    def _get_static_forward_context(vllm_config: VllmConfig):
+        compilation_config = getattr(vllm_config, "compilation_config", None)
+        static_forward_context = getattr(compilation_config, "static_forward_context", None)
+        if static_forward_context is None or not hasattr(static_forward_context, "clear"):
+            return None
+        return static_forward_context
+
+    @staticmethod
+    def _clear_static_forward_context(vllm_config: VllmConfig) -> None:
+        """Clear static layer registrations before rebuilding the model on fallback."""
+        candidates = [("vllm_config", vllm_config)]
+        if get_current_vllm_config is not None:
+            try:
+                candidates.append(("current_vllm_config", get_current_vllm_config()))
+            except Exception as e:
+                logger.debug("Failed to get current vLLM config while clearing static context: %s", e)
+
+        cleared_contexts = []
+        seen_context_ids = set()
+        for source, config in candidates:
+            static_forward_context = ModelNetLoaderElastic._get_static_forward_context(config)
+            if static_forward_context is None:
+                continue
+
+            context_id = id(static_forward_context)
+            if context_id in seen_context_ids:
+                continue
+            seen_context_ids.add(context_id)
+
+            try:
+                context_size = str(len(static_forward_context))
+            except TypeError:
+                context_size = "unknown"
+            static_forward_context.clear()
+            cleared_contexts.append(f"{source}:{context_size}")
+
+        if cleared_contexts:
+            logger.info("Cleared static_forward_context before fallback: %s", cleared_contexts)
+
     def load_model(self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = "") -> nn.Module:
         """
         Loads the model using the specified configuration.
@@ -240,6 +301,9 @@ def load_model(self, vllm_config: VllmConfig, model_config: ModelConfig, prefix:
                         logger.info("Empty CUDA cache")
                         torch.cuda.empty_cache()
 
+                    # Clear registrations from the failed initialize_model
+                    self._clear_static_forward_context(vllm_config)
+
                     model, need_process_weights_after_loading = self.revert_to_default(
                         model_config, vllm_config, device_config, prefix
                     )
@@ -293,6 +357,7 @@ def load_model(self, vllm_config: VllmConfig, model_config: ModelConfig, prefix:
                         logger.error("Unknown error: %s", e)
 
                 try:
+                    server_int8_cache = "hbm" if is_draft and self.int8_cache != "no" else self.int8_cache
                     elastic_server = ElasticServer(
                         driver_ip,
                         listen_port,
@@ -301,7 +366,7 @@ def load_model(self, vllm_config: VllmConfig, model_config: ModelConfig, prefix:
                         model_config.model,
                         parallel_config.tensor_parallel_size,
                         parallel_config.pipeline_parallel_size,
-                        self.int8_cache,
+                        server_int8_cache,
                         self.int8_cache_name,
                         group_name=group_name,
                     )
@@ -321,9 +386,12 @@ def load_model(self, vllm_config: VllmConfig, model_config: ModelConfig, prefix:
         if need_process_weights_after_loading:
             process_weights_after_loading(model, model_config, torch.device(device_config.device))
 
+        if not is_draft:
+            self._sync_target_netloader_before_draft(vllm_config)
+
         if model is None:
             logger.error("NetLoader elastic loads model fails")
-            return None
+            raise RuntimeError("NetLoader elastic loads model fails")
 
         return model.eval()