diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..c1d6fe4 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,3 @@ +# Unified diffs contain a one-character context marker on blank context lines. +# Validate whitespace after applying the patch to its source tree instead. +patch/*.patch -whitespace diff --git a/Dockerfile.sm120-v024 b/Dockerfile.sm120-v024 index 5e0e4f7..27aa928 100644 --- a/Dockerfile.sm120-v024 +++ b/Dockerfile.sm120-v024 @@ -72,8 +72,9 @@ RUN pip install --no-cache-dir --break-system-packages flashinfer-python==0.6.14 --index-url https://flashinfer.ai/whl/cu130 \ "flashinfer-jit-cache==0.6.14+cu130" -# 3. our patch, applied to the installed package (pure-Python delta; -# includes the DSpark speculative-decoding backport) +# 3. Combined v0.24.0 patch: W2 streaming/recovery, the step-driven tier +# manager, the DSpark speculative backport, and SM120 fixes. One canonical +# patch keeps the image label an exact source identity. COPY patch/vllm-moet-v0.24.0.patch /tmp/vllm-moet.patch RUN cd "$(python3 -c 'import vllm, os; print(os.path.dirname(os.path.dirname(vllm.__file__)))')" \ && git apply --verbose /tmp/vllm-moet.patch \ diff --git a/patch/FILES.txt b/patch/FILES.txt index cbf9443..179ab44 100644 --- a/patch/FILES.txt +++ b/patch/FILES.txt @@ -4,6 +4,9 @@ # patch but not in the vllm fork branch the patch is generated from — # merge it into the branch first; do not ship the regeneration. csrc/nvfp4_ds_mla/concat_and_cache_nvfp4_ds_mla.cu +tests/model_executor/layers/quantization/test_moe_w2_cgroup_memory.py +tests/model_executor/layers/quantization/test_moe_w2_padded_routes.py +tests/model_executor/layers/quantization/test_moe_w2_step_pins.py tools/nvfp4_flashinfer_sm120/README.md tools/nvfp4_flashinfer_sm120/nvfp4_expand.cuh tools/nvfp4_flashinfer_sm120/patch_flashinfer.py @@ -12,6 +15,7 @@ vllm/compilation/cuda_graph.py vllm/config/speculative.py vllm/config/vllm.py vllm/envs.py +vllm/forward_context.py vllm/model_executor/layers/attention/mla_attention.py vllm/model_executor/layers/quantization/fp8.py vllm/model_executor/layers/quantization/modelopt.py @@ -27,6 +31,9 @@ vllm/model_executor/layers/quantization/utils/moe_w2_store.py vllm/model_executor/layers/quantization/utils/prefill_timers.py vllm/model_executor/layers/quantization/utils/skinny_fp8_cubit.py vllm/model_executor/layers/sparse_attn_indexer.py +vllm/model_executor/model_loader/__init__.py +vllm/model_executor/model_loader/default_loader.py +vllm/model_executor/model_loader/weight_utils.py vllm/model_executor/models/deepseek_mtp.py vllm/model_executor/models/qwen3_dflash.py vllm/model_executor/models/qwen3_dspark.py diff --git a/patch/vllm-moet-v0.24.0.patch b/patch/vllm-moet-v0.24.0.patch index 5fe2a51..05e230e 100644 --- a/patch/vllm-moet-v0.24.0.patch +++ b/patch/vllm-moet-v0.24.0.patch @@ -1,6 +1,6 @@ diff --git a/csrc/nvfp4_ds_mla/concat_and_cache_nvfp4_ds_mla.cu b/csrc/nvfp4_ds_mla/concat_and_cache_nvfp4_ds_mla.cu new file mode 100644 -index 000000000..df62c57cc +index 0000000..df62c57 --- /dev/null +++ b/csrc/nvfp4_ds_mla/concat_and_cache_nvfp4_ds_mla.cu @@ -0,0 +1,142 @@ @@ -146,9 +146,978 @@ index 000000000..df62c57cc + m.def("concat_and_cache_nvfp4_ds_mla", &concat_and_cache_nvfp4_ds_mla, + "Write MLA KV into the packed nvfp4_ds_mla layout (352 B/token)"); +} +diff --git a/tests/model_executor/layers/quantization/test_moe_w2_cgroup_memory.py b/tests/model_executor/layers/quantization/test_moe_w2_cgroup_memory.py +new file mode 100644 +index 0000000..97a0b6a +--- /dev/null ++++ b/tests/model_executor/layers/quantization/test_moe_w2_cgroup_memory.py +@@ -0,0 +1,124 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import importlib.util ++import logging ++import os ++import sys ++import tempfile ++import unittest ++from pathlib import Path ++from types import ModuleType ++from unittest import mock ++ ++ROOT = Path(__file__).resolve().parents[4] ++STORE_PATH = ROOT / "vllm/model_executor/layers/quantization/utils/moe_w2_store.py" ++GIB = 1 << 30 ++SAFETY_ENV = { ++ "VLLM_MOE_W2_MIN_MEM_AVAILABLE_GB": "16", ++ "VLLM_MOE_W2_MIN_CGROUP_HEADROOM_GB": "4", ++} ++ ++ ++def _load_store_module(): ++ """Load the store in isolation; these checks need only CPU torch.""" ++ vllm_module = ModuleType("vllm") ++ vllm_module.__path__ = [str(ROOT / "vllm")] ++ logger_module = ModuleType("vllm.logger") ++ logger_module.init_logger = logging.getLogger ++ spec = importlib.util.spec_from_file_location("_test_moe_w2_store", STORE_PATH) ++ assert spec is not None and spec.loader is not None ++ module = importlib.util.module_from_spec(spec) ++ with mock.patch.dict( ++ sys.modules, ++ {"vllm": vllm_module, "vllm.logger": logger_module}, ++ ): ++ spec.loader.exec_module(module) ++ return module ++ ++ ++def _cgroup_status(max_available: int, high_available: int | None) -> dict: ++ return { ++ "known": True, ++ "version": 2, ++ "path": "/test", ++ "limited": True, ++ "max_available": max_available, ++ "high_available": high_available, ++ "current": 1 * GIB, ++ "events": {}, ++ } ++ ++ ++class TestMoeW2CgroupMemory(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ cls.store = _load_store_module() ++ ++ def _v2_status(self, current: int, high: str, maximum: str) -> dict: ++ with tempfile.TemporaryDirectory() as root: ++ files = { ++ "memory.current": str(current), ++ "memory.high": high, ++ "memory.max": maximum, ++ "memory.stat": "anon 1\nfile 2\nfile_mapped 3\n", ++ "memory.events": "high 0\nmax 0\noom 0\n", ++ "memory.swap.current": "0", ++ "memory.swap.max": "max", ++ } ++ for name, value in files.items(): ++ Path(root, name).write_text(value) ++ with mock.patch.object( ++ self.store, "_active_cgroup_v2_dirs", return_value=[root] ++ ): ++ return self.store._cgroup_memory_status() ++ ++ def test_crossed_soft_high_is_separate_from_hard_max_headroom(self): ++ status = self._v2_status(11 * GIB, str(10 * GIB), str(20 * GIB)) ++ ++ self.assertTrue(status["limited"]) ++ self.assertEqual(status["max_available"], 9 * GIB) ++ self.assertEqual(status["high_available"], -1 * GIB) ++ ++ def test_crossed_soft_high_does_not_refuse_safe_hard_headroom(self): ++ with ( ++ mock.patch.dict(os.environ, SAFETY_ENV, clear=False), ++ mock.patch.object( ++ self.store, "_mem_available_bytes", return_value=64 * GIB ++ ), ++ mock.patch.object( ++ self.store, ++ "_cgroup_memory_status", ++ return_value=_cgroup_status(8 * GIB, -1 * GIB), ++ ), ++ ): ++ report = self.store._memory_preflight("soft-high", 2 * GIB) ++ ++ self.assertEqual(report["cgroup_max_available"], 8 * GIB) ++ self.assertEqual(report["cgroup_high_available"], -1 * GIB) ++ ++ def test_hard_max_headroom_still_refuses_unsafe_transient(self): ++ with ( ++ mock.patch.dict(os.environ, SAFETY_ENV, clear=False), ++ mock.patch.object( ++ self.store, "_mem_available_bytes", return_value=64 * GIB ++ ), ++ mock.patch.object( ++ self.store, ++ "_cgroup_memory_status", ++ return_value=_cgroup_status(6 * GIB, 20 * GIB), ++ ), ++ self.assertRaisesRegex(RuntimeError, "memory.max headroom"), ++ ): ++ self.store._memory_preflight("hard-max", 3 * GIB) ++ ++ def test_finite_soft_high_does_not_make_unlimited_max_limited(self): ++ status = self._v2_status(5 * GIB, str(10 * GIB), "max") ++ ++ self.assertFalse(status["limited"]) ++ self.assertIsNone(status["max_available"]) ++ self.assertEqual(status["high_available"], 5 * GIB) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/tests/model_executor/layers/quantization/test_moe_w2_padded_routes.py b/tests/model_executor/layers/quantization/test_moe_w2_padded_routes.py +new file mode 100644 +index 0000000..2b4a328 +--- /dev/null ++++ b/tests/model_executor/layers/quantization/test_moe_w2_padded_routes.py +@@ -0,0 +1,363 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import ast ++import importlib.util ++import logging ++import sys ++import unittest ++from pathlib import Path ++from types import ModuleType, SimpleNamespace ++from unittest import mock ++ ++import numpy as np ++import torch ++ ++ROOT = Path(__file__).resolve().parents[4] ++CUBIT_PATH = ROOT / "vllm/model_executor/layers/quantization/utils/moe_w2_cubit.py" ++DELTA_PATH = ROOT / "vllm/model_executor/layers/quantization/utils/moe_w2_delta.py" ++FORWARD_CONTEXT_PATH = ROOT / "vllm/forward_context.py" ++RUNNER_PATH = ROOT / "vllm/v1/worker/gpu_model_runner.py" ++UBATCH_PATH = ROOT / "vllm/v1/worker/gpu_ubatch_wrapper.py" ++ ++ ++def _load_function(path: Path, name: str): ++ tree = ast.parse(path.read_text()) ++ function = next( ++ node ++ for node in tree.body ++ if isinstance(node, ast.FunctionDef) and node.name == name ++ ) ++ module = ast.fix_missing_locations(ast.Module(body=[function], type_ignores=[])) ++ namespace = {"_BULK_PREFILL_TOKENS": 96, "np": np, "torch": torch} ++ exec(compile(module, str(path), "exec"), namespace) ++ return namespace[name] ++ ++ ++def _load_delta_module(): ++ vllm_module = ModuleType("vllm") ++ vllm_module.__path__ = [str(ROOT / "vllm")] ++ logger_module = ModuleType("vllm.logger") ++ logger_module.init_logger = logging.getLogger ++ spec = importlib.util.spec_from_file_location("_test_moe_w2_delta_pad", DELTA_PATH) ++ assert spec is not None and spec.loader is not None ++ module = importlib.util.module_from_spec(spec) ++ with ( ++ mock.patch.dict( ++ sys.modules, ++ {"vllm": vllm_module, "vllm.logger": logger_module}, ++ ), ++ mock.patch.dict("os.environ", {"VLLM_MOE_W2_DELTA_TRACE": "0"}, clear=False), ++ ): ++ spec.loader.exec_module(module) ++ return module ++ ++ ++class TestMoeW2PaddedRoutes(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ cls.route_metadata = staticmethod( ++ _load_function(CUBIT_PATH, "_masked_route_metadata") ++ ) ++ cls.get_slot_mapping = staticmethod( ++ _load_function(CUBIT_PATH, "_get_token_slot_mapping") ++ ) ++ cls.get_has_prefill = staticmethod( ++ _load_function(CUBIT_PATH, "_get_has_prefill") ++ ) ++ cls.batch_has_prefill = staticmethod( ++ _load_function(RUNNER_PATH, "_batch_has_prefill") ++ ) ++ cls.delta = _load_delta_module() ++ ++ def test_runner_prefill_classification_uses_prompt_progress(self): ++ self.assertTrue( ++ self.batch_has_prefill( ++ np.array([0], dtype=np.int32), np.array([14], dtype=np.int32) ++ ) ++ ) ++ self.assertTrue( ++ self.batch_has_prefill( ++ np.array([20, 7], dtype=np.int32), ++ np.array([20, 8], dtype=np.int32), ++ ) ++ ) ++ self.assertFalse( ++ self.batch_has_prefill( ++ np.array([14, 20], dtype=np.int32), ++ np.array([14, 8], dtype=np.int32), ++ ) ++ ) ++ ++ def test_forward_context_prefill_overrides_bulk_fallback(self): ++ vllm_module = ModuleType("vllm") ++ vllm_module.__path__ = [str(ROOT / "vllm")] ++ context_module = ModuleType("vllm.forward_context") ++ context = SimpleNamespace(has_prefill=True) ++ context_module.get_forward_context = lambda: context ++ with mock.patch.dict( ++ sys.modules, ++ {"vllm": vllm_module, "vllm.forward_context": context_module}, ++ ): ++ self.assertTrue(self.get_has_prefill(14)) ++ context.has_prefill = False ++ self.assertFalse(self.get_has_prefill(128)) ++ context.has_prefill = None ++ self.assertFalse(self.get_has_prefill(14)) ++ self.assertTrue(self.get_has_prefill(128)) ++ ++ def test_fixed_shape_route_metadata_masks_padding(self): ++ sorted_ids = torch.arange(8) ++ slot_mapping = torch.tensor([20, 21, 22, -1]) ++ ++ token_valid, route_valid, pair_live, rows = self.route_metadata( ++ sorted_ids, slot_mapping, top_k=2, mblock=1, pad_row=4 ++ ) ++ ++ torch.testing.assert_close(token_valid, torch.tensor([True, True, True, False])) ++ torch.testing.assert_close( ++ route_valid, ++ torch.tensor([True, True, True, True, True, True, False, False]), ++ ) ++ torch.testing.assert_close(pair_live, route_valid) ++ torch.testing.assert_close(rows, torch.tensor([0, 0, 1, 1, 2, 2, 4, 4])) ++ ++ slot_mapping.copy_(torch.tensor([20, -1, -1, -1])) ++ _, route_valid, pair_live, rows = self.route_metadata( ++ sorted_ids, slot_mapping, top_k=2, mblock=1, pad_row=4 ++ ) ++ torch.testing.assert_close( ++ route_valid, ++ torch.tensor([True, True, False, False, False, False, False, False]), ++ ) ++ torch.testing.assert_close(pair_live, route_valid) ++ torch.testing.assert_close(rows, torch.tensor([0, 0, 4, 4, 4, 4, 4, 4])) ++ ++ _, route_valid, pair_live, _ = self.route_metadata( ++ sorted_ids, slot_mapping, top_k=2, mblock=4, pad_row=4 ++ ) ++ torch.testing.assert_close( ++ route_valid, ++ torch.tensor([True, True, False, False, False, False, False, False]), ++ ) ++ torch.testing.assert_close(pair_live, torch.tensor([True, False])) ++ ++ def test_seen_mask_excludes_disjoint_padding_experts_in_both_tiers(self): ++ ids = torch.tensor([[0, 1], [2, 3], [6, 7]]) ++ token_valid = torch.tensor([True, True, False]) ++ base_seen = torch.zeros(8, dtype=torch.uint8) ++ fp4_seen = torch.zeros_like(base_seen) ++ ++ self.delta.mark_seen(base_seen, ids, token_valid) ++ self.delta.mark_seen(fp4_seen, ids, token_valid) ++ ++ expected = torch.tensor([1, 1, 1, 1, 0, 0, 0, 0], dtype=torch.uint8) ++ torch.testing.assert_close(base_seen, expected) ++ torch.testing.assert_close(fp4_seen, expected) ++ ++ def test_slot_mapping_shape_mismatch_fails_loud(self): ++ vllm_module = ModuleType("vllm") ++ vllm_module.__path__ = [str(ROOT / "vllm")] ++ context_module = ModuleType("vllm.forward_context") ++ context_module.get_forward_context = lambda: SimpleNamespace( ++ token_slot_mapping=torch.zeros(3, dtype=torch.int64) ++ ) ++ with ( ++ mock.patch.dict( ++ sys.modules, ++ {"vllm": vllm_module, "vllm.forward_context": context_module}, ++ ), ++ self.assertRaisesRegex(RuntimeError, "match padded T"), ++ ): ++ self.get_slot_mapping(4) ++ ++ def test_static_capture_and_topology_contracts(self): ++ cubit_tree = ast.parse(CUBIT_PATH.read_text()) ++ functions = { ++ node.name: node ++ for node in cubit_tree.body ++ if isinstance(node, ast.FunctionDef) ++ } ++ for name in ( ++ "_desc_build_kernel", ++ "_desc_build_kernel_w4s", ++ "_desc_build_kernel_basecache", ++ "_desc_build_kernel_base_delta", ++ "_desc_build_kernel_base_delta_split", ++ ): ++ function = functions[name] ++ self.assertIn("pair_live_ptr", [arg.arg for arg in function.args.args]) ++ self.assertTrue( ++ any( ++ isinstance(node, ast.Name) and node.id == "pair_live_ptr" ++ for node in ast.walk(function) ++ ) ++ ) ++ ++ real_args = [arg.arg for arg in functions["_moe_w2_forward"].args.args] ++ fake_args = [arg.arg for arg in functions["_moe_w2_forward_fake"].args.args] ++ self.assertEqual(real_args, fake_args) ++ self.assertEqual(real_args, ["x", "topk_weights", "topk_ids", "layer_key"]) ++ ++ runner_source = RUNNER_PATH.read_text() ++ self.assertIn("force_eager=force_w2_prefill_eager", runner_source) ++ self.assertIn("and not has_prefill", runner_source) ++ self.assertIn("self._get_attention_kv_cache_gid()", runner_source) ++ self.assertIn("decode_context_parallel_size != 1", runner_source) ++ self.assertIn("token_slot_mapping[ubatch.token_slice]", runner_source) ++ self.assertIn("_w2_profile_token_slot_mapping", runner_source) ++ self.assertIn( ++ "profile_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1)", ++ runner_source, ++ ) ++ runner_tree = ast.parse(runner_source) ++ context_calls = [ ++ node ++ for node in ast.walk(runner_tree) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Name) ++ and node.func.id == "set_forward_context" ++ and any(keyword.arg == "slot_mapping" for keyword in node.keywords) ++ ] ++ self.assertGreaterEqual(len(context_calls), 5) ++ for call in context_calls: ++ self.assertIn( ++ "token_slot_mapping", {keyword.arg for keyword in call.keywords} ++ ) ++ self.assertIn("has_prefill", {keyword.arg for keyword in call.keywords}) ++ ++ context_source = FORWARD_CONTEXT_PATH.read_text() ++ self.assertIn("token_slot_mapping=token_slot_mapping", context_source) ++ self.assertIn("has_prefill=has_prefill", context_source) ++ ubatch_source = UBATCH_PATH.read_text() ++ self.assertIn("token_slot_mapping[i]", ubatch_source) ++ self.assertIn("has_prefill=has_prefill", ubatch_source) ++ ++ timed = functions["_moe_w2_forward_timed"] ++ align_call = next( ++ node ++ for node in ast.walk(timed) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Name) ++ and node.func.id == "moe_align_block_size" ++ ) ++ pad_kw = next( ++ keyword ++ for keyword in align_call.keywords ++ if keyword.arg == "pad_sorted_ids" ++ ) ++ self.assertIsInstance(pad_kw.value, ast.Constant) ++ self.assertIs(pad_kw.value.value, True) ++ alignment_guard = next( ++ node ++ for node in ast.walk(timed) ++ if isinstance(node, ast.If) ++ and isinstance(node.test, ast.BinOp) ++ and isinstance(node.test.op, ast.Mod) ++ ) ++ self.assertTrue( ++ any(isinstance(node, ast.Raise) for node in ast.walk(alignment_guard)) ++ ) ++ ++ ++@unittest.skipUnless( ++ torch.cuda.is_available() and importlib.util.find_spec("triton") is not None, ++ "CUDA and Triton are required", ++) ++class TestMoeW2PaddedRoutesCUDA(unittest.TestCase): ++ def test_graph_replay_masks_seen_misses_and_output(self): ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_cubit, ++ moe_w2_delta, ++ ) ++ ++ device = torch.device("cuda") ++ T, top_k, mblock, n_experts = 4, 2, 1, 8 ++ sorted_ids = torch.arange(T * top_k, device=device, dtype=torch.int64) ++ expert_blocks = torch.arange(n_experts, device=device, dtype=torch.int32) ++ topk_ids = torch.arange(n_experts, device=device).view(T, top_k) ++ slot_mapping = torch.tensor([10, 11, 12, -1], device=device) ++ slot_row = torch.full((n_experts,), -1, dtype=torch.int32, device=device) ++ num_post = torch.tensor([T * top_k], dtype=torch.int32, device=device) ++ seen = torch.zeros(n_experts, dtype=torch.int32, device=device) ++ miss = torch.zeros(1, dtype=torch.int32, device=device) ++ desc = torch.zeros((2, n_experts, 6), dtype=torch.int64, device=device) ++ scratch = torch.zeros(1, dtype=torch.uint8, device=device) ++ routed_output = torch.zeros(T + 1, dtype=torch.float32, device=device) ++ route_values = torch.arange(1, n_experts + 1, device=device).float() ++ ++ def run_masked_routes(): ++ seen.zero_() ++ miss.zero_() ++ routed_output.zero_() ++ token_valid, route_valid, pair_live, rows = ( ++ moe_w2_cubit._masked_route_metadata( ++ sorted_ids, slot_mapping, top_k, mblock, T ++ ) ++ ) ++ moe_w2_delta.mark_seen(seen, topk_ids, token_valid) ++ moe_w2_cubit._desc_build_kernel_basecache[(1,)]( ++ expert_blocks, ++ num_post, ++ pair_live, ++ slot_row, ++ miss, ++ desc, ++ scratch.data_ptr(), ++ scratch.data_ptr(), ++ scratch.data_ptr(), ++ scratch.data_ptr(), ++ scratch.data_ptr(), ++ scratch.data_ptr(), ++ scratch.data_ptr(), ++ 1, ++ 0, ++ 0, ++ 0, ++ 1, ++ 1, ++ 1, ++ 1, ++ 1, ++ 1, ++ n_experts, ++ n_experts, ++ n_experts * 6, ++ mblock, ++ BLOCK=256, ++ ) ++ routed_output.index_add_( ++ 0, rows, torch.where(route_valid, route_values, 0.0) ++ ) ++ ++ side = torch.cuda.Stream() ++ side.wait_stream(torch.cuda.current_stream()) ++ with torch.cuda.stream(side): ++ run_masked_routes() ++ torch.cuda.current_stream().wait_stream(side) ++ ++ graph = torch.cuda.CUDAGraph() ++ with torch.cuda.graph(graph): ++ run_masked_routes() ++ ++ cases = ( ++ ([10, 11, 12, -1], 6, [3.0, 7.0, 11.0, 0.0]), ++ ([10, -1, -1, -1], 2, [3.0, 0.0, 0.0, 0.0]), ++ ) ++ for mapping, expected_routes, expected_output in cases: ++ slot_mapping.copy_(torch.tensor(mapping, device=device)) ++ graph.replay() ++ torch.cuda.synchronize() ++ ++ self.assertEqual(miss.item(), expected_routes) ++ self.assertEqual(seen.count_nonzero().item(), expected_routes) ++ promotion_candidates = ((seen > 0) & (slot_row < 0)).count_nonzero() ++ self.assertEqual(promotion_candidates.item(), expected_routes) ++ torch.testing.assert_close( ++ routed_output[:T], torch.tensor(expected_output, device=device) ++ ) ++ self.assertEqual(routed_output[T].item(), 0.0) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/tests/model_executor/layers/quantization/test_moe_w2_step_pins.py b/tests/model_executor/layers/quantization/test_moe_w2_step_pins.py +new file mode 100644 +index 0000000..711fb77 +--- /dev/null ++++ b/tests/model_executor/layers/quantization/test_moe_w2_step_pins.py +@@ -0,0 +1,464 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import ast ++import contextlib ++import importlib.util ++import logging ++import sys ++import threading ++import unittest ++from pathlib import Path ++from types import ModuleType ++from unittest import mock ++ ++import torch ++ ++ROOT = Path(__file__).resolve().parents[4] ++DELTA_PATH = ROOT / "vllm/model_executor/layers/quantization/utils/moe_w2_delta.py" ++GATE_PATH = ROOT / "vllm/model_executor/layers/quantization/utils/moe_w2_gate.py" ++RUNNER_PATH = ROOT / "vllm/v1/worker/gpu_model_runner.py" ++WORKER_PATH = ROOT / "vllm/v1/worker/gpu_worker.py" ++ ++ ++def _load_delta_module(): ++ """Load the tier in isolation; this regression needs only CPU torch.""" ++ vllm_module = ModuleType("vllm") ++ vllm_module.__path__ = [str(ROOT / "vllm")] ++ logger_module = ModuleType("vllm.logger") ++ logger_module.init_logger = logging.getLogger ++ spec = importlib.util.spec_from_file_location("_test_moe_w2_delta", DELTA_PATH) ++ assert spec is not None and spec.loader is not None ++ module = importlib.util.module_from_spec(spec) ++ with ( ++ mock.patch.dict( ++ sys.modules, ++ {"vllm": vllm_module, "vllm.logger": logger_module}, ++ ), ++ mock.patch.dict("os.environ", {"VLLM_MOE_W2_DELTA_TRACE": "0"}, clear=False), ++ ): ++ spec.loader.exec_module(module) ++ return module ++ ++ ++class TestMoeW2StepPins(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ cls.delta = _load_delta_module() ++ ++ def _saturated_lru_tier(self): ++ tier = object.__new__(self.delta.DeltaTier) ++ tier._lock = threading.Lock() ++ tier.n_slots = 3 ++ tier._free = [] ++ tier._alloc_owner(tier.n_slots) ++ tier._owner_li[:] = torch.tensor([0, 0, 0]) ++ tier._owner_ei[:] = torch.tensor([0, 1, 2]) ++ tier._owner_tick[:] = torch.tensor([0, 1, 2]) ++ tier._policy = "lru" ++ tier._tick = 10 ++ tier._step_pins = {0, 1, 2} ++ tier._coupled_fp4 = None ++ tier._seen_host = torch.zeros((1, 32), dtype=torch.uint8) ++ tier._freq = torch.zeros((1, 32), dtype=torch.float32) ++ tier._need = torch.zeros((1, 32), dtype=torch.float32) ++ tier.slot_table = torch.full((1, 32), -1, dtype=torch.int32) ++ tier._mirror = torch.full((1, 32), -1, dtype=torch.int32) ++ for slot, (_, expert, _) in enumerate(tier._owner): ++ tier.slot_table[0, expert] = slot ++ tier._mirror[0, expert] = slot ++ tier._n_evicted = 0 ++ tier._win_evicted = 0 ++ return tier ++ ++ def test_step_scope_reset_keeps_saturated_lru_evicting(self): ++ tier = self._saturated_lru_tier() ++ ++ # This is the production failure mode before the runner reset: once ++ # every saturated slot has accumulated in _step_pins, even emergency ++ # promotion cannot find a victim. ++ self.assertEqual(tier._take_slots_batch(1, emergency=True), []) ++ ++ for expert in range(3, 19): ++ tier._tick += 3 ++ tier.step_begin() ++ expected = min(range(tier.n_slots), key=lambda slot: tier._owner[slot][2]) ++ self.assertEqual(tier._take_slots_batch(1, emergency=True), [expected]) ++ ++ # Mirror force_promote's ownership handoff and current-step pin. ++ tier._own(expected, 0, expert) ++ tier.slot_table[0, expert] = expected ++ tier._mirror[0, expert] = expected ++ tier._step_pins.add(expected) ++ self.assertEqual(tier._free, []) ++ ++ self.assertEqual(tier._n_evicted, 16) ++ ++ def test_explicit_seen_snapshot_wins_over_shared_snapshot_overwrite(self): ++ tier = self._saturated_lru_tier() ++ tier._step_pins.clear() ++ # Simulate another caller overwriting the shared host buffer after this ++ # caller captured expert 0. The immutable set must still protect slot 0. ++ tier._seen_host.zero_() ++ tier._seen_host[0, 2] = 1 ++ self.assertEqual(tier._take_slot({(0, 0)}), 1) ++ ++ def test_target_boundary_clears_seen_but_preserves_pins(self): ++ tier = object.__new__(self.delta.DeltaTier) ++ tier.dev = torch.device("cpu") ++ tier._snap_lock = threading.Lock() ++ tier._stream = mock.Mock() ++ tier.seen = torch.ones((2, 4), dtype=torch.uint8) ++ tier._step_pins = {1, 2} ++ main = object() ++ ++ with ( ++ mock.patch.object( ++ self.delta.torch.cuda, "current_stream", return_value=main ++ ), ++ mock.patch.object( ++ self.delta.torch.cuda, "stream", return_value=contextlib.nullcontext() ++ ), ++ ): ++ tier.routing_step_begin() ++ ++ self.assertEqual(tier.seen.count_nonzero().item(), 0) ++ self.assertEqual(tier._step_pins, {1, 2}) ++ tier._stream.wait_stream.assert_called_once_with(main) ++ ++ def test_module_boundaries_reset_both_tiers(self): ++ base = mock.Mock() ++ fp4 = mock.Mock() ++ with ( ++ mock.patch.object(self.delta, "_BASE_TIER", base), ++ mock.patch.object(self.delta, "_TIER", fp4), ++ ): ++ self.delta.begin_target_step() ++ self.delta.begin_replay_step() ++ self.delta.finish_forward_step() ++ ++ base.pause_for_forward.assert_called_once_with() ++ fp4.pause_for_forward.assert_called_once_with() ++ base.routing_step_begin.assert_called_once_with() ++ fp4.routing_step_begin.assert_called_once_with() ++ base.step_begin.assert_called_once_with() ++ fp4.step_begin.assert_called_once_with() ++ base.wake.assert_called_once_with() ++ fp4.wake.assert_called_once_with() ++ ++ def test_manager_pass_cannot_overlap_forward_window(self): ++ tier = object.__new__(self.delta.DeltaTier) ++ tier._forward_lock = threading.Lock() ++ tier._forward_paused = False ++ tier._wake = mock.Mock() ++ tier._wake_driven = False ++ ++ tier.pause_for_forward() ++ entered = threading.Event() ++ ++ def manager_pass(): ++ with tier._forward_lock: ++ entered.set() ++ ++ thread = threading.Thread(target=manager_pass) ++ thread.start() ++ self.assertFalse(entered.wait(0.05)) ++ tier.wake() ++ self.assertTrue(entered.wait(1.0)) ++ thread.join(timeout=1.0) ++ self.assertFalse(thread.is_alive()) ++ self.assertFalse(tier._forward_paused) ++ self.assertTrue(tier._wake_driven) ++ tier._wake.set.assert_called_once_with() ++ ++ def test_ensure_resident_drains_then_uses_exact_layer_snapshot(self): ++ tree = ast.parse(DELTA_PATH.read_text()) ++ ensure = next( ++ node ++ for node in ast.walk(tree) ++ if isinstance(node, ast.FunctionDef) and node.name == "ensure_resident" ++ ) ++ snap_scope = next( ++ node ++ for node in ast.walk(ensure) ++ if isinstance(node, ast.With) ++ and any( ++ isinstance(item.context_expr, ast.Attribute) ++ and item.context_expr.attr == "_snap_lock" ++ for item in node.items ++ ) ++ ) ++ lock_scope = next( ++ node ++ for node in ast.walk(snap_scope) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "synchronize" ++ ) ++ pool_lock = next( ++ node ++ for node in ast.walk(ensure) ++ if isinstance(node, ast.With) ++ and any( ++ isinstance(item.context_expr, ast.Attribute) ++ and item.context_expr.attr == "_lock" ++ for item in node.items ++ ) ++ ) ++ self.assertLess(lock_scope.lineno, pool_lock.lineno) ++ layer_snapshot = next( ++ node ++ for node in ast.walk(ensure) ++ if isinstance(node, ast.Assign) ++ and any( ++ isinstance(target, ast.Name) and target.id == "layer_seen_set" ++ for target in node.targets ++ ) ++ ) ++ self.assertIsInstance(layer_snapshot.value, ast.SetComp) ++ take = next( ++ node ++ for node in ast.walk(pool_lock) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_take_slots_batch" ++ ) ++ seen_kw = next( ++ keyword for keyword in take.keywords if keyword.arg == "seen_set" ++ ) ++ self.assertIsInstance(seen_kw.value, ast.Name) ++ self.assertEqual(seen_kw.value.id, "layer_seen_set") ++ pin_clears = [ ++ node ++ for node in ast.walk(pool_lock) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "clear" ++ and isinstance(node.func.value, ast.Attribute) ++ and node.func.value.attr == "_step_pins" ++ ] ++ pin_mutations = [ ++ node ++ for node in ast.walk(pool_lock) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr in {"add", "update"} ++ and isinstance(node.func.value, ast.Attribute) ++ and node.func.value.attr == "_step_pins" ++ ] ++ self.assertEqual(len(pin_clears), 1) ++ self.assertGreaterEqual(len(pin_mutations), 2) ++ self.assertLess(pin_clears[0].lineno, take.lineno) ++ ++ def test_manager_and_force_promote_pass_immutable_seen_sets(self): ++ tree = ast.parse(DELTA_PATH.read_text()) ++ functions = { ++ node.name: node ++ for node in ast.walk(tree) ++ if isinstance(node, ast.FunctionDef) ++ } ++ tick_take = next( ++ node ++ for node in ast.walk(functions["_tick_once"]) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_take_slot" ++ ) ++ self.assertEqual(len(tick_take.args), 1) ++ self.assertIsInstance(tick_take.args[0], ast.Name) ++ self.assertEqual(tick_take.args[0].id, "seen_set") ++ ++ force_take = next( ++ node ++ for node in ast.walk(functions["force_promote"]) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_take_slots_batch" ++ ) ++ seen_kw = next( ++ keyword for keyword in force_take.keywords if keyword.arg == "seen_set" ++ ) ++ self.assertIsInstance(seen_kw.value, ast.Name) ++ self.assertEqual(seen_kw.value.id, "seen_set") ++ ++ wrapper_take = next( ++ node ++ for node in ast.walk(functions["_take_slot"]) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_take_slots_batch" ++ ) ++ wrapper_kw = next( ++ keyword for keyword in wrapper_take.keywords if keyword.arg == "seen_set" ++ ) ++ self.assertIsInstance(wrapper_kw.value, ast.Name) ++ self.assertEqual(wrapper_kw.value.id, "seen_set") ++ ++ for name in ("_tick_once", "force_promote"): ++ assignments = [ ++ node ++ for node in ast.walk(functions[name]) ++ if isinstance(node, ast.Assign) ++ and any( ++ isinstance(target, ast.Name) and target.id == "seen_set" ++ for target in node.targets ++ ) ++ ] ++ self.assertEqual(len(assignments), 1, name) ++ ++ def test_current_prefetch_hit_is_repinned_for_replay(self): ++ tier = object.__new__(self.delta.DeltaTier) ++ tier.dev = torch.device("cpu") ++ tier._store = {0: object()} ++ tier.n_layers = 1 ++ tier._store_mask_cache = None ++ tier._store_mask_n = -1 ++ tier._snap_lock = threading.Lock() ++ tier._lock = threading.Lock() ++ tier._stream = mock.Mock() ++ tier.seen = torch.zeros((1, 4), dtype=torch.uint8) ++ tier.seen[0, 2] = 1 ++ tier._seen_host = torch.zeros_like(tier.seen) ++ tier._mirror = torch.full((1, 4), -1, dtype=torch.int32) ++ tier._mirror[0, 2] = 0 ++ tier._alloc_owner(1) ++ tier._owner_li[0] = 0 ++ tier._owner_ei[0] = 2 ++ tier._owner_tick[0] = 3 ++ tier._tick = 10 ++ tier._step_pins = {0} # draft_prefetch pinned this current-step hit ++ tier._need = torch.zeros((1, 4), dtype=torch.float32) ++ main = object() ++ event = mock.Mock() ++ ++ tier.step_begin() ++ self.assertEqual(tier._step_pins, set()) ++ with ( ++ mock.patch.object( ++ self.delta.torch.cuda, "current_stream", return_value=main ++ ), ++ mock.patch.object( ++ self.delta.torch.cuda, "stream", return_value=contextlib.nullcontext() ++ ), ++ mock.patch.object(self.delta.torch.cuda, "Event", return_value=event), ++ ): ++ self.assertEqual(tier.force_promote(), 0) ++ ++ self.assertEqual(tier._step_pins, {0}) ++ self.assertEqual(tier._owner[0], (0, 2, tier._tick)) ++ tier._stream.wait_stream.assert_called_once_with(main) ++ event.synchronize.assert_called_once_with() ++ ++ def test_runner_uses_routing_then_post_forward_pin_boundaries(self): ++ tree = ast.parse(RUNNER_PATH.read_text()) ++ execute_model = next( ++ node ++ for node in ast.walk(tree) ++ if isinstance(node, ast.FunctionDef) and node.name == "execute_model" ++ ) ++ target_begins = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "begin_target_step" ++ and isinstance(node.func.value, ast.Name) ++ and node.func.value.id == "_w2d" ++ ] ++ replay_begins = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "begin_replay_step" ++ and isinstance(node.func.value, ast.Name) ++ and node.func.value.id == "_w2d" ++ ] ++ self.assertEqual(len(target_begins), 1) ++ self.assertEqual(len(replay_begins), 1) ++ ++ swallowing_handlers = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Try) ++ and ( ++ target_begins[0] in ast.walk(node) or replay_begins[0] in ast.walk(node) ++ ) ++ ] ++ self.assertEqual(swallowing_handlers, []) ++ ++ target_forwards = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_model_forward" ++ ] ++ first_forward = min(node.lineno for node in target_forwards) ++ post_forward_reset = replay_begins[0] ++ self.assertLess(target_begins[0].lineno, first_forward) ++ self.assertLess(first_forward, post_forward_reset.lineno) ++ ++ direct_tier_resets = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "step_begin" ++ ] ++ self.assertEqual(direct_tier_resets, []) ++ ++ def test_worker_releases_manager_only_after_gate_barrier(self): ++ tree = ast.parse(WORKER_PATH.read_text()) ++ execute_model = next( ++ node ++ for node in ast.walk(tree) ++ if isinstance(node, ast.FunctionDef) and node.name == "execute_model" ++ ) ++ finishes = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_finish_w2_manager_step" ++ ] ++ barriers = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "_gate_pp_barrier" ++ ] ++ self.assertEqual(len(finishes), 2) ++ self.assertEqual(len(barriers), 2) ++ finally_finishes = [ ++ node ++ for node in ast.walk(execute_model) ++ if isinstance(node, ast.Try) ++ and any(finish in ast.walk(node) for finish in finishes) ++ and any( ++ finish in ast.walk(final) ++ for final in node.finalbody ++ for finish in finishes ++ ) ++ ] ++ self.assertEqual(len(finally_finishes), 2) ++ ++ gate_tree = ast.parse(GATE_PATH.read_text()) ++ should_reforward = next( ++ node ++ for node in ast.walk(gate_tree) ++ if isinstance(node, ast.FunctionDef) and node.name == "should_reforward" ++ ) ++ unsafe_wakes = [ ++ node ++ for node in ast.walk(should_reforward) ++ if isinstance(node, ast.Call) ++ and isinstance(node.func, ast.Attribute) ++ and node.func.attr == "wake_all" ++ ] ++ self.assertEqual(unsafe_wakes, []) ++ ++ ++if __name__ == "__main__": ++ unittest.main() diff --git a/tools/nvfp4_flashinfer_sm120/README.md b/tools/nvfp4_flashinfer_sm120/README.md new file mode 100644 -index 000000000..a147137db +index 0000000..a147137 --- /dev/null +++ b/tools/nvfp4_flashinfer_sm120/README.md @@ -0,0 +1,49 @@ @@ -203,7 +1172,7 @@ index 000000000..a147137db +- Microbench (isolated KV gather, RTX 5090): 1.86× tokens/s vs 656 B. diff --git a/tools/nvfp4_flashinfer_sm120/nvfp4_expand.cuh b/tools/nvfp4_flashinfer_sm120/nvfp4_expand.cuh new file mode 100644 -index 000000000..e27203c56 +index 0000000..e27203c --- /dev/null +++ b/tools/nvfp4_flashinfer_sm120/nvfp4_expand.cuh @@ -0,0 +1,211 @@ @@ -420,7 +1389,7 @@ index 000000000..e27203c56 +} diff --git a/tools/nvfp4_flashinfer_sm120/patch_flashinfer.py b/tools/nvfp4_flashinfer_sm120/patch_flashinfer.py new file mode 100644 -index 000000000..cfa92d323 +index 0000000..cfa92d3 --- /dev/null +++ b/tools/nvfp4_flashinfer_sm120/patch_flashinfer.py @@ -0,0 +1,412 @@ @@ -837,7 +1806,7 @@ index 000000000..cfa92d323 +print("SYNTAX-OK python") +print("PATCH-FLASHINFER-DONE") diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py -index 6da3ec717..c84958ac0 100644 +index 6da3ec7..c84958a 100644 --- a/vllm/compilation/breakable_cudagraph.py +++ b/vllm/compilation/breakable_cudagraph.py @@ -175,10 +175,13 @@ class BreakableCUDAGraphCapture: @@ -857,7 +1826,7 @@ index 6da3ec717..c84958ac0 100644 self._capturing = True diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py -index b63d86199..3cad1f058 100644 +index b63d861..3cad1f0 100644 --- a/vllm/compilation/cuda_graph.py +++ b/vllm/compilation/cuda_graph.py @@ -314,6 +314,10 @@ class CUDAGraphWrapper: @@ -872,7 +1841,7 @@ index b63d86199..3cad1f058 100644 # `output` is managed by pytorch's cudagraph pool output = self.runnable(*args, **kwargs) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py -index de505e122..f0b49e294 100644 +index de505e1..f0b49e2 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -54,6 +54,7 @@ MTPModelTypes = Literal[ @@ -1074,7 +2043,7 @@ index de505e122..f0b49e294 100644 return self.num_speculative_tokens_per_batch_size is not None diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py -index ba7d26c93..67d3b5ea2 100644 +index ba7d26c..67d3b5e 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -524,6 +524,16 @@ class VllmConfig: @@ -1184,23 +2153,44 @@ index ba7d26c93..67d3b5ea2 100644 "nvfp4 KV cache is not supported with MLA (Multi-head Latent " "Attention) backends. Please use a different --kv-cache-dtype " diff --git a/vllm/envs.py b/vllm/envs.py -index 27a85bb3d..88f310c4f 100755 +index 27a85bb..4a657ca 100755 --- a/vllm/envs.py +++ b/vllm/envs.py -@@ -179,6 +179,9 @@ if TYPE_CHECKING: +@@ -179,6 +179,17 @@ if TYPE_CHECKING: VLLM_MOE_USE_DEEP_GEMM: bool = True VLLM_USE_DEEP_GEMM_E8M0: bool = True VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES: bool = True ++ VLLM_MOE_W2: bool = False ++ VLLM_MOE_W2_STORE_DIR: str = "" ++ VLLM_MOE_W2_PACK_ID: str = "" ++ VLLM_MOE_W2_CACHE_CONTROL: Literal[ ++ "required", "best-effort", "off" ++ ] = "required" ++ VLLM_MOE_W2_MIN_MEM_AVAILABLE_GB: float = 16.0 ++ VLLM_MOE_W2_MIN_CGROUP_HEADROOM_GB: float = 4.0 + # Opt-in: fused hand-written SASS (cubit) sparse-MLA on SM120 (experimental, + # eager only). See vllm/v1/attention/backends/mla/cubit_sparse_mla.py. + VLLM_SPARSE_MLA_CUBIT: bool | None = None VLLM_DEEP_GEMM_WARMUP: Literal[ "skip", "full", -@@ -1444,6 +1447,16 @@ environment_variables: dict[str, Callable[[], Any]] = { +@@ -1444,6 +1455,29 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES": lambda: bool( int(os.getenv("VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "1")) ), ++ # W2 pack-store cold-build safety. The implementation reads these at the ++ # point of use; registration here prevents the generic unknown-VLLM-env ++ # guard from misclassifying the supported controls. ++ "VLLM_MOE_W2": lambda: bool(int(os.getenv("VLLM_MOE_W2", "0"))), ++ "VLLM_MOE_W2_STORE_DIR": lambda: os.getenv("VLLM_MOE_W2_STORE_DIR", ""), ++ "VLLM_MOE_W2_PACK_ID": lambda: os.getenv("VLLM_MOE_W2_PACK_ID", ""), ++ "VLLM_MOE_W2_CACHE_CONTROL": env_with_choices( ++ "VLLM_MOE_W2_CACHE_CONTROL", "required", ++ ["required", "best-effort", "off"]), ++ "VLLM_MOE_W2_MIN_MEM_AVAILABLE_GB": lambda: float( ++ os.getenv("VLLM_MOE_W2_MIN_MEM_AVAILABLE_GB", "16")), ++ "VLLM_MOE_W2_MIN_CGROUP_HEADROOM_GB": lambda: float( ++ os.getenv("VLLM_MOE_W2_MIN_CGROUP_HEADROOM_GB", "4")), + # Opt-in: fused hand-written SASS (cubit) sparse-MLA decode on SM120, + # replacing the Triton accumulate+finish pair for supported decode shapes. + # Experimental; requires eager mode. See @@ -1214,8 +2204,65 @@ index 27a85bb3d..88f310c4f 100755 # DeepGemm JITs the kernels on-demand. The warmup attempts to make DeepGemm # JIT all the required kernels before model execution so there is no # JIT'ing in the hot-path. However, this warmup increases the engine +diff --git a/vllm/forward_context.py b/vllm/forward_context.py +index 5527ec1..a7cb4f4 100644 +--- a/vllm/forward_context.py ++++ b/vllm/forward_context.py +@@ -131,6 +131,10 @@ class ForwardContext: + no_compile_layers: dict[str, Any] + attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] + slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] ++ token_slot_mapping: torch.Tensor | list[torch.Tensor] | None = None ++ # True when any real request is still consuming prompt tokens. None is ++ # reserved for dummy/profile forwards that lack request lifecycle state. ++ has_prefill: bool | None = None + """ + Type Dict[str, AttentionMetadata] for v1, map from layer_name of each + attention layer to its attention metadata +@@ -209,6 +213,8 @@ def create_forward_context( + batch_descriptor: BatchDescriptor | None = None, + ubatch_slices: UBatchSlices | None = None, + slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, ++ token_slot_mapping: torch.Tensor | list[torch.Tensor] | None = None, ++ has_prefill: bool | None = None, + additional_kwargs: dict[str, Any] | None = None, + skip_compiled: bool = False, + ): +@@ -222,6 +228,8 @@ def create_forward_context( + all_moe_layers=all_moe_layers, + attn_metadata=attn_metadata, + slot_mapping=slot_mapping or {}, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, + dp_metadata=dp_metadata, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_descriptor, +@@ -256,6 +264,8 @@ def set_forward_context( + batch_descriptor: BatchDescriptor | None = None, + ubatch_slices: UBatchSlices | None = None, + slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, ++ token_slot_mapping: torch.Tensor | list[torch.Tensor] | None = None, ++ has_prefill: bool | None = None, + skip_compiled: bool = False, + ): + """A context manager that stores the current forward context, +@@ -313,9 +323,11 @@ def set_forward_context( + cudagraph_runtime_mode, + batch_descriptor, + ubatch_slices, +- slot_mapping, +- additional_kwargs, +- skip_compiled, ++ slot_mapping=slot_mapping, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, ++ additional_kwargs=additional_kwargs, ++ skip_compiled=skip_compiled, + ) + + try: diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py -index 051468ed1..33902056c 100644 +index 051468e..3390205 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -188,6 +188,7 @@ return curr_o @ W_O @@ -1272,7 +2319,7 @@ index 051468ed1..33902056c 100644 # Enforce that we enough for at least 1 page per request diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py -index 7cdb04cfb..64813139b 100644 +index 7cdb04c..6481313 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -671,6 +671,21 @@ class Fp8MoEMethod(FusedMoEMethodBase): @@ -1336,7 +2383,7 @@ index 7cdb04cfb..64813139b 100644 assert self.moe_kernel is not None return self.moe_kernel.apply( diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py -index d51a2dd31..2fe74cc55 100644 +index d51a2dd..2fe74cc 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1,6 +1,7 @@ @@ -1551,23 +2598,36 @@ index d51a2dd31..2fe74cc55 100644 return UnquantizedLinearMethod() return None diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py -index 1b2a8a74b..3afdb16fd 100644 +index 1b2a8a7..de83142 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py -@@ -617,6 +617,21 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): +@@ -617,6 +617,34 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) layer.register_parameter("w2_bias", w2_bias) set_weight_attrs(w2_bias, extra_weight_attrs) + # VLLM_MOE_W2: the raw checkpoint experts of all layers do not fit a -+ # single GPU; stage them in host RAM until the 2-bit planes are built -+ # in process_weights_after_loading. ++ # single GPU or a constrained host cgroup. Skip checkpoint staging ++ # entirely on a valid pack/cache hit; otherwise materialize one ++ # layer lazily and build its planes as soon as its last tensor loads. ++ # The all-layers CPU staging path remains only as the explicit ++ # VLLM_MOE_W2_STREAM_BUILD=0 fallback. + from vllm.model_executor.layers.quantization.utils import moe_w2_cubit -+ if moe_w2_cubit.is_w2_layer(getattr(layer, "layer_name", "")): -+ for pname in ("w13_weight", "w13_weight_scale", "w2_weight", -+ "w2_weight_scale"): ++ ++ if ( ++ moe_w2_cubit.is_w2_layer(getattr(layer, "layer_name", "")) ++ and not moe_w2_cubit.plan_pack_skip(layer) ++ and not moe_w2_cubit.arm_stream_build(layer, checkpoint_format="mxfp4") ++ ): ++ for pname in ( ++ "w13_weight", ++ "w13_weight_scale", ++ "w2_weight", ++ "w2_weight_scale", ++ ): + p_ = getattr(layer, pname) -+ attrs = {k: getattr(p_, k) for k in ("weight_loader",) -+ if hasattr(p_, k)} ++ attrs = { ++ k: getattr(p_, k) for k in ("weight_loader",) if hasattr(p_, k) ++ } + newp = torch.nn.Parameter(p_.data.cpu(), requires_grad=False) + layer.register_parameter(pname, newp) + set_weight_attrs(newp, attrs) @@ -1576,22 +2636,24 @@ index 1b2a8a74b..3afdb16fd 100644 def _setup_kernel( self, -@@ -722,6 +737,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): +@@ -722,6 +750,16 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) def process_weights_after_loading(self, layer): + # VLLM_MOE_W2: build 2-bit tensor-sym planes; skip Marlin/other backends. + from vllm.model_executor.layers.quantization.utils import moe_w2_cubit ++ + if moe_w2_cubit.is_w2_layer(getattr(layer, "layer_name", "")): -+ key = len(moe_w2_cubit._LAYERS) -+ moe_w2_cubit.build_layer_planes(layer, key) ++ key = getattr(layer, "_moe_w2_create_key", len(moe_w2_cubit._LAYERS)) ++ if not getattr(layer, "_moe_w2_stream_built", False): ++ moe_w2_cubit.build_layer_planes(layer, key) + layer._moe_w2_key = key + return + w13 = layer.w13_weight w2 = layer.w2_weight w13_scale = layer.w13_weight_scale -@@ -779,6 +802,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): +@@ -779,6 +817,17 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: @@ -1601,18 +2663,16 @@ index 1b2a8a74b..3afdb16fd 100644 + # routed-expert output, matching the other non-modular applies here. + w2_key = getattr(layer, "_moe_w2_key", None) + if w2_key is not None: -+ from vllm.model_executor.layers.quantization.utils import ( -+ moe_w2_cubit) -+ assert layer.expert_map is None and \ -+ not layer.apply_router_weight_on_input -+ return moe_w2_cubit.moe_w2_forward(x, topk_weights, topk_ids, -+ w2_key) ++ from vllm.model_executor.layers.quantization.utils import moe_w2_cubit ++ ++ assert layer.expert_map is None and not layer.apply_router_weight_on_input ++ return moe_w2_cubit.moe_w2_forward(x, topk_weights, topk_ids, w2_key) + assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py -index 32a2d8689..3e83d5a87 100644 +index 32a2d86..3e83d5a 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -1122,6 +1122,11 @@ def deepgemm_post_process_fp8_weight_block( @@ -1629,10 +2689,10 @@ index 32a2d8689..3e83d5a87 100644 mn=r, diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_cubit.py b/vllm/model_executor/layers/quantization/utils/moe_w2_cubit.py new file mode 100644 -index 000000000..43bad76a2 +index 0000000..cf35fde --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_cubit.py -@@ -0,0 +1,1702 @@ +@@ -0,0 +1,2275 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Routed experts on 2-bit tensor-sym planes (cubit moe_w2) for the @@ -1677,8 +2737,9 @@ index 000000000..43bad76a2 + +_KERN = b"moe_w2_mm" +_DIR = os.getenv("VLLM_MOE_W2_CUBIT_DIR", "/cubit-share") -+_BLOCK = 4 # tokens per pair == kernel M limit -+_NTHR = 256 # NWARP=8 (K>=1024) ++_BLOCK = 4 # tokens per pair == kernel M limit ++_NTHR = 256 # NWARP=8 (K>=1024) ++_BULK_PREFILL_TOKENS = 96 + + +def _nwarp_for_k(k: int) -> int: @@ -1693,6 +2754,7 @@ index 000000000..43bad76a2 + return n + return 1 + ++ +_cu = None +_fns: dict = {} +_state = "uninit" @@ -1720,9 +2782,52 @@ index 000000000..43bad76a2 + v = v.permute(0, 3, 2, 5, 4, 1, 6).reshape(pairs * 16, K) + return v.contiguous().view(a.dtype) + ++ ++def _masked_route_metadata( ++ sorted_ids: torch.Tensor, ++ token_slot_mapping: torch.Tensor, ++ top_k: int, ++ mblock: int, ++ pad_row: int, ++) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ++ """Build fixed-shape route masks from the runner's persistent slot map.""" ++ T = token_slot_mapping.shape[0] ++ token_valid = token_slot_mapping >= 0 ++ aligned_valid = sorted_ids < T * top_k ++ token_rows = (sorted_ids // top_k).clamp(max=T - 1) ++ route_valid = aligned_valid & token_valid[token_rows] ++ pair_live = route_valid.view(-1, mblock).any(dim=1) ++ rows = torch.where(route_valid, token_rows, torch.full_like(sorted_ids, pad_row)) ++ return token_valid, route_valid, pair_live, rows ++ ++ ++def _get_token_slot_mapping(T: int) -> torch.Tensor: ++ from vllm.forward_context import get_forward_context ++ ++ token_slot_mapping = get_forward_context().token_slot_mapping ++ if not isinstance(token_slot_mapping, torch.Tensor): ++ raise RuntimeError( ++ "moe_w2 requires a persistent token slot mapping in ForwardContext" ++ ) ++ if token_slot_mapping.ndim != 1 or token_slot_mapping.shape[0] != T: ++ raise RuntimeError( ++ "moe_w2 slot mapping must be one-dimensional and match padded T" ++ ) ++ return token_slot_mapping ++ ++ ++def _get_has_prefill(T: int) -> bool: ++ from vllm.forward_context import get_forward_context ++ ++ has_prefill = get_forward_context().has_prefill ++ if has_prefill is None: ++ return T > _BULK_PREFILL_TOKENS ++ return has_prefill ++ ++ +# layer_key -> dict(planes13, sc13, planes2, sc2, top_k, inter) +_LAYERS: dict[int, dict] = {} -+_WS: dict = {} # shared workspaces, sized lazily ++_WS: dict = {} # shared workspaces, sized lazily + +# ---- adaptive expert top-p (VLLM_MOE_W2_TOPP, colibri's --topp) ---------- +# Keep each token's routed experts only up to cumulative router weight p: @@ -1793,6 +2898,7 @@ index 000000000..43bad76a2 + return int(v) + try: + from vllm.config import get_current_vllm_config ++ + cfg = get_current_vllm_config().model_config.hf_config + cfg = cfg.get_text_config() + n = cfg.num_hidden_layers @@ -1814,6 +2920,7 @@ index 000000000..43bad76a2 + if "mtp" in name: + return False + import re ++ + m = re.search(r"\.layers\.(\d+)\.", name) + if m is None: + return False @@ -1824,13 +2931,22 @@ index 000000000..43bad76a2 + global _cu + if _cu is None: + cu = ctypes.CDLL("libcuda.so.1") -+ cu.cuLaunchKernel.argtypes = [ctypes.c_void_p] + [ctypes.c_uint] * 6 + [ -+ ctypes.c_uint, ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p), -+ ctypes.c_void_p] -+ cu.cuModuleLoad.argtypes = [ctypes.POINTER(ctypes.c_void_p), -+ ctypes.c_char_p] -+ cu.cuModuleGetFunction.argtypes = [ctypes.POINTER(ctypes.c_void_p), -+ ctypes.c_void_p, ctypes.c_char_p] ++ cu.cuLaunchKernel.argtypes = ( ++ [ctypes.c_void_p] ++ + [ctypes.c_uint] * 6 ++ + [ ++ ctypes.c_uint, ++ ctypes.c_void_p, ++ ctypes.POINTER(ctypes.c_void_p), ++ ctypes.c_void_p, ++ ] ++ ) ++ cu.cuModuleLoad.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_char_p] ++ cu.cuModuleGetFunction.argtypes = [ ++ ctypes.POINTER(ctypes.c_void_p), ++ ctypes.c_void_p, ++ ctypes.c_char_p, ++ ] + _cu = cu + return _cu + @@ -1850,9 +2966,13 @@ index 000000000..43bad76a2 + torch.cuda.init() + torch.zeros(1, device="cuda") + cu = _driver() -+ for tier, kern in (("w2", b"moe_w2_mm"), ("w4", b"moe_w4_mm"), -+ ("w4s", b"moe_w4s_mm"), -+ ("w2mc2", b"moe_w2_mm"), ("w2mc4", b"moe_w2_mm")): ++ for tier, kern in ( ++ ("w2", b"moe_w2_mm"), ++ ("w4", b"moe_w4_mm"), ++ ("w4s", b"moe_w4s_mm"), ++ ("w2mc2", b"moe_w2_mm"), ++ ("w2mc4", b"moe_w2_mm"), ++ ): + # GEMM contraction K: gate-up needs K=hidden (4096 DS4-Flash, + # 6144 GLM-5.x, 7168 Kimi-K2.x); down needs K=I/TP (2048 @ TP1, + # 1024 @ TP2, 512 @ TP4). Cubins are loaded opportunistically -- @@ -1867,11 +2987,15 @@ index 000000000..43bad76a2 + if not os.path.exists(path): + continue + mod = ctypes.c_void_p() -+ _ck(cu.cuModuleLoad(ctypes.byref(mod), path.encode()), -+ f"cuModuleLoad {path}") ++ _ck( ++ cu.cuModuleLoad(ctypes.byref(mod), path.encode()), ++ f"cuModuleLoad {path}", ++ ) + fn = ctypes.c_void_p() -+ _ck(cu.cuModuleGetFunction(ctypes.byref(fn), mod, kern), -+ "cuModuleGetFunction") ++ _ck( ++ cu.cuModuleGetFunction(ctypes.byref(fn), mod, kern), ++ "cuModuleGetFunction", ++ ) + _fns[(tier, k)] = fn + global _afrag_ok + if _AFRAG: @@ -1881,11 +3005,15 @@ index 000000000..43bad76a2 + if not os.path.exists(path): + continue + mod = ctypes.c_void_p() -+ _ck(cu.cuModuleLoad(ctypes.byref(mod), path.encode()), -+ f"cuModuleLoad {path}") ++ _ck( ++ cu.cuModuleLoad(ctypes.byref(mod), path.encode()), ++ f"cuModuleLoad {path}", ++ ) + fn = ctypes.c_void_p() -+ _ck(cu.cuModuleGetFunction(ctypes.byref(fn), mod, b"moe_w2_mm"), -+ "cuModuleGetFunction afrag") ++ _ck( ++ cu.cuModuleGetFunction(ctypes.byref(fn), mod, b"moe_w2_mm"), ++ "cuModuleGetFunction afrag", ++ ) + _fns[("w2mc4afrag", k)] = fn + _afrag_ok = True + logger.info("moe_w2_cubit: AFRAG prefill cubins loaded") @@ -1905,6 +3033,7 @@ index 000000000..43bad76a2 +# Load-time plane building +# -------------------------------------------------------------------------- + ++ +def _require_kernels(K13: int, K2: int, need_w4: bool) -> None: + """Fail loudly at weight load when the cubins this model's shapes need are + missing from _DIR (they are loaded opportunistically in _ensure_ready).""" @@ -1916,7 +3045,8 @@ index 000000000..43bad76a2 + missing = [f"{t}_k{k}" for t, k in need if (t, k) not in _fns] + assert not missing, ( + f"moe_w2_cubit: missing cubins for K13={K13}/K2={K2}: {missing} " -+ f"(dir {_DIR}; set VLLM_MOE_W2_CUBIT_DIR)") ++ f"(dir {_DIR}; set VLLM_MOE_W2_CUBIT_DIR)" ++ ) + + +def _fp4_tier_for_build(E: int, dev, n13k13: int, n2k2: int): @@ -1930,12 +3060,18 @@ index 000000000..43bad76a2 + the base slot the refinement is residency-coupled to.""" + from vllm.model_executor.layers.quantization.utils import moe_w2_delta + split = moe_w2_delta.split_enabled() -+ sc13, sc2 = ((n13k13 // 32, n2k2 // 32) -+ if moe_w2_delta.base_enabled() and not split else (0, 0)) ++ sc13, sc2 = ( ++ (n13k13 // 32, n2k2 // 32) ++ if moe_w2_delta.base_enabled() and not split ++ else (0, 0) ++ ) + div = 4 if split else 2 -+ return moe_w2_delta.get_tier(n_experts=E, dev=dev, -+ w13_bytes=n13k13 // div + sc13, -+ w2_bytes=n2k2 // div + sc2) ++ return moe_w2_delta.get_tier( ++ n_experts=E, ++ dev=dev, ++ w13_bytes=n13k13 // div + sc13, ++ w2_bytes=n2k2 // div + sc2, ++ ) + + +def _stage_fp4_host(tier, layer_key: int, fp13, sc13, fp2, sc2) -> None: @@ -1990,10 +3126,11 @@ index 000000000..43bad76a2 + their loaders. Returns True when the layer boots with zero checkpoint + staging.""" + from vllm.model_executor.layers.quantization.utils import moe_w2_delta -+ from vllm.model_executor.layers.quantization.utils import ( -+ moe_w2_planes_cache as _pc) ++ from vllm.model_executor.layers.quantization.utils import moe_w2_planes_cache as _pc + from vllm.model_executor.layers.quantization.utils.moe_w2_store import ( -+ pack_has_layer) ++ pack_has_layer, ++ ) ++ + global _n_created + key = _n_created + _n_created += 1 @@ -2012,8 +3149,7 @@ index 000000000..43bad76a2 + # host-resident: the pack store serves the base (and the FP4 + # need-pool when configured) — probe the sidecars. + n_keys = _layer_cutoff() + 1 -+ if not pack_has_layer("base", key, n_keys, E, -+ c13len + s13len + c2len + s2len): ++ if not pack_has_layer("base", key, n_keys, E, c13len + s13len + c2len + s2len): + return False + if moe_w2_delta.enabled(): + # over-base FP4 need-pool: mirror _fp4_tier_for_build's sizing @@ -2032,11 +3168,11 @@ index 000000000..43bad76a2 + # layer index, sized exactly like process-time try_load). + lidx = _pc.layer_idx_from_name(getattr(layer, "layer_name", "")) + if lidx is None or not _pc.cache_has_layer( -+ lidx, _pc.expected_sizes( -+ E, N13, K13, N2, K2, want_fp4=moe_w2_delta.enabled())): ++ lidx, ++ _pc.expected_sizes(E, N13, K13, N2, K2, want_fp4=moe_w2_delta.enabled()), ++ ): + return False -+ for pname in ("w13_weight", "w13_weight_scale", -+ "w2_weight", "w2_weight_scale"): ++ for pname in ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"): + p = getattr(layer, pname) + p.data = torch.empty(0, dtype=p.data.dtype, device="cpu") + p.weight_loader = _noop_loader @@ -2048,7 +3184,9 @@ index 000000000..43bad76a2 + logger.info( + "moe_w2 LOADER-SKIP armed: pack-resident expert layers are " + "neither host-staged nor copied from the checkpoint " -+ "(first: key %d)", key) ++ "(first: key %d)", ++ key, ++ ) + logger.debug("moe_w2: layer key %d loader-skipped", key) + return True + @@ -2086,11 +3224,28 @@ index 000000000..43bad76a2 + # layer's buffer the moment its first tensor arrives; with a + # layer-major checkpoint only a couple of layers are ever + # in-flight, an unordered one merely degrades to the old profile. ++ allocation_label = None + if param.data.numel() == 0: + shape = self._layer._moe_w2_stream_shapes[self._pname] -+ param.data = torch.empty(shape, dtype=param.data.dtype, -+ device="cpu") -+ ret = self._inner(param, loaded_weight, *args, **kwargs) ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_store, ++ ) ++ ++ allocation_bytes = param.data.element_size() ++ for dim in shape: ++ allocation_bytes *= dim ++ label = f"lazy expert staging {self._pname} {tuple(shape)}" ++ moe_w2_store.allocation_preflight(label, allocation_bytes) ++ param.data = torch.empty(shape, dtype=param.data.dtype, device="cpu") ++ allocation_label = label ++ try: ++ ret = self._inner(param, loaded_weight, *args, **kwargs) ++ finally: ++ # torch.empty is demand-paged; the copy in `_inner` is what ++ # commits RSS/cgroup memory. Check after that touch, including a ++ # partial-copy failure, rather than after the virtual allocation. ++ if allocation_label is not None: ++ moe_w2_store.allocation_postflight(allocation_label) + ok = (ret is True) if kwargs.get("return_success") else True + if not ok: + return ret @@ -2100,11 +3255,13 @@ index 000000000..43bad76a2 + f"moe_w2 stream-build: {self._pname} load arrived after " + f"the layer was already built — more (expert, shard) " + f"tensors than expected; set VLLM_MOE_W2_STREAM_BUILD=0 " -+ f"and report the checkpoint") ++ f"and report the checkpoint" ++ ) + pend[self._pname] -= 1 + if all(v == 0 for v in pend.values()): + key = self._layer._moe_w2_create_key -+ build_layer_planes_nvfp4(self._layer, key) ++ builder = self._layer._moe_w2_stream_builder ++ builder(self._layer, key) + # Drop the staging storage IN PLACE on the ORIGINAL Parameter + # objects. _finish_layer replaced the layer's attributes with + # stub Parameters, but load_weights' params_dict (built once, @@ -2116,19 +3273,18 @@ index 000000000..43bad76a2 + p.data = torch.empty(0, dtype=p.data.dtype, device="cpu") + self._layer._moe_w2_stream_orig = () + self._layer._moe_w2_stream_built = True -+ logger.debug("moe_w2: layer key %d stream-built during load", -+ key) ++ logger.debug("moe_w2: layer key %d stream-built during load", key) + return ret + + -+def arm_stream_build(layer) -> bool: ++def arm_stream_build(layer, checkpoint_format: str = "nvfp4") -> bool: + """Arm the streaming per-layer build on a layer plan_pack_skip missed + (first boot, or a store that does not yet hold it). Expected loads per -+ param: w13-side shards land twice per expert (w1, w3), w2-side once — -+ exact counts over ALL SIX params the requant reads (including scale_2: ++ param: w13-side shards land twice per expert (w1, w3), w2-side once. ++ NVFP4 counts all six params the requant reads (including scale_2: + building before they land would bake uninitialized per-tensor scales -+ into the planes AND the caches), so completeness needs no -+ checkpoint-ordering assumption. ++ into the planes AND the caches); MXFP4 counts its four weight/scale ++ params. Completeness therefore needs no checkpoint-ordering assumption. + + Staging is LAZY: the four big params become 0-byte stubs here and a + layer's buffers materialize on its FIRST loaded tensor (the up-front @@ -2145,25 +3301,44 @@ index 000000000..43bad76a2 + E = layer.w13_weight.shape[0] + except Exception: # noqa: BLE001 - unexpected layout: staged path + return False -+ expected = {"w13_weight": 2 * E, "w13_weight_scale": 2 * E, -+ "w13_weight_scale_2": 2 * E, -+ "w2_weight": E, "w2_weight_scale": E, -+ "w2_weight_scale_2": E} ++ if checkpoint_format == "nvfp4": ++ expected = { ++ "w13_weight": 2 * E, ++ "w13_weight_scale": 2 * E, ++ "w13_weight_scale_2": 2 * E, ++ "w2_weight": E, ++ "w2_weight_scale": E, ++ "w2_weight_scale_2": E, ++ } ++ builder = build_layer_planes_nvfp4 ++ elif checkpoint_format == "mxfp4": ++ expected = { ++ "w13_weight": 2 * E, ++ "w13_weight_scale": 2 * E, ++ "w2_weight": E, ++ "w2_weight_scale": E, ++ } ++ builder = build_layer_planes ++ else: ++ raise ValueError( ++ f"moe_w2 stream-build: unsupported checkpoint format {checkpoint_format!r}" ++ ) + big = ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale") + wrappers = {} + for pname in expected: + p = getattr(layer, pname, None) + inner = getattr(p, "weight_loader", None) + if p is None or inner is None: -+ return False # leave the layer fully staged ++ return False # leave the layer fully staged + wrappers[pname] = (p, _StreamLoader(layer, pname, inner)) + layer._moe_w2_pending = expected ++ layer._moe_w2_stream_builder = builder ++ layer._moe_w2_stream_format = checkpoint_format + # originals of the BIG params: their storage is dropped in place at + # build time, and their SHAPES feed the lazy materialization + layer._moe_w2_stream_orig = tuple(getattr(layer, p) for p in big) -+ layer._moe_w2_stream_shapes = { -+ p: tuple(getattr(layer, p).shape) for p in big} -+ for pname in big: # lazy: nothing staged until loaded ++ layer._moe_w2_stream_shapes = {p: tuple(getattr(layer, p).shape) for p in big} ++ for pname in big: # lazy: nothing staged until loaded + p = getattr(layer, pname) + p.data = torch.empty(0, dtype=p.data.dtype, device="cpu") + for pname, (p, wrap) in wrappers.items(): @@ -2174,12 +3349,14 @@ index 000000000..43bad76a2 + "moe_w2 STREAM-BUILD armed: layer staging materializes on its " + "first loaded tensor and requants on its last (peak staging = " + "layers in flight, not the checkpoint); " -+ "VLLM_MOE_W2_STREAM_BUILD=0 restores the old path") ++ "VLLM_MOE_W2_STREAM_BUILD=0 restores the old path" ++ ) + return True + + -+def _try_skip_requant(layer, layer_key: int, E: int, N13: int, K13: int, -+ N2: int, K2: int, param_names) -> bool: ++def _try_skip_requant( ++ layer, layer_key: int, E: int, N13: int, K13: int, N2: int, K2: int, param_names ++) -> bool: + """Boot-from-pack: when every host store this config serves from already + holds this layer's rows (valid pack written by a previous boot), the + dequant->requant of the checkpoint experts produces bytes NOBODY reads — @@ -2198,37 +3375,45 @@ index 000000000..43bad76a2 + but ITS pack misses the layer, we also requant (the fp4 sections can + only be rebuilt from the checkpoint bytes).""" + from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ + if not moe_w2_delta.base_enabled(): + return False + dev = torch.device("cuda") + c13len, s13len = N13 * K13 // 4, N13 * K13 // 32 + c2len, s2len = N2 * K2 // 4, N2 * K2 // 32 + btier = moe_w2_delta.get_base_tier( -+ _layer_cutoff() + 1, E, dev, -+ w13_bytes=c13len + s13len, w2_bytes=c2len + s2len) ++ _layer_cutoff() + 1, E, dev, w13_bytes=c13len + s13len, w2_bytes=c2len + s2len ++ ) + if layer_key not in btier._store: + return False + tier = _fp4_tier_for_build(E, dev, N13 * K13, N2 * K2) + if tier is not None and layer_key not in tier._store: + return False -+ from vllm.model_executor.layers.quantization.utils import ( -+ moe_w2_planes_cache as _pc) ++ from vllm.model_executor.layers.quantization.utils import moe_w2_planes_cache as _pc ++ + _LAYERS[layer_key] = dict( -+ N13=N13, K13=K13, N2=N2, K2=K2, E=E, base=True, ++ N13=N13, ++ K13=K13, ++ N2=N2, ++ K2=K2, ++ E=E, ++ base=True, + tl_idx=_pc.layer_idx_from_name(getattr(layer, "layer_name", "")), -+ off_s13=c13len, off_c2=c13len + s13len, ++ off_s13=c13len, ++ off_c2=c13len + s13len, + off_s2=c13len + s13len + c2len, -+ off4_s13=2 * c13len, off4_c2=2 * c13len + s13len, ++ off4_s13=2 * c13len, ++ off4_c2=2 * c13len + s13len, + off4_s2=2 * c13len + s13len + 2 * c2len, + ) + stub = torch.empty(0, dtype=torch.uint8, device=dev) + for name in param_names: -+ layer.register_parameter( -+ name, torch.nn.Parameter(stub, requires_grad=False)) ++ layer.register_parameter(name, torch.nn.Parameter(stub, requires_grad=False)) + logger.info( -+ "moe_w2: layer %d requant SKIPPED — %s serving from pack " -+ "(boot-from-pack)", layer_key, -+ "base+fp4" if tier is not None else "base") ++ "moe_w2: layer %d requant SKIPPED — %s serving from pack (boot-from-pack)", ++ layer_key, ++ "base+fp4" if tier is not None else "base", ++ ) + return True + + @@ -2241,18 +3426,56 @@ index 000000000..43bad76a2 + """ + assert _ensure_ready(), "moe_w2 cubins missing" + dev = torch.device("cuda") -+ w13 = layer.w13_weight.data # [E, 2I, H/2] u8 (cpu) -+ s13 = layer.w13_weight_scale.data # [E, 2I, H/32] u8 -+ w2 = layer.w2_weight.data # [E, H, I/2] u8 -+ s2 = layer.w2_weight_scale.data # [E, H, I/32] u8 ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ ++ if getattr(layer, "_moe_w2_pack_skip", False): ++ # Loader-level skip leaves 0-byte parameter stubs, so shapes must ++ # come from the create-time stash. The store probed by ++ # plan_pack_skip must still serve the layer; checkpoint bytes no ++ # longer exist as a fallback. ++ E, N13, K13, N2, K2 = layer._moe_w2_shapes ++ _require_kernels(K13, K2, need_w4=moe_w2_delta.enabled()) ++ if moe_w2_delta.base_enabled(): ++ assert _try_skip_requant( ++ layer, ++ layer_key, ++ E, ++ N13, ++ K13, ++ N2, ++ K2, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ), ( ++ f"moe_w2: layer {layer_key} was loader-skipped on a pack " ++ f"sidecar hit but the pack no longer serves it " ++ f"(dir/sidecar changed mid-load?) — restart without the " ++ f"stale VLLM_MOE_W2_STORE_DIR state" ++ ) ++ return ++ assert _consume_planes_cache(layer, layer_key, dev, E, N13, K13, N2, K2), ( ++ f"moe_w2: layer {layer_key} was loader-skipped on a planes-" ++ f"cache hit but the cache no longer serves it — restart " ++ f"without the stale VLLM_MOE_W2_PLANES_CACHE state" ++ ) ++ return ++ w13 = layer.w13_weight.data # [E, 2I, H/2] u8 (cpu) ++ s13 = layer.w13_weight_scale.data # [E, 2I, H/32] u8 ++ w2 = layer.w2_weight.data # [E, H, I/2] u8 ++ s2 = layer.w2_weight_scale.data # [E, H, I/32] u8 + E, N13, _ = w13.shape + _, N2, _ = w2.shape -+ K13, K2 = N2, N13 // 2 # H, I (4096/2048 on DS4-Flash TP1) -+ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ K13, K2 = N2, N13 // 2 # H, I (4096/2048 on DS4-Flash TP1) + _require_kernels(K13, K2, need_w4=moe_w2_delta.enabled()) -+ if _try_skip_requant(layer, layer_key, E, N13, K13, N2, K2, -+ ("w13_weight", "w13_weight_scale", "w2_weight", -+ "w2_weight_scale")): ++ if _try_skip_requant( ++ layer, ++ layer_key, ++ E, ++ N13, ++ K13, ++ N2, ++ K2, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ): + return + + planes13 = torch.empty(E, N13 * K13 // 4, dtype=torch.uint8, device=dev) @@ -2261,7 +3484,10 @@ index 000000000..43bad76a2 + sc2 = torch.empty(E, N2 * K2 // 32, dtype=torch.uint8, device=dev) + + from vllm.model_executor.layers.quantization.utils.moe_w2_planes import ( -+ mxfp4_to_nibbles, pack_fp4_fragment_major) ++ mxfp4_to_nibbles, ++ pack_fp4_fragment_major, ++ ) ++ + # Pass the PER-RANK FP4 plane sizes (N*K//2 bytes/expert) so the delta tier's + # slots, host store, and pool indexing match the (TP-sharded) planes. On TP1 + # these equal the module constants -> the single-GPU path is unchanged. @@ -2300,14 +3526,26 @@ index 000000000..43bad76a2 + # created; the old "start on layer NUM_LAYERS-1" trigger never fired + # under PP, where layer_keys are local per rank and never reach 42) + -+ _finish_layer(layer, layer_key, dev, planes13, sc13, planes2, sc2, -+ N13, K13, N2, K2, E, -+ ("w13_weight", "w13_weight_scale", "w2_weight", -+ "w2_weight_scale")) ++ _finish_layer( ++ layer, ++ layer_key, ++ dev, ++ planes13, ++ sc13, ++ planes2, ++ sc2, ++ N13, ++ K13, ++ N2, ++ K2, ++ E, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ) + + -+def build_layer_planes_fp8(layer, layer_key: int, -+ scale_suffix: str = "weight_scale_inv") -> None: ++def build_layer_planes_fp8( ++ layer, layer_key: int, scale_suffix: str = "weight_scale_inv" ++) -> None: + """FP8 block-quant checkpoint variant of build_layer_planes (Fp8MoEMethod: + DS4-Flash-FP8, GLM-5.2-FP8 — models without an FP4 release). + @@ -2320,21 +3558,30 @@ index 000000000..43bad76a2 + """ + from vllm.model_executor.layers.quantization.utils import moe_w2_delta + from vllm.model_executor.layers.quantization.utils.moe_w2_planes import ( -+ fp8_block_to_codes_scales, pack_fp4_fragment_major) ++ fp8_block_to_codes_scales, ++ pack_fp4_fragment_major, ++ ) + + assert _ensure_ready(), "moe_w2 cubins missing" + dev = torch.device("cuda") -+ w13 = layer.w13_weight.data # [E, 2I, H] e4m3 (cpu) ++ w13 = layer.w13_weight.data # [E, 2I, H] e4m3 (cpu) + s13 = getattr(layer, f"w13_{scale_suffix}").data # [E, 2I/128, H/128] f32 -+ w2 = layer.w2_weight.data # [E, H, I] e4m3 -+ s2 = getattr(layer, f"w2_{scale_suffix}").data # [E, H/128, I/128] f32 ++ w2 = layer.w2_weight.data # [E, H, I] e4m3 ++ s2 = getattr(layer, f"w2_{scale_suffix}").data # [E, H/128, I/128] f32 + assert w13.dtype == torch.float8_e4m3fn, w13.dtype + E, N13, K13 = w13.shape + _, N2, K2 = w2.shape + _require_kernels(K13, K2, need_w4=moe_w2_delta.enabled()) -+ if _try_skip_requant(layer, layer_key, E, N13, K13, N2, K2, -+ ("w13_weight", f"w13_{scale_suffix}", "w2_weight", -+ f"w2_{scale_suffix}")): ++ if _try_skip_requant( ++ layer, ++ layer_key, ++ E, ++ N13, ++ K13, ++ N2, ++ K2, ++ ("w13_weight", f"w13_{scale_suffix}", "w2_weight", f"w2_{scale_suffix}"), ++ ): + return + + planes13 = torch.empty(E, N13 * K13 // 4, dtype=torch.uint8, device=dev) @@ -2359,7 +3606,8 @@ index 000000000..43bad76a2 + sg = s13[e0:e1].to(dev, non_blocking=True) + for i in range(e1 - e0): + codes, sbytes, nib = fp8_block_to_codes_scales( -+ wg[i], sg[i], want_nibbles=fp13 is not None) ++ wg[i], sg[i], want_nibbles=fp13 is not None ++ ) + planes13[e0 + i] = pack_fragment_major(codes) + sc13[e0 + i] = pack_scales(sbytes) + if fp13 is not None: @@ -2368,7 +3616,8 @@ index 000000000..43bad76a2 + sg = s2[e0:e1].to(dev, non_blocking=True) + for i in range(e1 - e0): + codes, sbytes, nib = fp8_block_to_codes_scales( -+ wg[i], sg[i], want_nibbles=fp2 is not None) ++ wg[i], sg[i], want_nibbles=fp2 is not None ++ ) + planes2[e0 + i] = pack_fragment_major(codes) + sc2[e0 + i] = pack_scales(sbytes) + if fp2 is not None: @@ -2378,15 +3627,26 @@ index 000000000..43bad76a2 + _stage_fp4_host(tier, layer_key, fp13, sc13, fp2, sc2) + del fp13, fp2 + -+ _finish_layer(layer, layer_key, dev, planes13, sc13, planes2, sc2, -+ N13, K13, N2, K2, E, -+ ("w13_weight", f"w13_{scale_suffix}", "w2_weight", -+ f"w2_{scale_suffix}")) ++ _finish_layer( ++ layer, ++ layer_key, ++ dev, ++ planes13, ++ sc13, ++ planes2, ++ sc2, ++ N13, ++ K13, ++ N2, ++ K2, ++ E, ++ ("w13_weight", f"w13_{scale_suffix}", "w2_weight", f"w2_{scale_suffix}"), ++ ) + + -+def _consume_planes_cache(layer, layer_key: int, dev, -+ E: int, N13: int, K13: int, N2: int, -+ K2: int) -> bool: ++def _consume_planes_cache( ++ layer, layer_key: int, dev, E: int, N13: int, K13: int, N2: int, K2: int ++) -> bool: + """Serve one layer's planes from the planes cache (GPU-resident + configs). CPU tensors from the cache feed the same _stage_fp4_host/ + _finish_layer sinks as a fresh requant (their copy_ calls are @@ -2394,13 +3654,17 @@ index 000000000..43bad76a2 + requant) and the loader-skip path (stubs; the cache is the ONLY + source). Returns True on a hit.""" + from vllm.model_executor.layers.quantization.utils import ( -+ moe_w2_planes_cache as planes_cache) ++ moe_w2_planes_cache as planes_cache, ++ ) ++ + lidx = planes_cache.layer_idx_from_name(getattr(layer, "layer_name", "")) + if not planes_cache.enabled() or lidx is None: + return False + tier = _fp4_tier_for_build(E, dev, N13 * K13, N2 * K2) -+ cached = planes_cache.try_load(lidx, planes_cache.expected_sizes( -+ E, N13, K13, N2, K2, want_fp4=tier is not None)) ++ cached = planes_cache.try_load( ++ lidx, ++ planes_cache.expected_sizes(E, N13, K13, N2, K2, want_fp4=tier is not None), ++ ) + if cached is None: + return False + planes13 = cached["planes13"].view(E, -1).to(dev) @@ -2408,12 +3672,29 @@ index 000000000..43bad76a2 + planes2 = cached["planes2"].view(E, -1).to(dev) + sc2 = cached["sc2"].view(E, -1).to(dev) + if tier is not None: -+ _stage_fp4_host(tier, layer_key, cached["fp13"].view(E, -1), -+ sc13, cached["fp2"].view(E, -1), sc2) -+ _finish_layer(layer, layer_key, dev, planes13, sc13, planes2, -+ sc2, N13, K13, N2, K2, E, -+ ("w13_weight", "w13_weight_scale", "w2_weight", -+ "w2_weight_scale")) ++ _stage_fp4_host( ++ tier, ++ layer_key, ++ cached["fp13"].view(E, -1), ++ sc13, ++ cached["fp2"].view(E, -1), ++ sc2, ++ ) ++ _finish_layer( ++ layer, ++ layer_key, ++ dev, ++ planes13, ++ sc13, ++ planes2, ++ sc2, ++ N13, ++ K13, ++ N2, ++ K2, ++ E, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ) + logger.info("moe_w2: layer %d planes from cache", lidx) + return True + @@ -2432,7 +3713,9 @@ index 000000000..43bad76a2 + """ + from vllm.model_executor.layers.quantization.utils import moe_w2_delta + from vllm.model_executor.layers.quantization.utils.moe_w2_planes import ( -+ nvfp4_to_codes_scales, pack_fp4_fragment_major) ++ nvfp4_to_codes_scales, ++ pack_fp4_fragment_major, ++ ) + + assert _ensure_ready(), "moe_w2 cubins missing" + dev = torch.device("cuda") @@ -2445,40 +3728,56 @@ index 000000000..43bad76a2 + _require_kernels(K13, K2, need_w4=moe_w2_delta.enabled()) + if moe_w2_delta.base_enabled(): + assert _try_skip_requant( -+ layer, layer_key, E, N13, K13, N2, K2, -+ ("w13_weight", "w13_weight_scale", "w2_weight", -+ "w2_weight_scale")), ( ++ layer, ++ layer_key, ++ E, ++ N13, ++ K13, ++ N2, ++ K2, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ), ( + f"moe_w2: layer {layer_key} was loader-skipped on a pack " + f"sidecar hit but the pack no longer serves it " + f"(dir/sidecar changed mid-load?) — restart without the " -+ f"stale VLLM_MOE_W2_STORE_DIR state") ++ f"stale VLLM_MOE_W2_STORE_DIR state" ++ ) + return + # GPU-resident: materialize the planes from the planes cache + # (probed at create time; a miss here means the cache dir changed + # under a live load). -+ assert _consume_planes_cache(layer, layer_key, dev, -+ E, N13, K13, N2, K2), ( ++ assert _consume_planes_cache(layer, layer_key, dev, E, N13, K13, N2, K2), ( + f"moe_w2: layer {layer_key} was loader-skipped on a planes-" + f"cache hit but the cache no longer serves it — restart " -+ f"without the stale VLLM_MOE_W2_PLANES_CACHE state") ++ f"without the stale VLLM_MOE_W2_PLANES_CACHE state" ++ ) + return -+ w13 = layer.w13_weight.data # [E, 2I, H/2] u8 (cpu) -+ s13 = layer.w13_weight_scale.data # [E, 2I, H/16] e4m3 -+ s13_2 = layer.w13_weight_scale_2.data # [E, 2] f32 (w1, w3) -+ w2 = layer.w2_weight.data # [E, H, I/2] u8 -+ s2 = layer.w2_weight_scale.data # [E, H, I/16] e4m3 -+ s2_2 = layer.w2_weight_scale_2.data # [E] f32 ++ w13 = layer.w13_weight.data # [E, 2I, H/2] u8 (cpu) ++ s13 = layer.w13_weight_scale.data # [E, 2I, H/16] e4m3 ++ s13_2 = layer.w13_weight_scale_2.data # [E, 2] f32 (w1, w3) ++ w2 = layer.w2_weight.data # [E, H, I/2] u8 ++ s2 = layer.w2_weight_scale.data # [E, H, I/16] e4m3 ++ s2_2 = layer.w2_weight_scale_2.data # [E] f32 + assert w13.dtype == torch.uint8 and s13.dtype == torch.float8_e4m3fn, ( -+ w13.dtype, s13.dtype) ++ w13.dtype, ++ s13.dtype, ++ ) + E, N13, K13h = w13.shape + K13 = K13h * 2 + _, N2, K2h = w2.shape + K2 = K2h * 2 -+ group = K13 // s13.shape[2] # 16 for NVFP4 ++ group = K13 // s13.shape[2] # 16 for NVFP4 + _require_kernels(K13, K2, need_w4=moe_w2_delta.enabled()) -+ if _try_skip_requant(layer, layer_key, E, N13, K13, N2, K2, -+ ("w13_weight", "w13_weight_scale", "w2_weight", -+ "w2_weight_scale")): ++ if _try_skip_requant( ++ layer, ++ layer_key, ++ E, ++ N13, ++ K13, ++ N2, ++ K2, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ): + return + + # Planes cache (VLLM_MOE_W2_PLANES_CACHE): the requant below is @@ -2491,7 +3790,9 @@ index 000000000..43bad76a2 + if _consume_planes_cache(layer, layer_key, dev, E, N13, K13, N2, K2): + return + from vllm.model_executor.layers.quantization.utils import ( -+ moe_w2_planes_cache as planes_cache) ++ moe_w2_planes_cache as planes_cache, ++ ) ++ + lidx = planes_cache.layer_idx_from_name(getattr(layer, "layer_name", "")) + tier = _fp4_tier_for_build(E, dev, N13 * K13, N2 * K2) + @@ -2515,12 +3816,12 @@ index 000000000..43bad76a2 + wg = w13[e0:e1].to(dev, non_blocking=True) + sg = s13[e0:e1].to(dev, non_blocking=True) + s2g = s13_2[e0:e1].to(dev, non_blocking=True) -+ half = N13 // 2 # rows [0:I]=w1, [I:2I]=w3 ++ half = N13 // 2 # rows [0:I]=w1, [I:2I]=w3 + for i in range(e1 - e0): + s2_row = torch.cat((s2g[i, 0].expand(half), s2g[i, 1].expand(half))) + codes, sbytes, nib = nvfp4_to_codes_scales( -+ wg[i], sg[i], s2_row, group=group, -+ want_nibbles=fp13 is not None) ++ wg[i], sg[i], s2_row, group=group, want_nibbles=fp13 is not None ++ ) + planes13[e0 + i] = pack_fragment_major(codes) + sc13[e0 + i] = pack_scales(sbytes) + if fp13 is not None: @@ -2530,36 +3831,69 @@ index 000000000..43bad76a2 + s2g = s2_2[e0:e1].to(dev, non_blocking=True) + for i in range(e1 - e0): + codes, sbytes, nib = nvfp4_to_codes_scales( -+ wg[i], sg[i], s2g[i], group=group, -+ want_nibbles=fp2 is not None) ++ wg[i], sg[i], s2g[i], group=group, want_nibbles=fp2 is not None ++ ) + planes2[e0 + i] = pack_fragment_major(codes) + sc2[e0 + i] = pack_scales(sbytes) + if fp2 is not None: + fp2[e0 + i] = _pack_fp4_plane(nib) + + if planes_cache.enabled() and lidx is not None: -+ planes_cache.store(lidx, dict(planes13=planes13, sc13=sc13, -+ planes2=planes2, sc2=sc2, -+ fp13=fp13, fp2=fp2)) ++ planes_cache.store( ++ lidx, ++ dict( ++ planes13=planes13, ++ sc13=sc13, ++ planes2=planes2, ++ sc2=sc2, ++ fp13=fp13, ++ fp2=fp2, ++ ), ++ ) + + if tier is not None: + _stage_fp4_host(tier, layer_key, fp13, sc13, fp2, sc2) + del fp13, fp2 + -+ _finish_layer(layer, layer_key, dev, planes13, sc13, planes2, sc2, -+ N13, K13, N2, K2, E, -+ ("w13_weight", "w13_weight_scale", "w2_weight", -+ "w2_weight_scale")) ++ _finish_layer( ++ layer, ++ layer_key, ++ dev, ++ planes13, ++ sc13, ++ planes2, ++ sc2, ++ N13, ++ K13, ++ N2, ++ K2, ++ E, ++ ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), ++ ) + + -+def _finish_layer(layer, layer_key, dev, planes13, sc13, planes2, sc2, -+ N13, K13, N2, K2, E, param_names) -> None: ++def _finish_layer( ++ layer, ++ layer_key, ++ dev, ++ planes13, ++ sc13, ++ planes2, ++ sc2, ++ N13, ++ K13, ++ N2, ++ K2, ++ E, ++ param_names, ++) -> None: + from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ + # transformer layer index of this layer_key (dense-offset models: GLM's + # first sparse layer 3 -> key 0). LOOKA uses it to pair each key with + # its transformer layer's router (mlp.gate) weights. -+ from vllm.model_executor.layers.quantization.utils import ( -+ moe_w2_planes_cache as _pc) ++ from vllm.model_executor.layers.quantization.utils import moe_w2_planes_cache as _pc ++ + _tl = _pc.layer_idx_from_name(getattr(layer, "layer_name", "")) + if moe_w2_delta.base_enabled(): + # BASE cache (inverted delta): the 2-bit planes go to PINNED HOST RAM @@ -2570,61 +3904,97 @@ index 000000000..43bad76a2 + c13len, s13len = planes13.shape[1], sc13.shape[1] + c2len, s2len = planes2.shape[1], sc2.shape[1] + btier = moe_w2_delta.get_base_tier( -+ _layer_cutoff() + 1, E, dev, -+ w13_bytes=c13len + s13len, w2_bytes=c2len + s2len) ++ _layer_cutoff() + 1, ++ E, ++ dev, ++ w13_bytes=c13len + s13len, ++ w2_bytes=c2len + s2len, ++ ) + btier.add_layer_host_planes( + layer_key, + torch.cat((planes13, sc13), dim=1), -+ torch.cat((planes2, sc2), dim=1)) ++ torch.cat((planes2, sc2), dim=1), ++ ) + _LAYERS[layer_key] = dict( -+ N13=N13, K13=K13, N2=N2, K2=K2, E=E, base=True, tl_idx=_tl, -+ off_s13=c13len, off_c2=c13len + s13len, ++ N13=N13, ++ K13=K13, ++ N2=N2, ++ K2=K2, ++ E=E, ++ base=True, ++ tl_idx=_tl, ++ off_s13=c13len, ++ off_c2=c13len + s13len, + off_s2=c13len + s13len + c2len, + # FP4 need-pool slot sections ([fp4_13|sc13|fp4_2|sc2]; fp4 codes + # are 2x the 2-bit codes, scale sections identical) — read by the + # base+delta desc kernel when the FP4 tier coexists. -+ off4_s13=2 * c13len, off4_c2=2 * c13len + s13len, ++ off4_s13=2 * c13len, ++ off4_c2=2 * c13len + s13len, + off4_s2=2 * c13len + s13len + 2 * c2len, + ) + del planes13, sc13, planes2, sc2 + stub = torch.empty(0, dtype=torch.uint8, device=dev) + for name in param_names: + layer.register_parameter( -+ name, torch.nn.Parameter(stub, requires_grad=False)) -+ logger.info("moe_w2: layer %d planes HOST-staged (base cache, " -+ "%.2f GiB pinned)", layer_key, -+ E * btier.slot_bytes / 2**30) ++ name, torch.nn.Parameter(stub, requires_grad=False) ++ ) ++ logger.info( ++ "moe_w2: layer %d planes HOST-staged (base cache, %.2f GiB pinned)", ++ layer_key, ++ E * btier.slot_bytes / 2**30, ++ ) + return + + _LAYERS[layer_key] = dict( -+ planes13=planes13, sc13=sc13, planes2=planes2, sc2=sc2, -+ N13=N13, K13=K13, N2=N2, K2=K2, E=E, tl_idx=_tl, ++ planes13=planes13, ++ sc13=sc13, ++ planes2=planes2, ++ sc2=sc2, ++ N13=N13, ++ K13=K13, ++ N2=N2, ++ K2=K2, ++ E=E, ++ tl_idx=_tl, + ) + # Release checkpoint copies; keep CUDA stubs so device probes stay happy. + stub = torch.empty(0, dtype=torch.uint8, device=dev) + for name in param_names: -+ layer.register_parameter( -+ name, torch.nn.Parameter(stub, requires_grad=False)) -+ logger.info("moe_w2: layer %d planes built (%.2f GiB)", layer_key, -+ (planes13.nbytes + sc13.nbytes + planes2.nbytes + sc2.nbytes) -+ / 2**30) ++ layer.register_parameter(name, torch.nn.Parameter(stub, requires_grad=False)) ++ logger.info( ++ "moe_w2: layer %d planes built (%.2f GiB)", ++ layer_key, ++ (planes13.nbytes + sc13.nbytes + planes2.nbytes + sc2.nbytes) / 2**30, ++ ) + + +# -------------------------------------------------------------------------- +# Forward +# -------------------------------------------------------------------------- + -+def _workspaces(slots: int, tokens: int, dev, inter: int = 2048, -+ hidden: int = 4096, n_experts: int = 256) -> dict: ++ ++def _workspaces( ++ slots: int, ++ tokens: int, ++ dev, ++ inter: int = 2048, ++ hidden: int = 4096, ++ n_experts: int = 256, ++) -> dict: + # `inter` = per-rank expert intermediate size I (2048 on 1 GPU; 1024 @ TP2, + # 512 @ TP4 as the experts shard). The hidden H (4096 DS4, 6144 GLM-5.x) is + # NOT sharded, so the A-side (a1), x-quant (xq) and w2 output (c2) buffers + # stay H-wide; only the gate/up output (c13 = 2I), the intermediate + # activation (act/a2 = I) and its group-128 scales (as2 = I/128) follow the + # shard. -+ if (_WS.get("slots", 0) < slots or _WS.get("tokens", 0) < tokens -+ or _WS.get("inter") != inter or _WS.get("hidden") != hidden -+ or _WS.get("n_experts", 0) < n_experts): ++ if ( ++ _WS.get("slots", 0) < slots ++ or _WS.get("tokens", 0) < tokens ++ or _WS.get("inter") != inter ++ or _WS.get("hidden") != hidden ++ or _WS.get("n_experts", 0) < n_experts ++ ): + slots = max(slots, _WS.get("slots", 0)) + tokens = max(tokens, _WS.get("tokens", 0)) + n_experts = max(n_experts, _WS.get("n_experts", 0)) @@ -2637,46 +4007,43 @@ index 000000000..43bad76a2 + # token-side quant buffers; the LAST row is the permanent zero + # pad row (gather source for filler slots) — quant only ever + # writes rows [:T]. -+ xq=torch.zeros(tokens + 1, hidden, dtype=torch.float8_e4m3fn, -+ device=dev), -+ xs=torch.zeros(tokens + 1, hidden // 128, dtype=torch.float32, -+ device=dev), -+ a1=torch.zeros(slots + 4, hidden, dtype=torch.float8_e4m3fn, -+ device=dev), -+ as1=torch.zeros(slots + 4, hidden // 128, dtype=torch.float32, -+ device=dev), ++ xq=torch.zeros(tokens + 1, hidden, dtype=torch.float8_e4m3fn, device=dev), ++ xs=torch.zeros(tokens + 1, hidden // 128, dtype=torch.float32, device=dev), ++ a1=torch.zeros(slots + 4, hidden, dtype=torch.float8_e4m3fn, device=dev), ++ as1=torch.zeros(slots + 4, hidden // 128, dtype=torch.float32, device=dev), + # zeros, not empty: pad-pair rows are never written by the kernel + # (early EXIT) yet flow through silu/scatter math with weight 0; + # uninitialized inf/nan would poison 0*x. -+ c13=torch.zeros(slots + 4, 2 * inter, dtype=torch.bfloat16, -+ device=dev), ++ c13=torch.zeros(slots + 4, 2 * inter, dtype=torch.bfloat16, device=dev), + act=torch.zeros(slots + 4, inter, dtype=torch.bfloat16, device=dev), -+ a2=torch.zeros(slots + 4, inter, dtype=torch.float8_e4m3fn, -+ device=dev), -+ as2=torch.zeros(slots + 4, max(inter // 128, 1), -+ dtype=torch.float32, device=dev), -+ c2=torch.zeros(slots + 4, hidden, dtype=torch.bfloat16, -+ device=dev), -+ desc=torch.empty(4, slots // _BLOCK, 6, dtype=torch.int64, -+ device=dev), ++ a2=torch.zeros(slots + 4, inter, dtype=torch.float8_e4m3fn, device=dev), ++ as2=torch.zeros( ++ slots + 4, max(inter // 128, 1), dtype=torch.float32, device=dev ++ ), ++ c2=torch.zeros(slots + 4, hidden, dtype=torch.bfloat16, device=dev), ++ desc=torch.empty(4, slots // _BLOCK, 6, dtype=torch.int64, device=dev), + # split-FP4 (moe_w4s_mm) desc tables: 8 u64 per pair, 64 B ABI -+ desc4s=torch.empty(2, slots // _BLOCK, 8, dtype=torch.int64, -+ device=dev), ++ desc4s=torch.empty( ++ 2, slots // _BLOCK, 8, dtype=torch.int64, device=dev ++ ), + # -1 slot row for the tier-less desc path; sized to the MODEL's + # expert count (256 = DS4 default; 384 Kimi-K2.x reads past a + # fixed 256-row table). -+ no_slots=torch.full((max(n_experts, 256),), -1, -+ dtype=torch.int32, device=dev), ++ no_slots=torch.full( ++ (max(n_experts, 256),), -1, dtype=torch.int32, device=dev ++ ), + ) + if _afrag_ok: + # AFRAG destination buffers: the triton repack streams row-major + # a1/a2 into these (single pass, no copy-back); the desc tables + # point the GEMM at them instead of a1/a2. + _WS.update( -+ a1f=torch.zeros(slots + 4, hidden, dtype=torch.float8_e4m3fn, -+ device=dev), -+ a2f=torch.zeros(slots + 4, inter, dtype=torch.float8_e4m3fn, -+ device=dev), ++ a1f=torch.zeros( ++ slots + 4, hidden, dtype=torch.float8_e4m3fn, device=dev ++ ), ++ a2f=torch.zeros( ++ slots + 4, inter, dtype=torch.float8_e4m3fn, device=dev ++ ), + ) + return _WS + @@ -2715,13 +4082,38 @@ index 000000000..43bad76a2 + +@triton.jit +def _desc_build_kernel( -+ eids_ptr, npost_ptr, slot_ptr, d_ptr, -+ a1b, as1b, c13b, a2b, as2b, c2b, -+ p13b, s13b, p2b, s2b, poolb, -+ p13s, s13s, p2s, s2s, -+ slot_bytes, w13_bytes, -+ a1_rb, as1_rb, c13_rb, a2_rb, as2_rb, c2_rb, -+ n_experts, pairs, cap6, mblock, ++ eids_ptr, ++ npost_ptr, ++ pair_live_ptr, ++ slot_ptr, ++ d_ptr, ++ a1b, ++ as1b, ++ c13b, ++ a2b, ++ as2b, ++ c2b, ++ p13b, ++ s13b, ++ p2b, ++ s2b, ++ poolb, ++ p13s, ++ s13s, ++ p2s, ++ s2s, ++ slot_bytes, ++ w13_bytes, ++ a1_rb, ++ as1_rb, ++ c13_rb, ++ a2_rb, ++ as2_rb, ++ c2_rb, ++ n_experts, ++ pairs, ++ cap6, ++ mblock, + BLOCK: tl.constexpr, +): + """All four moe desc tables in one launch (24 columns per pair). @@ -2738,7 +4130,8 @@ index 000000000..43bad76a2 + e = tl.minimum(tl.maximum(e, 0), n_experts - 1) + slot = tl.load(slot_ptr + e, mask=mask, other=-1).to(tl.int64) + npost = tl.load(npost_ptr).to(tl.int64) -+ live = p < npost // mblock ++ has_real = tl.load(pair_live_ptr + p, mask=mask, other=0) != 0 ++ live = (p < npost // mblock) & has_real + is4 = slot >= 0 + m2 = tl.where(live & ~is4, mblock, 0).to(tl.int64) + m4 = tl.where(live & is4, mblock, 0).to(tl.int64) @@ -2759,11 +4152,16 @@ index 000000000..43bad76a2 + elif gi == 1: + b, s, a, as_, c, m = p2b + e * p2s, bs2, a2, as2, c2, m2 + elif gi == 2: -+ b, s, a, as_, c, m = (poolb + slot_c * slot_bytes, bs13, -+ a1, as1, c13, m4) ++ b, s, a, as_, c, m = (poolb + slot_c * slot_bytes, bs13, a1, as1, c13, m4) + else: -+ b, s, a, as_, c, m = (poolb + slot_c * slot_bytes + w13_bytes, -+ bs2, a2, as2, c2, m4) ++ b, s, a, as_, c, m = ( ++ poolb + slot_c * slot_bytes + w13_bytes, ++ bs2, ++ a2, ++ as2, ++ c2, ++ m4, ++ ) + tl.store(d + 0, a, mask=mask) + tl.store(d + 1, as_, mask=mask) + tl.store(d + 2, b, mask=mask) @@ -2774,7 +4172,7 @@ index 000000000..43bad76a2 + +@triton.jit +def _desc_build_kernel_w4s( -+ eids_ptr, npost_ptr, slot_ptr, d_ptr, ++ eids_ptr, npost_ptr, pair_live_ptr, slot_ptr, d_ptr, + a1b, as1b, c13b, a2b, as2b, c2b, + p13b, s13b, p2b, s2b, poolb, + p13s, s13s, p2s, s2s, @@ -2795,7 +4193,8 @@ index 000000000..43bad76a2 + e = tl.minimum(tl.maximum(e, 0), n_experts - 1) + slot = tl.load(slot_ptr + e, mask=mask, other=-1).to(tl.int64) + npost = tl.load(npost_ptr).to(tl.int64) -+ live = p < npost // mblock ++ has_real = tl.load(pair_live_ptr + p, mask=mask, other=0) != 0 ++ live = (p < npost // mblock) & has_real + is4 = slot >= 0 + m4 = tl.where(live & is4, mblock, 0).to(tl.int64) + base = p.to(tl.int64) * mblock @@ -2826,11 +4225,33 @@ index 000000000..43bad76a2 + +@triton.jit +def _desc_build_kernel_basecache( -+ eids_ptr, npost_ptr, slot_ptr, miss_ptr, d_ptr, -+ a1b, as1b, c13b, a2b, as2b, c2b, -+ poolb, slot_bytes, off_s13, off_c2, off_s2, -+ a1_rb, as1_rb, c13_rb, a2_rb, as2_rb, c2_rb, -+ n_experts, pairs, cap6, mblock, ++ eids_ptr, ++ npost_ptr, ++ pair_live_ptr, ++ slot_ptr, ++ miss_ptr, ++ d_ptr, ++ a1b, ++ as1b, ++ c13b, ++ a2b, ++ as2b, ++ c2b, ++ poolb, ++ slot_bytes, ++ off_s13, ++ off_c2, ++ off_s2, ++ a1_rb, ++ as1_rb, ++ c13_rb, ++ a2_rb, ++ as2_rb, ++ c2_rb, ++ n_experts, ++ pairs, ++ cap6, ++ mblock, + BLOCK: tl.constexpr, +): + """Base-cache variant of _desc_build_kernel: the 2-bit BASE planes live in @@ -2847,7 +4268,8 @@ index 000000000..43bad76a2 + e = tl.minimum(tl.maximum(e, 0), n_experts - 1) + slot = tl.load(slot_ptr + e, mask=mask, other=-1).to(tl.int64) + npost = tl.load(npost_ptr).to(tl.int64) -+ live = p < npost // mblock ++ has_real = tl.load(pair_live_ptr + p, mask=mask, other=0) != 0 ++ live = (p < npost // mblock) & has_real + hit = slot >= 0 + m = tl.where(live & hit, mblock, 0).to(tl.int64) + n_miss = tl.sum(tl.where(mask & live & ~hit, 1, 0)) @@ -2877,12 +4299,39 @@ index 000000000..43bad76a2 + +@triton.jit +def _desc_build_kernel_base_delta( -+ eids_ptr, npost_ptr, bslot_ptr, fslot_ptr, miss_ptr, d_ptr, -+ a1b, as1b, c13b, a2b, as2b, c2b, -+ bpoolb, bslot_bytes, off_s13, off_c2, off_s2, -+ fpoolb, fslot_bytes, off4_s13, off4_c2, off4_s2, -+ a1_rb, as1_rb, c13_rb, a2_rb, as2_rb, c2_rb, -+ n_experts, pairs, cap6, mblock, ++ eids_ptr, ++ npost_ptr, ++ pair_live_ptr, ++ bslot_ptr, ++ fslot_ptr, ++ miss_ptr, ++ d_ptr, ++ a1b, ++ as1b, ++ c13b, ++ a2b, ++ as2b, ++ c2b, ++ bpoolb, ++ bslot_bytes, ++ off_s13, ++ off_c2, ++ off_s2, ++ fpoolb, ++ fslot_bytes, ++ off4_s13, ++ off4_c2, ++ off4_s2, ++ a1_rb, ++ as1_rb, ++ c13_rb, ++ a2_rb, ++ as2_rb, ++ c2_rb, ++ n_experts, ++ pairs, ++ cap6, ++ mblock, + BLOCK: tl.constexpr, +): + """Base cache + FP4 need-pool coexistence variant: TWO slot tables with @@ -2900,7 +4349,8 @@ index 000000000..43bad76a2 + bslot = tl.load(bslot_ptr + e, mask=mask, other=-1).to(tl.int64) + fslot = tl.load(fslot_ptr + e, mask=mask, other=-1).to(tl.int64) + npost = tl.load(npost_ptr).to(tl.int64) -+ live = p < npost // mblock ++ has_real = tl.load(pair_live_ptr + p, mask=mask, other=0) != 0 ++ live = (p < npost // mblock) & has_real + is4 = fslot >= 0 + bhit = bslot >= 0 + m2 = tl.where(live & bhit & ~is4, mblock, 0).to(tl.int64) @@ -2936,12 +4386,39 @@ index 000000000..43bad76a2 + +@triton.jit +def _desc_build_kernel_base_delta_split( -+ eids_ptr, npost_ptr, bslot_ptr, fslot_ptr, miss_ptr, d_ptr, d4s_ptr, -+ a1b, as1b, c13b, a2b, as2b, c2b, -+ bpoolb, bslot_bytes, off_s13, off_c2, off_s2, -+ fpoolb, fslot_bytes, w13r_bytes, -+ a1_rb, as1_rb, c13_rb, a2_rb, as2_rb, c2_rb, -+ n_experts, pairs, cap6, cap8, mblock, ++ eids_ptr, ++ npost_ptr, ++ pair_live_ptr, ++ bslot_ptr, ++ fslot_ptr, ++ miss_ptr, ++ d_ptr, ++ d4s_ptr, ++ a1b, ++ as1b, ++ c13b, ++ a2b, ++ as2b, ++ c2b, ++ bpoolb, ++ bslot_bytes, ++ off_s13, ++ off_c2, ++ off_s2, ++ fpoolb, ++ fslot_bytes, ++ w13r_bytes, ++ a1_rb, ++ as1_rb, ++ c13_rb, ++ a2_rb, ++ as2_rb, ++ c2_rb, ++ n_experts, ++ pairs, ++ cap6, ++ cap8, ++ mblock, + BLOCK: tl.constexpr, +): + """Base cache + SPLIT FP4 need-pool: refinement slots are read AGAINST @@ -2961,9 +4438,10 @@ index 000000000..43bad76a2 + bslot = tl.load(bslot_ptr + e, mask=mask, other=-1).to(tl.int64) + fslot = tl.load(fslot_ptr + e, mask=mask, other=-1).to(tl.int64) + npost = tl.load(npost_ptr).to(tl.int64) -+ live = p < npost // mblock ++ has_real = tl.load(pair_live_ptr + p, mask=mask, other=0) != 0 ++ live = (p < npost // mblock) & has_real + bhit = bslot >= 0 -+ is4 = (fslot >= 0) & bhit # split serve needs BOTH resident ++ is4 = (fslot >= 0) & bhit # split serve needs BOTH resident + m2 = tl.where(live & bhit & ~is4, mblock, 0).to(tl.int64) + m4 = tl.where(live & is4, mblock, 0).to(tl.int64) + n_miss = tl.sum(tl.where(mask & live & ~bhit, 1, 0)) @@ -2977,7 +4455,7 @@ index 000000000..43bad76a2 + a2 = a2b + base * a2_rb + as2 = as2b + base * as2_rb + c2 = c2b + base * c2_rb -+ for gi in tl.static_range(2): # w2 tables (base pool sections) ++ for gi in tl.static_range(2): # w2 tables (base pool sections) + d = d_ptr + gi * cap6 + p * 6 + if gi == 0: + b, s, a, as_, c = bs, bs + off_s13, a1, as1, c13 @@ -2989,13 +4467,19 @@ index 000000000..43bad76a2 + tl.store(d + 3, s, mask=mask) + tl.store(d + 4, c, mask=mask) + tl.store(d + 5, m2, mask=mask) -+ for gi in tl.static_range(2): # w4s tables (base + refinement) ++ for gi in tl.static_range(2): # w4s tables (base + refinement) + d = d4s_ptr + gi * cap8 + p * 8 + if gi == 0: + bb, rr, ss, a, as_, c = bs, fs, bs + off_s13, a1, as1, c13 + else: -+ bb, rr, ss, a, as_, c = (bs + off_c2, fs + w13r_bytes, -+ bs + off_s2, a2, as2, c2) ++ bb, rr, ss, a, as_, c = ( ++ bs + off_c2, ++ fs + w13r_bytes, ++ bs + off_s2, ++ a2, ++ as2, ++ c2, ++ ) + tl.store(d + 0, a, mask=mask) + tl.store(d + 1, as_, mask=mask) + tl.store(d + 2, bb, mask=mask) @@ -3006,19 +4490,34 @@ index 000000000..43bad76a2 + tl.store(d + 7, tl.zeros_like(m4), mask=mask) + + -+def _launch(tier: str, K: int, desc: torch.Tensor, n_rows: int, pairs: int, -+ stream): ++def _launch(tier: str, K: int, desc: torch.Tensor, n_rows: int, pairs: int, stream): + fn = _fns[(tier, K)] -+ args = [ctypes.c_uint64(desc.data_ptr()), -+ ctypes.c_uint32(K), -+ ctypes.c_uint32(K // 64), -+ ctypes.c_uint32(n_rows * 2), -+ ctypes.c_uint32(K // 128)] ++ args = [ ++ ctypes.c_uint64(desc.data_ptr()), ++ ctypes.c_uint32(K), ++ ctypes.c_uint32(K // 64), ++ ctypes.c_uint32(n_rows * 2), ++ ctypes.c_uint32(K // 128), ++ ] + argv = (ctypes.c_void_p * len(args))( -+ *[ctypes.cast(ctypes.byref(x), ctypes.c_void_p) for x in args]) -+ _ck(_driver().cuLaunchKernel(fn, n_rows // 16, pairs, 1, -+ _nwarp_for_k(K) * 32, 1, 1, 0, -+ stream, argv, None), "launch") ++ *[ctypes.cast(ctypes.byref(x), ctypes.c_void_p) for x in args] ++ ) ++ _ck( ++ _driver().cuLaunchKernel( ++ fn, ++ n_rows // 16, ++ pairs, ++ 1, ++ _nwarp_for_k(K) * 32, ++ 1, ++ 1, ++ 0, ++ stream, ++ argv, ++ None, ++ ), ++ "launch", ++ ) + + +def _moe_w2_forward( @@ -3028,6 +4527,7 @@ index 000000000..43bad76a2 + layer_key: int, +) -> torch.Tensor: + from vllm.model_executor.layers.quantization.utils import prefill_timers ++ + with prefill_timers.span("moe_w2"): + return _moe_w2_forward_timed(x, topk_weights, topk_ids, layer_key) + @@ -3047,6 +4547,8 @@ index 000000000..43bad76a2 + + st = _LAYERS[layer_key] + T, H = x.shape ++ token_slot_mapping = _get_token_slot_mapping(T) ++ logical_prefill = _get_has_prefill(T) + # adaptive expert top-p (env-gated; identity when off). Must run before + # moe_align/mark_seen/route_log so dropped experts are neither fetched + # nor counted as routed. @@ -3054,23 +4556,26 @@ index 000000000..43bad76a2 + top_k = topk_ids.shape[1] + dev = x.device + stream = ctypes.c_void_p(torch.cuda.current_stream(dev).cuda_stream) -+ -+ # decode-sized calls use the proven 4-token kernel + delta tier; -+ # prefill-sized calls use the MC4 kernel (16 tokens per pair-entry = full -+ # QMMA-M, plane reads amortized 4x, ~1.5x over MC2) on the 2-bit base only. -+ # 96 = the largest cudagraph capture size: anything above is necessarily a -+ # prefill chunk; short tail chunks keep the delta-quality path. -+ prefill = T > 96 -+ mblock = 16 if prefill else _BLOCK ++ capturing = torch.cuda.is_current_stream_capturing() ++ if logical_prefill and capturing: ++ raise RuntimeError("moe_w2 prefill residency cannot run under CUDA capture") ++ ++ # Bulk prefills use the MC4 kernel (16 tokens per pair-entry = full ++ # QMMA-M) on the 2-bit base. Short prefill tails retain the 4-token kernel ++ # and delta-quality path, but still use eager per-layer base residency. ++ bulk_prefill = logical_prefill and T > _BULK_PREFILL_TOKENS ++ mblock = 16 if bulk_prefill else _BLOCK + sorted_ids, expert_blocks, num_post = moe_align_block_size( -+ topk_ids, mblock, st["E"]) ++ topk_ids, mblock, st["E"], pad_sorted_ids=True ++ ) + slots = sorted_ids.numel() ++ if slots % mblock: ++ raise RuntimeError("moe_w2 aligned route capacity must be mblock-divisible") + pairs = slots // mblock + # st["K2"] = per-rank expert intermediate I (w2 contraction), st["K13"] = + # hidden H (w13 contraction) -> size the workspaces for the model's shapes + # (and correctly under tensor parallelism). -+ ws = _workspaces(slots, T, dev, inter=st["K2"], hidden=st["K13"], -+ n_experts=st["E"]) ++ ws = _workspaces(slots, T, dev, inter=st["K2"], hidden=st["K13"], n_experts=st["E"]) + + # ---- activation quant (group-128) into the padded buffer; the buffer's + # last row is the permanent zero pad row for filler slots. @@ -3078,26 +4583,32 @@ index 000000000..43bad76a2 + pad_row = xq.shape[0] - 1 + _, xs = per_token_group_quant_fp8(x, 128, out_q=xq[:T]) + ws["xs"][:T] = xs -+ valid = sorted_ids < T * top_k -+ rows = torch.where(valid, sorted_ids // top_k, -+ torch.full_like(sorted_ids, pad_row)) -+ torch.index_select(xq.view(torch.uint8), 0, rows, -+ out=ws["a1"][:slots].view(torch.uint8)) ++ # Runner updates this persistent buffer every step and writes -1 to the ++ # cudagraph padding tail. Reading it on-device keeps replay dynamic while ++ # all tensor shapes and addresses stay capture-stable. ++ token_valid, valid, pair_live, rows = _masked_route_metadata( ++ sorted_ids, token_slot_mapping, top_k, mblock, pad_row ++ ) ++ torch.index_select( ++ xq.view(torch.uint8), 0, rows, out=ws["a1"][:slots].view(torch.uint8) ++ ) + torch.index_select(ws["xs"], 0, rows, out=ws["as1"][:slots]) + + # ---- desc tables in ONE triton launch + from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ + base_mode = st.get("base", False) + # AFRAG (prefill): the GEMM reads fragment-major activations from the + # dedicated a1f/a2f buffers (filled by the single-pass triton repack + # below); point the desc 'a' fields there. w4 tables are decode-only, + # so redirecting the shared base in prefill is safe. -+ use_afrag = prefill and _afrag_ok ++ use_afrag = bulk_prefill and _afrag_ok + a1_base = ws["a1f"] if use_afrag else ws["a1"] + a2_base = ws["a2f"] if use_afrag else ws["a2"] + d = ws["desc"] + cap = d.shape[1] + miss_rows = None ++ use_w4s = False + if base_mode: + # BASE cache: 2-bit planes come from the base tier's GPU pool; a live + # pair with a non-resident expert contributes zero and bumps the miss @@ -3107,29 +4618,30 @@ index 000000000..43bad76a2 + # (delta tier over the base cache, gate-filled) coexists on the + # decode path: FP4-resident pairs divert to the w4 tier. + btier = moe_w2_delta._BASE_TIER -+ tier = moe_w2_delta._TIER # FP4 need-pool (None unless opted in) -+ if torch.cuda.is_current_stream_capturing(): ++ tier = moe_w2_delta._TIER # FP4 need-pool (None unless opted in) ++ if capturing: + btier.notify_capture() + if tier is not None: + tier.notify_capture() -+ elif prefill: -+ btier.ensure_resident(layer_key, topk_ids.view(-1)) -+ moe_w2_delta.mark_seen(btier.seen[layer_key], topk_ids.view(-1).long()) ++ elif logical_prefill: ++ btier.ensure_resident(layer_key, topk_ids[token_valid].reshape(-1)) ++ moe_w2_delta.mark_seen(btier.seen[layer_key], topk_ids.long(), token_valid) + if tier is not None: + # the gate's force_promote reads the FP4 tier's own seen scatter -+ moe_w2_delta.mark_seen(tier.seen[layer_key], -+ topk_ids.view(-1).long()) -+ if not prefill: ++ moe_w2_delta.mark_seen(tier.seen[layer_key], topk_ids.long(), token_valid) ++ if not logical_prefill: + # LOOKA/PILOT (router-lookahead): score predictors + write the + # next layer's prediction. Must run BEFORE the route_log + # overwrite below (predictor [0] reads last step's ids from it). + # In-graph safe (persistent buffers, static shapes); no-op + # unless armed. -+ from vllm.model_executor.layers.quantization.utils import ( -+ moe_w2_looka) ++ from vllm.model_executor.layers.quantization.utils import moe_w2_looka ++ + if moe_w2_looka.enabled(): -+ moe_w2_looka.record(layer_key, x, topk_ids, btier.route_log) -+ if not prefill and btier.route_log is not None: ++ moe_w2_looka.record( ++ layer_key, x, topk_ids, btier.route_log, token_valid ++ ) ++ if not logical_prefill and btier.route_log is not None: + # per-(token,layer) routing log for the draft-prefetch predictor: + # a static [n_layers, T_cap, k_cap] buffer the runner reads back + # post-step (~KBs). In-graph safe: fixed shapes per captured @@ -3137,53 +4649,127 @@ index 000000000..43bad76a2 + # count hold stale ids — the host slices by the true T. + _t = min(topk_ids.shape[0], btier.route_log.shape[1]) + _k = min(topk_ids.shape[1], btier.route_log.shape[2]) ++ logged_ids = torch.where( ++ token_valid[:, None], topk_ids, torch.full_like(topk_ids, -1) ++ ) + btier.route_log[layer_key, :_t, :_k].copy_( -+ topk_ids[:_t, :_k], non_blocking=True) ++ logged_ids[:_t, :_k], non_blocking=True ++ ) + if layer_key == 0: + # per-step counter reset, in-graph (layer 0 runs first each step) + btier.miss_count.zero_() + slot_row = btier.slot_table[layer_key] -+ use_fp4 = tier is not None and not prefill ++ use_fp4 = tier is not None and not bulk_prefill + use_w4s_base = use_fp4 and moe_w2_delta.split_enabled() ++ use_w4s = use_w4s_base + if use_w4s_base: + fslot_row = tier.slot_table[layer_key] + d4s = ws["desc4s"] + _desc_build_kernel_base_delta_split[(triton.cdiv(pairs, 256),)]( -+ expert_blocks, num_post, slot_row, fslot_row, -+ btier.miss_count, d, d4s, -+ a1_base.data_ptr(), ws["as1"].data_ptr(), ws["c13"].data_ptr(), -+ a2_base.data_ptr(), ws["as2"].data_ptr(), ws["c2"].data_ptr(), -+ btier.pool.data_ptr(), btier.slot_bytes, -+ st["off_s13"], st["off_c2"], st["off_s2"], -+ tier.pool.data_ptr(), tier.slot_bytes, tier.w13_bytes, -+ st["K13"], (st["K13"] // 128) * 4, 4 * st["K2"], st["K2"], -+ (st["K2"] // 128) * 4, 2 * st["K13"], -+ st["E"], pairs, cap * 6, d4s.shape[1] * 8, mblock, BLOCK=256) ++ expert_blocks, ++ num_post, ++ pair_live, ++ slot_row, ++ fslot_row, ++ btier.miss_count, ++ d, ++ d4s, ++ a1_base.data_ptr(), ++ ws["as1"].data_ptr(), ++ ws["c13"].data_ptr(), ++ a2_base.data_ptr(), ++ ws["as2"].data_ptr(), ++ ws["c2"].data_ptr(), ++ btier.pool.data_ptr(), ++ btier.slot_bytes, ++ st["off_s13"], ++ st["off_c2"], ++ st["off_s2"], ++ tier.pool.data_ptr(), ++ tier.slot_bytes, ++ tier.w13_bytes, ++ st["K13"], ++ (st["K13"] // 128) * 4, ++ 4 * st["K2"], ++ st["K2"], ++ (st["K2"] // 128) * 4, ++ 2 * st["K13"], ++ st["E"], ++ pairs, ++ cap * 6, ++ d4s.shape[1] * 8, ++ mblock, ++ BLOCK=256, ++ ) + elif use_fp4: + fslot_row = tier.slot_table[layer_key] + _desc_build_kernel_base_delta[(triton.cdiv(pairs, 256),)]( -+ expert_blocks, num_post, slot_row, fslot_row, -+ btier.miss_count, d, -+ a1_base.data_ptr(), ws["as1"].data_ptr(), ws["c13"].data_ptr(), -+ a2_base.data_ptr(), ws["as2"].data_ptr(), ws["c2"].data_ptr(), -+ btier.pool.data_ptr(), btier.slot_bytes, -+ st["off_s13"], st["off_c2"], st["off_s2"], -+ tier.pool.data_ptr(), tier.slot_bytes, -+ st["off4_s13"], st["off4_c2"], st["off4_s2"], -+ st["K13"], (st["K13"] // 128) * 4, 4 * st["K2"], st["K2"], -+ (st["K2"] // 128) * 4, 2 * st["K13"], -+ st["E"], pairs, cap * 6, mblock, BLOCK=256) ++ expert_blocks, ++ num_post, ++ pair_live, ++ slot_row, ++ fslot_row, ++ btier.miss_count, ++ d, ++ a1_base.data_ptr(), ++ ws["as1"].data_ptr(), ++ ws["c13"].data_ptr(), ++ a2_base.data_ptr(), ++ ws["as2"].data_ptr(), ++ ws["c2"].data_ptr(), ++ btier.pool.data_ptr(), ++ btier.slot_bytes, ++ st["off_s13"], ++ st["off_c2"], ++ st["off_s2"], ++ tier.pool.data_ptr(), ++ tier.slot_bytes, ++ st["off4_s13"], ++ st["off4_c2"], ++ st["off4_s2"], ++ st["K13"], ++ (st["K13"] // 128) * 4, ++ 4 * st["K2"], ++ st["K2"], ++ (st["K2"] // 128) * 4, ++ 2 * st["K13"], ++ st["E"], ++ pairs, ++ cap * 6, ++ mblock, ++ BLOCK=256, ++ ) + else: + _desc_build_kernel_basecache[(triton.cdiv(pairs, 256),)]( -+ expert_blocks, num_post, slot_row, -+ btier.miss_count, d, -+ a1_base.data_ptr(), ws["as1"].data_ptr(), ws["c13"].data_ptr(), -+ a2_base.data_ptr(), ws["as2"].data_ptr(), ws["c2"].data_ptr(), -+ btier.pool.data_ptr(), btier.slot_bytes, -+ st["off_s13"], st["off_c2"], st["off_s2"], -+ st["K13"], (st["K13"] // 128) * 4, 4 * st["K2"], st["K2"], -+ (st["K2"] // 128) * 4, 2 * st["K13"], -+ st["E"], pairs, cap * 6, mblock, BLOCK=256) ++ expert_blocks, ++ num_post, ++ pair_live, ++ slot_row, ++ btier.miss_count, ++ d, ++ a1_base.data_ptr(), ++ ws["as1"].data_ptr(), ++ ws["c13"].data_ptr(), ++ a2_base.data_ptr(), ++ ws["as2"].data_ptr(), ++ ws["c2"].data_ptr(), ++ btier.pool.data_ptr(), ++ btier.slot_bytes, ++ st["off_s13"], ++ st["off_c2"], ++ st["off_s2"], ++ st["K13"], ++ (st["K13"] // 128) * 4, ++ 4 * st["K2"], ++ st["K2"], ++ (st["K2"] // 128) * 4, ++ 2 * st["K13"], ++ st["E"], ++ pairs, ++ cap * 6, ++ mblock, ++ BLOCK=256, ++ ) + # Miss pairs get scatter weight 0: the GEMMs early-EXIT on m=0 and + # never write their c13/c2 rows, but those workspace rows hold STALE + # values from a previous forward — zeroing the WEIGHT (not the rows) @@ -3192,84 +4778,130 @@ index 000000000..43bad76a2 + # — except under split, where serving needs the BASE slot too (an + # FP4-mapped/base-missing pair contributed zero and must replay). + e_pair = expert_blocks.to(torch.long).clamp_(0, st["E"] - 1) -+ resident = (slot_row[e_pair] >= 0) ++ resident = slot_row[e_pair] >= 0 + if use_fp4 and not use_w4s_base: -+ resident |= (fslot_row[e_pair] >= 0) ++ resident |= fslot_row[e_pair] >= 0 + miss_rows = resident.repeat_interleave(mblock)[:slots] + if not use_fp4: -+ tier = None # downstream w4 launches key off `tier` ++ tier = None # downstream w4 launches key off `tier` + else: -+ tier = moe_w2_delta._TIER # peek only; created by the plane builder -+ if tier is not None and not prefill: -+ if torch.cuda.is_current_stream_capturing(): ++ tier = moe_w2_delta._TIER # peek only; created by the plane builder ++ if tier is not None and not bulk_prefill: ++ if capturing: + tier.notify_capture() + slot_row = tier.slot_table[layer_key] + pool_ptr = tier.pool.data_ptr() -+ moe_w2_delta.mark_seen(tier.seen[layer_key], -+ topk_ids.view(-1).long()) ++ moe_w2_delta.mark_seen(tier.seen[layer_key], topk_ids.long(), token_valid) + else: + if tier is not None: -+ moe_w2_delta.mark_seen(tier.seen[layer_key], -+ topk_ids.view(-1).long()) ++ moe_w2_delta.mark_seen( ++ tier.seen[layer_key], topk_ids.long(), token_valid ++ ) + slot_row = ws["no_slots"] -+ pool_ptr = ws["a1"].data_ptr() # never dereferenced (m4=0) ++ pool_ptr = ws["a1"].data_ptr() # never dereferenced (m4=0) + _desc_build_kernel[(triton.cdiv(pairs, 256),)]( -+ expert_blocks, num_post, slot_row, d, -+ a1_base.data_ptr(), ws["as1"].data_ptr(), ws["c13"].data_ptr(), -+ a2_base.data_ptr(), ws["as2"].data_ptr(), ws["c2"].data_ptr(), -+ st["planes13"].data_ptr(), st["sc13"].data_ptr(), -+ st["planes2"].data_ptr(), st["sc2"].data_ptr(), pool_ptr, -+ st["planes13"].shape[1], st["sc13"].shape[1], -+ st["planes2"].shape[1], st["sc2"].shape[1], ++ expert_blocks, ++ num_post, ++ pair_live, ++ slot_row, ++ d, ++ a1_base.data_ptr(), ++ ws["as1"].data_ptr(), ++ ws["c13"].data_ptr(), ++ a2_base.data_ptr(), ++ ws["as2"].data_ptr(), ++ ws["c2"].data_ptr(), ++ st["planes13"].data_ptr(), ++ st["sc13"].data_ptr(), ++ st["planes2"].data_ptr(), ++ st["sc2"].data_ptr(), ++ pool_ptr, ++ st["planes13"].shape[1], ++ st["sc13"].shape[1], ++ st["planes2"].shape[1], ++ st["sc2"].shape[1], + (tier.slot_bytes if tier is not None else moe_w2_delta.SLOT_BYTES), + (tier.w13_bytes if tier is not None else moe_w2_delta.W13_BYTES), + # row strides (bytes). H-side: a1 fp8 [H], as1 f32 [H/128], c2 bf16 + # [H]. per-rank intermediate side: c13 bf16 [2I], a2 fp8 [I], as2 + # f32 [I/128]. K13 = H, K2 = I -> identical to the old literals on + # DS4 TP1 (H=4096, I=2048); GLM-5.x gets H=6144, TP shards shrink I. -+ st["K13"], (st["K13"] // 128) * 4, 4 * st["K2"], st["K2"], -+ (st["K2"] // 128) * 4, 2 * st["K13"], -+ st["E"], pairs, cap * 6, mblock, BLOCK=256) -+ if tier is not None and not prefill and moe_w2_delta.split_enabled(): ++ st["K13"], ++ (st["K13"] // 128) * 4, ++ 4 * st["K2"], ++ st["K2"], ++ (st["K2"] // 128) * 4, ++ 2 * st["K13"], ++ st["E"], ++ pairs, ++ cap * 6, ++ mblock, ++ BLOCK=256, ++ ) ++ use_w4s = ( ++ tier is not None ++ and not bulk_prefill ++ and not base_mode ++ and moe_w2_delta.split_enabled() ++ ) ++ if use_w4s: + # split-FP4: the extra 8-field tables for moe_w4s_mm (base/bs = + # the resident plane rows, ref = the slot's refinement sections) + d4s = ws["desc4s"] + _desc_build_kernel_w4s[(triton.cdiv(pairs, 256),)]( -+ expert_blocks, num_post, slot_row, d4s, -+ a1_base.data_ptr(), ws["as1"].data_ptr(), -+ ws["c13"].data_ptr(), a2_base.data_ptr(), -+ ws["as2"].data_ptr(), ws["c2"].data_ptr(), -+ st["planes13"].data_ptr(), st["sc13"].data_ptr(), -+ st["planes2"].data_ptr(), st["sc2"].data_ptr(), pool_ptr, -+ st["planes13"].shape[1], st["sc13"].shape[1], -+ st["planes2"].shape[1], st["sc2"].shape[1], -+ tier.slot_bytes, tier.w13_bytes, -+ st["K13"], (st["K13"] // 128) * 4, 4 * st["K2"], st["K2"], -+ (st["K2"] // 128) * 4, 2 * st["K13"], -+ st["E"], pairs, d4s.shape[1] * 8, mblock, BLOCK=256) ++ expert_blocks, ++ num_post, ++ pair_live, ++ slot_row, ++ d4s, ++ a1_base.data_ptr(), ++ ws["as1"].data_ptr(), ++ ws["c13"].data_ptr(), ++ a2_base.data_ptr(), ++ ws["as2"].data_ptr(), ++ ws["c2"].data_ptr(), ++ st["planes13"].data_ptr(), ++ st["sc13"].data_ptr(), ++ st["planes2"].data_ptr(), ++ st["sc2"].data_ptr(), ++ pool_ptr, ++ st["planes13"].shape[1], ++ st["sc13"].shape[1], ++ st["planes2"].shape[1], ++ st["sc2"].shape[1], ++ tier.slot_bytes, ++ tier.w13_bytes, ++ st["K13"], ++ (st["K13"] // 128) * 4, ++ 4 * st["K2"], ++ st["K2"], ++ (st["K2"] // 128) * 4, ++ 2 * st["K13"], ++ st["E"], ++ pairs, ++ d4s.shape[1] * 8, ++ mblock, ++ BLOCK=256, ++ ) + + # ---- w13 GEMMs (both tiers) -> fused silu*up -> quant -> w2 GEMMs + # AFRAG prefill: single-pass triton repack row-major a1/a2 -> fragment-major + # a1f/a2f (desc built against a1f/a2f above) so the GEMM loads each m16k32 + # A-fragment in one LDG.128. Numerics bit-identical to mc4. -+ w2tier = ("w2mc4afrag" if use_afrag else "w2mc4") if prefill else "w2" -+ # AFRAG repacks COMPLETE 16-row tiles. `slots` is moe_align's OVER-ALLOCATED -+ # row count (sorted_ids.numel() = topk*T + E*15), NOT a multiple of 16; the -+ # desc/kernel only ever touch the first `pairs*16` rows (num_post <= pairs*16), -+ # so repack exactly that tile-aligned region. Rows [pairs*16:slots] are unused -+ # filler (never read). Capacity is fine: pairs*16 <= slots <= a1.shape[0]-4. ++ w2tier = ("w2mc4afrag" if use_afrag else "w2mc4") if bulk_prefill else "w2" ++ # AFRAG repacks complete 16-row tiles. moe_align pads sorted_ids to mblock, ++ # so the bulk path's entire slot region is tile-aligned. + if use_afrag: + _afrag_repack(ws["a1"], ws["a1f"], pairs, st["K13"]) + _launch(w2tier, st["K13"], d[0], st["N13"], pairs, stream) + # split-FP4 dispatch: both residency modes fill ws["desc4s"] (classic: + # _desc_build_kernel_w4s against resident planes; base cache: + # _desc_build_kernel_base_delta_split against the coupled base slots) -+ use_w4s = (tier is not None and not prefill -+ and moe_w2_delta.split_enabled()) -+ if tier is not None and not prefill: ++ if tier is not None and not bulk_prefill: + if use_w4s: -+ _launch("w4s", st["K13"], ws["desc4s"][0], st["N13"], pairs, -+ stream) ++ _launch( ++ "w4s", st["K13"], ws["desc4s"][0], st["N13"], pairs, stream ++ ) + else: + _launch("w4", st["K13"], d[2], st["N13"], pairs, stream) + act = ws["act"][:slots] @@ -3279,9 +4911,11 @@ index 000000000..43bad76a2 + if use_afrag: + _afrag_repack(ws["a2"], ws["a2f"], pairs, st["K2"]) + _launch(w2tier, st["K2"], d[1], st["N2"], pairs, stream) -+ if tier is not None and not prefill: ++ if tier is not None and not bulk_prefill: + if use_w4s: -+ _launch("w4s", st["K2"], ws["desc4s"][1], st["N2"], pairs, stream) ++ _launch( ++ "w4s", st["K2"], ws["desc4s"][1], st["N2"], pairs, stream ++ ) + else: + _launch("w4", st["K2"], d[3], st["N2"], pairs, stream) + @@ -3304,9 +4938,8 @@ index 000000000..43bad76a2 + # (their GEMMs early-EXITed) — zero their scatter weight so a miss + # contributes exactly nothing (the replay recomputes them properly). + w = w * miss_rows.to(torch.float32) -+ dump = T * top_k # collision row for filler slots -+ dst = torch.where(valid, sorted_ids, -+ torch.full_like(sorted_ids, dump)).long() ++ dump = T * top_k # collision row for filler slots ++ dst = torch.where(valid, sorted_ids, torch.full_like(sorted_ids, dump)).long() + gath = torch.zeros(dump + 1, H, dtype=torch.float32, device=dev) + gath.index_copy_(0, dst, ws["c2"][:slots].float() * w.unsqueeze(1)) + return gath[:dump].view(T, top_k, H).sum(dim=1).to(x.dtype) @@ -3337,10 +4970,10 @@ index 000000000..43bad76a2 + return enabled() and _ensure_ready() diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_delta.py b/vllm/model_executor/layers/quantization/utils/moe_w2_delta.py new file mode 100644 -index 000000000..f7905fd40 +index 0000000..f9b12c8 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_delta.py -@@ -0,0 +1,1783 @@ +@@ -0,0 +1,1899 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FP4 delta tier for the 1-GPU 2-bit MoE path (quality restoration). @@ -3357,13 +4990,13 @@ index 000000000..f7905fd40 + (w13 8.4 MiB + w2 4.2 MiB packed back-to-back per slot); + - slot table: int32 [layers, 256] on GPU (-1 = base tier), read by the + desc-build kernel inside CUDA graphs; -+ - manager thread: consumes the forward's last-seen expert flags ++ - manager thread: between forwards, consumes the last-seen expert flags + (event-synced D2H), promotes seen-but-uncached experts (H2D on a side + stream, capped per pass), evicts only experts cold for >= 2 passes. -+ Passes are EVENT-DRIVEN: the runner signals step boundaries (wake_all -+ via step_begin / the gate decision) and the manager runs at most one -+ pass per signal, rate-limited by VLLM_MOE_W2_DELTA_TICK_MS; a wall -+ clock timeout only provides liveness for configs that never signal. ++ Passes are EVENT-DRIVEN: the worker opens a manager window after each ++ completed forward and closes it before the next one. The manager runs at ++ most one pass per signal, rate-limited by VLLM_MOE_W2_DELTA_TICK_MS; a ++ wall-clock timeout only provides liveness for configs that never signal. + +Consistency model (deliberate): the table update is racy versus graph +replay — the worst case is one step reading the OLD tier for an expert, @@ -3634,6 +5267,12 @@ index 000000000..f7905fd40 + self._tick = 0 + self._stop = False + self._thread = None ++ # Manager passes may rewrite pool slots that CUDA graphs read without ++ # taking the Python tier lock. Keep the whole pass outside the forward ++ # window: pause_for_forward() acquires this lock before a target starts; ++ # wake() releases it only after every fixed-point/gate replay finishes. ++ self._forward_lock = threading.Lock() ++ self._forward_paused = False + # Step-boundary wakeup (wake()/wake_all()): coalescing event; the + # _wake_driven latch flips the loop from legacy polling to pure + # event cadence the moment the first signal arrives. @@ -3671,8 +5310,8 @@ index 000000000..f7905fd40 + self._kpi_2nd = 0 # extra replays for second-order misses + self._kpi_fp_giveup = 0 # steps that accepted second-order residue + self._kpi_fp_resid = 0 # residual missing pairs in those steps -+ # Slots touched since step_begin(): promoted or hit by any pass of -+ # the CURRENT step. Never evictable (even in the emergency pass) — ++ # Slots touched since step_begin(), or by the current eager prefill ++ # layer. Never evictable (even in the emergency pass) — + # without this, a fixed-point iteration can evict pass-k's fetches + # to serve pass-k+1 (their seen marks are zeroed after each + # snapshot) and ping-pong past the replay cap. @@ -3857,21 +5496,39 @@ index 000000000..f7905fd40 + # runner): keep the legacy fixed-period poll. + self._wake.wait(timeout=_TICK_S) + self._wake.clear() -+ gap = _TICK_S - (time.monotonic() - last) -+ if gap > 0: -+ time.sleep(gap) -+ last = time.monotonic() -+ try: -+ torch.cuda.set_device(self.dev) -+ self._tick_once() -+ except Exception as e: # noqa: BLE001 - never kill serving -+ logger.warning("delta tick failed: %s", e) -+ time.sleep(1.0) ++ # A target/replay sequence holds this lock. If the next target ++ # starts before an event-driven pass does, the pass waits; if the ++ # pass already started, the target waits for it to finish. ++ with self._forward_lock: ++ gap = _TICK_S - (time.monotonic() - last) ++ if gap > 0: ++ time.sleep(gap) ++ last = time.monotonic() ++ try: ++ torch.cuda.set_device(self.dev) ++ self._tick_once() ++ except Exception as e: # noqa: BLE001 - never kill serving ++ logger.warning("delta tick failed: %s", e) ++ time.sleep(1.0) ++ ++ def pause_for_forward(self) -> None: ++ """Block until any manager pass finishes, then exclude new passes. ++ ++ Called on the runner thread before routing marks are cleared and a new ++ target forward starts. Idempotence keeps a failed prior forward safe: ++ the next call remains paused until a later worker completion releases ++ the window. ++ """ ++ if self._forward_paused: ++ return ++ self._forward_lock.acquire() ++ self._forward_paused = True + + def wake(self): -+ """Step-boundary signal (runner thread): run one manager pass as -+ soon as the rate limiter allows. Lock-free and coalescing — a burst -+ of calls between two passes collapses into one.""" ++ """Finish the forward window and signal one manager pass.""" ++ if self._forward_paused: ++ self._forward_paused = False ++ self._forward_lock.release() + self._wake_driven = True + self._wake.set() + @@ -3896,6 +5553,10 @@ index 000000000..f7905fd40 + # token counts for the hit-rate below — read under the snap lock + # so a concurrent snapshot can't swap the values underneath. + cnt_raw = self._seen_host[seen[:, 0], seen[:, 1]] ++ # Keep eviction protection bound to this exact snapshot. Other ++ # snapshot users may overwrite the shared host buffer after the ++ # lock is released. ++ seen_set = {tuple(pair) for pair in seen.tolist()} + if seen.numel() == 0: + return + if _CAPTURE and not self._cap_done: @@ -3953,7 +5614,7 @@ index 000000000..f7905fd40 + for li, ei in cand: + if promoted >= _PROMOTE_PER_TICK: + break -+ slot = self._take_slot() ++ slot = self._take_slot(seen_set) + if slot is None: + break + self._promote(li, ei, slot) @@ -4114,8 +5775,13 @@ index 000000000..f7905fd40 + logger.warning("moe_w2 pool-heat preload failed: %s", e) + return 0 + -+ def _take_slots_batch(self, k: int, emergency: bool = False, -+ min_cold: int = 2) -> list[int]: ++ def _take_slots_batch( ++ self, ++ k: int, ++ emergency: bool = False, ++ min_cold: int = 2, ++ seen_set: set[tuple[int, int]] | None = None, ++ ) -> list[int]: + """Take up to k slots (lock held by caller): free list first, then ONE + vectorized eviction pass over all slots. Replaces the old per-slot + python scan per promotion — O(n_slots) per TAKEN slot — which at GLM @@ -4129,9 +5795,9 @@ index 000000000..f7905fd40 + + Eviction policy is unchanged: least-valuable slot by _POLICY key + (need / freq / lru), restricted to slots whose owner is not active in -+ the current seen window (read directly from the _seen_host snapshot — -+ the call sites always passed a set built from exactly that) and cold -+ >= 2 ticks, so in-flight graph reads never hit a rewritten slot. ++ the caller's immutable seen snapshot (or _seen_host for best-effort ++ callers without one) and cold >= 2 ticks, so in-flight graph reads ++ never hit a rewritten slot. + Victims are unmapped here (graphs stop dispatching w4 before bytes + change); the caller reserves _owner for each returned slot.""" + out: list[int] = [] @@ -4148,14 +5814,22 @@ index 000000000..f7905fd40 + elif self._policy == "freq": + key = self._freq[lic, eic].double() + else: -+ # double, not clone: the tensorized owner-tick is int64 and the -+ # inf sentinel below cannot be represented there (lru is not a -+ # production policy for these tiers; surfaced by unit tests) ++ # LRU ticks are int64, but masked candidates need an infinity ++ # sentinel below. Promote to float so a saturated LRU pool can ++ # exclude in-flight slots instead of raising on the assignment. + key = tk.double() + # Hard exclusions: free markers, owners active in the current seen + # window (their slots may be read by this step's graph/replay), and + # step-pinned slots (touched by any pass of the current step). -+ blocked = (li < 0) | self._seen_host[lic, eic].to(torch.bool) ++ if seen_set is None: ++ active_seen = self._seen_host[lic, eic].to(torch.bool) ++ else: ++ seen_mask = torch.zeros_like(self._seen_host, dtype=torch.bool) ++ if seen_set: ++ seen_li, seen_ei = zip(*seen_set) ++ seen_mask[list(seen_li), list(seen_ei)] = True ++ active_seen = seen_mask[lic, eic] ++ blocked = (li < 0) | active_seen + if self._step_pins: + blocked[list(self._step_pins)] = True + # Residency coupling (split-FP4 over the base cache): never evict a @@ -4212,7 +5886,7 @@ index 000000000..f7905fd40 + + def _take_slot(self, seen_set=None): + """Single-slot wrapper (kept for the unit tests / external callers).""" -+ slots = self._take_slots_batch(1) ++ slots = self._take_slots_batch(1, seen_set=seen_set) + return slots[0] if slots else None + + def _promote(self, li, ei, slot): @@ -4235,16 +5909,22 @@ index 000000000..f7905fd40 + + def step_begin(self) -> None: + """Open a new step's pin scope (runner: before the first miss read; -+ prefill: at each ensure_resident). Slots touched after this call are ++ prefill: before the first ensure_resident). Slots touched after this call are + pinned against eviction until the next step_begin — the fixed-point -+ replay's passes must never cannibalize each other's fetches. -+ -+ Doubles as the manager's step-boundary signal: every live tier -+ (base AND fp4) gets one pass per step instead of a free-running -+ poll.""" ++ replay's passes must never cannibalize each other's fetches. The worker ++ signals the manager only after every replay has completed.""" + with self._lock: + self._step_pins.clear() -+ wake_all() ++ ++ def routing_step_begin(self) -> None: ++ """Discard routing marks left by warmup or the previous target step.""" ++ main = torch.cuda.current_stream(self.dev) ++ with self._snap_lock: ++ self.seen.zero_() ++ # The manager snapshots on this side stream without otherwise ++ # waiting for the runner. Order future snapshots after the zero. ++ with torch.cuda.stream(self._stream): ++ self._stream.wait_stream(main) + + # ---- draft-affinity prefetch (VLLM_MOE_W2_PREFETCH=1) ------------------ + @@ -4395,6 +6075,9 @@ index 000000000..f7905fd40 + seen = self._seen_host.nonzero() + if seen.numel() == 0: + return 0 ++ # Eviction must use the immutable routing snapshot captured above, ++ # not the shared host buffer that another caller can overwrite. ++ seen_set = {tuple(pair) for pair in seen.tolist()} + # Bound the working set to RECENT steps: `seen` otherwise accumulates + # up to 4 manager ticks of routings (the manager zeroes it lazily), so + # on deep/wide models a single fire tried to force-promote every @@ -4453,7 +6136,9 @@ index 000000000..f7905fd40 + # emergency=True: this is the synchronous runner-thread path with + # no forward in flight — leaving a miss UNRESTORED is worse than + # evicting a warm-but-idle slot (see _take_slots_batch). -+ slots = self._take_slots_batch(len(cand), emergency=True) ++ slots = self._take_slots_batch( ++ len(cand), emergency=True, seen_set=seen_set ++ ) + plan = [((li, ei), slot) for (li, ei), slot in zip(cand, slots)] + if not plan: + return 0 @@ -4504,17 +6189,25 @@ index 000000000..f7905fd40 + if layer_key not in self._store: + return 0 + ids = ids.unique().long() ++ ids_cpu = ids.cpu() ++ layer_seen_set = {(layer_key, int(e)) for e in ids_cpu} + mark_seen(self.seen[layer_key], ids.to(self.dev)) -+ # snapshot seen (protects eviction) exactly like force_promote ++ # Wait for prior-layer work before selecting eviction victims. Keep ++ # the device `seen` aggregate intact for telemetry, but protect only ++ # this layer's experts during the layer-at-a-time eager prefill scan; ++ # retaining prior layers in the eviction mask union-saturates the pool. + main = torch.cuda.current_stream(self.dev) + with self._snap_lock: + with torch.cuda.stream(self._stream): + self._stream.wait_stream(main) -+ self._seen_host.copy_(self.seen, non_blocking=True) + ev = torch.cuda.Event() + ev.record(self._stream) + ev.synchronize() + with self._lock: ++ # Each eager layer starts after the prior layer drained. Scope ++ # pins to this layer so the manager cannot recycle a current ++ # hit or freshly loaded slot before its GEMMs consume it. ++ self._step_pins.clear() + slots = self._mirror[layer_key].long()[ids.cpu()] + hit = slots >= 0 + if bool(hit.any()): @@ -4522,13 +6215,16 @@ index 000000000..f7905fd40 + # force_promote: protect them from a racing manager + # eviction while this layer's eager GEMMs read them) + self._owner_tick[slots[hit]] = self._tick ++ self._step_pins.update(int(s) for s in slots[hit].tolist()) + cand = [(layer_key, int(e)) for e in ids.cpu()[~hit].tolist()] + if not cand: + return 0 + # emergency=True: prefill MUST have its whole layer resident — + # an unfetched expert here zeroes contributions for EVERY token + # of the chunk (the existing pool-too-small warning path). -+ slots = self._take_slots_batch(len(cand), emergency=True) ++ slots = self._take_slots_batch( ++ len(cand), emergency=True, seen_set=layer_seen_set ++ ) + plan = [((li, ei), slot) for (li, ei), slot in zip(cand, slots)] + if not plan: + return 0 @@ -4538,6 +6234,7 @@ index 000000000..f7905fd40 + rows = self._store.rows_for([p for p, _ in plan], scan=True) + for ((li, ei), slot), row in zip(plan, rows): + self._own(slot, li, ei) ++ self._step_pins.add(slot) + with torch.cuda.stream(self._stream): + self.pool[slot].copy_(row, non_blocking=True) + with torch.cuda.stream(self._stream): @@ -4836,14 +6533,34 @@ index 000000000..f7905fd40 + self._cap_frames = [] + + -+def mark_seen(seen_row, ids): ++def mark_seen(seen_row, ids, token_valid=None): + """Record routed experts into a layer's seen row from the forward. Token + COUNTS when observability is on (token-weighted hit-rate / capture), else a -+ cheap binary flag. `ids` = flattened topk_ids (int64). Graph-capture-safe.""" ++ cheap binary flag. ``token_valid`` masks padded token rows without changing ++ the captured shape.""" ++ if token_valid is not None: ++ if ids.ndim != 2 or token_valid.ndim != 1: ++ raise RuntimeError( ++ "masked moe_w2 seen recording requires [T, K] ids and [T] validity" ++ ) ++ if ids.shape[0] != token_valid.shape[0]: ++ raise RuntimeError( ++ "moe_w2 token-validity length does not match routed token rows" ++ ) ++ updates = token_valid[:, None].expand_as(ids).reshape(-1).to(seen_row.dtype) ++ ids = ids.reshape(-1) ++ else: ++ updates = None + if _COUNT: -+ seen_row.index_add_(0, ids, torch.ones_like(ids, dtype=seen_row.dtype)) ++ if updates is None: ++ updates = torch.ones_like(ids, dtype=seen_row.dtype) ++ seen_row.index_add_(0, ids, updates) + else: -+ seen_row.index_fill_(0, ids, 1) ++ if updates is None: ++ seen_row.index_fill_(0, ids, 1) ++ else: ++ seen_row.scatter_reduce_(0, ids, updates, ++ reduce="amax", include_self=True) + + +_TIER: DeltaTier | None = None @@ -4869,6 +6586,26 @@ index 000000000..f7905fd40 +_BASE_GB = float(os.getenv("VLLM_MOE_W2_BASE_CACHE_GB", "0")) +_BASE_TIER: DeltaTier | None = None + ++ ++def begin_target_step() -> None: ++ """Exclude manager rewrites, then clear stale marks before the target.""" ++ for tier in (_BASE_TIER, _TIER): ++ if tier is not None: ++ tier.pause_for_forward() ++ tier.routing_step_begin() ++ ++ ++def begin_replay_step() -> None: ++ """Open both tiers' pin scopes after the target forward, before replay.""" ++ for tier in (_BASE_TIER, _TIER): ++ if tier is not None: ++ tier.step_begin() ++ ++ ++def finish_forward_step() -> None: ++ """Open the between-forward manager window and signal one pass.""" ++ wake_all() ++ +# Miss tolerance: a decode step with <= TOL missing routed (layer, expert) +# pairs keeps its logits (the missing pairs contributed zero) instead of +# replaying the graph. Rationale: at 99.9% token hit-rate a 600-pair step @@ -5007,6 +6744,25 @@ index 000000000..f7905fd40 + return _BASE_GB > 0 + + ++def _arm_split_coupling() -> None: ++ """Couple split-FP4 refinement residency to the base tier. ++ ++ Fresh requant builds create the FP4 tier before ``_finish_layer`` creates ++ the base tier, while pack-hit restores create them in the opposite order. ++ Invoke this after either singleton lookup so coupling cannot depend on ++ which path constructed its tier first. ++ """ ++ if not split_enabled() or _BASE_TIER is None or _TIER is None: ++ return ++ if _BASE_TIER._coupled_fp4 is _TIER: ++ return ++ _BASE_TIER._coupled_fp4 = _TIER ++ logger.info( ++ "moe_w2 delta: split-FP4 residency coupling armed " ++ "(base evictions exclude FP4-mapped experts)" ++ ) ++ ++ +def spec_suppressed() -> bool: + """Spec-guard latch (VLLM_MOE_W2_SPEC_GUARD): True while the base pool + is too cold for speculation to pay — the runner then skips scheduling @@ -5016,10 +6772,7 @@ index 000000000..f7905fd40 + + +def wake_all() -> None: -+ """Step-boundary broadcast: nudge every live tier's manager (base + -+ fp4/delta) to run one pass. Called once per decode step from -+ step_begin (base-cache configs) and the gate decision (delta-only -+ configs); lock-free, coalescing, safe before tiers exist.""" ++ """Post-forward broadcast: release and nudge each live tier manager.""" + t = _BASE_TIER + if t is not None: + t.wake() @@ -5073,6 +6826,7 @@ index 000000000..f7905fd40 + _BASE_TIER.n_slots, _BASE_TIER.slot_bytes / 2**20, + _BASE_TIER.n_slots * _BASE_TIER.slot_bytes / 2**30, + cov, n_layers * n_experts) ++ _arm_split_coupling() + return _BASE_TIER + + @@ -5109,12 +6863,6 @@ index 000000000..f7905fd40 + w2_bytes=W2_BYTES if w2_bytes is None else w2_bytes, + policy=policy, tag=fp4_tag if base_enabled() else "delta", + host_pinned=not base_enabled()) -+ if base_enabled() and split_enabled() and _BASE_TIER is not None: -+ # residency coupling: the base tier must not evict slots the -+ # FP4 tier's refinement rows are mapped against -+ _BASE_TIER._coupled_fp4 = _TIER -+ logger.info("moe_w2 delta: split-FP4 residency coupling armed " -+ "(base evictions exclude FP4-mapped experts)") + # Start the background manager as soon as the tier exists. It idles until + # experts are actually routed (seen empty -> early return) and only + # promotes layers whose host planes are already staged, so an early start @@ -5123,13 +6871,14 @@ index 000000000..f7905fd40 + # "start on the last layer built" trigger never ran and the tier sat + # inactive (pool allocated but no promotions). + _TIER.start() ++ _arm_split_coupling() + return _TIER diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_gate.py b/vllm/model_executor/layers/quantization/utils/moe_w2_gate.py new file mode 100644 -index 000000000..fcf18c484 +index 0000000..8891cf5 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_gate.py -@@ -0,0 +1,188 @@ +@@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Confidence-gated FP4 re-forward for the 2-bit MoE path (directive 2 / Step B). @@ -5262,12 +7011,6 @@ index 000000000..fcf18c484 + """ + global _n_steps, _n_fired + _n_steps += 1 -+ # Step-boundary signal for the tier managers: this is the one gate call -+ # guaranteed once per decode step, so delta-only configs (no base tier, -+ # hence no runner step_begin) still get event-driven manager passes -+ # instead of the legacy free-running poll. -+ from vllm.model_executor.layers.quantization.utils import moe_w2_delta -+ moe_w2_delta.wake_all() + if logits is None or logits.numel() == 0: + return False + if logits.dim() == 1: @@ -5320,10 +7063,10 @@ index 000000000..fcf18c484 + fire_rate=(_n_fired / _n_steps if _n_steps else 0.0)) diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_looka.py b/vllm/model_executor/layers/quantization/utils/moe_w2_looka.py new file mode 100644 -index 000000000..f5de680cf +index 0000000..6e1c3ea --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_looka.py -@@ -0,0 +1,256 @@ +@@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Router-lookahead measurement (LOOKA) + prefetch (PILOT) for the moe_w2 @@ -5470,7 +7213,8 @@ index 000000000..f5de680cf + + +def record(layer_key: int, x: torch.Tensor, topk_ids: torch.Tensor, -+ route_log: torch.Tensor | None) -> None: ++ route_log: torch.Tensor | None, ++ token_valid: torch.Tensor) -> None: + """In-graph hook, called from the moe_w2 decode forward of every BASE + layer BEFORE the route_log is overwritten. Pure tensor ops on persistent + buffers (capture-safe); `layer_key` is a python int, so the branches @@ -5482,13 +7226,14 @@ index 000000000..f5de680cf + return + k_true = topk_ids.shape[1] + true = topk_ids[:, :k_true].int() ++ valid_routes = token_valid[:, None] + # [0] previous decode step, same layer (the affinity-class baseline). + # route_log still holds LAST step's ids for this layer here. + if route_log is not None: + prev = route_log[layer_key, :T, :k_true] + m0 = (true.unsqueeze(2) == prev.unsqueeze(1)).any(dim=2) -+ _hit[0] += m0.sum() -+ _tot[0] += true.numel() ++ _hit[0] += (m0 & valid_routes).sum() ++ _tot[0] += token_valid.sum() * k_true + # [1] router-lookahead: the prediction targeting THIS layer, written + # into _pred_buf by the previous layer's record() call (same step; the + # keys run in order inside one forward). Exists iff this layer's gate @@ -5496,8 +7241,8 @@ index 000000000..f5de680cf + if layer_key > 0 and layer_key in _gate_w: + pred = _pred_buf[:T] + m1 = (true.unsqueeze(2) == pred.unsqueeze(1)).any(dim=2) -+ _hit[1] += m1.sum() -+ _tot[1] += true.numel() ++ _hit[1] += (m1 & valid_routes).sum() ++ _tot[1] += token_valid.sum() * k_true + # predict layer_key+1's routing from THIS layer's expert input + w = _gate_w.get(layer_key + 1) + if w is not None: @@ -5507,6 +7252,8 @@ index 000000000..f5de680cf + if b is not None: + scores = scores + b + pred_ids = torch.topk(scores, _PILOT_K, dim=-1).indices.int() ++ pred_ids = torch.where(token_valid[:, None], pred_ids, ++ torch.full_like(pred_ids, -1)) + _pred_buf[:T].copy_(pred_ids) + if _pilot_log is not None: + _pilot_log[layer_key + 1, :T].copy_(pred_ids) @@ -5582,7 +7329,7 @@ index 000000000..f5de680cf + f"(top-{_PILOT_K}, n={tot[1]})") diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_planes.py b/vllm/model_executor/layers/quantization/utils/moe_w2_planes.py new file mode 100644 -index 000000000..634454b8d +index 0000000..634454b --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_planes.py @@ -0,0 +1,286 @@ @@ -5874,10 +7621,10 @@ index 000000000..634454b8d + return vals * s diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_planes_cache.py b/vllm/model_executor/layers/quantization/utils/moe_w2_planes_cache.py new file mode 100644 -index 000000000..752f88ded +index 0000000..9bffa48 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_planes_cache.py -@@ -0,0 +1,226 @@ +@@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Disk cache for built 2-bit expert planes (+ optional FP4 delta planes). @@ -5912,6 +7659,7 @@ index 000000000..752f88ded +import queue +import re +import threading ++from contextlib import suppress + +import numpy as np +import torch @@ -5948,15 +7696,33 @@ index 000000000..752f88ded + + +def _ckpt_id() -> str: -+ """Identity of the checkpoint the planes derive from: model path + -+ sha1 of the safetensors index (covers shard layout and tensor set).""" ++ """Cheap identity of the checkpoint the planes derive from. ++ ++ Hash the canonical model path and index contents, then the name, size and ++ nanosecond mtime of every referenced shard. This catches an in-place ++ checkpoint replacement without hashing a DS4-sized checkpoint at boot. ++ """ + from vllm.config import get_current_vllm_config + model = get_current_vllm_config().model_config.model + h = hashlib.sha1(model.encode()) + idx = os.path.join(model, "model.safetensors.index.json") ++ shards: set[str] = set() + if os.path.exists(idx): + with open(idx, "rb") as f: -+ h.update(f.read()) ++ raw = f.read() ++ h.update(raw) ++ with suppress(AttributeError, TypeError, ValueError, ++ json.JSONDecodeError): ++ shards.update(json.loads(raw).get("weight_map", {}).values()) ++ if not shards and os.path.isdir(model): ++ shards.update(name for name in os.listdir(model) ++ if name.endswith(".safetensors")) ++ for name in sorted(shards): ++ path = os.path.join(model, name) ++ stat = os.stat(path) ++ h.update(name.encode()) ++ h.update(str(stat.st_size).encode()) ++ h.update(str(stat.st_mtime_ns).encode()) + return h.hexdigest() + + @@ -5975,6 +7741,11 @@ index 000000000..752f88ded + ) + + ++def cache_identity() -> dict: ++ """Public deterministic identity shared by every derived W2 cache.""" ++ return _meta() ++ ++ +def _rank_dir() -> str: + world, rank = _tp_ids() + return os.path.join(os.environ["VLLM_MOE_W2_PLANES_CACHE"], @@ -6007,7 +7778,11 @@ index 000000000..752f88ded + try: + d = _rank_dir() + mp = os.path.join(d, "meta.json") -+ if not os.path.exists(mp) or json.load(open(mp)) != _meta(): ++ if not os.path.exists(mp): ++ return False ++ with open(mp) as f: ++ meta = json.load(f) ++ if meta != _meta(): + return False + for part, nbytes in sizes.items(): + p = os.path.join(d, f"layer{layer_idx}.{part}.bin") @@ -6028,7 +7803,9 @@ index 000000000..752f88ded + mp = os.path.join(d, "meta.json") + if not os.path.exists(mp): + return None -+ if json.load(open(mp)) != _meta(): ++ with open(mp) as f: ++ meta = json.load(f) ++ if meta != _meta(): + _mark_broken(f"meta mismatch in {d} (stale cache?) — rebuilding") + return None + out = {} @@ -6083,10 +7860,13 @@ index 000000000..752f88ded + os.makedirs(d, exist_ok=True) + if not _meta_written: + mp = os.path.join(d, "meta.json") -+ if os.path.exists(mp) and json.load(open(mp)) != _meta(): -+ # stale cache from another checkpoint/config: start over -+ for f in os.listdir(d): -+ os.unlink(os.path.join(d, f)) ++ if os.path.exists(mp): ++ with open(mp) as f: ++ old_meta = json.load(f) ++ if old_meta != _meta(): ++ # stale cache from another checkpoint/config: start over ++ for name in os.listdir(d): ++ os.unlink(os.path.join(d, name)) + with open(mp + ".tmp", "w") as f: + json.dump(_meta(), f) + os.replace(mp + ".tmp", mp) @@ -6106,67 +7886,628 @@ index 000000000..752f88ded + _mark_broken(f"store failed: {e}") diff --git a/vllm/model_executor/layers/quantization/utils/moe_w2_store.py b/vllm/model_executor/layers/quantization/utils/moe_w2_store.py new file mode 100644 -index 000000000..c0b937425 +index 0000000..a64cb58 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/moe_w2_store.py -@@ -0,0 +1,656 @@ +@@ -0,0 +1,1405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Host-side expert stores for the 2-bit MoE tiers (moe_w2_delta.DeltaTier). + -+Three backends behind one tiny interface: ++Three backends behind one tiny interface: ++ ++ - PinnedHostStore: today's behaviour — per-layer [E, slot_bytes] host ++ tensors (pinned or pageable), rows handed to cudaMemcpyAsync directly. ++ Default; byte-identical to the pre-store code path. ++ - MmapPackStore (VLLM_MOE_W2_STORE_DIR=): rows live in a per-rank ++ PACK FILE on disk; reads are buffered preads -> pinned stage -> H2D. ++ The kernel page cache is the RAM tier (LRU for free), so host RAM holds ++ only the hot part of the base instead of the whole 73-190 GiB store — ++ and the pack doubles as a persistent quantization cache across boots (a ++ layer already in the pack skips D2H staging entirely). ++ - TieredPackStore (additionally VLLM_MOE_W2_BASE_RAM_GB=, base ++ tier only): a PINNED arena of the N most-recently-used rows over the ++ same pack. An arena hit is a zero-copy pinned view (H2D DMAs straight ++ from it — the exact PinnedHostStore hot path, no syscall, no memcpy); ++ a miss preadv's the row into the arena slot (the arena IS the bounce ++ buffer), buffered by default so the page cache serves as an ++ opportunistic L3 under the arena (VLLM_MOE_W2_TIER_DIRECT=1 for ++ O_DIRECT misses). The arena itself can never be reclaimed under ++ memory pressure — the hot fetch set stays RAM-fast even on hosts with ++ zero spare page cache. Policy is recency (LRU): freq-pinning lost on ++ live GLM traces (routing too flat — see GLM_RAMTIER_FINDINGS). ++ ++Pack layout (per tier tag, per TP rank): ++ /.rankof.pack raw rows, offset = (li*E+ei)*stride ++ /.rankof.json sidecar: shapes + layers written ++ ++`stride` is slot_bytes rounded up to 4 KiB so the SAME pack serves the ++O_DIRECT reader without a repack (O_DIRECT needs 4K-aligned offset/length/ ++buffer; the pinned arena is page-aligned and stride-strided, so every row ++satisfies all three). Rows of layers not listed in the sidecar are holes ++(sparse file) and are never read. ++ ++Concurrency: every read/write caller already holds the owning DeltaTier's ++lock (manager tick, force_promote, ensure_resident are serialized there), ++so the shared pinned stage buffer / arena bookkeeping need no lock of ++their own. ++""" ++ ++import json ++import os ++import threading ++import time ++from collections import deque ++from concurrent.futures import ThreadPoolExecutor ++from contextlib import suppress ++ ++import torch ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++_ALIGN = 4096 ++_PACK_VERSION = 2 ++_GIB = 1 << 30 ++_CACHE_CONTROL_MODES = {"required", "best-effort", "off"} ++_pending_checkpoint_drops: set[str] = set() ++_pending_checkpoint_lock = threading.Lock() ++ ++ ++def _env_true(name: str) -> bool: ++ return os.getenv(name, "").strip().lower() in ("1", "true", "yes", "on") ++ ++ ++def checkpoint_cache_safety_enabled() -> bool: ++ """Whether checkpoint reads belong to a live W2 pack-store load. ++ ++ Keep this predicate narrow: weight_utils calls the hooks for every ++ safetensors model, while only W2 pack builds need the fail-closed cache ++ discipline implemented here. ++ """ ++ return _env_true("VLLM_MOE_W2") and bool( ++ os.getenv("VLLM_MOE_W2_STORE_DIR", "").strip() ++ ) ++ ++ ++def _cache_control_mode() -> str: ++ mode = os.getenv("VLLM_MOE_W2_CACHE_CONTROL", "required").strip().lower() ++ if mode not in _CACHE_CONTROL_MODES: ++ raise ValueError( ++ "VLLM_MOE_W2_CACHE_CONTROL must be one of " ++ f"{sorted(_CACHE_CONTROL_MODES)}, got {mode!r}" ++ ) ++ if mode == "off": ++ logger.warning_once( ++ "moe_w2 SAFETY OVERRIDE: page-cache eviction is OFF; the " ++ "MemAvailable guard remains armed, but a cold pack rebuild may " ++ "abort before completion" ++ ) ++ return mode ++ ++ ++def _cache_control_failure(message: str, error: Exception | None = None) -> bool: ++ mode = _cache_control_mode() ++ detail = f" ({error})" if error is not None else "" ++ if mode == "required": ++ raise RuntimeError(message + detail) from error ++ logger.warning_once("%s%s", message, detail) ++ return False ++ ++ ++def _require_cache_control() -> bool: ++ """Fail before a cold build when DONTNEED cannot be issued. ++ ++ `best-effort` and `off` are explicit operator overrides. The default is ++ deliberately fail-closed because the fallback has hard-wedged 128 GiB ++ hosts while restaging DS4-class checkpoints. ++ """ ++ mode = _cache_control_mode() ++ if mode == "off": ++ return False ++ if not hasattr(os, "posix_fadvise") or not hasattr(os, "POSIX_FADV_DONTNEED"): ++ return _cache_control_failure( ++ "moe_w2 pack build requires POSIX_FADV_DONTNEED; set " ++ "VLLM_MOE_W2_CACHE_CONTROL=best-effort or off only as an " ++ "explicit unsafe override" ++ ) ++ return True ++ ++ ++def _fadvise_dontneed(fd: int, offset: int, length: int, label: str) -> bool: ++ """Discard clean file-backed pages for one completed staging extent.""" ++ if not _require_cache_control(): ++ return False ++ try: ++ os.posix_fadvise(fd, offset, length, os.POSIX_FADV_DONTNEED) ++ return True ++ except OSError as e: ++ return _cache_control_failure( ++ f"moe_w2 could not evict page cache for {label}", e ++ ) ++ ++ ++def _drop_path_page_cache(path: str, label: str) -> bool: ++ flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) ++ try: ++ fd = os.open(path, flags) ++ except OSError as e: ++ return _cache_control_failure( ++ f"moe_w2 could not open {label} for page-cache eviction", e ++ ) ++ try: ++ return _fadvise_dontneed(fd, 0, 0, label) ++ finally: ++ os.close(fd) ++ ++ ++def _configured_gib(name: str, default: float) -> int: ++ raw = os.getenv(name, str(default)).strip() ++ try: ++ value = float(raw) ++ except ValueError as e: ++ raise ValueError(f"{name} must be a non-negative GiB value, got {raw!r}") from e ++ if value < 0: ++ raise ValueError(f"{name} must be non-negative, got {value}") ++ return int(value * _GIB) ++ ++ ++def _mem_available_bytes() -> int: ++ try: ++ with open("/proc/meminfo") as f: ++ for line in f: ++ if line.startswith("MemAvailable:"): ++ return int(line.split()[1]) * 1024 ++ except OSError as e: ++ raise RuntimeError("moe_w2 memory preflight cannot read /proc/meminfo") from e ++ raise RuntimeError("moe_w2 memory preflight found no MemAvailable value") ++ ++ ++def _read_text(path: str) -> str: ++ with open(path) as f: ++ return f.read().strip() ++ ++ ++def _read_kv_ints(path: str) -> dict[str, int]: ++ try: ++ return { ++ parts[0]: int(parts[1]) ++ for line in _read_text(path).splitlines() ++ if len(parts := line.split()) == 2 ++ } ++ except (OSError, ValueError): ++ return {} ++ ++ ++def _finite_limit(raw: str) -> int | None: ++ if raw == "max": ++ return None ++ value = int(raw) ++ return None if value >= 1 << 60 else value ++ ++ ++def _fmt_gib(value: int | None) -> str: ++ return "n/a" if value is None else f"{value / _GIB:.1f} GiB" ++ ++ ++def _active_cgroup_v2_dirs() -> list[str]: ++ root = "/sys/fs/cgroup" ++ candidates: list[str] = [] ++ try: ++ for line in _read_text("/proc/self/cgroup").splitlines(): ++ fields = line.split(":", 2) ++ if len(fields) == 3 and fields[0] == "0": ++ rel = fields[2].lstrip("/") ++ if rel: ++ candidates.append(os.path.join(root, rel)) ++ break ++ except OSError: ++ pass ++ # A cgroup namespace commonly mounts the process's cgroup as root. ++ candidates.append(root) ++ leaf = next( ++ (p for p in candidates if os.path.exists(os.path.join(p, "memory.current"))), ++ None, ++ ) ++ if leaf is None: ++ return [] ++ dirs = [] ++ current = os.path.realpath(leaf) ++ root_real = os.path.realpath(root) ++ while current.startswith(root_real): ++ if os.path.exists(os.path.join(current, "memory.current")): ++ dirs.append(current) ++ if current == root_real: ++ break ++ parent = os.path.dirname(current) ++ if parent == current: ++ break ++ current = parent ++ return dirs ++ ++ ++def _active_cgroup_v1_dirs() -> list[str]: ++ root = "/sys/fs/cgroup/memory" ++ candidates: list[str] = [] ++ try: ++ for line in _read_text("/proc/self/cgroup").splitlines(): ++ fields = line.split(":", 2) ++ if len(fields) != 3 or "memory" not in fields[1].split(","): ++ continue ++ rel = fields[2].lstrip("/") ++ if rel: ++ candidates.append(os.path.join(root, rel)) ++ break ++ except OSError: ++ pass ++ # A cgroup namespace can expose the process leaf as the mount root. ++ candidates.append(root) ++ leaf = next( ++ ( ++ p ++ for p in candidates ++ if os.path.exists(os.path.join(p, "memory.usage_in_bytes")) ++ ), ++ None, ++ ) ++ if leaf is None: ++ return [] ++ dirs = [] ++ current = os.path.realpath(leaf) ++ root_real = os.path.realpath(root) ++ while current.startswith(root_real): ++ if os.path.exists(os.path.join(current, "memory.usage_in_bytes")): ++ dirs.append(current) ++ if current == root_real: ++ break ++ parent = os.path.dirname(current) ++ if parent == current: ++ break ++ current = parent ++ return dirs ++ ++ ++def _cgroup_memory_status() -> dict: ++ """Resolve hard max and soft high headroom independently.""" ++ dirs = _active_cgroup_v2_dirs() ++ if dirs: ++ max_headrooms: list[int] = [] ++ high_headrooms: list[int] = [] ++ limits: list[tuple[str, str, int]] = [] ++ try: ++ for directory in dirs: ++ current = int(_read_text(os.path.join(directory, "memory.current"))) ++ for filename in ("memory.high", "memory.max"): ++ path = os.path.join(directory, filename) ++ if not os.path.exists(path): ++ continue ++ limit = _finite_limit(_read_text(path)) ++ if limit is not None: ++ headroom = limit - current ++ if filename == "memory.max": ++ max_headrooms.append(max(0, headroom)) ++ else: ++ high_headrooms.append(headroom) ++ limits.append((directory, filename, limit)) ++ leaf = dirs[0] ++ stat = _read_kv_ints(os.path.join(leaf, "memory.stat")) ++ events = _read_kv_ints(os.path.join(leaf, "memory.events")) ++ swap_current = None ++ swap_limit = None ++ swap_current_path = os.path.join(leaf, "memory.swap.current") ++ swap_max_path = os.path.join(leaf, "memory.swap.max") ++ if os.path.exists(swap_current_path): ++ swap_current = int(_read_text(swap_current_path)) ++ if os.path.exists(swap_max_path): ++ swap_limit = _finite_limit(_read_text(swap_max_path)) ++ return dict( ++ known=True, ++ version=2, ++ path=leaf, ++ limited=bool(max_headrooms), ++ max_available=min(max_headrooms) if max_headrooms else None, ++ high_available=min(high_headrooms) if high_headrooms else None, ++ current=int(_read_text(os.path.join(leaf, "memory.current"))), ++ limits=limits, ++ file=stat.get("file"), ++ file_mapped=stat.get("file_mapped"), ++ anon=stat.get("anon"), ++ swap_current=swap_current, ++ swap_limit=swap_limit, ++ events=events, ++ ) ++ except (OSError, ValueError) as e: ++ return dict(known=False, version=2, path=dirs[0], error=str(e)) ++ ++ # cgroup v1 fallback. Resolve the process leaf and every visible ancestor; ++ # the mount root is often unlimited while the Docker/systemd leaf is not. ++ # `memory.limit_in_bytes` uses a huge sentinel for unlimited. ++ dirs = _active_cgroup_v1_dirs() ++ if dirs: ++ try: ++ headrooms: list[int] = [] ++ limits: list[tuple[str, str, int]] = [] ++ for directory in dirs: ++ current = int( ++ _read_text(os.path.join(directory, "memory.usage_in_bytes")) ++ ) ++ limit = _finite_limit( ++ _read_text(os.path.join(directory, "memory.limit_in_bytes")) ++ ) ++ if limit is not None: ++ headrooms.append(max(0, limit - current)) ++ limits.append((directory, "memory.limit_in_bytes", limit)) ++ leaf = dirs[0] ++ stat = _read_kv_ints(os.path.join(leaf, "memory.stat")) ++ leaf_current = int(_read_text(os.path.join(leaf, "memory.usage_in_bytes"))) ++ failcnt = None ++ failcnt_path = os.path.join(leaf, "memory.failcnt") ++ if os.path.exists(failcnt_path): ++ failcnt = int(_read_text(failcnt_path)) ++ memsw_current = None ++ memsw_limit = None ++ memsw_current_path = os.path.join(leaf, "memory.memsw.usage_in_bytes") ++ memsw_limit_path = os.path.join(leaf, "memory.memsw.limit_in_bytes") ++ if os.path.exists(memsw_current_path): ++ memsw_current = int(_read_text(memsw_current_path)) ++ if os.path.exists(memsw_limit_path): ++ memsw_limit = _finite_limit(_read_text(memsw_limit_path)) ++ return dict( ++ known=True, ++ version=1, ++ path=leaf, ++ limited=bool(headrooms), ++ max_available=min(headrooms) if headrooms else None, ++ high_available=None, ++ current=leaf_current, ++ limits=limits, ++ file=stat.get("cache"), ++ file_mapped=stat.get("mapped_file"), ++ anon=stat.get("rss"), ++ swap_current=( ++ None ++ if memsw_current is None ++ else max(0, memsw_current - leaf_current) ++ ), ++ swap_limit=memsw_limit, ++ events={} if failcnt is None else {"failcnt": failcnt}, ++ ) ++ except (OSError, ValueError) as e: ++ return dict(known=False, version=1, path=dirs[0], error=str(e)) ++ return dict( ++ known=False, ++ version=None, ++ path=None, ++ error="no readable cgroup memory controller", ++ ) ++ ++ ++def _memory_preflight(label: str, transient_bytes: int = 0) -> dict: ++ """Refuse an allocation/read before it can cross the safety floor. ++ ++ `transient_bytes` is the largest additional anonymous or page-cache ++ extent the next indivisible operation can create. Checks occur before ++ every checkpoint shard, pack-layer write, and pinned-arena allocation, ++ so peak cache growth is bounded by one shard plus one layer rather than ++ the whole checkpoint plus pack. ++ """ ++ transient_bytes = max(0, int(transient_bytes)) ++ host_reserve = _configured_gib("VLLM_MOE_W2_MIN_MEM_AVAILABLE_GB", 16) ++ cgroup_reserve = _configured_gib("VLLM_MOE_W2_MIN_CGROUP_HEADROOM_GB", 4) ++ available = _mem_available_bytes() ++ host_need = host_reserve + transient_bytes ++ if host_reserve and available < host_need: ++ raise RuntimeError( ++ f"moe_w2 memory preflight REFUSED {label}: MemAvailable " ++ f"{available / _GIB:.1f} GiB < required " ++ f"{host_need / _GIB:.1f} GiB (reserve " ++ f"{host_reserve / _GIB:.1f} + transient " ++ f"{transient_bytes / _GIB:.1f}); no checkpoint/pack I/O began" ++ ) ++ cgroup = _cgroup_memory_status() ++ if cgroup_reserve and not cgroup.get("known"): ++ raise RuntimeError( ++ f"moe_w2 memory preflight REFUSED {label}: cannot determine " ++ f"the active cgroup memory.max limit/headroom " ++ f"({cgroup.get('error', 'unknown error')}); set " ++ "VLLM_MOE_W2_MIN_CGROUP_HEADROOM_GB=0 only as an explicit " ++ "unsafe override" ++ ) ++ cgroup_max_available = cgroup.get("max_available") ++ cgroup_high_available = cgroup.get("high_available") ++ cgroup_need = cgroup_reserve + transient_bytes ++ if ( ++ cgroup_reserve ++ and cgroup.get("limited") ++ and cgroup_max_available is not None ++ and cgroup_max_available < cgroup_need ++ ): ++ raise RuntimeError( ++ f"moe_w2 memory preflight REFUSED {label}: cgroup memory.max " ++ f"headroom {cgroup_max_available / _GIB:.1f} GiB < required " ++ f"{cgroup_need / _GIB:.1f} GiB (reserve " ++ f"{cgroup_reserve / _GIB:.1f} + transient " ++ f"{transient_bytes / _GIB:.1f}); no checkpoint/pack I/O began" ++ ) ++ max_headroom = ( ++ "unknown" ++ if not cgroup.get("known") ++ else "unlimited" ++ if not cgroup.get("limited") ++ else f"{cgroup_max_available / _GIB:.1f} GiB" ++ ) ++ high_headroom = ( ++ "unknown" ++ if not cgroup.get("known") ++ else "unlimited" ++ if cgroup.get("version") == 2 and cgroup_high_available is None ++ else "n/a" ++ if cgroup_high_available is None ++ else f"{cgroup_high_available / _GIB:.1f} GiB" ++ ) ++ logger.info( ++ "moe_w2 safety preflight[%s]: MemAvailable %.1f GiB, cgroup " ++ "memory.max headroom %s, memory.high headroom %s, transient %.1f GiB, " ++ "floors host %.1f / cgroup %.1f GiB", ++ label, ++ available / _GIB, ++ max_headroom, ++ high_headroom, ++ transient_bytes / _GIB, ++ host_reserve / _GIB, ++ cgroup_reserve / _GIB, ++ ) ++ logger.info( ++ "moe_w2 cgroup[%s]: current %s, anon %s, file %s, mapped %s, " ++ "swap %s/%s, events %s", ++ cgroup.get("path"), ++ _fmt_gib(cgroup.get("current")), ++ _fmt_gib(cgroup.get("anon")), ++ _fmt_gib(cgroup.get("file")), ++ _fmt_gib(cgroup.get("file_mapped")), ++ _fmt_gib(cgroup.get("swap_current")), ++ _fmt_gib(cgroup.get("swap_limit")), ++ cgroup.get("events", {}), ++ ) ++ return dict( ++ available=available, ++ cgroup_max_available=cgroup_max_available, ++ cgroup_high_available=cgroup_high_available, ++ cgroup=cgroup, ++ transient=transient_bytes, ++ host_reserve=host_reserve, ++ cgroup_reserve=cgroup_reserve, ++ ) ++ + -+ - PinnedHostStore: today's behaviour — per-layer [E, slot_bytes] host -+ tensors (pinned or pageable), rows handed to cudaMemcpyAsync directly. -+ Default; byte-identical to the pre-store code path. -+ - MmapPackStore (VLLM_MOE_W2_STORE_DIR=): rows live in a per-rank -+ PACK FILE on disk; reads are buffered preads -> pinned stage -> H2D. -+ The kernel page cache is the RAM tier (LRU for free), so host RAM holds -+ only the hot part of the base instead of the whole 73-190 GiB store — -+ and the pack doubles as a persistent quantization cache across boots (a -+ layer already in the pack skips D2H staging entirely). -+ - TieredPackStore (additionally VLLM_MOE_W2_BASE_RAM_GB=, base -+ tier only): a PINNED arena of the N most-recently-used rows over the -+ same pack. An arena hit is a zero-copy pinned view (H2D DMAs straight -+ from it — the exact PinnedHostStore hot path, no syscall, no memcpy); -+ a miss preadv's the row into the arena slot (the arena IS the bounce -+ buffer), buffered by default so the page cache serves as an -+ opportunistic L3 under the arena (VLLM_MOE_W2_TIER_DIRECT=1 for -+ O_DIRECT misses). The arena itself can never be reclaimed under -+ memory pressure — the hot fetch set stays RAM-fast even on hosts with -+ zero spare page cache. Policy is recency (LRU): freq-pinning lost on -+ live GLM traces (routing too flat — see GLM_RAMTIER_FINDINGS). ++def checkpoint_file_preflight(path: str, extra_bytes: int = 0) -> None: ++ """Guard one safetensors shard before its mmap/read can populate cache.""" ++ if not checkpoint_cache_safety_enabled(): ++ return ++ _require_cache_control() ++ # The previous shard's generator finally can run while its consumer ++ # still owns the last yielded tensor or safetensors handle. Its immediate ++ # DONTNEED is therefore deliberately queued. By the time the consumer ++ # asks for the next shard, the W2 iterator has cloned its loop-variable ++ # tail and more mappings are usually releasable; model-specific aliases ++ # can lag longer, so retry every queued path at every boundary. Keep paths ++ # queued for the final post-consumer retry as a second fail-closed guard. ++ checkpoint_retry_pending() ++ try: ++ file_bytes = os.path.getsize(path) ++ except OSError as e: ++ raise RuntimeError(f"moe_w2 cannot stat checkpoint shard {path!r}") from e ++ _memory_preflight( ++ f"checkpoint shard {os.path.basename(path)}", ++ file_bytes + max(0, int(extra_bytes)), ++ ) + -+Pack layout (per tier tag, per TP rank): -+ /.rankof.pack raw rows, offset = (li*E+ei)*stride -+ /.rankof.json sidecar: shapes + layers written + -+`stride` is slot_bytes rounded up to 4 KiB so the SAME pack serves the -+O_DIRECT reader without a repack (O_DIRECT needs 4K-aligned offset/length/ -+buffer; the pinned arena is page-aligned and stride-strided, so every row -+satisfies all three). Rows of layers not listed in the sidecar are holes -+(sparse file) and are never read. ++def checkpoint_file_done(path: str) -> None: ++ """Evict a consumed shard now and queue a post-consumer retry. + -+Concurrency: every read/write caller already holds the owning DeltaTier's -+lock (manager tick, force_promote, ensure_resident are serialized there), -+so the shared pinned stage buffer / arena bookkeeping need no lock of -+their own. -+""" ++ The source generator can finish while its consumer still holds the last ++ yielded tensor. The W2-safe iterator clones that tensor, but retaining a ++ retry until `model.load_weights` unwinds also covers model-specific loader ++ references and exception paths. ++ """ ++ if not checkpoint_cache_safety_enabled(): ++ return ++ with _pending_checkpoint_lock: ++ _pending_checkpoint_drops.add(path) ++ _drop_path_page_cache(path, f"checkpoint shard {path}") ++ _memory_preflight(f"after checkpoint shard {os.path.basename(path)}") + -+import json -+import os -+import time -+from collections import deque -+from concurrent.futures import ThreadPoolExecutor + -+import torch ++def _retry_pending_checkpoint_drops(*, clear: bool) -> None: ++ if not checkpoint_cache_safety_enabled(): ++ return ++ with _pending_checkpoint_lock: ++ paths = sorted(_pending_checkpoint_drops) ++ if not paths: ++ return ++ for path in paths: ++ _drop_path_page_cache(path, f"released checkpoint shard {path}") ++ _memory_preflight( ++ f"after retrying {len(paths)} pending checkpoint shard cache drops" ++ ) ++ if clear: ++ # Clear only the successful snapshot; a concurrent/newer path remains ++ # queued. If either eviction or the postflight raises, the full set is ++ # retained for the exception-unwind retry instead of silently losing ++ # a failed safety obligation. ++ with _pending_checkpoint_lock: ++ _pending_checkpoint_drops.difference_update(paths) ++ ++ ++def checkpoint_retry_pending() -> None: ++ """Retry queued shard eviction between sequential shard mappings. ++ ++ Do not clear the queue here: a model-specific loader may retain an older ++ tensor longer than one yield. Repeating at each shard boundary is bounded ++ by checkpoint shard count (about 1K cheap fadvise calls for 46 shards), ++ and final cleanup retries once more after model.load_weights has fully ++ unwound, then clears the successful snapshot. ++ """ ++ _retry_pending_checkpoint_drops(clear=False) + -+from vllm.logger import init_logger + -+logger = init_logger(__name__) ++def checkpoint_cleanup_pending() -> None: ++ """Retry queued shard evictions after the model consumer releases refs.""" ++ _retry_pending_checkpoint_drops(clear=True) + -+_ALIGN = 4096 -+_PACK_VERSION = 1 ++ ++def allocation_preflight(label: str, allocation_bytes: int) -> None: ++ """Guard lazy anonymous staging that occurs while a shard is mapped.""" ++ if checkpoint_cache_safety_enabled(): ++ _memory_preflight(label, allocation_bytes) ++ ++ ++def allocation_postflight(label: str) -> None: ++ """Prove a guarded lazy allocation left the configured floor intact.""" ++ if checkpoint_cache_safety_enabled(): ++ _memory_preflight(f"after {label}") ++ ++ ++def guarded_checkpoint_clone(label: str, tensor: torch.Tensor) -> torch.Tensor: ++ """Clone the consumer-retained shard tail so no mmap reference escapes.""" ++ if not checkpoint_cache_safety_enabled(): ++ return tensor ++ allocation_bytes = tensor.numel() * tensor.element_size() ++ allocation_preflight(label, allocation_bytes) ++ try: ++ return tensor.clone() ++ finally: ++ allocation_postflight(label) ++ ++ ++def _pack_build_identity() -> dict: ++ """Identity of checkpoint/config bytes represented by a persistent pack.""" ++ explicit = os.getenv("VLLM_MOE_W2_PACK_ID", "").strip() ++ if explicit: ++ return { ++ "operator_id": explicit, ++ "zero_mode": os.getenv("VLLM_MOE_W2_ZERO_MODE", "auto"), ++ } ++ try: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_planes_cache, ++ ) ++ ++ return moe_w2_planes_cache.cache_identity() ++ except Exception as e: # noqa: BLE001 ++ if checkpoint_cache_safety_enabled(): ++ raise RuntimeError( ++ "moe_w2 cannot identify the checkpoint/config for safe pack " ++ "reuse; set VLLM_MOE_W2_PACK_ID to an explicit immutable " ++ "deployment identity only if current vLLM config is " ++ "unavailable" ++ ) from e ++ return { ++ "unresolved": True, ++ "zero_mode": os.getenv("VLLM_MOE_W2_ZERO_MODE", "auto"), ++ } + + +def _rank_suffix() -> str: @@ -6176,9 +8517,12 @@ index 000000000..c0b937425 + sidecar). Graceful fallback when torch.distributed is uninitialized + (single-GPU, tests, offline tools).""" + try: -+ from vllm.distributed import (get_pp_group, -+ get_tensor_model_parallel_rank, -+ get_tensor_model_parallel_world_size) ++ from vllm.distributed import ( ++ get_pp_group, ++ get_tensor_model_parallel_rank, ++ get_tensor_model_parallel_world_size, ++ ) ++ + tp_rank = get_tensor_model_parallel_rank() + tp_world = get_tensor_model_parallel_world_size() + pp = get_pp_group() @@ -6209,11 +8553,12 @@ index 000000000..c0b937425 + + def add_layer(self, layer_key: int, parts) -> None: + E = parts[0].shape[0] -+ host = torch.empty(E, self.slot_bytes, dtype=torch.uint8, -+ pin_memory=self._pinned) ++ host = torch.empty( ++ E, self.slot_bytes, dtype=torch.uint8, pin_memory=self._pinned ++ ) + off = 0 + for t in parts: -+ host[:, off:off + t.shape[1]].copy_(t, non_blocking=False) ++ host[:, off : off + t.shape[1]].copy_(t, non_blocking=False) + off += t.shape[1] + assert off == self.slot_bytes, (off, self.slot_bytes) + self._layers[layer_key] = host @@ -6238,8 +8583,10 @@ index 000000000..c0b937425 + + resident = False + -+ def __init__(self, dir_: str, tag: str, n_layers: int, n_experts: int, -+ slot_bytes: int): ++ def __init__( ++ self, dir_: str, tag: str, n_layers: int, n_experts: int, slot_bytes: int ++ ): ++ _require_cache_control() + self.slot_bytes = slot_bytes + self.E = n_experts + self.n_layers = n_layers @@ -6248,52 +8595,100 @@ index 000000000..c0b937425 + base = f"{tag}.{_rank_suffix()}" + self.path = os.path.join(dir_, base + ".pack") + self._sidecar_path = os.path.join(dir_, base + ".json") -+ self._meta = dict(version=_PACK_VERSION, tag=tag, E=n_experts, -+ n_layers=n_layers, slot_bytes=slot_bytes, -+ stride=self.stride, layers=[]) ++ self._meta = dict( ++ version=_PACK_VERSION, ++ tag=tag, ++ E=n_experts, ++ n_layers=n_layers, ++ slot_bytes=slot_bytes, ++ stride=self.stride, ++ build_identity=_pack_build_identity(), ++ layers=[], ++ ) + if os.path.exists(self._sidecar_path): + try: + with open(self._sidecar_path) as f: + old = json.load(f) -+ match = all(old.get(k) == self._meta[k] for k in -+ ("version", "E", "slot_bytes", "stride")) ++ match = all( ++ old.get(k) == self._meta[k] ++ for k in ( ++ "version", ++ "tag", ++ "E", ++ "n_layers", ++ "slot_bytes", ++ "stride", ++ "build_identity", ++ ) ++ ) + if match: + self._meta["layers"] = sorted( -+ int(li) for li in old.get("layers", [])) ++ int(li) for li in old.get("layers", []) ++ ) + else: + logger.warning( + "moe_w2 store: pack %s shape mismatch " -+ "(have %s, want %s) — rebuilding", self.path, -+ {k: old.get(k) for k in ("E", "slot_bytes", "stride")}, -+ {k: self._meta[k] for k in ("E", "slot_bytes", -+ "stride")}) ++ "(have %s, want %s) — rebuilding", ++ self.path, ++ { ++ k: old.get(k) ++ for k in ( ++ "version", ++ "E", ++ "n_layers", ++ "slot_bytes", ++ "stride", ++ "build_identity", ++ ) ++ }, ++ { ++ k: self._meta[k] ++ for k in ( ++ "version", ++ "E", ++ "n_layers", ++ "slot_bytes", ++ "stride", ++ "build_identity", ++ ) ++ }, ++ ) + except (OSError, ValueError, json.JSONDecodeError) as e: -+ logger.warning("moe_w2 store: unreadable sidecar %s (%s) — " -+ "rebuilding", self._sidecar_path, e) ++ logger.warning( ++ "moe_w2 store: unreadable sidecar %s (%s) — rebuilding", ++ self._sidecar_path, ++ e, ++ ) + self._present = set(self._meta["layers"]) + size = self.n_layers * self.E * self.stride + flags = os.O_RDWR | os.O_CREAT + self._fd = os.open(self.path, flags, 0o644) + if os.fstat(self._fd).st_size < size: -+ os.ftruncate(self._fd, size) # sparse until layers are written ++ os.ftruncate(self._fd, size) # sparse until layers are written + # reusable pinned stage for reads (grown on demand; callers hold the + # tier lock, and the tier syncs its H2D copies before the next call, + # so reuse is safe). -+ self._stage = torch.empty(0, slot_bytes, dtype=torch.uint8, -+ pin_memory=True) ++ self._stage = torch.empty(0, slot_bytes, dtype=torch.uint8, pin_memory=True) + # write-side staging reused across layers (pageable [E, stride]) + self._wbuf: torch.Tensor | None = None + self._pool = ThreadPoolExecutor( + max_workers=int(os.getenv("VLLM_MOE_W2_STORE_THREADS", "8")), -+ thread_name_prefix="moe-w2-store") ++ thread_name_prefix="moe-w2-store", ++ ) + self._reads = 0 + self._read_bytes = 0 + self._read_s = 0.0 ++ self._write_cache_drop_calls = 0 ++ self._write_cache_drop_bytes = 0 + if self._present: + logger.info( + "moe_w2 store[%s]: pack %s has %d/%d layers — staging for " + "those layers will be SKIPPED (persistent quant cache)", -+ tag, self.path, len(self._present), n_layers) ++ tag, ++ self.path, ++ len(self._present), ++ n_layers, ++ ) + + # ---- staging (load time) ---------------------------------------- + @@ -6305,23 +8700,94 @@ index 000000000..c0b937425 + + def add_layer(self, layer_key: int, parts) -> None: + if layer_key in self._present: -+ return # already packed on a previous boot ++ return # already packed on a previous boot + E = parts[0].shape[0] + assert E == self.E, (E, self.E) + if self._wbuf is None: ++ _memory_preflight( ++ f"allocating {os.path.basename(self.path)} write staging", ++ self.E * self.stride, ++ ) + self._wbuf = torch.zeros(self.E, self.stride, dtype=torch.uint8) + off = 0 + for t in parts: -+ self._wbuf[:, off:off + t.shape[1]].copy_(t, non_blocking=False) ++ self._wbuf[:, off : off + t.shape[1]].copy_(t, non_blocking=False) + off += t.shape[1] + assert off == self.slot_bytes, (off, self.slot_bytes) + mv = memoryview(self._wbuf.numpy()).cast("B") + base_off = layer_key * self.E * self.stride ++ _memory_preflight( ++ f"writing {os.path.basename(self.path)} layer {layer_key}", len(mv) ++ ) + written = 0 -+ while written < len(mv): # pwrite may be partial (>2 GiB rows) -+ written += os.pwrite(self._fd, mv[written:written + (1 << 30)], -+ base_off + written) -+ os.fdatasync(self._fd) ++ try: ++ while written < len(mv): # pwrite may be partial (>2 GiB rows) ++ n = os.pwrite( ++ self._fd, mv[written : written + (1 << 30)], base_off + written ++ ) ++ if n <= 0: ++ raise OSError( ++ f"moe_w2 pack short write @ {base_off + written} ({self.path})" ++ ) ++ written += n ++ os.fdatasync(self._fd) ++ except BaseException: ++ # A failed layer is never published in the sidecar. Best-effort ++ # writeback + cache cleanup avoids retaining a failed attempt's ++ # dirty cache while the boot unwinds. DONTNEED operates on whole ++ # pages and cannot discard dirty pages, so retry fdatasync first ++ # and round a partial write up to the pack's 4 KiB alignment. ++ # Preserve the primary write/sync error if cleanup also fails. ++ if written: ++ try: ++ os.fdatasync(self._fd) ++ except OSError: ++ logger.exception( ++ "moe_w2 pack writeback cleanup also failed after " ++ "write error for %s layer %d", ++ self.path, ++ layer_key, ++ ) ++ cleanup_bytes = min(len(mv), (written + _ALIGN - 1) // _ALIGN * _ALIGN) ++ try: ++ _fadvise_dontneed( ++ self._fd, ++ base_off, ++ cleanup_bytes, ++ f"failed pack {self.path} layer {layer_key}", ++ ) ++ except Exception: # noqa: BLE001 ++ logger.exception( ++ "moe_w2 pack cleanup also failed after write error " ++ "for %s layer %d", ++ self.path, ++ layer_key, ++ ) ++ self._wbuf = None ++ try: ++ _memory_preflight( ++ f"after failed {os.path.basename(self.path)} layer {layer_key}" ++ ) ++ except Exception: # noqa: BLE001 ++ logger.exception( ++ "moe_w2 memory floor also failed after pack error for %s layer %d", ++ self.path, ++ layer_key, ++ ) ++ raise ++ try: ++ cache_dropped = _fadvise_dontneed( ++ self._fd, base_off, len(mv), f"pack {self.path} layer {layer_key}" ++ ) ++ except BaseException: ++ self._wbuf = None ++ raise ++ if cache_dropped: ++ self._write_cache_drop_calls += 1 ++ self._write_cache_drop_bytes += len(mv) ++ _memory_preflight( ++ f"after writing {os.path.basename(self.path)} layer {layer_key}" ++ ) + self._present.add(layer_key) + self._meta["layers"] = sorted(self._present) + tmp = self._sidecar_path + ".tmp" @@ -6329,7 +8795,7 @@ index 000000000..c0b937425 + json.dump(self._meta, f) + os.replace(tmp, self._sidecar_path) + if len(self._present) == self.n_layers: -+ self._wbuf = None # all layers packed; drop write staging ++ self._wbuf = None # all layers packed; drop write staging + + # ---- reads (serve time) ----------------------------------------- + @@ -6348,28 +8814,31 @@ index 000000000..c0b937425 + RAM tier.""" + n = len(pairs) + if self._stage.shape[0] < n: -+ self._stage = torch.empty(max(n, 2 * self._stage.shape[0]), -+ self.slot_bytes, dtype=torch.uint8, -+ pin_memory=True) ++ self._stage = torch.empty( ++ max(n, 2 * self._stage.shape[0]), ++ self.slot_bytes, ++ dtype=torch.uint8, ++ pin_memory=True, ++ ) + t0 = time.perf_counter() + offs = [(li * self.E + ei) * self.stride for li, ei in pairs] + for off in offs: + try: -+ os.posix_fadvise(self._fd, off, self.slot_bytes, -+ os.POSIX_FADV_WILLNEED) ++ os.posix_fadvise(self._fd, off, self.slot_bytes, os.POSIX_FADV_WILLNEED) + except OSError: -+ break # advisory only ++ break # advisory only + stage_mv = memoryview(self._stage.numpy()).cast("B") + + def _read_one(i_off): + i, off = i_off -+ row = stage_mv[i * self.slot_bytes:(i + 1) * self.slot_bytes] ++ row = stage_mv[i * self.slot_bytes : (i + 1) * self.slot_bytes] + done = 0 + while done < self.slot_bytes: + got = os.preadv(self._fd, [row[done:]], off + done) + if got <= 0: -+ raise IOError(f"moe_w2 pack short read @ {off + done} " -+ f"({self.path})") ++ raise OSError( ++ f"moe_w2 pack short read @ {off + done} ({self.path})" ++ ) + done += got + + # preadv releases the GIL: a pool turns both page-cache memcpys and @@ -6387,17 +8856,20 @@ index 000000000..c0b937425 + + def release(self) -> None: + self._pool.shutdown(wait=False) -+ try: ++ with suppress(OSError): + os.close(self._fd) -+ except OSError: -+ pass + self._present = set() + self._stage = torch.empty(0, self.slot_bytes, dtype=torch.uint8) + self._wbuf = None + + def stats(self) -> dict: -+ return dict(reads=self._reads, read_bytes=self._read_bytes, -+ read_s=self._read_s) ++ return dict( ++ reads=self._reads, ++ read_bytes=self._read_bytes, ++ read_s=self._read_s, ++ write_cache_drop_calls=self._write_cache_drop_calls, ++ write_cache_drop_bytes=self._write_cache_drop_bytes, ++ ) + + +class TieredPackStore(MmapPackStore): @@ -6451,17 +8923,30 @@ index 000000000..c0b937425 + logged, expected only for absurdly small arenas). + """ + -+ def __init__(self, dir_: str, tag: str, n_layers: int, n_experts: int, -+ slot_bytes: int, ram_gb: float): ++ def __init__( ++ self, ++ dir_: str, ++ tag: str, ++ n_layers: int, ++ n_experts: int, ++ slot_bytes: int, ++ ram_gb: float, ++ ): + super().__init__(dir_, tag, n_layers, n_experts, slot_bytes) + self.n_arena = max(int(ram_gb * 2**30) // self.stride, 16) -+ self._arena = torch.empty(self.n_arena, self.stride, -+ dtype=torch.uint8, pin_memory=True) ++ arena_bytes = self.n_arena * self.stride ++ _memory_preflight( ++ f"allocating {tag} pinned arena ({self.n_arena} rows)", arena_bytes ++ ) ++ self._arena = torch.empty( ++ self.n_arena, self.stride, dtype=torch.uint8, pin_memory=True ++ ) ++ _memory_preflight(f"after allocating {tag} pinned arena") + assert self._arena.data_ptr() % _ALIGN == 0, "pinned base unaligned?" + self._arena_mv = memoryview(self._arena.numpy()).cast("B") -+ self._pos: dict[tuple[int, int], int] = {} # (li,ei) -> slot ++ self._pos: dict[tuple[int, int], int] = {} # (li,ei) -> slot + self._owner_pair: list = [None] * self.n_arena -+ self._last = [0] * self.n_arena # recency clock stamps ++ self._last = [0] * self.n_arena # recency clock stamps + self._clock = 0 + self._free = list(range(self.n_arena)) + # Miss-read mode: buffered (default; page cache = opportunistic L3) @@ -6469,15 +8954,14 @@ index 000000000..c0b937425 + # is separate; the parent's buffered fd keeps serving writes + # (add_layer) and stage-overflow reads. + self.direct = os.getenv("VLLM_MOE_W2_TIER_DIRECT", "0") == "1" -+ self._dfd = (os.open(self.path, os.O_RDONLY | os.O_DIRECT) -+ if self.direct else -1) ++ self._dfd = os.open(self.path, os.O_RDONLY | os.O_DIRECT) if self.direct else -1 + self.scan_enabled = os.getenv("VLLM_MOE_W2_TIER_SCAN", "1") == "1" + # fetch metrics (read by DeltaTier._log_summary/_dump) + self._hit_rows = 0 + self._miss_rows = 0 + self._miss_bytes = 0 -+ self._lat_hit_ms = deque(maxlen=2048) # pure arena-hit calls -+ self._lat_miss_ms = deque(maxlen=2048) # calls with >=1 NVMe row ++ self._lat_hit_ms = deque(maxlen=2048) # pure arena-hit calls ++ self._lat_miss_ms = deque(maxlen=2048) # calls with >=1 NVMe row + self._calls = 0 + self._heat_path = self.path + ".heat.json" + if os.getenv("VLLM_MOE_W2_TIER_PREHEAT", "1") == "1": @@ -6490,15 +8974,15 @@ index 000000000..c0b937425 + body). O_DIRECT mode reads the full stride (offset/length/buffer + all 4K-aligned by construction); buffered mode reads just + slot_bytes through the page cache.""" -+ row = self._arena_mv[slot * self.stride:(slot + 1) * self.stride] -+ fd, want = ((self._dfd, self.stride) if self.direct -+ else (self._fd, self.slot_bytes)) ++ row = self._arena_mv[slot * self.stride : (slot + 1) * self.stride] ++ fd, want = ( ++ (self._dfd, self.stride) if self.direct else (self._fd, self.slot_bytes) ++ ) + done = 0 + while done < want: + got = os.preadv(fd, [row[done:want]], off + done) + if got <= 0: -+ raise IOError(f"moe_w2 pack short read @ {off + done} " -+ f"({self.path})") ++ raise OSError(f"moe_w2 pack short read @ {off + done} ({self.path})") + done += got + + def _evict_order(self, busy: set) -> list: @@ -6515,8 +8999,7 @@ index 000000000..c0b937425 + json.dump(dict(version=1, keys=keys), f) + os.replace(tmp, self._heat_path) + except OSError as e: -+ logger.warning_once("moe_w2 tiered store: heat dump failed: %s", -+ e) ++ logger.warning_once("moe_w2 tiered store: heat dump failed: %s", e) + + def _preheat(self) -> None: + """Refill the arena with the previous run's hot set (boot time, @@ -6527,35 +9010,59 @@ index 000000000..c0b937425 + keys = [tuple(k) for k in json.load(f).get("keys", [])] + except (OSError, ValueError, json.JSONDecodeError): + return -+ keys = [k for k in dict.fromkeys(keys) # dedupe, keep order -+ if k[0] in self._present and 0 <= k[1] < self.E] -+ keys = keys[:self.n_arena] ++ keys = [ ++ k ++ for k in dict.fromkeys(keys) # dedupe, keep order ++ if k[0] in self._present and 0 <= k[1] < self.E ++ ] ++ keys = keys[: self.n_arena] + if not keys: + return + t0 = time.perf_counter() ++ _memory_preflight( ++ f"preheating {os.path.basename(self.path)} arena", ++ len(keys) * self.slot_bytes, ++ ) ++ fills = [ ++ (i, (li * self.E + ei) * self.stride) for i, (li, ei) in enumerate(keys) ++ ] + try: -+ fills = [(i, (li * self.E + ei) * self.stride) -+ for i, (li, ei) in enumerate(keys)] + list(self._pool.map(lambda p: self._read_row(p[0], p[1]), fills)) -+ for i, k in enumerate(keys): -+ self._pos[k] = i -+ self._owner_pair[i] = k -+ self._last[i] = len(keys) - i # heat order = recency -+ self._free = list(range(len(keys), self.n_arena)) -+ self._clock = len(keys) + 1 -+ logger.info( -+ "moe_w2 tiered store: arena PREHEATED — %d rows " -+ "(%.1f GiB) from %s in %.1f s", -+ len(keys), len(keys) * self.stride / 2**30, -+ self._heat_path, time.perf_counter() - t0) + except Exception as e: # noqa: BLE001 - preheat must not kill boot -+ logger.warning("moe_w2 tiered store: preheat failed (%s) — " -+ "starting cold", e) ++ # The pack read is an optional optimization and may degrade to a ++ # cold arena. Cache eviction and the postflight are outside this ++ # recoverable block: required-mode safety failures must abort. ++ _fadvise_dontneed( ++ self._fd, 0, 0, f"pack {self.path} after failed pinned-arena preheat" ++ ) ++ _memory_preflight(f"after failed preheat of {os.path.basename(self.path)}") ++ logger.warning( ++ "moe_w2 tiered store: preheat failed (%s) — starting cold", e ++ ) + self._pos = {} + self._owner_pair = [None] * self.n_arena + self._last = [0] * self.n_arena + self._free = list(range(self.n_arena)) + self._clock = 0 ++ return ++ _fadvise_dontneed( ++ self._fd, 0, 0, f"pack {self.path} after pinned-arena preheat" ++ ) ++ _memory_preflight(f"after preheating {os.path.basename(self.path)} arena") ++ for i, k in enumerate(keys): ++ self._pos[k] = i ++ self._owner_pair[i] = k ++ self._last[i] = len(keys) - i # heat order = recency ++ self._free = list(range(len(keys), self.n_arena)) ++ self._clock = len(keys) + 1 ++ logger.info( ++ "moe_w2 tiered store: arena PREHEATED — %d rows " ++ "(%.1f GiB) from %s in %.1f s", ++ len(keys), ++ len(keys) * self.stride / 2**30, ++ self._heat_path, ++ time.perf_counter() - t0, ++ ) + + # -- reads ---------------------------------------------------------- + @@ -6584,27 +9091,26 @@ index 000000000..c0b937425 + evict_order = None + ev_i = 0 + overflow: list[int] = [] -+ placed: list[tuple[int, int, int]] = [] # (idx, slot, offset) ++ placed: list[tuple[int, int, int]] = [] # (idx, slot, offset) + for i in miss_idx: + li, ei = pairs[i] + s = self._pos.get((li, ei)) -+ if s is not None: # duplicate earlier in this batch ++ if s is not None: # duplicate earlier in this batch + out[i] = s + busy.add(s) + continue + if self._free: + slot = self._free.pop() + elif scan: -+ overflow.append(i) # scans never evict ++ overflow.append(i) # scans never evict + continue + else: + if evict_order is None: + evict_order = self._evict_order(busy) -+ while ev_i < len(evict_order) \ -+ and evict_order[ev_i] in busy: ++ while ev_i < len(evict_order) and evict_order[ev_i] in busy: + ev_i += 1 + if ev_i >= len(evict_order): -+ overflow.append(i) # batch > arena; stage fallback ++ overflow.append(i) # batch > arena; stage fallback + continue + slot = evict_order[ev_i] + ev_i += 1 @@ -6624,13 +9130,13 @@ index 000000000..c0b937425 + if not self.direct: + for _, _, off in placed: + try: -+ os.posix_fadvise(self._fd, off, self.slot_bytes, -+ os.POSIX_FADV_WILLNEED) ++ os.posix_fadvise( ++ self._fd, off, self.slot_bytes, os.POSIX_FADV_WILLNEED ++ ) + except OSError: -+ break # advisory only ++ break # advisory only + if len(placed) > 2: -+ list(self._pool.map( -+ lambda p: self._read_row(p[1], p[2]), placed)) ++ list(self._pool.map(lambda p: self._read_row(p[1], p[2]), placed)) + else: + for p in placed: + self._read_row(p[1], p[2]) @@ -6642,10 +9148,12 @@ index 000000000..c0b937425 + logger.warning( + "moe_w2 tiered store: batch of %d rows exceeds the " + "arena (%d slots) — %d rows served via buffered stage; " -+ "raise VLLM_MOE_W2_BASE_RAM_GB", n, self.n_arena, -+ len(overflow)) -+ srows = MmapPackStore.rows_for( -+ self, [pairs[i] for i in overflow]) ++ "raise VLLM_MOE_W2_BASE_RAM_GB", ++ n, ++ self.n_arena, ++ len(overflow), ++ ) ++ srows = MmapPackStore.rows_for(self, [pairs[i] for i in overflow]) + stage_rows = dict(zip(overflow, srows)) + # metrics + result assembly + n_miss = len(placed) + len(overflow) @@ -6655,19 +9163,19 @@ index 000000000..c0b937425 + dt_ms = (time.perf_counter() - t0) * 1e3 + (self._lat_miss_ms if n_miss else self._lat_hit_ms).append(dt_ms) + if placed and self._calls % 1024 == 0: -+ keys = sorted(self._pos, key=lambda k: self._last[self._pos[k]], -+ reverse=True) ++ keys = sorted( ++ self._pos, key=lambda k: self._last[self._pos[k]], reverse=True ++ ) + self._pool.submit(self._dump_heat, [list(k) for k in keys]) -+ return [stage_rows[i] if out[i] is None -+ else self._arena[out[i], :self.slot_bytes] -+ for i in range(n)] ++ return [ ++ stage_rows[i] if out[i] is None else self._arena[out[i], : self.slot_bytes] ++ for i in range(n) ++ ] + + def release(self) -> None: + if self._dfd >= 0: -+ try: ++ with suppress(OSError): + os.close(self._dfd) -+ except OSError: -+ pass + self._pos = {} + self._owner_pair = [] + self._free = [] @@ -6681,6 +9189,7 @@ index 000000000..c0b937425 + return 0.0 + v = sorted(d) + return v[min(int(len(v) * q), len(v) - 1)] ++ + st = super().stats() + st.update( + arena_slots=self.n_arena, @@ -6696,8 +9205,9 @@ index 000000000..c0b937425 + return st + + -+def pack_has_layer(tag: str, layer_key: int, n_layers: int, n_experts: int, -+ slot_bytes: int) -> bool: ++def pack_has_layer( ++ tag: str, layer_key: int, n_layers: int, n_experts: int, slot_bytes: int ++) -> bool: + """Sidecar-only presence probe: does the pack this config would serve + from already hold `layer_key`? Used at WEIGHT-CREATE time (before any + store exists) to decide the loader-level skip — a pack-resident layer's @@ -6714,8 +9224,15 @@ index 000000000..c0b937425 + return False + with open(sidecar) as f: + meta = json.load(f) -+ want = dict(version=_PACK_VERSION, E=n_experts, -+ slot_bytes=slot_bytes, stride=stride) ++ want = dict( ++ version=_PACK_VERSION, ++ tag=tag, ++ E=n_experts, ++ n_layers=n_layers, ++ slot_bytes=slot_bytes, ++ stride=stride, ++ build_identity=_pack_build_identity(), ++ ) + if any(meta.get(k) != v for k, v in want.items()): + return False + return int(layer_key) in {int(li) for li in meta.get("layers", [])} @@ -6723,8 +9240,7 @@ index 000000000..c0b937425 + return False + + -+def make_store(tag: str, n_layers: int, n_experts: int, slot_bytes: int, -+ pinned: bool): ++def make_store(tag: str, n_layers: int, n_experts: int, slot_bytes: int, pinned: bool): + """Store factory: pack-file backends when VLLM_MOE_W2_STORE_DIR is set + (plus a pinned arena for the BASE tier when VLLM_MOE_W2_BASE_RAM_GB + is set), else the classic pinned/pageable host store. Env read at call @@ -6735,7 +9251,8 @@ index 000000000..c0b937425 + "experiment, moe_w2_nvme) is superseded by the pack store and " + "IGNORED. Equivalent config: VLLM_MOE_W2_STORE_DIR= + " + "VLLM_MOE_W2_BASE_RAM_GB= (the arena fraction is " -+ "the RAM share; it also persists quantization across boots).") ++ "the RAM share; it also persists quantization across boots)." ++ ) + dir_ = os.getenv("VLLM_MOE_W2_STORE_DIR", "").strip() + if not dir_: + return PinnedHostStore(slot_bytes, pinned=pinned) @@ -6744,31 +9261,43 @@ index 000000000..c0b937425 + stride = (slot_bytes + _ALIGN - 1) // _ALIGN * _ALIGN + pack_gib = n_layers * n_experts * stride / 2**30 + ram_gb = 0.25 * pack_gib if ram_raw == "auto" else float(ram_raw) -+ store = TieredPackStore(dir_, tag, n_layers, n_experts, slot_bytes, -+ ram_gb) ++ store = TieredPackStore(dir_, tag, n_layers, n_experts, slot_bytes, ram_gb) + logger.info( + "moe_w2 store[%s]: TIERED backend %s — pinned arena %.1f GiB " + "(%d slots, %.0f%% of the %.1f GiB pack) + %s NVMe misses", -+ tag, store.path, store.n_arena * store.stride / 2**30, -+ store.n_arena, 100.0 * store.n_arena / (n_layers * n_experts), -+ pack_gib, "O_DIRECT" if store.direct else "buffered") ++ tag, ++ store.path, ++ store.n_arena * store.stride / 2**30, ++ store.n_arena, ++ 100.0 * store.n_arena / (n_layers * n_experts), ++ pack_gib, ++ "O_DIRECT" if store.direct else "buffered", ++ ) + if store.n_arena < 2 * n_experts: + logger.warning( + "moe_w2 store[%s]: arena of %d slots is smaller than one " + "prefill layer's worst case (2*E=%d) — expect stage " + "overflows; raise VLLM_MOE_W2_BASE_RAM_GB", -+ tag, store.n_arena, 2 * n_experts) ++ tag, ++ store.n_arena, ++ 2 * n_experts, ++ ) + return store + store = MmapPackStore(dir_, tag, n_layers, n_experts, slot_bytes) + logger.info( + "moe_w2 store[%s]: PACK-FILE backend %s (slot %.2f MiB, " + "stride %d, %d layers x %d experts; host RAM tier = page cache)", -+ tag, store.path, slot_bytes / 2**20, store.stride, n_layers, -+ n_experts) ++ tag, ++ store.path, ++ slot_bytes / 2**20, ++ store.stride, ++ n_layers, ++ n_experts, ++ ) + return store diff --git a/vllm/model_executor/layers/quantization/utils/prefill_timers.py b/vllm/model_executor/layers/quantization/utils/prefill_timers.py new file mode 100644 -index 000000000..fd83f500b +index 0000000..fd83f50 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/prefill_timers.py @@ -0,0 +1,53 @@ @@ -6827,7 +9356,7 @@ index 000000000..fd83f500b + name, _total_ms[name], _count[name]) diff --git a/vllm/model_executor/layers/quantization/utils/skinny_fp8_cubit.py b/vllm/model_executor/layers/quantization/utils/skinny_fp8_cubit.py new file mode 100644 -index 000000000..adf4754de +index 0000000..adf4754 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/skinny_fp8_cubit.py @@ -0,0 +1,238 @@ @@ -7070,7 +9599,7 @@ index 000000000..adf4754de + torch.zeros(1, device="cuda") + return _ensure_ready() diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py -index fe2b268cd..d56825ef8 100644 +index fe2b268..d56825e 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -340,6 +340,9 @@ def sparse_attn_indexer( @@ -7083,8 +9612,313 @@ index fe2b268cd..d56825ef8 100644 ) use_persistent_topk = current_platform.is_cuda() and topk_tokens in ( 512, +diff --git a/vllm/model_executor/model_loader/__init__.py b/vllm/model_executor/model_loader/__init__.py +index 1ae78b7..254ce7f 100644 +--- a/vllm/model_executor/model_loader/__init__.py ++++ b/vllm/model_executor/model_loader/__init__.py +@@ -122,6 +122,15 @@ def register_model_loader(load_format: str): + def get_model_loader(load_config: LoadConfig) -> BaseModelLoader: + """Get a model loader based on the load format.""" + load_format = load_config.load_format ++ from vllm.model_executor.layers.quantization.utils import moe_w2_store ++ ++ guarded_formats = {"auto", "hf", "mistral", "safetensors"} ++ if (moe_w2_store.checkpoint_cache_safety_enabled() ++ and load_format not in guarded_formats): ++ raise RuntimeError( ++ f"load format {load_format!r} bypasses the W2 pack-store " ++ "checkpoint cache-safety hooks; use one of " ++ f"{sorted(guarded_formats)} with sequential safetensors") + if load_format not in _LOAD_FORMAT_TO_MODEL_LOADER: + raise ValueError(f"Load format `{load_format}` is not supported") + return _LOAD_FORMAT_TO_MODEL_LOADER[load_format](load_config) +diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py +index 3ea76f4..b2c329e 100644 +--- a/vllm/model_executor/model_loader/default_loader.py ++++ b/vllm/model_executor/model_loader/default_loader.py +@@ -253,6 +253,25 @@ class DefaultModelLoader(BaseModelLoader): + source.fall_back_to_pt, + source.allow_patterns_overrides, + ) ++ from vllm.model_executor.layers.quantization.utils import moe_w2_store ++ ++ if (moe_w2_store.checkpoint_cache_safety_enabled() ++ and not use_safetensors): ++ raise RuntimeError( ++ f"load format {self.load_config.load_format!r} selected a " ++ "non-safetensors checkpoint path that bypasses the W2 " ++ "pack-store cache-safety hooks; use the default sequential " ++ "safetensors loader for a guarded restage") ++ if (moe_w2_store.checkpoint_cache_safety_enabled() ++ and extra_config.get("enable_multithread_load") ++ and self.load_config.safetensors_load_strategy not in ( ++ None, "lazy")): ++ raise RuntimeError( ++ "multi-thread safetensors loading cannot safely honor " ++ "safetensors_load_strategy=" ++ f"{self.load_config.safetensors_load_strategy!r} while the " ++ "W2 pack-store guard is active; use the default sequential " ++ "lazy strategy") + if self.load_config.load_format == "npcache": + # Currently np_cache only support *.bin checkpoints + assert use_safetensors is False +@@ -264,6 +283,13 @@ class DefaultModelLoader(BaseModelLoader): + self.load_config.use_tqdm_on_load, + ) + elif use_safetensors: ++ if (moe_w2_store.checkpoint_cache_safety_enabled() ++ and self.load_config.load_format in ( ++ "fastsafetensors", "instanttensor")): ++ raise RuntimeError( ++ f"load format {self.load_config.load_format!r} bypasses " ++ "the W2 pack-store checkpoint cache-safety hooks; use " ++ "the default safetensors loader for a guarded restage") + if self.load_config.load_format == "fastsafetensors": + weights_iterator = fastsafetensors_weights_iterator( + hf_weights_files, +@@ -316,7 +342,16 @@ class DefaultModelLoader(BaseModelLoader): + if self.counter_before_loading_weights == 0.0: + self.counter_before_loading_weights = time.perf_counter() + # Apply the prefix. +- return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) ++ def prefixed_weights_iterator(): ++ try: ++ for name, tensor in weights_iterator: ++ yield source.prefix + name, tensor ++ finally: ++ close = getattr(weights_iterator, "close", None) ++ if close is not None: ++ close() ++ ++ return prefixed_weights_iterator() + + def get_all_weights( + self, +@@ -424,7 +459,23 @@ class DefaultModelLoader(BaseModelLoader): + + self._init_ep_weight_filter(model_config) + +- loaded_weights = model.load_weights(self.get_all_weights(model_config, model)) ++ from vllm.model_executor.layers.quantization.utils import moe_w2_store ++ ++ weights_iterator = self.get_all_weights(model_config, model) ++ try: ++ loaded_weights = model.load_weights(weights_iterator) ++ finally: ++ try: ++ # A model loader may stop consuming in the middle of a shard. ++ # Explicitly close the iterator so its per-shard finally runs ++ # and registers that shard before the post-consumer retry. ++ close = getattr(weights_iterator, "close", None) ++ if close is not None: ++ close() ++ finally: ++ # The consumer's final `loaded_weight` local is gone only ++ # after model.load_weights unwinds. Retry every shard here. ++ moe_w2_store.checkpoint_cleanup_pending() + + self.counter_after_loading_weights = time.perf_counter() + logger.info_once( +diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py +index 47c6c02..6737f5d 100644 +--- a/vllm/model_executor/model_loader/weight_utils.py ++++ b/vllm/model_executor/model_loader/weight_utils.py +@@ -837,6 +837,15 @@ def safetensors_weights_iterator( + loading_desc += " (eager)" + + sorted_files = sorted(hf_weights_files, key=_natural_sort_key) ++ from vllm.model_executor.layers.quantization.utils import moe_w2_store ++ ++ w2_cache_safety = moe_w2_store.checkpoint_cache_safety_enabled() ++ if w2_cache_safety and safetensors_load_strategy == "torchao": ++ raise RuntimeError( ++ "torchao safetensors reconstruction can retain mmap-backed " ++ "tensors across shards and is not supported by the W2 pack-store " ++ "cache-safety path; use the default safetensors loader" ++ ) + + fs_type = _get_fs_type(sorted_files) + is_net_fs = fs_type in ("nfs", "nfs4", "lustre") +@@ -895,6 +904,19 @@ def safetensors_weights_iterator( + avail_bytes / 1024**3, + ) + ++ if w2_cache_safety and should_prefetch: ++ if safetensors_load_strategy == "prefetch": ++ raise RuntimeError( ++ "safetensors checkpoint prefetch is unsafe while a W2 pack " ++ "store is active: it can populate the entire checkpoint in " ++ "uncapped host page cache; use the default sequential loader" ++ ) ++ logger.warning_once( ++ "Disabling automatic safetensors prefetch while the W2 pack-store " ++ "cache-safety guard is active" ++ ) ++ should_prefetch = False ++ + if should_prefetch: + _prefetch_all_checkpoints( + sorted_files, +@@ -909,49 +931,78 @@ def safetensors_weights_iterator( + disable=not enable_tqdm(use_tqdm_on_load), + bar_format=_BAR_FORMAT, + ): +- if safetensors_load_strategy == "eager": +- with open(st_file, "rb") as f: +- state_dict = load(f.read()) +- for name, param in state_dict.items(): +- if not should_skip_weight(name, local_expert_ids): +- yield name, param +- elif safetensors_load_strategy == "torchao": +- # we can't load flattened torchao tensor subclasses directly into the model +- # instead we reconstruct the subclasses here before returning +- if not torchao_version_at_least("0.15.0"): +- raise ValueError( +- "Please use torchao version >= 0.15.0 " +- "to load torchao safetensors checkpoint" ++ extra_bytes = ( ++ os.path.getsize(st_file) ++ if w2_cache_safety and safetensors_load_strategy == "eager" ++ else 0 ++ ) ++ moe_w2_store.checkpoint_file_preflight(st_file, extra_bytes) ++ state_dict = None ++ unflattened_state_dict = None ++ param = None ++ try: ++ if safetensors_load_strategy == "eager": ++ with open(st_file, "rb") as f: ++ state_dict = load(f.read()) ++ for name, param in state_dict.items(): ++ if not should_skip_weight(name, local_expert_ids): ++ yield name, param ++ elif safetensors_load_strategy == "torchao": ++ # we can't load flattened torchao tensor subclasses directly ++ # into the model ++ # instead we reconstruct the subclasses here before returning ++ if not torchao_version_at_least("0.15.0"): ++ raise ValueError( ++ "Please use torchao version >= 0.15.0 " ++ "to load torchao safetensors checkpoint" ++ ) ++ from torchao.prototype.safetensors.safetensors_support import ( ++ unflatten_tensor_state_dict, + ) +- from torchao.prototype.safetensors.safetensors_support import ( +- unflatten_tensor_state_dict, +- ) + +- with safe_open(st_file, framework="pt") as f: +- state_dict = {} +- for name in f.keys(): # noqa: SIM118 +- if should_skip_weight(name, local_expert_ids): +- continue +- state_dict[name] = f.get_tensor(name) +- +- # update with leftover tensor data from previous iteration, if any +- state_dict.update(leftover_state_dict) +- metadata = f.metadata() +- # due to sharded checkpoints, we are not guaranteed that we have all +- # tensor subclass data on one file +- # state_dict has the leftover data from this step and we wait for +- # missing information to be provided in a future iteration +- unflattened_state_dict, leftover_state_dict = ( +- unflatten_tensor_state_dict(state_dict, metadata) +- ) +- yield from unflattened_state_dict.items() +- else: +- with safe_open(st_file, framework="pt") as f: +- for name in f.keys(): # noqa: SIM118 +- if should_skip_weight(name, local_expert_ids): +- continue +- param = f.get_tensor(name) +- yield name, param ++ with safe_open(st_file, framework="pt") as f: ++ state_dict = {} ++ for name in f.keys(): # noqa: SIM118 ++ if should_skip_weight(name, local_expert_ids): ++ continue ++ state_dict[name] = f.get_tensor(name) ++ ++ # update with leftover tensor data from previous iteration, if any ++ state_dict.update(leftover_state_dict) ++ metadata = f.metadata() ++ # due to sharded checkpoints, we are not guaranteed that we have all ++ # tensor subclass data on one file ++ # state_dict has the leftover data from this step and we wait for ++ # missing information to be provided in a future iteration ++ unflattened_state_dict, leftover_state_dict = ( ++ unflatten_tensor_state_dict(state_dict, metadata) ++ ) ++ yield from unflattened_state_dict.items() ++ else: ++ with safe_open(st_file, framework="pt") as f: ++ names = [ ++ name ++ for name in f.keys() # noqa: SIM118 ++ if not should_skip_weight(name, local_expert_ids) ++ ] ++ for i, name in enumerate(names): ++ param = f.get_tensor(name) ++ if w2_cache_safety and i == len(names) - 1: ++ param = moe_w2_store.guarded_checkpoint_clone( ++ f"checkpoint shard tail {name}", param ++ ) ++ yield name, param ++ finally: ++ # Release mmap-backed tensors before DONTNEED. This path also ++ # runs when the consumer closes or throws into the generator. ++ param = None ++ if state_dict is not None: ++ state_dict.clear() ++ state_dict = None ++ if unflattened_state_dict is not None: ++ unflattened_state_dict.clear() ++ unflattened_state_dict = None ++ moe_w2_store.checkpoint_file_done(st_file) + + + def multi_thread_safetensors_weights_iterator( +@@ -961,6 +1012,37 @@ def multi_thread_safetensors_weights_iterator( + ) -> Generator[tuple[str, torch.Tensor], None, None]: + """Multi-Thread iterate over the weights in the model safetensor files.""" + ++ from vllm.model_executor.layers.quantization.utils import moe_w2_store ++ ++ if moe_w2_store.checkpoint_cache_safety_enabled(): ++ logger.warning_once( ++ "Serializing multi-thread safetensors loading while the W2 " ++ "pack-store cache-safety guard is active" ++ ) ++ for st_file in tqdm( ++ sorted(hf_weights_files, key=_natural_sort_key), ++ desc="Loading safetensors checkpoint shards (W2 safe sequential)", ++ disable=not enable_tqdm(use_tqdm_on_load), ++ bar_format=_BAR_FORMAT, ++ ): ++ moe_w2_store.checkpoint_file_preflight(st_file) ++ state_dict = None ++ try: ++ state_dict = load_file(st_file, device="cpu") ++ keys = list(state_dict) ++ for i, key in enumerate(keys): ++ tensor = state_dict.pop(key) ++ if i == len(keys) - 1: ++ tensor = moe_w2_store.guarded_checkpoint_clone( ++ f"checkpoint shard tail {key}", tensor ++ ) ++ yield key, tensor ++ finally: ++ if state_dict is not None: ++ state_dict.clear() ++ moe_w2_store.checkpoint_file_done(st_file) ++ return ++ + def _load_file(st_file: str): + result = load_file(st_file, device="cpu") + return result diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py -index 88f33ac02..546f92265 100644 +index 88f33ac..546f922 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -35,6 +35,7 @@ from .deepseek_v2 import ( @@ -7112,7 +9946,7 @@ index 88f33ac02..546f92265 100644 super().__init__() self.config = vllm_config.model_config.hf_config diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py -index 820260f79..6340e22f3 100644 +index 820260f..6340e22 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -345,7 +345,7 @@ class DFlashQwen3Model(nn.Module): @@ -7158,7 +9992,7 @@ index 820260f79..6340e22f3 100644 self.model.precompute_and_store_context_kv( diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py new file mode 100644 -index 000000000..b387bae3d +index 0000000..b387bae --- /dev/null +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -0,0 +1,220 @@ @@ -7383,7 +10217,7 @@ index 000000000..b387bae3d + loader.load_weights(model_weights.items()) + self.model._build_fused_kv_buffers() diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py -index a18d54acd..92e3bede8 100644 +index a18d54a..92e3bed 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -604,6 +604,8 @@ _SPECULATIVE_DECODING_MODELS = { @@ -7396,7 +10230,7 @@ index a18d54acd..92e3bede8 100644 "PeagleLlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), diff --git a/vllm/models/deepseek_v4/__init__.py b/vllm/models/deepseek_v4/__init__.py -index 44e486db7..d03eba3d4 100644 +index 44e486d..d03eba3 100644 --- a/vllm/models/deepseek_v4/__init__.py +++ b/vllm/models/deepseek_v4/__init__.py @@ -17,14 +17,23 @@ from .quant_config import DeepseekV4FP8Config @@ -7425,7 +10259,7 @@ index 44e486db7..d03eba3d4 100644 "DeepseekV4ForCausalLM", diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py new file mode 100644 -index 000000000..be4a87b32 +index 0000000..be4a87b --- /dev/null +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -0,0 +1,488 @@ @@ -7918,7 +10752,7 @@ index 000000000..be4a87b32 + return f"model.{rest}" + return f"model.layers.{stage}.{rest}" diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py -index aa60ad34c..fa6da8cbe 100644 +index aa60ad3..fa6da8c 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -46,7 +46,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -8004,7 +10838,7 @@ index aa60ad34c..fa6da8cbe 100644 # Default mapper assumes the original FP4-expert checkpoint layout. diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py -index 64715deae..8c456a69a 100644 +index 64715de..8c456a6 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -40,6 +40,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -8030,7 +10864,7 @@ index 64715deae..8c456a69a 100644 super().__init__() self.config = vllm_config.model_config.hf_config diff --git a/vllm/models/deepseek_v4/nvidia/ops/o_proj.py b/vllm/models/deepseek_v4/nvidia/ops/o_proj.py -index 18e3b1056..d53669202 100644 +index 18e3b10..d536692 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/o_proj.py +++ b/vllm/models/deepseek_v4/nvidia/ops/o_proj.py @@ -15,13 +15,16 @@ def compute_fp8_einsum_recipe() -> tuple[tuple[int, int, int], bool]: @@ -8053,7 +10887,7 @@ index 18e3b1056..d53669202 100644 diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py -index 0dc3ccce0..fb75a278c 100644 +index 0dc3ccc..fb75a27 100644 --- a/vllm/transformers_utils/configs/speculators/algos.py +++ b/vllm/transformers_utils/configs/speculators/algos.py @@ -118,3 +118,47 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None: @@ -8105,7 +10939,7 @@ index 0dc3ccce0..fb75a278c 100644 + if config_dict.get(key) is not None: + pre_trained_config[key] = config_dict[key] diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py -index 80319003d..40193a3a6 100755 +index 8031900..40193a3 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -408,7 +408,9 @@ class FlashInferBackend(AttentionBackend): @@ -8121,7 +10955,7 @@ index 80319003d..40193a3a6 100755 return torch.float8_e5m2 diff --git a/vllm/v1/attention/backends/mla/cubit_sparse_mla.py b/vllm/v1/attention/backends/mla/cubit_sparse_mla.py new file mode 100644 -index 000000000..16d22c7cd +index 0000000..16d22c7 --- /dev/null +++ b/vllm/v1/attention/backends/mla/cubit_sparse_mla.py @@ -0,0 +1,1171 @@ @@ -9297,7 +12131,7 @@ index 000000000..16d22c7cd + tcap, (T + tcap - 1) // tcap, int(capturing), _pfc_launches) + return True diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py -index 2a944d061..d4807a5f7 100644 +index 2a944d0..d4807a5 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -148,6 +148,7 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase): @@ -9328,7 +12162,7 @@ index 2a944d061..d4807a5f7 100644 # fp8_ds_mla packed layout: 512 NoPE + 16 scales + 128 RoPE. return (num_blocks, block_size, 656) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py -index d802f5688..d192d3aa3 100644 +index d802f56..d192d3a 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -62,10 +62,10 @@ class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMet @@ -9392,7 +12226,7 @@ index d802f5688..d192d3aa3 100644 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], diff --git a/vllm/v1/attention/backends/mla/nvfp4_ds_mla_cache.py b/vllm/v1/attention/backends/mla/nvfp4_ds_mla_cache.py new file mode 100644 -index 000000000..c37d3729e +index 0000000..c37d372 --- /dev/null +++ b/vllm/v1/attention/backends/mla/nvfp4_ds_mla_cache.py @@ -0,0 +1,53 @@ @@ -9451,7 +12285,7 @@ index 000000000..c37d3729e + slot_mapping) diff --git a/vllm/v1/attention/backends/mla/sparse_mla_env.py b/vllm/v1/attention/backends/mla/sparse_mla_env.py new file mode 100644 -index 000000000..931614177 +index 0000000..9316141 --- /dev/null +++ b/vllm/v1/attention/backends/mla/sparse_mla_env.py @@ -0,0 +1,216 @@ @@ -9672,7 +12506,7 @@ index 000000000..931614177 + except ValueError: + return 8192 diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py -index df23f3437..962c098f8 100644 +index df23f34..962c098 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -306,16 +306,19 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): @@ -9880,7 +12714,7 @@ index df23f3437..962c098f8 100644 + mask=offset < index_width, + ) diff --git a/vllm/v1/attention/ops/merge_attn_states.py b/vllm/v1/attention/ops/merge_attn_states.py -index cf4338fb1..f0577501e 100644 +index cf4338f..f057750 100644 --- a/vllm/v1/attention/ops/merge_attn_states.py +++ b/vllm/v1/attention/ops/merge_attn_states.py @@ -46,6 +46,19 @@ def merge_attn_states( @@ -9904,7 +12738,7 @@ index cf4338fb1..f0577501e 100644 # does not support FP8 dtype for inputs, fallback to use Triton kernel. # However, when output_scale is provided, the inputs are still BF16/FP16 diff --git a/vllm/v1/attention/ops/triton_decode_attention.py b/vllm/v1/attention/ops/triton_decode_attention.py -index dbe3c5705..f12d719fb 100644 +index dbe3c57..f12d719 100644 --- a/vllm/v1/attention/ops/triton_decode_attention.py +++ b/vllm/v1/attention/ops/triton_decode_attention.py @@ -530,6 +530,13 @@ def _decode_grouped_att_m_fwd( @@ -9922,7 +12756,7 @@ index dbe3c5705..f12d719fb 100644 _fwd_grouped_kernel_stage1[grid]( q, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py -index 90d93a110..418f33346 100644 +index 90d93a1..418f333 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -57,7 +57,10 @@ from vllm.v1.metrics.perf import ModelMetrics, PerfStats @@ -9982,7 +12816,7 @@ index 90d93a110..418f33346 100644 self, spec_decoding_stats: SpecDecodingStats | None, diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py -index f97f697de..df3827905 100644 +index f97f697..df38279 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -147,6 +147,43 @@ class EngineCore: @@ -10041,7 +12875,7 @@ index f97f697de..df3827905 100644 self.check_for_draft_tokens = ( self.use_spec_decode or vllm_config.model_config.is_diffusion diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py -index 5a2a5c5e2..33c628e8e 100644 +index 5a2a5c5..33c628e 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -382,6 +382,10 @@ class MLAAttentionSpec(FullAttentionSpec): @@ -10056,7 +12890,7 @@ index 5a2a5c5e2..33c628e8e 100644 if self.model_version == "deepseek_v4": # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token. diff --git a/vllm/v1/spec_decode/dynamic/utils.py b/vllm/v1/spec_decode/dynamic/utils.py -index de869b19a..348eaf2e9 100644 +index de869b1..348eaf2 100644 --- a/vllm/v1/spec_decode/dynamic/utils.py +++ b/vllm/v1/spec_decode/dynamic/utils.py @@ -1,8 +1,16 @@ @@ -10194,7 +13028,7 @@ index de869b19a..348eaf2e9 100644 + vllm_num_speculative_tokens=self._num_spec_tokens, + ) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py -index dad1777b4..358acf464 100644 +index dad1777..358acf4 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -109,6 +109,13 @@ def get_uniform_token_count( @@ -10324,7 +13158,7 @@ index dad1777b4..358acf464 100644 self.hidden_states: torch.Tensor | None = None self.aux_hidden_states: list[torch.Tensor] = [] diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py -index 6b750fe7e..167ad6851 100644 +index 6b750fe..167ad68 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -93,6 +93,11 @@ class InputBatch: @@ -10400,7 +13234,7 @@ index 6b750fe7e..167ad6851 100644 # last sampled token in addition to all draft tokens. BLOCK_SIZE=triton.next_power_of_2( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py -index 30ca2ddc5..894e258bf 100644 +index 30ca2dd..894e258 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -47,6 +47,7 @@ from vllm.tasks import SupportedTask @@ -10716,7 +13550,7 @@ index 30ca2ddc5..894e258bf 100644 # Post-step KV connector related operations. diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py -index fab53fef7..7bacdbc51 100644 +index fab53fe..7bacdbc 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -213,6 +213,11 @@ def gumbel_sample( @@ -10732,7 +13566,7 @@ index fab53fef7..7bacdbc51 100644 BLOCK_SIZE = 1024 num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE) diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py -index 09153dd20..c70f169f7 100644 +index 09153dd..c70f169 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -14,6 +14,12 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device): @@ -10749,7 +13583,7 @@ index 09153dd20..c70f169f7 100644 from vllm.v1.worker.gpu.spec_decode.gemma4.speculator import ( Gemma4Speculator, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py -index 1bd130838..9014fc0e7 100644 +index 1bd1308..9014fc0 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -29,6 +29,8 @@ logger = init_logger(__name__) @@ -11099,7 +13933,7 @@ index 1bd130838..9014fc0e7 100644 ) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/__init__.py b/vllm/v1/worker/gpu/spec_decode/dspark/__init__.py new file mode 100644 -index 000000000..208f01a7c +index 0000000..208f01a --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dspark/__init__.py @@ -0,0 +1,2 @@ @@ -11107,7 +13941,7 @@ index 000000000..208f01a7c +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/scheduler.py b/vllm/v1/worker/gpu/spec_decode/dspark/scheduler.py new file mode 100644 -index 000000000..28a341755 +index 0000000..28a3417 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dspark/scheduler.py @@ -0,0 +1,373 @@ @@ -11486,7 +14320,7 @@ index 000000000..28a341755 + self._prev_sched_l = length diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py new file mode 100644 -index 000000000..88841d5f4 +index 0000000..88841d5 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -0,0 +1,302 @@ @@ -11794,7 +14628,7 @@ index 000000000..88841d5f4 + self._sample_sequential(num_reqs, head_hidden) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py new file mode 100644 -index 000000000..b5f2a540a +index 0000000..b5f2a54 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -0,0 +1,76 @@ @@ -11875,7 +14709,7 @@ index 000000000..b5f2a540a + + return draft_model diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py -index 360f64921..66d0ba8b4 100644 +index 360f649..66d0ba8 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py @@ -44,6 +44,15 @@ def get_eagle3_aux_layers_from_config( @@ -11895,7 +14729,7 @@ index 360f64921..66d0ba8b4 100644 return tuple(layer_ids) return None diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py -index 3868604d3..5e7cadc68 100644 +index 3868604..5e7cadc 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -100,12 +100,22 @@ class RejectionSampler: @@ -11922,7 +14756,7 @@ index 3868604d3..5e7cadc68 100644 processed_logits = self.sampler.apply_sampling_params( logits, diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py -index 4ab45b2ae..35a89634e 100644 +index 4ab45b2..35a8963 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -16,35 +16,58 @@ class DraftTokensHandler: @@ -12023,7 +14857,7 @@ index 4ab45b2ae..35a89634e 100644 + " `ptd_token_id` for parallel drafting." ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py -index 74938a823..27135b752 100644 +index 74938a8..909ce63 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4,6 +4,7 @@ @@ -12049,7 +14883,22 @@ index 74938a823..27135b752 100644 from vllm.compilation.breakable_cudagraph import ( BreakableCUDAGraphWrapper, is_breakable_cudagraph_enabled, -@@ -898,6 +907,12 @@ class GPUModelRunner( +@@ -233,6 +242,14 @@ if TYPE_CHECKING: + + logger = init_logger(__name__) + ++ ++def _batch_has_prefill( ++ num_computed_tokens: np.ndarray, num_prompt_tokens: np.ndarray ++) -> bool: ++ """Return whether any scheduled request is still consuming its prompt.""" ++ return bool(np.any(num_computed_tokens < num_prompt_tokens)) ++ ++ + AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata] + # list when ubatching is enabled + PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict +@@ -898,6 +915,12 @@ class GPUModelRunner( # Ephemeral state transferred between execute_model() and sample_tokens(). self.execute_model_state: ExecuteModelState | None = None self.kv_connector_output: KVConnectorOutput | None = None @@ -12062,7 +14911,7 @@ index 74938a823..27135b752 100644 self.mamba_state_idx: dict[str, int] = {} self._mamba_bufs: mamba_utils.MambaBuffers | None = None self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None -@@ -1355,10 +1370,19 @@ class GPUModelRunner( +@@ -1355,10 +1378,19 @@ class GPUModelRunner( req_state.output_token_ids.extend( new_token_ids[-num_new_tokens:] ) @@ -12086,7 +14935,7 @@ index 74938a823..27135b752 100644 del req_state.output_token_ids[num_output_tokens:] if req_index is not None: end_idx = ( -@@ -2473,8 +2497,11 @@ class GPUModelRunner( +@@ -2473,8 +2505,11 @@ class GPUModelRunner( cm.slot_mapping = slot_mappings[kv_cache_gid] if self.speculative_config and spec_decode_common_attn_metadata is None: @@ -12099,7 +14948,7 @@ index 74938a823..27135b752 100644 ( EagleProposer, DFlashProposer, -@@ -2482,16 +2509,20 @@ class GPUModelRunner( +@@ -2482,16 +2517,20 @@ class GPUModelRunner( ExtractHiddenStatesProposer, ), ): @@ -12123,7 +14972,7 @@ index 74938a823..27135b752 100644 self.drafter.set_per_group_block_table( kv_cache_gid, cm.block_table_tensor ) -@@ -3560,6 +3591,18 @@ class GPUModelRunner( +@@ -3560,6 +3599,18 @@ class GPUModelRunner( intermediate_tensors = self.sync_and_gather_intermediate_tensors( num_input_tokens, intermediate_tensors, True ) @@ -12142,10 +14991,165 @@ index 74938a823..27135b752 100644 if is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: # Run the encoder, just like we do with other multimodal inputs. -@@ -4306,6 +4349,30 @@ class GPUModelRunner( +@@ -3978,6 +4029,7 @@ class GPUModelRunner( + ) -> tuple[ + dict[int, torch.Tensor] | None, + dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None, ++ torch.Tensor | list[torch.Tensor] | None, + ]: + """ + Build slot mappings in both formats needed by the system. +@@ -3992,13 +4044,40 @@ class GPUModelRunner( + A tuple of: + - slot_mappings_by_gid: dict[int, torch.Tensor] for attention metadata + - slot_mappings_by_layer: dict[str, torch.Tensor] or list for ForwardContext ++ - token_slot_mapping: persistent FullAttention mapping, optionally sliced + """ + if not ( + hasattr(self, "kv_cache_config") + and self.kv_cache_config is not None + and len(self.kv_cache_config.kv_cache_groups) > 0 + ): +- return None, None ++ if os.getenv("VLLM_MOE_W2", "0") != "1": ++ return None, None, None ++ # Memory profiling runs a real W2 dummy forward before KV caches ++ # (and therefore attention slot mappings) exist. Give that phase ++ # the same fixed-address validity contract using runner-owned ++ # storage; graph capture after cache initialization uses the real ++ # FullAttention slot mapping below. ++ profile_storage = getattr(self, "_w2_profile_token_slot_mapping", None) ++ if profile_storage is None: ++ if torch.cuda.is_current_stream_capturing(): ++ raise RuntimeError( ++ "cannot allocate the moe_w2 profile token mask during capture" ++ ) ++ profile_storage = torch.empty( ++ (self.max_num_tokens,), dtype=torch.int64, device=self.device ++ ) ++ self._w2_profile_token_slot_mapping = profile_storage ++ profile_mapping = profile_storage[:num_tokens_padded] ++ profile_mapping[:num_tokens_unpadded].zero_() ++ profile_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1) ++ if ubatch_slices is not None: ++ return ( ++ None, ++ None, ++ [profile_mapping[ubatch.token_slice] for ubatch in ubatch_slices], ++ ) ++ return None, None, profile_mapping + + def _get_slot_mapping(kv_cache_gid: int): + assert num_reqs_padded is not None and num_tokens_padded is not None +@@ -4025,6 +4104,14 @@ class GPUModelRunner( + gid: _get_slot_mapping(gid) + for gid, _ in enumerate(self.kv_cache_config.kv_cache_groups) + } ++ if ( ++ os.getenv("VLLM_MOE_W2", "0") == "1" ++ and self.parallel_config.decode_context_parallel_size != 1 ++ ): ++ raise RuntimeError( ++ "moe_w2 padded-route masking requires decode context parallel size 1" ++ ) ++ token_slot_mapping = slot_mappings_by_gid[self._get_attention_kv_cache_gid()] + + slot_mappings_by_layer: dict[str, torch.Tensor] = {} + for gid, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups): +@@ -4034,14 +4121,16 @@ class GPUModelRunner( + + if ubatch_slices is not None: + result: list[dict[str, torch.Tensor]] = [] ++ token_result: list[torch.Tensor] = [] + for ubatch in ubatch_slices: + sliced_mappings: dict[str, torch.Tensor] = {} + for layer_name, slot_mapping in slot_mappings_by_layer.items(): + sliced_mappings[layer_name] = slot_mapping[ubatch.token_slice] + result.append(sliced_mappings) +- return slot_mappings_by_gid, result ++ token_result.append(token_slot_mapping[ubatch.token_slice]) ++ return slot_mappings_by_gid, result, token_result + +- return slot_mappings_by_gid, slot_mappings_by_layer ++ return slot_mappings_by_gid, slot_mappings_by_layer, token_slot_mapping + + def _is_all_reqs_chunked_prefill(self) -> bool: + """Check if all scheduled requests are marked to discard sampled tokens. +@@ -4136,6 +4225,15 @@ class GPUModelRunner( + num_scheduled_tokens_np = np.array(tokens, dtype=np.int32) + max_num_scheduled_tokens = int(num_scheduled_tokens_np.max()) + num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens ++ has_prefill = _batch_has_prefill( ++ self.input_batch.num_computed_tokens_cpu[:num_reqs], ++ self.input_batch.num_prompt_tokens[:num_reqs], ++ ) ++ force_w2_prefill_eager = ( ++ os.getenv("VLLM_MOE_W2", "0") == "1" ++ and not self.is_pooling_model ++ and has_prefill ++ ) + + logits_indices, spec_decode_metadata = self._prepare_inputs( + scheduler_output, +@@ -4164,9 +4262,13 @@ class GPUModelRunner( + num_scheduled_tokens_np=num_scheduled_tokens_np, + max_num_scheduled_tokens=max_num_scheduled_tokens, + use_cascade_attn=cascade_attn_prefix_lens is not None, ++ force_eager=force_w2_prefill_eager, + num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), + ) + ++ if force_w2_prefill_eager and cudagraph_mode != CUDAGraphMode.NONE: ++ raise RuntimeError("moe_w2 prefill must execute without CUDA graphs") ++ + logger.debug( + "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " + "should_ubatch: %s, num_tokens_across_dp: %s", +@@ -4253,15 +4355,19 @@ class GPUModelRunner( + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices + +- slot_mappings_by_group, slot_mappings = self._get_slot_mappings( +- num_tokens_padded=num_tokens_padded +- if pad_attn or has_separate_kv_update +- else num_tokens_unpadded, +- num_reqs_padded=( +- num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs +- ), +- num_tokens_unpadded=num_tokens_unpadded, +- ubatch_slices=ubatch_slices_padded, ++ slot_mappings_by_group, slot_mappings, token_slot_mapping = ( ++ self._get_slot_mappings( ++ num_tokens_padded=num_tokens_padded ++ if pad_attn or has_separate_kv_update ++ else num_tokens_unpadded, ++ num_reqs_padded=( ++ num_reqs_padded ++ if pad_attn or has_separate_kv_update ++ else num_reqs ++ ), ++ num_tokens_unpadded=num_tokens_unpadded, ++ ubatch_slices=ubatch_slices_padded, ++ ) + ) + + attn_metadata, spec_decode_common_attn_metadata = ( +@@ -4306,6 +4412,45 @@ class GPUModelRunner( self.model_config.is_encoder_decoder and num_encoder_reqs > 0 ) ++ # CUDA graph capture/warmup and the previous target can leave routing ++ # marks in the tier singletons. Clear both tiers before the next target ++ # forward so its replay snapshot contains only this logical step. Keep ++ # prior-step pins until the target completes: they protect slots from a ++ # racing background manager while a graph or eager prefill starts. ++ if os.getenv("VLLM_MOE_W2", "0") == "1" and not self.is_pooling_model: ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_delta as _w2d, ++ ) ++ ++ _w2d.begin_target_step() ++ + # moe_w2 draft-affinity PREFETCH (VLLM_MOE_W2_PREFETCH=1): before the + # forward, fold the PREVIOUS step's in-graph routing log into the + # token->experts table and prefetch this step's predicted experts on @@ -12153,10 +15157,12 @@ index 74938a823..27135b752 100644 + # step's sampled+draft tokens — the draft signal. Decode-shaped + # steps only (a prefill chunk would poison the table; prefill + # prefetches via ensure_resident anyway). -+ if moe_w2_gate is not None and not self.is_pooling_model: ++ if moe_w2_gate is not None and not self.is_pooling_model and not has_prefill: + try: + from vllm.model_executor.layers.quantization.utils import ( -+ moe_w2_delta as _w2d) ++ moe_w2_delta as _w2d, ++ ) ++ + _btier = _w2d._BASE_TIER + if ( + _btier is not None @@ -12165,7 +15171,8 @@ index 74938a823..27135b752 100644 + ): + _n_real = min( + scheduler_output.total_num_scheduled_tokens, -+ _btier.route_log.shape[1]) ++ _btier.route_log.shape[1], ++ ) + _btier.draft_prefetch(input_ids[:_n_real]) + except Exception as e: # noqa: BLE001 - never crash serving + logger.warning_once("moe_w2 draft prefetch skipped: %s", e) @@ -12173,10 +15180,31 @@ index 74938a823..27135b752 100644 # Run the model. # Use persistent buffers for CUDA graphs. # When spec decode is enabled, defer connector finalization -@@ -4337,6 +4404,128 @@ class GPUModelRunner( +@@ -4321,6 +4466,8 @@ class GPUModelRunner( + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, + skip_compiled=has_encoder_input, + ), + record_function_or_nullcontext("gpu_model_runner: forward"), +@@ -4337,6 +4484,151 @@ class GPUModelRunner( **model_kwargs, ) ++ # Open both tiers' pin scopes exactly once after the target pass and ++ # before any fixed-point or confidence-gated replay can promote experts. ++ if os.getenv("VLLM_MOE_W2", "0") == "1" and not self.is_pooling_model: ++ # This is a quality invariant, not a best-effort optimization. If ++ # the scope cannot be opened, continuing would silently preserve ++ # stale pins and eventually freeze a saturated pool. ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_delta as _w2d, ++ ) ++ ++ _w2d.begin_replay_step() ++ + # moe_w2 BASE cache under PIPELINE parallelism: a miss is LOCAL to the + # stage — its desc kernels zeroed the missing pairs' contributions and + # bumped ITS miss counter, and the stage still holds this step's @@ -12199,26 +15227,26 @@ index 74938a823..27135b752 100644 + ): + try: + from vllm.model_executor.layers.quantization.utils import ( -+ moe_w2_delta as _w2d) ++ moe_w2_delta as _w2d, ++ ) ++ + _btier = _w2d._BASE_TIER + if _btier is not None: -+ # open this step's pin scope: slots touched by any pass -+ # below are ineligible for eviction until the next step -+ _btier.step_begin() + _miss = int(_btier.miss_count.item()) + _tp = get_tp_group() + if _tp.world_size > 1: + _t = torch.tensor([_miss], device=_btier.dev) + torch.distributed.all_reduce( -+ _t, op=torch.distributed.ReduceOp.MAX, -+ group=_tp.device_group) ++ _t, ++ op=torch.distributed.ReduceOp.MAX, ++ group=_tp.device_group, ++ ) + _max_miss = int(_t.item()) + else: + _max_miss = _miss + if _miss > 0: + _btier.force_promote(max_promote=None) -+ _btier.kpi_step( -+ _max_miss, _max_miss > _w2d.base_miss_tol()) ++ _btier.kpi_step(_max_miss, _max_miss > _w2d.base_miss_tol()) + # Replay to a FIXED POINT (bounded): the corrected early + # layers can re-route later layers to experts the first + # pass never fetched — those SECOND-ORDER misses zero @@ -12239,6 +15267,8 @@ index 74938a823..27135b752 100644 + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, + skip_compiled=has_encoder_input, + ): + model_output = self._model_forward( @@ -12252,8 +15282,10 @@ index 74938a823..27135b752 100644 + if _tp.world_size > 1: + _t = torch.tensor([_miss], device=_btier.dev) + torch.distributed.all_reduce( -+ _t, op=torch.distributed.ReduceOp.MAX, -+ group=_tp.device_group) ++ _t, ++ op=torch.distributed.ReduceOp.MAX, ++ group=_tp.device_group, ++ ) + _max_miss = int(_t.item()) + else: + _max_miss = _miss @@ -12262,8 +15294,7 @@ index 74938a823..27135b752 100644 + if _replays: + _btier.kpi_fp(_replays, _max_miss) + except Exception as e: # noqa: BLE001 - never crash serving -+ logger.warning( -+ "moe_w2 base-cache PP stage replay skipped: %s", e) ++ logger.warning("moe_w2 base-cache PP stage replay skipped: %s", e) + + # moe_w2 confidence gate under PP: cache this step's forward context so + # the worker can drive a FULL second pipeline pass (gate_reforward) when @@ -12284,25 +15315,33 @@ index 74938a823..27135b752 100644 + and moe_w2_gate.enabled() + and get_pp_group().world_size > 1 + and not self.is_pooling_model ++ and not has_prefill + and spec_decode_metadata is None + and max_num_scheduled_tokens <= 1 + ): + self._gate_ctx = dict( -+ input_ids=input_ids, positions=positions, ++ input_ids=input_ids, ++ positions=positions, + intermediate_tensors=intermediate_tensors, -+ inputs_embeds=inputs_embeds, model_kwargs=model_kwargs, -+ attn_metadata=attn_metadata, num_tokens_padded=num_tokens_padded, ++ inputs_embeds=inputs_embeds, ++ model_kwargs=model_kwargs, ++ attn_metadata=attn_metadata, ++ num_tokens_padded=num_tokens_padded, + num_tokens_across_dp=num_tokens_across_dp, -+ cudagraph_mode=cudagraph_mode, batch_desc=batch_desc, ++ cudagraph_mode=cudagraph_mode, ++ batch_desc=batch_desc, + ubatch_slices_padded=ubatch_slices_padded, -+ slot_mappings=slot_mappings, has_encoder_input=has_encoder_input, ++ slot_mappings=slot_mappings, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, ++ has_encoder_input=has_encoder_input, + logits_indices=logits_indices, + ) + with record_function_or_nullcontext("gpu_model_runner: postprocess"): if self.use_aux_hidden_state_outputs: # True when EAGLE 3 is used. -@@ -4395,6 +4584,189 @@ class GPUModelRunner( +@@ -4395,6 +4687,196 @@ class GPUModelRunner( assert broadcasted is not None logits = broadcasted["logits"] @@ -12316,26 +15355,27 @@ index 74938a823..27135b752 100644 + # error the zero-contribution logits are kept (quality blip, not a + # crash). + if ( -+ moe_w2_gate is not None # module family deployed ++ moe_w2_gate is not None # module family deployed + and logits is not None + and not self.is_pooling_model + and get_pp_group().world_size == 1 + ): + try: + from vllm.model_executor.layers.quantization.utils import ( -+ moe_w2_delta as _w2d) ++ moe_w2_delta as _w2d, ++ ) ++ + _btier = _w2d._BASE_TIER + if _btier is not None: -+ # open this step's pin scope: slots touched by any pass -+ # below are ineligible for eviction until the next step -+ _btier.step_begin() + _miss = int(_btier.miss_count.item()) + _tp = get_tp_group() + if _tp.world_size > 1: + _t = torch.tensor([_miss], device=logits.device) + torch.distributed.all_reduce( -+ _t, op=torch.distributed.ReduceOp.MAX, -+ group=_tp.device_group) ++ _t, ++ op=torch.distributed.ReduceOp.MAX, ++ group=_tp.device_group, ++ ) + _max_miss = int(_t.item()) + else: + _max_miss = _miss @@ -12350,8 +15390,7 @@ index 74938a823..27135b752 100644 + _btier.force_promote(max_promote=None) + # KPI: per-step replay rate + missing pairs (windowed + # INFO line) — the pool-sizing signal. -+ _btier.kpi_step( -+ _max_miss, _max_miss > _w2d.base_miss_tol()) ++ _btier.kpi_step(_max_miss, _max_miss > _w2d.base_miss_tol()) + # Replay to a FIXED POINT (bounded): corrected early + # layers can re-route later layers onto experts the + # first pass never fetched; those second-order misses @@ -12372,6 +15411,8 @@ index 74938a823..27135b752 100644 + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, + skip_compiled=has_encoder_input, + ): + _re_out = self._model_forward( @@ -12391,8 +15432,10 @@ index 74938a823..27135b752 100644 + if _tp.world_size > 1: + _t = torch.tensor([_miss], device=logits.device) + torch.distributed.all_reduce( -+ _t, op=torch.distributed.ReduceOp.MAX, -+ group=_tp.device_group) ++ _t, ++ op=torch.distributed.ReduceOp.MAX, ++ group=_tp.device_group, ++ ) + _max_miss = int(_t.item()) + else: + _max_miss = _miss @@ -12419,6 +15462,7 @@ index 74938a823..27135b752 100644 + and moe_w2_gate.enabled() + and logits is not None + and not self.is_pooling_model ++ and not has_prefill + and get_pp_group().is_last_rank + and (spec_decode_metadata is not None or max_num_scheduled_tokens <= 1) + ): @@ -12449,8 +15493,8 @@ index 74938a823..27135b752 100644 + return flag + t = torch.tensor([1 if flag else 0], device=logits.device) + torch.distributed.all_reduce( -+ t, op=torch.distributed.ReduceOp.MAX, -+ group=_tp.device_group) ++ t, op=torch.distributed.ReduceOp.MAX, group=_tp.device_group ++ ) + return bool(t.item()) + + fire = _or_tp(fire) @@ -12471,6 +15515,8 @@ index 74938a823..27135b752 100644 + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, + skip_compiled=has_encoder_input, + ): + regated_output = self._model_forward( @@ -12492,7 +15538,7 @@ index 74938a823..27135b752 100644 self.execute_model_state = ExecuteModelState( scheduler_output, logits, -@@ -4416,6 +4788,95 @@ class GPUModelRunner( +@@ -4416,6 +4898,101 @@ class GPUModelRunner( return None @@ -12542,6 +15588,8 @@ index 74938a823..27135b752 100644 + batch_descriptor=ctx["batch_desc"], + ubatch_slices=ctx["ubatch_slices_padded"], + slot_mapping=ctx["slot_mappings"], ++ token_slot_mapping=ctx["token_slot_mapping"], ++ has_prefill=ctx["has_prefill"], + skip_compiled=ctx["has_encoder_input"], + ): + out = self._model_forward( @@ -12566,8 +15614,10 @@ index 74938a823..27135b752 100644 + } + pp.send_tensor_dict(send, all_gather_group=tp) + if _trace: -+ logger.info("[gate-pp] rank=%d replayed stage at FP4 -> sent", -+ pp.rank_in_group) ++ logger.info( ++ "[gate-pp] rank=%d replayed stage at FP4 -> sent", ++ pp.rank_in_group, ++ ) + else: + sample_hidden_states = hidden_states[ctx["logits_indices"]] + new_logits = self.model.compute_logits(sample_hidden_states) @@ -12580,15 +15630,17 @@ index 74938a823..27135b752 100644 + aux_hidden_states=aux_hidden_states, + ) + if _trace: -+ logger.info("[gate-pp] rank=%d (last) replayed -> FP4 logits", -+ pp.rank_in_group) ++ logger.info( ++ "[gate-pp] rank=%d (last) replayed -> FP4 logits", ++ pp.rank_in_group, ++ ) + except Exception as e: # noqa: BLE001 - gate must never crash serving + logger.warning("moe_w2 PP gate re-forward skipped: %s", e) + def _input_fits_in_drafter( self, common_attn_metadata: CommonAttentionMetadata | None ) -> bool: -@@ -4441,6 +4902,12 @@ class GPUModelRunner( +@@ -4441,6 +5018,12 @@ class GPUModelRunner( # receive sampled token ids from the last PP rank. if self.use_async_scheduling and not get_pp_group().is_last_rank: self._pp_receive_prev_sampled_token_ids_to_input_batch() @@ -12601,7 +15653,7 @@ index 74938a823..27135b752 100644 # In case of PP with kv transfer, we need to pass through the # kv_connector_output return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) -@@ -4605,6 +5072,23 @@ class GPUModelRunner( +@@ -4605,6 +5188,23 @@ class GPUModelRunner( # tokens on the CPU, so they are run after bookkeeping. propose_draft_token_ids(valid_sampled_token_ids) @@ -12625,7 +15677,7 @@ index 74938a823..27135b752 100644 # Finalize KV connector (wait_for_save + clear metadata) after # draft model runs. Deferred from target model forward to allow # draft model to also save its KV cache. -@@ -4696,13 +5180,35 @@ class GPUModelRunner( +@@ -4696,13 +5296,35 @@ class GPUModelRunner( def _pp_broadcast_prev_sampled_token_ids( self, sampled_token_ids: torch.Tensor ) -> None: @@ -12666,7 +15718,7 @@ index 74938a823..27135b752 100644 # Skip for chunked prefill: sampled tokens are dummy # and will be discarded, no need to broadcast. if not self._is_all_reqs_chunked_prefill(): -@@ -4711,16 +5217,33 @@ class GPUModelRunner( +@@ -4711,16 +5333,33 @@ class GPUModelRunner( ) def _pp_receive_prev_sampled_token_ids_to_input_batch(self) -> None: @@ -12704,7 +15756,7 @@ index 74938a823..27135b752 100644 # construct `prev_req_id_to_index` here so `_prepare_input_ids` # can map req_id -> previous batch row -@@ -4731,18 +5254,155 @@ class GPUModelRunner( +@@ -4731,18 +5370,157 @@ class GPUModelRunner( if i in discard_req_indices_set: continue prev_req_id_to_index[req_id] = i @@ -12857,7 +15909,9 @@ index 74938a823..27135b752 100644 + if moe_w2_gate is not None: + try: + from vllm.model_executor.layers.quantization.utils import ( -+ moe_w2_delta as _w2d) ++ moe_w2_delta as _w2d, ++ ) ++ + if _w2d.spec_suppressed(): + return None + except Exception: # noqa: BLE001 - guard must never crash @@ -12865,7 +15919,7 @@ index 74938a823..27135b752 100644 draft_token_ids, req_ids = self._get_draft_token_ids_cpu() return DraftTokenIds(req_ids, draft_token_ids) -@@ -5209,6 +5869,40 @@ class GPUModelRunner( +@@ -5209,6 +5987,45 @@ class GPUModelRunner( ) eplb_models += 1 @@ -12892,54 +15946,113 @@ index 74938a823..27135b752 100644 + # any cudagraph capture. No-op unless armed via env. + if moe_w2_gate is not None: + try: -+ from vllm.model_executor.layers.quantization.utils \ -+ import moe_w2_delta as _w2d -+ from vllm.model_executor.layers.quantization.utils \ -+ import moe_w2_looka as _w2l ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_delta as _w2d, ++ ) ++ from vllm.model_executor.layers.quantization.utils import ( ++ moe_w2_looka as _w2l, ++ ) ++ + if _w2d._BASE_TIER is not None: -+ _w2l.arm(self.model, -+ _w2d._BASE_TIER.n_layers, -+ _w2d._BASE_TIER.dev) ++ _w2l.arm( ++ self.model, ++ _w2d._BASE_TIER.n_layers, ++ _w2d._BASE_TIER.dev, ++ ) + except Exception as e: # noqa: BLE001 - never fatal + logger.warning("moe_w2 LOOKA arm failed: %s", e) + self._setup_eagle3_aux_hidden_state_outputs() # Resolve the MoE model, unwrapping VLM wrappers if needed. -@@ -5969,7 +6663,8 @@ class GPUModelRunner( +@@ -5823,11 +6640,13 @@ class GPUModelRunner( + + attn_metadata: PerLayerAttnMetadata | None = None + +- slot_mappings_by_group, slot_mappings = self._get_slot_mappings( +- num_tokens_padded=num_tokens_padded, +- num_reqs_padded=num_reqs_padded, +- num_tokens_unpadded=num_tokens_unpadded, +- ubatch_slices=ubatch_slices_padded, ++ slot_mappings_by_group, slot_mappings, token_slot_mapping = ( ++ self._get_slot_mappings( ++ num_tokens_padded=num_tokens_padded, ++ num_reqs_padded=num_reqs_padded, ++ num_tokens_unpadded=num_tokens_unpadded, ++ ubatch_slices=ubatch_slices_padded, ++ ) + ) + + # Dummy runs have no real slot assignments — fill with -1 so +@@ -5954,6 +6773,8 @@ class GPUModelRunner( + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=None, + ), + ): + outputs = self.model( +@@ -5969,10 +6790,15 @@ class GPUModelRunner( else: hidden_states = outputs - if self.speculative_config and ( +- self.speculative_config.use_eagle() +- or self.speculative_config.uses_draft_model() +- or self.speculative_config.uses_extract_hidden_states() + # The drafter lives on the last PP rank only. -+ if self.speculative_config and get_pp_group().is_last_rank and ( - self.speculative_config.use_eagle() - or self.speculative_config.uses_draft_model() - or self.speculative_config.uses_extract_hidden_states() -@@ -6875,8 +7570,9 @@ class GPUModelRunner( ++ if ( ++ self.speculative_config ++ and get_pp_group().is_last_rank ++ and ( ++ self.speculative_config.use_eagle() ++ or self.speculative_config.uses_draft_model() ++ or self.speculative_config.uses_extract_hidden_states() ++ ) + ): + assert isinstance( + self.drafter, +@@ -6875,10 +7701,15 @@ class GPUModelRunner( # because some of them change the threshold at init time. self.calculate_reorder_batch_threshold() - # Initialize drafter attention backend - if self.speculative_config and ( +- self.speculative_config.use_eagle() +- or self.speculative_config.uses_draft_model() + # Initialize drafter attention backend (drafter lives on the last PP + # rank only). -+ if self.speculative_config and get_pp_group().is_last_rank and ( - self.speculative_config.use_eagle() - or self.speculative_config.uses_draft_model() ++ if ( ++ self.speculative_config ++ and get_pp_group().is_last_rank ++ and ( ++ self.speculative_config.use_eagle() ++ or self.speculative_config.uses_draft_model() ++ ) ): -@@ -6929,7 +7625,8 @@ class GPUModelRunner( + assert isinstance( + self.drafter, +@@ -6929,9 +7760,14 @@ class GPUModelRunner( ) # Initialize drafter's cudagraph dispatcher if using spec decode. - if self.speculative_config and ( +- self.speculative_config.use_eagle() +- or self.speculative_config.uses_extract_hidden_states() + # The drafter lives on the last PP rank only. -+ if self.speculative_config and get_pp_group().is_last_rank and ( - self.speculative_config.use_eagle() - or self.speculative_config.uses_extract_hidden_states() ++ if ( ++ self.speculative_config ++ and get_pp_group().is_last_rank ++ and ( ++ self.speculative_config.use_eagle() ++ or self.speculative_config.uses_extract_hidden_states() ++ ) ): + assert isinstance( + self.drafter, diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py -index 657fc8267..dee74ffb4 100644 +index 657fc82..9dfdb20 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -278,6 +278,9 @@ class UBatchWrapper: @@ -12952,8 +16065,65 @@ index 657fc8267..dee74ffb4 100644 ): ubatch_metadata[0].context.cpu_wait_event.set() for thread in ubatch_threads: +@@ -335,6 +338,8 @@ class UBatchWrapper: + ubatch_slices, + attn_metadata, + slot_mapping, ++ token_slot_mapping, ++ has_prefill, + input_ids, + positions, + inputs_embeds, +@@ -349,6 +354,9 @@ class UBatchWrapper: + # slot_mapping can be None, an empty dict (from create_forward_context + # converting None to {}), or a list of dicts (one per ubatch) + has_slot_mapping = slot_mapping and isinstance(slot_mapping, list) ++ has_token_slot_mapping = isinstance(token_slot_mapping, list) ++ if token_slot_mapping is not None and not has_token_slot_mapping: ++ raise RuntimeError("ubatched token slot mapping must be a list of views") + for i, ubatch_slice in enumerate(ubatch_slices): + forward_contexts.append( + create_forward_context( +@@ -358,6 +366,10 @@ class UBatchWrapper: + batch_descriptor=batch_descriptor, + cudagraph_runtime_mode=cudagraph_runtime_mode, + slot_mapping=slot_mapping[i] if has_slot_mapping else None, ++ token_slot_mapping=( ++ token_slot_mapping[i] if has_token_slot_mapping else None ++ ), ++ has_prefill=has_prefill, + ) + ) + +@@ -455,6 +467,8 @@ class UBatchWrapper: + + attn_metadata = forward_context.attn_metadata + slot_mapping = forward_context.slot_mapping ++ token_slot_mapping = forward_context.token_slot_mapping ++ has_prefill = forward_context.has_prefill + num_tokens = sum(ubatch_slice.num_tokens for ubatch_slice in ubatch_slices) + input_ids = kwargs["input_ids"] + positions = kwargs["positions"] +@@ -488,6 +502,8 @@ class UBatchWrapper: + ubatch_slices=ubatch_slices, + attn_metadata=attn_metadata, + slot_mapping=slot_mapping, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, +@@ -514,6 +530,8 @@ class UBatchWrapper: + ubatch_slices=ubatch_slices, + attn_metadata=attn_metadata, + slot_mapping=slot_mapping, ++ token_slot_mapping=token_slot_mapping, ++ has_prefill=has_prefill, + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py -index 5e266a313..75eb17b84 100644 +index 5e266a3..d8fa6bb 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -616,6 +616,31 @@ class Worker(WorkerBase): @@ -13005,16 +16175,19 @@ index 5e266a313..75eb17b84 100644 def get_supported_tasks(self) -> tuple[SupportedTask, ...]: return self.model_runner.get_supported_tasks() -@@ -905,6 +940,8 @@ class Worker(WorkerBase): +@@ -905,6 +940,11 @@ class Worker(WorkerBase): if isinstance( output, ModelRunnerOutput | AsyncModelRunnerOutput | NoneType ): + # Last PP rank: re-decide at FP4 (full pipeline) before sampling. -+ self._gate_pp_barrier(forward_pass) ++ try: ++ self._gate_pp_barrier(forward_pass) ++ finally: ++ self._finish_w2_manager_step(forward_pass) return output assert isinstance(output, IntermediateTensors) -@@ -914,15 +951,69 @@ class Worker(WorkerBase): +@@ -914,15 +954,80 @@ class Worker(WorkerBase): and not get_pp_group().is_last_rank ) @@ -13042,9 +16215,20 @@ index 5e266a313..75eb17b84 100644 + # Non-last PP rank: participate in the gate barrier + (if fired) re-run + # this stage at FP4 for the full-pipeline re-decide. -+ self._gate_pp_barrier(forward_pass) ++ try: ++ self._gate_pp_barrier(forward_pass) ++ finally: ++ self._finish_w2_manager_step(forward_pass) return None ++ def _finish_w2_manager_step(self, forward_pass: bool) -> None: ++ """Release tier managers only after every target/replay has drained.""" ++ if not forward_pass or os.getenv("VLLM_MOE_W2", "0") != "1": ++ return ++ from vllm.model_executor.layers.quantization.utils import moe_w2_delta ++ ++ moe_w2_delta.finish_forward_step() ++ + def _gate_pp_barrier(self, forward_pass: bool) -> None: + """Confidence-gate full re-forward under PP (opt-in, VLLM_MOE_W2_GATE). + diff --git a/tools/test_moe_w2_forward.py b/tools/test_moe_w2_forward.py index 2f8f294..4cf4773 100644 --- a/tools/test_moe_w2_forward.py +++ b/tools/test_moe_w2_forward.py @@ -8,6 +8,7 @@ Run (inside the vllm image): python3 test_moe_w2_forward.py """ + import os import sys @@ -15,48 +16,89 @@ os.environ.setdefault("VLLM_MOE_W2", "1") -from vllm.model_executor.layers.quantization.utils import moe_w2_cubit # noqa: E402 +from vllm.model_executor.layers.quantization.utils import ( # noqa: E402 + moe_w2_cubit, + moe_w2_delta, +) from vllm.model_executor.layers.quantization.utils.moe_w2_planes import ( # noqa: E402 - mxfp4_to_codes, pack_fragment_major, pack_scales, + mxfp4_to_codes, + mxfp4_to_nibbles, + nibbles_to_refinement, + pack_fp4_fragment_major, + pack_fragment_major, + pack_scales, + split_fp4_dequant, ) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( # noqa: E402 per_token_group_quant_fp8, ) +from vllm.forward_context import ( # noqa: E402 + ForwardContext, + override_forward_context, +) assert moe_w2_cubit._ensure_ready(), "cubins not found" dev = torch.device("cuda") torch.manual_seed(11) E = int(os.environ.get("E", "32")) -H = int(os.environ.get("H", "4096")) # 4096 DS4, 6144 GLM-5.x, 7168 Kimi-K2.x -I = int(os.environ.get("I", "2048")) # per-rank I under TP (1024 TP2, 512 TP4) -T = int(os.environ.get("T", "9")) # T>96 exercises the PREFILL tier (mc4/afrag) +H = int(os.environ.get("H", "4096")) # 4096 DS4, 6144 GLM-5.x, 7168 Kimi-K2.x +INTERMEDIATE = int(os.environ.get("I", "2048")) # 1024 TP2, 512 TP4 +T = int(os.environ.get("T", "9")) # T>96 exercises the PREFILL tier (mc4/afrag) TOPK = 6 LEVELS = torch.tensor([-4.0, -1.0, 1.0, 4.0], device=dev) -w13_pack = torch.randint(0, 256, (E, 2 * I, H // 2), dtype=torch.uint8, device=dev) -s13 = torch.randint(118, 124, (E, 2 * I, H // 32), dtype=torch.uint8, device=dev) -w2_pack = torch.randint(0, 256, (E, H, I // 2), dtype=torch.uint8, device=dev) -s2 = torch.randint(118, 124, (E, H, I // 32), dtype=torch.uint8, device=dev) +w13_pack = torch.randint( + 0, 256, (E, 2 * INTERMEDIATE, H // 2), dtype=torch.uint8, device=dev +) +s13 = torch.randint( + 118, 124, (E, 2 * INTERMEDIATE, H // 32), dtype=torch.uint8, device=dev +) +w2_pack = torch.randint( + 0, 256, (E, H, INTERMEDIATE // 2), dtype=torch.uint8, device=dev +) +s2 = torch.randint(118, 124, (E, H, INTERMEDIATE // 32), dtype=torch.uint8, device=dev) -st = dict(N13=2 * I, K13=H, N2=H, K2=I, E=E) -st["planes13"] = torch.stack([pack_fragment_major(mxfp4_to_codes(w13_pack[e])) for e in range(E)]) +st = dict(N13=2 * INTERMEDIATE, K13=H, N2=H, K2=INTERMEDIATE, E=E) +st["planes13"] = torch.stack( + [pack_fragment_major(mxfp4_to_codes(w13_pack[e])) for e in range(E)] +) st["sc13"] = torch.stack([pack_scales(s13[e]) for e in range(E)]) -st["planes2"] = torch.stack([pack_fragment_major(mxfp4_to_codes(w2_pack[e])) for e in range(E)]) +st["planes2"] = torch.stack( + [pack_fragment_major(mxfp4_to_codes(w2_pack[e])) for e in range(E)] +) st["sc2"] = torch.stack([pack_scales(s2[e]) for e in range(E)]) moe_w2_cubit._LAYERS[0] = st def dequant(pack, sc): codes = mxfp4_to_codes(pack) - return LEVELS[codes.long()] * torch.exp2(sc.float() - 127.0).repeat_interleave(32, -1) + return LEVELS[codes.long()] * torch.exp2(sc.float() - 127.0).repeat_interleave( + 32, -1 + ) x = (torch.randn(T, H, device=dev) * 0.3).to(torch.bfloat16) -topk_ids = torch.stack([torch.randperm(E, device=dev)[:TOPK] for _ in range(T)]).to(torch.int32) +topk_ids = torch.stack([torch.randperm(E, device=dev)[:TOPK] for _ in range(T)]).to( + torch.int32 +) topk_w = torch.rand(T, TOPK, device=dev) * 0.5 -got = moe_w2_cubit._moe_w2_forward(x, topk_w, topk_ids, 0) + +def run_forward(): + """Exercise the kernel under the same persistent-slot contract as serving.""" + context = ForwardContext( + no_compile_layers={}, + attn_metadata={}, + slot_mapping={}, + token_slot_mapping=torch.arange(T, device=dev), + has_prefill=T > 96, + ) + with override_forward_context(context): + return moe_w2_cubit._moe_w2_forward(x, topk_w, topk_ids, 0) + + +got = run_forward() # ---- reference with the same activation-quant numerics a8, as8 = per_token_group_quant_fp8(x, 128) @@ -67,7 +109,7 @@ def dequant(pack, sc): e = int(topk_ids[t, j]) w13d = dequant(w13_pack[e], s13[e]) c13 = a_deq[t] @ w13d.T - act = torch.nn.functional.silu(c13[:I]) * c13[I:] + act = torch.nn.functional.silu(c13[:INTERMEDIATE]) * c13[INTERMEDIATE:] q2, qs2 = per_token_group_quant_fp8(act.to(torch.bfloat16).unsqueeze(0), 128) act_deq = q2.float() * qs2.repeat_interleave(128, 1) w2d = dequant(w2_pack[e], s2[e]) @@ -75,7 +117,8 @@ def dequant(pack, sc): rel = (got.float() - ref).abs().max().item() / ref.abs().max().item() cos = torch.nn.functional.cosine_similarity( - got.float().flatten(), ref.flatten(), dim=0).item() + got.float().flatten(), ref.flatten(), dim=0 +).item() print(f"T={T} E={E}: max_rel={rel:.3e} cos={cos:.6f}") ok = rel < 0.06 and cos > 0.999 @@ -87,21 +130,19 @@ def dequant(pack, sc): print("DELTA mixed: skipped (prefill tier is 2-bit-only by design)") print("RESULT:", "PASS" if ok else "FAIL") sys.exit(0 if ok else 1) -from vllm.model_executor.layers.quantization.utils import moe_w2_delta -from vllm.model_executor.layers.quantization.utils.moe_w2_planes import ( - mxfp4_to_nibbles, pack_fp4_fragment_major) - os.environ["VLLM_MOE_W2_DELTA_GB"] = "1" # per-expert FP4 plane bytes for THIS model's shapes (the module defaults are # the DS4 TP1 sizes; production passes these via _fp4_tier_for_build) -tier = moe_w2_delta.DeltaTier(1, E, dev, - w13_bytes=2 * I * H // 2, - w2_bytes=H * I // 2) +tier = moe_w2_delta.DeltaTier( + 1, E, dev, w13_bytes=2 * INTERMEDIATE * H // 2, w2_bytes=H * INTERMEDIATE // 2 +) moe_w2_delta._TIER = tier -fp13 = torch.stack([pack_fp4_fragment_major(mxfp4_to_nibbles(w13_pack[e])) - for e in range(E)]) -fp2 = torch.stack([pack_fp4_fragment_major(mxfp4_to_nibbles(w2_pack[e])) - for e in range(E)]) +fp13 = torch.stack( + [pack_fp4_fragment_major(mxfp4_to_nibbles(w13_pack[e])) for e in range(E)] +) +fp2 = torch.stack( + [pack_fp4_fragment_major(mxfp4_to_nibbles(w2_pack[e])) for e in range(E)] +) tier.add_layer_host_planes(0, fp13, fp2) promoted = list(range(0, E, 2)) for e in promoted: @@ -109,7 +150,7 @@ def dequant(pack, sc): tier._promote(0, e, slot) torch.cuda.synchronize() -E2M1 = torch.tensor([0, .5, 1, 1.5, 2, 3, 4, 6] * 2, device=dev) +E2M1 = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6] * 2, device=dev) E2M1[8:] *= -1 @@ -118,16 +159,16 @@ def dequant_fp4(pack, sc): return E2M1[nib.long()] * torch.exp2(sc.float() - 127.0).repeat_interleave(32, -1) -got2 = moe_w2_cubit._moe_w2_forward(x, topk_w, topk_ids, 0) +got2 = run_forward() ref2 = torch.zeros(T, H, device=dev) for t in range(T): for j in range(TOPK): e = int(topk_ids[t, j]) dq13 = dequant_fp4 if e in promoted else dequant dq2 = dequant_fp4 if e in promoted else dequant - w13d = (dq13(w13_pack[e], s13[e])) + w13d = dq13(w13_pack[e], s13[e]) c13 = a_deq[t] @ w13d.T - act = torch.nn.functional.silu(c13[:I]) * c13[I:] + act = torch.nn.functional.silu(c13[:INTERMEDIATE]) * c13[INTERMEDIATE:] q2, qs2 = per_token_group_quant_fp8(act.to(torch.bfloat16).unsqueeze(0), 128) act_deq = q2.float() * qs2.repeat_interleave(128, 1) w2d = dq2(w2_pack[e], s2[e]) @@ -135,28 +176,34 @@ def dequant_fp4(pack, sc): rel2 = (got2.float() - ref2).abs().max().item() / ref2.abs().max().item() cos2 = torch.nn.functional.cosine_similarity( - got2.float().flatten(), ref2.flatten(), dim=0).item() -print(f"DELTA mixed ({len(promoted)}/{E} promoted): max_rel={rel2:.3e} " - f"cos={cos2:.6f}") + got2.float().flatten(), ref2.flatten(), dim=0 +).item() +print(f"DELTA mixed ({len(promoted)}/{E} promoted): max_rel={rel2:.3e} cos={cos2:.6f}") ok = ok and rel2 < 0.06 and cos2 > 0.999 # ---- SPLIT FP4 (VLLM_MOE_W2_DELTA_SPLIT): the delta slots hold 2-bit # REFINEMENT planes and moe_w4s_mm reads them alongside the resident base. # Reference: split_fp4_dequant (true FP4 modulo the mag-0 -> 0.5 merge). -from vllm.model_executor.layers.quantization.utils.moe_w2_planes import ( - nibbles_to_refinement, split_fp4_dequant) - -moe_w2_delta._SPLIT = True # env is read at import; force for test +moe_w2_delta._SPLIT = True # env is read at import; force for test assert moe_w2_delta.split_enabled() -tier_s = moe_w2_delta.DeltaTier(1, E, dev, - w13_bytes=2 * I * H // 4, - w2_bytes=H * I // 4) +tier_s = moe_w2_delta.DeltaTier( + 1, E, dev, w13_bytes=2 * INTERMEDIATE * H // 4, w2_bytes=H * INTERMEDIATE // 4 +) moe_w2_delta._TIER = tier_s -rf13 = torch.stack([pack_fragment_major( - nibbles_to_refinement(mxfp4_to_nibbles(w13_pack[e]))) for e in range(E)]) -rf2 = torch.stack([pack_fragment_major( - nibbles_to_refinement(mxfp4_to_nibbles(w2_pack[e]))) for e in range(E)]) -assert rf13.shape[1] == 2 * I * H // 4 and rf2.shape[1] == H * I // 4 +rf13 = torch.stack( + [ + pack_fragment_major(nibbles_to_refinement(mxfp4_to_nibbles(w13_pack[e]))) + for e in range(E) + ] +) +rf2 = torch.stack( + [ + pack_fragment_major(nibbles_to_refinement(mxfp4_to_nibbles(w2_pack[e]))) + for e in range(E) + ] +) +assert rf13.shape[1] == 2 * INTERMEDIATE * H // 4 +assert rf2.shape[1] == H * INTERMEDIATE // 4 tier_s.add_layer_host_planes(0, rf13, rf2) for e in promoted: slot = tier_s._take_slot(set()) @@ -166,11 +213,12 @@ def dequant_fp4(pack, sc): def dequant_split(pack, sc): nib = mxfp4_to_nibbles(pack) - return (split_fp4_dequant(nib) - * torch.exp2(sc.float() - 127.0).repeat_interleave(32, -1)) + return split_fp4_dequant(nib) * torch.exp2(sc.float() - 127.0).repeat_interleave( + 32, -1 + ) -got3 = moe_w2_cubit._moe_w2_forward(x, topk_w, topk_ids, 0) +got3 = run_forward() ref3 = torch.zeros(T, H, device=dev) for t in range(T): for j in range(TOPK): @@ -178,7 +226,7 @@ def dequant_split(pack, sc): dq = dequant_split if e in promoted else dequant w13d = dq(w13_pack[e], s13[e]) c13 = a_deq[t] @ w13d.T - act = torch.nn.functional.silu(c13[:I]) * c13[I:] + act = torch.nn.functional.silu(c13[:INTERMEDIATE]) * c13[INTERMEDIATE:] q2, qs2 = per_token_group_quant_fp8(act.to(torch.bfloat16).unsqueeze(0), 128) act_deq = q2.float() * qs2.repeat_interleave(128, 1) w2d = dq(w2_pack[e], s2[e]) @@ -186,9 +234,12 @@ def dequant_split(pack, sc): rel3 = (got3.float() - ref3).abs().max().item() / ref3.abs().max().item() cos3 = torch.nn.functional.cosine_similarity( - got3.float().flatten(), ref3.flatten(), dim=0).item() -print(f"SPLIT mixed ({len(promoted)}/{E} promoted, slots {rf13.shape[1]}+" - f"{rf2.shape[1]} B): max_rel={rel3:.3e} cos={cos3:.6f}") + got3.float().flatten(), ref3.flatten(), dim=0 +).item() +print( + f"SPLIT mixed ({len(promoted)}/{E} promoted, slots {rf13.shape[1]}+" + f"{rf2.shape[1]} B): max_rel={rel3:.3e} cos={cos3:.6f}" +) ok = ok and rel3 < 0.06 and cos3 > 0.999 moe_w2_delta._SPLIT = False print("RESULT:", "PASS" if ok else "FAIL") diff --git a/tools/test_three_tier_split.py b/tools/test_three_tier_split.py index de96087..d170e68 100644 --- a/tools/test_three_tier_split.py +++ b/tools/test_three_tier_split.py @@ -18,6 +18,7 @@ Run (inside the vllm image): python3 test_three_tier_split.py """ + import os import sys @@ -25,16 +26,25 @@ os.environ.setdefault("VLLM_MOE_W2", "1") os.environ["VLLM_MOE_W2_DELTA_SPLIT"] = "1" +os.environ["VLLM_MOE_W2_DELTA_GB"] = "0.1" from vllm.model_executor.layers.quantization.utils import moe_w2_cubit # noqa: E402 from vllm.model_executor.layers.quantization.utils import moe_w2_delta # noqa: E402 from vllm.model_executor.layers.quantization.utils.moe_w2_planes import ( # noqa: E402 - mxfp4_to_codes, mxfp4_to_nibbles, nibbles_to_refinement, - pack_fragment_major, pack_scales, split_fp4_dequant, + mxfp4_to_codes, + mxfp4_to_nibbles, + nibbles_to_refinement, + pack_fragment_major, + pack_scales, + split_fp4_dequant, ) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( # noqa: E402 per_token_group_quant_fp8, ) +from vllm.forward_context import ( # noqa: E402 + ForwardContext, + override_forward_context, +) assert moe_w2_cubit._ensure_ready(), "cubins not found" assert moe_w2_delta.split_enabled() @@ -43,55 +53,99 @@ E = int(os.environ.get("E", "32")) H = int(os.environ.get("H", "4096")) -I = int(os.environ.get("I", "2048")) +INTERMEDIATE = int(os.environ.get("I", "2048")) T = int(os.environ.get("T", "9")) TOPK = 6 LEVELS = torch.tensor([-4.0, -1.0, 1.0, 4.0], device=dev) -w13_pack = torch.randint(0, 256, (E, 2 * I, H // 2), dtype=torch.uint8, device=dev) -s13 = torch.randint(118, 124, (E, 2 * I, H // 32), dtype=torch.uint8, device=dev) -w2_pack = torch.randint(0, 256, (E, H, I // 2), dtype=torch.uint8, device=dev) -s2 = torch.randint(118, 124, (E, H, I // 32), dtype=torch.uint8, device=dev) +w13_pack = torch.randint( + 0, 256, (E, 2 * INTERMEDIATE, H // 2), dtype=torch.uint8, device=dev +) +s13 = torch.randint( + 118, 124, (E, 2 * INTERMEDIATE, H // 32), dtype=torch.uint8, device=dev +) +w2_pack = torch.randint( + 0, 256, (E, H, INTERMEDIATE // 2), dtype=torch.uint8, device=dev +) +s2 = torch.randint(118, 124, (E, H, INTERMEDIATE // 32), dtype=torch.uint8, device=dev) -planes13 = torch.stack([pack_fragment_major(mxfp4_to_codes(w13_pack[e])) for e in range(E)]) +planes13 = torch.stack( + [pack_fragment_major(mxfp4_to_codes(w13_pack[e])) for e in range(E)] +) sc13p = torch.stack([pack_scales(s13[e]) for e in range(E)]) -planes2 = torch.stack([pack_fragment_major(mxfp4_to_codes(w2_pack[e])) for e in range(E)]) +planes2 = torch.stack( + [pack_fragment_major(mxfp4_to_codes(w2_pack[e])) for e in range(E)] +) sc2p = torch.stack([pack_scales(s2[e]) for e in range(E)]) c13len, s13len = planes13.shape[1], sc13p.shape[1] c2len, s2len = planes2.shape[1], sc2p.shape[1] # base-cache layer state, exactly _finish_layer's base branch moe_w2_cubit._LAYERS[0] = dict( - N13=2 * I, K13=H, N2=H, K2=I, E=E, base=True, - off_s13=c13len, off_c2=c13len + s13len, + N13=2 * INTERMEDIATE, + K13=H, + N2=H, + K2=INTERMEDIATE, + E=E, + base=True, + off_s13=c13len, + off_c2=c13len + s13len, off_s2=c13len + s13len + c2len, - off4_s13=2 * c13len, off4_c2=2 * c13len + s13len, + off4_s13=2 * c13len, + off4_c2=2 * c13len + s13len, off4_s2=2 * c13len + s13len + 2 * c2len, ) -# BASE tier: small pool (forces eviction pressure in check 3) -moe_w2_delta._BASE_GB = 0.25 # base_enabled() -> True -btier = moe_w2_delta.DeltaTier(1, E, dev, - w13_bytes=c13len + s13len, - w2_bytes=c2len + s2len, - pool_gb=0.25, policy="lru", tag="base") +# Exercise the production fresh-requant creation order: the split-FP4 tier +# is constructed first, then _finish_layer creates/looks up the base tier. +moe_w2_delta._BASE_GB = 0.25 # base_enabled() -> True +assert moe_w2_delta._TIER is None and moe_w2_delta._BASE_TIER is None +tier = moe_w2_delta.get_tier( + n_layers=1, + n_experts=E, + dev=dev, + w13_bytes=2 * INTERMEDIATE * H // 4, + w2_bytes=H * INTERMEDIATE // 4, +) +assert tier is not None +assert tier._tag == "fp4s" + +# BASE tier: small pool (forces eviction pressure in check 3). Assign it as +# if get_base_tier had just built it, then call the production lookup to prove +# the tier-first path arms coupling too. +btier = moe_w2_delta.DeltaTier( + 1, + E, + dev, + w13_bytes=c13len + s13len, + w2_bytes=c2len + s2len, + pool_gb=0.25, + policy="lru", + tag="base", +) btier.miss_count = torch.zeros(1, dtype=torch.int32, device=dev) moe_w2_delta._BASE_TIER = btier -btier.add_layer_host_planes(0, torch.cat((planes13, sc13p), dim=1), - torch.cat((planes2, sc2p), dim=1)) +assert moe_w2_delta.get_base_tier(1, E, dev, c13len + s13len, c2len + s2len) is btier +assert btier._coupled_fp4 is tier +btier.add_layer_host_planes( + 0, + torch.cat((planes13, sc13p), dim=1), + torch.cat((planes2, sc2p), dim=1), +) # FP4 need-pool in SPLIT mode: refinement planes, no scale sections -tier = moe_w2_delta.DeltaTier(1, E, dev, - w13_bytes=2 * I * H // 4, - w2_bytes=H * I // 4, - pool_gb=0.1, policy="freq", tag="fp4", - host_pinned=True) -moe_w2_delta._TIER = tier -btier._coupled_fp4 = tier # the residency coupling under test -rf13 = torch.stack([pack_fragment_major( - nibbles_to_refinement(mxfp4_to_nibbles(w13_pack[e]))) for e in range(E)]) -rf2 = torch.stack([pack_fragment_major( - nibbles_to_refinement(mxfp4_to_nibbles(w2_pack[e]))) for e in range(E)]) +rf13 = torch.stack( + [ + pack_fragment_major(nibbles_to_refinement(mxfp4_to_nibbles(w13_pack[e]))) + for e in range(E) + ] +) +rf2 = torch.stack( + [ + pack_fragment_major(nibbles_to_refinement(mxfp4_to_nibbles(w2_pack[e]))) + for e in range(E) + ] +) tier.add_layer_host_planes(0, rf13, rf2) # make every expert BASE-resident, half of them FP4-resident @@ -105,21 +159,27 @@ assert all(int(btier._mirror[0, e]) >= 0 for e in range(E)) x = (torch.randn(T, H, device=dev) * 0.3).to(torch.bfloat16) -topk_ids = torch.stack([torch.randperm(E, device=dev)[:TOPK] for _ in range(T)]).to(torch.int32) +topk_ids = torch.stack([torch.randperm(E, device=dev)[:TOPK] for _ in range(T)]).to( + torch.int32 +) topk_w = torch.rand(T, TOPK, device=dev) * 0.5 -E2M1 = torch.tensor([0, .5, 1, 1.5, 2, 3, 4, 6] * 2, device=dev) +E2M1 = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6] * 2, device=dev) E2M1[8:] *= -1 def dequant2(pack, sc): codes = mxfp4_to_codes(pack) - return LEVELS[codes.long()] * torch.exp2(sc.float() - 127.0).repeat_interleave(32, -1) + return LEVELS[codes.long()] * torch.exp2(sc.float() - 127.0).repeat_interleave( + 32, -1 + ) def dequant_split(pack, sc): nib = mxfp4_to_nibbles(pack) - return split_fp4_dequant(nib) * torch.exp2(sc.float() - 127.0).repeat_interleave(32, -1) + return split_fp4_dequant(nib) * torch.exp2(sc.float() - 127.0).repeat_interleave( + 32, -1 + ) def reference(zero_experts=()): @@ -134,22 +194,39 @@ def reference(zero_experts=()): dq = dequant_split if e in promoted else dequant2 w13d = dq(w13_pack[e], s13[e]) c13 = a_deq[t] @ w13d.T - act = torch.nn.functional.silu(c13[:I]) * c13[I:] - q2, qs2 = per_token_group_quant_fp8(act.to(torch.bfloat16).unsqueeze(0), 128) + act = torch.nn.functional.silu(c13[:INTERMEDIATE]) * c13[INTERMEDIATE:] + q2, qs2 = per_token_group_quant_fp8( + act.to(torch.bfloat16).unsqueeze(0), 128 + ) act_deq = q2.float() * qs2.repeat_interleave(128, 1) w2d = dq(w2_pack[e], s2[e]) ref[t] += float(topk_w[t, j]) * (act_deq[0] @ w2d.T) return ref +def run_forward(): + """Exercise the kernel under serving's persistent padded-slot contract.""" + context = ForwardContext( + no_compile_layers={}, + attn_metadata={}, + slot_mapping={}, + token_slot_mapping=torch.arange(T, device=dev), + has_prefill=False, + ) + with override_forward_context(context): + return moe_w2_cubit._moe_w2_forward(x, topk_w, topk_ids, 0) + + # ---- 1. mixed dispatch -------------------------------------------------- -got = moe_w2_cubit._moe_w2_forward(x, topk_w, topk_ids, 0) +got = run_forward() ref = reference() rel = (got.float() - ref).abs().max().item() / ref.abs().max().item() cos = torch.nn.functional.cosine_similarity( - got.float().flatten(), ref.flatten(), dim=0).item() -print(f"three-tier SPLIT mixed ({len(promoted)}/{E} FP4): max_rel={rel:.3e} " - f"cos={cos:.6f}") + got.float().flatten(), ref.flatten(), dim=0 +).item() +print( + f"three-tier SPLIT mixed ({len(promoted)}/{E} FP4): max_rel={rel:.3e} cos={cos:.6f}" +) ok = rel < 0.06 and cos > 0.999 assert int(btier.miss_count.item()) == 0, "unexpected misses in check 1" @@ -158,26 +235,30 @@ def reference(zero_experts=()): vb = int(btier._mirror[0, victim]) btier.slot_table[0, victim] = -1 btier._mirror[0, victim] = -1 -got2 = moe_w2_cubit._moe_w2_forward(x, topk_w, topk_ids, 0) +got2 = run_forward() miss = int(btier.miss_count.item()) ref2 = reference(zero_experts={victim}) rel2 = (got2.float() - ref2).abs().max().item() / ref2.abs().max().item() routed = int((topk_ids == victim).sum()) -print(f"transient (expert {victim} base-unmapped, routed {routed}x): " - f"miss_count={miss} max_rel={rel2:.3e}") +print( + f"transient (expert {victim} base-unmapped, routed {routed}x): " + f"miss_count={miss} max_rel={rel2:.3e}" +) ok = ok and miss > 0 and rel2 < 0.06 -btier.slot_table[0, victim] = vb # restore +btier.slot_table[0, victim] = vb # restore btier._mirror[0, victim] = vb # ---- 3. eviction coupling: FP4-mapped experts are never victims --------- btier.seen.zero_() btier._seen_host.zero_() -btier._tick += 10 # everything cold +btier._tick += 10 # everything cold with btier._lock: taken = btier._take_slots_batch(E, emergency=True) still = [e for e in promoted if int(btier._mirror[0, e]) >= 0] -print(f"eviction pressure: took {len(taken)} slots; FP4-mapped intact " - f"{len(still)}/{len(promoted)}") +print( + f"eviction pressure: took {len(taken)} slots; FP4-mapped intact " + f"{len(still)}/{len(promoted)}" +) ok = ok and len(still) == len(promoted) and len(taken) > 0 print("RESULT:", "PASS" if ok else "FAIL")