[Feature] Support NZ static buffers for prefetch offload (vllm-project#11945)

### What this PR does / why we need it?

This PR supports Ascend NZ-format static buffers for prefetch CPU
offload.

- Marks parameters whose quantized weights are transformed to
`ACL_FORMAT_FRACTAL_NZ` during `process_weights_after_loading`, after
`maybe_trans_nz` has run.
- Casts the corresponding prefetch `StaticBufferPool` buffers to
`ACL_FORMAT_FRACTAL_NZ` during `post_init`, so offloaded weights are
copied back into buffers with the expected Ascend storage format.
- Refactors the Ascend prefetch offloader to reuse upstream vLLM
`PrefetchOffloader` behavior while installing Ascend-specific module
offloaders and keeping CUDA stream/event aliases mapped to NPU APIs
after model runner initialization.

In my setup, this makes it possible to run prefill for GLM-5.1-W8A8 with
context `max-model-len 202752`, `max-num-batched-tokens 128`,
`block-size 128` and `gpu-memory-utilization 0.92` on a single-node
16-card a3, although it is still very slow. I use prefetch CPU offload
for selected MLP weights with:

```bash
--offload-backend prefetch
--offload-group-size 4
--offload-num-in-group 2
--offload-prefetch-step 2
--offload-params mlp
```

### Previous works


[vllm-project#10251](vllm-project#10251)
added the initial Ascend prefetch offloader by porting the upstream vLLM
implementation and replacing CUDA stream/event APIs with NPU
equivalents. However, that implementation still reused the normal-layout
`StaticBufferPool` flow and did not preserve Ascend NZ weight format for
static buffers.

This matters because vLLM Ascend has enabled NZ conversion for quantized
weights by default since
[vllm-project#4878](vllm-project#4878), included
from `v0.13.0rc1`: `VLLM_ASCEND_ENABLE_NZ=1` means quantized weights are
transformed to NZ by `maybe_trans_nz` during
`process_weights_after_loading`. When prefetch offload copies those NZ
weights through normal-layout static buffers, the offloaded weights may
be restored with the wrong layout. I tested the vllm-project#10251 behavior on
Ascend A3 and reproduced accuracy issues from this mismatch.


- vLLM version: v0.24.0
- vLLM main:
vllm-project/vllm@85c09e9

---------

Signed-off-by: KINGFIOX <me@kingfiox.work>
Co-authored-by: kurumi5210 <Jaychou1620@Gmail.com>

Author: KINGFIOX

.github/workflows/scripts/test_config.yaml

@@ -268,6 +268,7 @@
     - vllm_ascend/worker/model_runner_v1.py
     - vllm_ascend/compilation/acl_graph.py
   tests:
+    - tests/ut/model_executor
     - tests/e2e/pull_request/one_card/test_cpu_weight_offload.py
 
 # Model Loader
@@ -717,7 +718,7 @@ estimated_times:
   tests/e2e/pull_request/one_card/test_camem.py: 140
   tests/e2e/pull_request/one_card/test_completion_with_prompt_embeds.py: 200
   tests/e2e/pull_request/one_card/test_cpu_offloading.py: 30
-  tests/e2e/pull_request/one_card/test_cpu_weight_offload.py: 830
+  tests/e2e/pull_request/one_card/test_cpu_weight_offload.py: 1400
   tests/e2e/pull_request/one_card/test_guided_decoding.py: 610
   tests/e2e/pull_request/one_card/test_minicpm.py: 330
   tests/e2e/pull_request/one_card/test_multi_instance.py: 140

tests/e2e/pull_request/one_card/test_cpu_weight_offload.py

@@ -6,63 +6,109 @@
 #
 """End-to-end tests for CPU weight offloading on Ascend NPU.
 
-Covers both the prefetch backend (NPUPrefetchOffloader) and the UVA
-backend (functional_call fallback path, since UVA is not available on
-NPU hardware).  Tests verify that offloading produces the same outputs
+Covers both the prefetch backend (AscendPrefetchOffloader).
+Tests verify that offloading produces the same outputs
 as the baseline (no offloading).
 """
 
-import os
-
 import pytest
 
-from tests.e2e.conftest import wait_until_npu_memory_free
+from tests.e2e.conftest import VllmRunner, wait_until_npu_memory_free
+from tests.e2e.pull_request import utils as e2e_utils
 from tests.e2e.pull_request.utils import PROMPTS_SHORT, compare_logprobs
 
 MODEL = "Qwen/Qwen3-0.6B"
 
+_OFFLOAD_KEYS = {
+    "offload_backend",
+    "offload_group_size",
+    "offload_num_in_group",
+    "offload_prefetch_step",
+    "offload_params",
+    "cpu_offload_gb",
+}
+
+
+def _compare_offload_logprobs(
+    runner_kwargs: dict,
+    prompts: list[str],
+    atol: float = 0.0689,
+    decode_atol: float | None = None,
+) -> None:
+    """Compare prefetch/offload run against a no-offload eager baseline.
+
+    Unlike ``compare_logprobs``, this keeps ``additional_config`` (e.g.
+    ``weight_nz_mode``) on both sides and strips offload-related kwargs from
+    the baseline so accuracy of the offloader itself is exercised.
+    """
+    if decode_atol is None:
+        decode_atol = 2 * atol
+
+    baseline_kwargs = {k: v for k, v in runner_kwargs.items() if k not in _OFFLOAD_KEYS}
+    baseline_kwargs.pop("cudagraph_capture_sizes", None)
+    baseline_kwargs["enforce_eager"] = True
+
+    # baseline(eager, no offload)
+    with VllmRunner(**baseline_kwargs) as runner:
+        baseline_outputs = runner.model.generate(
+            prompts=prompts,
+            sampling_params=e2e_utils._LOGPROB_SAMPLING_PARAMS,
+        )
+
+    # enabled offload
+    with VllmRunner(**runner_kwargs) as runner:
+        offload_outputs = runner.model.generate(
+            prompts=prompts,
+            sampling_params=e2e_utils._LOGPROB_SAMPLING_PARAMS,
+        )
+
+    for prompt_idx, (base_out, offload_out) in enumerate(zip(baseline_outputs, offload_outputs)):
+        base_seq = base_out.outputs[0]
+        offload_seq = offload_out.outputs[0]
+
+        assert base_seq.logprobs is not None and offload_seq.logprobs is not None, (
+            f"logprobs not returned for prompt {prompt_idx}"
+        )
+        assert len(base_seq.token_ids) == len(offload_seq.token_ids) == 3, (
+            f"Expected 3 tokens for prompt {prompt_idx}, "
+            f"got baseline={len(base_seq.token_ids)}, offload={len(offload_seq.token_ids)}"
+        )
+
+        e2e_utils._check_prefill_token(base_seq, offload_seq, prompt_idx, atol)
+        for token_idx in range(1, 3):
+            e2e_utils._check_decode_token(base_seq, offload_seq, token_idx, prompt_idx, decode_atol)
+
 
 # -------------------- Prefetch backend tests --------------------
 
 
+@pytest.mark.parametrize("enforce_eager", [True, False], ids=["eager", "graph"])
+@pytest.mark.parametrize("nz_mode", [0, 1], ids=["ND", "NZ"])
 @wait_until_npu_memory_free()
-def test_prefetch_offload_eager():
-    """Test prefetch CPU offloading in eager mode.
+def test_prefetch_offload_accuracy(enforce_eager, nz_mode):
+    """Test prefetch CPU offloading across eager/graph × ND/NZ.
 
     Compares outputs between:
-    1. Baseline (eager, no offloading)
+    1. Baseline (eager, no offloading, same weight_nz_mode)
     2. Prefetch offloading (group_size=4, num_in_group=1)
-       with enforce_eager=True (no ACL graph capture)
-    """
-    runner_kwargs = {
-        "model_name": MODEL,
-        "max_model_len": 512,
-        "enforce_eager": True,
-        "offload_backend": "prefetch",
-        "offload_group_size": 4,
-        "offload_num_in_group": 1,
-    }
-    compare_logprobs(runner_kwargs=runner_kwargs, prompts=PROMPTS_SHORT)
-
 
-@wait_until_npu_memory_free()
-def test_prefetch_offload_aclgraph():
-    """Test prefetch CPU offloading with ACL graph capture.
-
-    Compares outputs between:
-    1. Baseline (eager, no offloading)
-    2. Prefetch offloading (group_size=4, num_in_group=1)
-       with ACL graph capture enabled (default, non-eager)
+    NZ uses weight_nz_mode=2 so BF16 weights are converted to FRACTAL_NZ
+    (mode 1 only enables NZ for quantized weights).
     """
-    runner_kwargs = {
+    runner_kwargs: dict = {
         "model_name": MODEL,
         "max_model_len": 512,
-        "cudagraph_capture_sizes": [1, 2, 4, 8],
         "offload_backend": "prefetch",
         "offload_group_size": 4,
         "offload_num_in_group": 1,
+        "additional_config": {"weight_nz_mode": nz_mode},
     }
-    compare_logprobs(runner_kwargs=runner_kwargs, prompts=PROMPTS_SHORT)
+    if enforce_eager:
+        runner_kwargs["enforce_eager"] = True
+    else:
+        runner_kwargs["cudagraph_capture_sizes"] = [1, 2, 4, 8]
+
+    _compare_offload_logprobs(runner_kwargs=runner_kwargs, prompts=PROMPTS_SHORT)
 
 
 @wait_until_npu_memory_free()
@@ -83,48 +129,3 @@ def test_prefetch_offload_selective_params():
         "offload_params": {"gate_up_proj", "down_proj"},
     }
     compare_logprobs(runner_kwargs=runner_kwargs, prompts=PROMPTS_SHORT)
-
-
-# -------------------- UVA backend tests --------------------
-# UVA (Unified Virtual Addressing) is not available on Ascend NPU, so
-# the UVA offloader falls back to the functional_call path that moves
-# weights to device on-demand.  Tests below mirror the upstream
-# test_cpu_offload.py parametrization but only exercise the non-UVA
-# (functional_call) path with enforce_eager, as NPU does not support
-# UVA zero-copy.
-
-
-@pytest.mark.parametrize("disable_pin_memory", [False, True])
-@wait_until_npu_memory_free()
-def test_uva_offload_functional_call(disable_pin_memory):
-    """Test UVA offloader's functional_call fallback on NPU.
-
-    With UVA disabled (forced by env var), the UVA offloader falls back
-    to moving weights to device inside a functional_call wrapper.
-    enforce_eager is required because this fallback is incompatible
-    with graph capture.
-
-    Parametrized over pin_memory to cover both pinned and unpinned
-    CPU storage paths.
-    """
-    old_uva = os.environ.get("VLLM_WEIGHT_OFFLOADING_DISABLE_UVA")
-    old_pin = os.environ.get("VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY")
-    try:
-        os.environ["VLLM_WEIGHT_OFFLOADING_DISABLE_UVA"] = "1"
-        os.environ["VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY"] = str(int(disable_pin_memory))
-        runner_kwargs = {
-            "model_name": MODEL,
-            "max_model_len": 512,
-            "enforce_eager": True,
-            "cpu_offload_gb": 1,
-        }
-        compare_logprobs(runner_kwargs=runner_kwargs, prompts=PROMPTS_SHORT)
-    finally:
-        if old_uva is None:
-            os.environ.pop("VLLM_WEIGHT_OFFLOADING_DISABLE_UVA", None)
-        else:
-            os.environ["VLLM_WEIGHT_OFFLOADING_DISABLE_UVA"] = old_uva
-        if old_pin is None:
-            os.environ.pop("VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY", None)
-        else:
-            os.environ["VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY"] = old_pin

tests/ut/model_executor/test_offloader.py

@@ -0,0 +1,37 @@
+from types import SimpleNamespace
+
+from vllm.model_executor.offloader.base import NoopOffloader
+
+from vllm_ascend.model_executor.offloader.base import create_offloader
+from vllm_ascend.model_executor.offloader.prefetch import _is_using_nz_weight
+
+
+def test_create_offloader_without_config_returns_noop():
+    offloader = create_offloader(None)
+
+    assert isinstance(offloader, NoopOffloader)
+
+
+def test_create_offloader_for_non_prefetch_backend_returns_noop():
+    offload_config = SimpleNamespace(
+        offload_backend=None,
+        prefetch=None,
+    )
+
+    offloader = create_offloader(offload_config)
+
+    assert isinstance(offloader, NoopOffloader)
+
+
+def test_is_using_nz_weight_handles_invalid_npu_format(monkeypatch):
+    param = SimpleNamespace(
+        data=SimpleNamespace(device=SimpleNamespace(type="npu")),
+    )
+
+    monkeypatch.setattr(
+        "vllm_ascend.model_executor.offloader.prefetch.torch_npu.get_npu_format",
+        lambda _: object(),
+        raising=False,
+    )
+
+    assert not _is_using_nz_weight(param)

vllm_ascend/model_executor/offloader/__init__.py

@@ -0,0 +1,9 @@
+"""Ascend-specific model parameter offloading."""
+
+from vllm_ascend.model_executor.offloader.base import create_offloader
+from vllm_ascend.model_executor.offloader.prefetch import AscendPrefetchOffloader
+
+__all__ = [
+    "AscendPrefetchOffloader",
+    "create_offloader",
+]

vllm_ascend/model_executor/offloader/base.py

@@ -0,0 +1,28 @@
+from typing import TYPE_CHECKING
+
+from vllm.model_executor.offloader.base import BaseOffloader, NoopOffloader
+
+from vllm_ascend.model_executor.offloader.prefetch import AscendPrefetchOffloader
+
+if TYPE_CHECKING:
+    from vllm.config import OffloadConfig
+
+
+def create_offloader(offload_config: "OffloadConfig | None") -> BaseOffloader:
+    """Create an Ascend-aware offloader while preserving vLLM defaults."""
+
+    if offload_config is None:
+        return NoopOffloader()
+
+    backend = offload_config.offload_backend
+    prefetch = offload_config.prefetch
+
+    if backend == "prefetch":
+        return AscendPrefetchOffloader(
+            group_size=prefetch.offload_group_size,
+            num_in_group=prefetch.offload_num_in_group,
+            prefetch_step=prefetch.offload_prefetch_step,
+            offload_params=prefetch.offload_params,
+            mode="cpu",
+        )
+    return NoopOffloader()