perf(glm5next): size the persisting-L2 set-aside for the weight prefetcher

tests/models/test_glm5next_l2_prefetch_persist.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Persisting-L2 set-aside sizing for the GLM-5.3 L2 weight prefetcher."""
+
+import pytest
+import torch
+
+from vllm.models.glm5next.nvidia import l2_prefetch as l2pf
+
+MAX = 84_000_000
+
+
+@pytest.mark.parametrize(
+    ("raw", "max_bytes", "expected"),
+    [
+        ("max", MAX, MAX),
+        (" MAX ", MAX, MAX),
+        ("all", MAX, MAX),
+        ("40", MAX, 40_000_000),
+        ("40.5", MAX, 40_500_000),
+        ("200", MAX, MAX),  # clamped to the device maximum
+        ("0", MAX, 0),
+        ("off", MAX, 0),
+        ("", MAX, 0),
+        (None, MAX, 0),
+        ("-5", MAX, 0),
+        ("bogus", MAX, 0),
+        ("max", 0, 0),  # device without a persisting L2 set-aside
+    ],
+)
+def test_persisting_l2_request(raw, max_bytes, expected):
+    assert l2pf.persisting_l2_request(raw, max_bytes) == expected
+
+
+def test_default_request_is_device_maximum():
+    assert l2pf.PERSIST_L2 == "max"
+
+
+def _driver():
+    from cuda.bindings import driver as cu
+
+    return cu
+
+
+def _read_limit(cu, device: torch.device) -> int:
+    with torch.cuda.device(device):
+        torch.empty(1, device=device)  # make the primary context current
+        err, value = cu.cuCtxGetLimit(cu.CUlimit.CU_LIMIT_PERSISTING_L2_CACHE_SIZE)
+    assert err == cu.CUresult.CUDA_SUCCESS
+    return int(value)
+
+
+def _set_limit(cu, device: torch.device, value: int) -> None:
+    with torch.cuda.device(device):
+        torch.empty(1, device=device)
+        (err,) = cu.cuCtxSetLimit(cu.CUlimit.CU_LIMIT_PERSISTING_L2_CACHE_SIZE, value)
+    assert err == cu.CUresult.CUDA_SUCCESS
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA device required")
+def test_configure_persisting_l2_applies_to_the_primary_context():
+    cu = _driver()
+    device = torch.device("cuda", 0)
+    err, dev = cu.cuDeviceGet(0)
+    assert err == cu.CUresult.CUDA_SUCCESS
+    err, max_bytes = cu.cuDeviceGetAttribute(
+        cu.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_PERSISTING_L2_CACHE_SIZE, dev
+    )
+    assert err == cu.CUresult.CUDA_SUCCESS
+    if max_bytes <= 0:
+        pytest.skip("device has no persisting L2 set-aside")
+    before = _read_limit(cu, device)
+    try:
+        assert l2pf.configure_persisting_l2(device, request="max") == max_bytes
+        assert _read_limit(cu, device) == max_bytes
+
+        # A zero request never touches the driver state.
+        assert l2pf.configure_persisting_l2(device, request="0") == 0
+        assert _read_limit(cu, device) == max_bytes
+
+        # Over-sized requests are clamped to the device maximum.
+        assert l2pf.configure_persisting_l2(device, request="100000") == max_bytes
+    finally:
+        _set_limit(cu, device, before)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA device required")
+def test_prefetcher_applies_the_env_request_once(monkeypatch):
+    cu = _driver()
+    device = torch.device("cuda", 0)
+    before = _read_limit(cu, device)
+    calls: list[str | None] = []
+    real = l2pf.configure_persisting_l2
+
+    def spy(dev, request=None):
+        calls.append(request)
+        return real(dev, request=request)
+
+    monkeypatch.setattr(l2pf, "configure_persisting_l2", spy)
+    monkeypatch.setattr(l2pf.L2Prefetcher, "_instances", {})
+    monkeypatch.setattr(l2pf, "_persisting_l2_applied", {})
+    try:
+        first = l2pf.L2Prefetcher.get(device)
+        second = l2pf.L2Prefetcher.get(device)
+        assert first is second
+        assert calls == [None]
+        assert first.persisting_l2_bytes == _read_limit(cu, device)
+    finally:
+        _set_limit(cu, device, before)

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -487,6 +487,11 @@ def _maybe_reduce_final_output(
         if output_is_reduced is None:
             output_is_reduced = self._fused_output_is_reduced
 
+        # Optional model-installed callback fired before the final all-reduce
+        # (GLM-5.3 L2 weight prefetch: the reduction leaves device memory idle).
+        _hook = getattr(self, "_l2_prefetch_pre_reduce_hook", None)
+        if _hook is not None:
+            _hook(states.shape[0])
         if (
             not self.moe_config.is_sequence_parallel
             and not self.moe_config.skip_final_all_reduce

vllm/model_executor/layers/linear.py

@@ -1656,6 +1656,11 @@ def forward(
         bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
         output_parallel = self.quant_method.apply(self, input_parallel, bias_)
 
+        # Optional model-installed callback fired before the all-reduce (GLM-5.3
+        # L2 weight prefetch: the reduction leaves device memory idle).
+        _hook = getattr(self, "_l2_prefetch_pre_reduce_hook", None)
+        if _hook is not None:
+            _hook(output_parallel.shape[0])
         if self.reduce_results and self.tp_size > 1:
             output = tensor_model_parallel_all_reduce(output_parallel)
         else:

vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py

@@ -755,6 +755,11 @@ def forward(
     ) -> None:
         num_tokens = hidden_states.size(0)
         projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
+        # Optional model-installed callback (e.g. GLM-5.3 L2 weight prefetch of
+        # o_proj while the small projections and the recurrence run).
+        _hook = getattr(self, "_l2_prefetch_hook", None)
+        if _hook is not None:
+            _hook(hidden_states.shape[0])
         if self.use_full_rank_gate:
             split_sizes = [
                 3 * self.local_projection_size,

vllm/model_executor/layers/mla.py

@@ -194,6 +194,11 @@ def forward(
         k_pe = k_pe.unsqueeze(1)
 
         q = q_proj_layer(q_proj_input)[0]
+        # Optional model-installed callback (e.g. GLM-5.3 L2 weight prefetch of
+        # o_proj while the indexer, rope and attention core run).
+        _hook = getattr(self, "_l2_prefetch_hook", None)
+        if _hook is not None:
+            _hook(hidden_states.shape[0])
         heads = self.num_heads
         if self.dcp_q_replicate:
             heads *= q_proj_layer.group_size