From b9c65c63460f097b5570ea85467f0a6e3137dcfa Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:12:18 +0900 Subject: [PATCH 01/13] [MoE] Expert pool: one shared GPU bank of NVFP4 Marlin expert rows with a device-side planner (moe_expert_pool_rows) Opt-in via --moe-expert-pool-rows N (default 0: unchanged). MoE layers keep their expert tensors in pinned host memory; after loading, one VRAM bank shared by all layers holds N rows per layer, a device-side LRU step program plans promotions per forward (no host code in the forward, so the MoE op stays inside CUDA graphs), and a Marlin consumer runs on the bank with logical alignment and a physical-row remap. Wider batches take a bank + host-view partition path. The placement is frozen (gate closed) through profiling and graph capture and opened at the end of warm-up. Supported: ModelOpt NVFP4 Marlin MoE backend, no EP/DP. Rejected at layer construction otherwise. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- tests/config/test_moe_expert_pool_rows.py | 35 + tests/kernels/expert_pool/__init__.py | 0 tests/kernels/expert_pool/marlin_fixture.py | 137 ++++ .../expert_pool/test_install_guards.py | 68 ++ .../expert_pool/test_pool_layer_helpers.py | 46 ++ .../expert_pool/test_pool_marlin_cuda.py | 161 +++++ .../expert_pool/test_pool_shared_experts.py | 21 + tests/kernels/expert_pool/test_pool_tables.py | 308 +++++++++ vllm/config/offload.py | 9 + vllm/engine/arg_utils.py | 5 + .../layers/fused_moe/expert_pool/__init__.py | 3 + .../layers/fused_moe/expert_pool/copy.py | 435 +++++++++++++ .../layers/fused_moe/expert_pool/install.py | 270 ++++++++ .../layers/fused_moe/expert_pool/layer.py | 261 ++++++++ .../layers/fused_moe/expert_pool/pool.py | 118 ++++ .../layers/fused_moe/expert_pool/tables.py | 610 ++++++++++++++++++ .../layers/fused_moe/routed_experts.py | 52 ++ .../layers/quantization/modelopt.py | 52 +- vllm/model_executor/model_loader/utils.py | 22 +- vllm/v1/worker/gpu_worker.py | 7 + 20 files changed, 2615 insertions(+), 5 deletions(-) create mode 100644 tests/config/test_moe_expert_pool_rows.py create mode 100644 tests/kernels/expert_pool/__init__.py create mode 100644 tests/kernels/expert_pool/marlin_fixture.py create mode 100644 tests/kernels/expert_pool/test_install_guards.py create mode 100644 tests/kernels/expert_pool/test_pool_layer_helpers.py create mode 100644 tests/kernels/expert_pool/test_pool_marlin_cuda.py create mode 100644 tests/kernels/expert_pool/test_pool_shared_experts.py create mode 100644 tests/kernels/expert_pool/test_pool_tables.py create mode 100644 vllm/model_executor/layers/fused_moe/expert_pool/__init__.py create mode 100644 vllm/model_executor/layers/fused_moe/expert_pool/copy.py create mode 100644 vllm/model_executor/layers/fused_moe/expert_pool/install.py create mode 100644 vllm/model_executor/layers/fused_moe/expert_pool/layer.py create mode 100644 vllm/model_executor/layers/fused_moe/expert_pool/pool.py create mode 100644 vllm/model_executor/layers/fused_moe/expert_pool/tables.py diff --git a/tests/config/test_moe_expert_pool_rows.py b/tests/config/test_moe_expert_pool_rows.py new file mode 100644 index 000000000000..b070b42b134f --- /dev/null +++ b/tests/config/test_moe_expert_pool_rows.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""moe_expert_pool_rows: config field, CLI -> EngineArgs, hash.""" + +import pytest +from pydantic import ValidationError + +from vllm.config.offload import OffloadConfig +from vllm.engine.arg_utils import EngineArgs +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +def test_default_is_off(): + assert OffloadConfig().moe_expert_pool_rows == 0 + + +def test_hash_distinguishes_the_pool_size(): + assert ( + OffloadConfig().compute_hash() + != OffloadConfig(moe_expert_pool_rows=8).compute_hash() + ) + + +def test_negative_rows_are_rejected(): + with pytest.raises(ValidationError): + OffloadConfig(moe_expert_pool_rows=-1) + + +def test_cli_reaches_engine_args_and_offload_config(): + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + args = parser.parse_args(["--moe-expert-pool-rows", "16"]) + engine_args = EngineArgs.from_cli_args(args) + assert engine_args.moe_expert_pool_rows == 16 + offload = OffloadConfig(moe_expert_pool_rows=engine_args.moe_expert_pool_rows) + assert offload.moe_expert_pool_rows == 16 diff --git a/tests/kernels/expert_pool/__init__.py b/tests/kernels/expert_pool/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/kernels/expert_pool/marlin_fixture.py b/tests/kernels/expert_pool/marlin_fixture.py new file mode 100644 index 000000000000..c3ede294aded --- /dev/null +++ b/tests/kernels/expert_pool/marlin_fixture.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Small NVFP4 Marlin MoE layers for the expert pool CUDA tests.""" + +import pytest +import torch + +from tests.kernels.moe.modular_kernel_tools.parallel_utils import _set_vllm_config +from tests.kernels.moe.utils import _scaled_fp4_quant_emulated +from vllm.config import ( + CompilationConfig, + ParallelConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory +from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.worker.workspace import ( + init_workspace_manager, + is_workspace_manager_initialized, +) + +E, K, N, TOP_K, M = 8, 256, 128, 2, 8 + +EXPERT_TENSORS = ("w13_weight", "w2_weight", "w13_weight_scale", "w2_weight_scale") + + +def vllm_config(pool_rows: int) -> VllmConfig: + cfg = VllmConfig( + parallel_config=ParallelConfig(), compilation_config=CompilationConfig() + ) + cfg.kernel_config.moe_backend = "marlin" + cfg.offload_config.moe_expert_pool_rows = pool_rows + return cfg + + +@pytest.fixture(scope="module") +def dist_env(): + cfg = vllm_config(0) + _set_vllm_config(cfg, 1, rank=0, local_rank=0) + if not is_workspace_manager_initialized(): + init_workspace_manager(torch.accelerator.current_accelerator()) + return cfg + + +def _quantize_row_major(w: torch.Tensor): + qs, ss, gs = [], [], [] + for i in range(w.shape[0]): + amax = w[i].abs().max().to(torch.float32) + g = torch.tensor(448.0 * 6.0, device=w.device) / amax + q, s = _scaled_fp4_quant_emulated(w[i], g) + qs.append(q) + ss.append(s) + gs.append(g) + return torch.stack(qs), torch.stack(ss), torch.stack(gs) + + +def quantized_weights(device, n: int = N, seed_offset: int = 0): + set_random_seed(11 + seed_offset) + w1 = torch.randn(E, 2 * n, K, dtype=torch.bfloat16, device=device) + w2 = torch.randn(E, K, n, dtype=torch.bfloat16, device=device) + # Distinct per-expert magnitudes so the global scales differ per expert. + mag = torch.tensor([0.5 + i for i in range(E)], device=device).view(E, 1, 1) + w1 = (w1 * mag).to(torch.bfloat16) + w2 = (w2 * mag.flip(0)).to(torch.bfloat16) + # Row-major [E, rows, K/16] block scales as a checkpoint stores them. The + # CUDA quant op would return the swizzled 128x4 layout padded to 128 rows, + # which is neither the checkpoint layout nor what Marlin's permute reads. + w1q, w1s, w1gs = _quantize_row_major(w1) + w2q, w2s, w2gs = _quantize_row_major(w2) + params = { + "w13_weight": w1q, + "w2_weight": w2q, + "w13_weight_scale": w1s, + "w2_weight_scale": w2s, + "w13_weight_scale_2": (1.0 / w1gs).unsqueeze(1).expand(-1, 2).contiguous(), + "w2_weight_scale_2": 1.0 / w2gs, + "w13_input_scale": torch.ones((E, 2), dtype=torch.float32, device=device), + "w2_input_scale": torch.ones(E, dtype=torch.float32, device=device), + } + assert params["w13_weight_scale"].shape == (E, 2 * n, K // 16) + assert params["w2_weight_scale"].shape == (E, K, n // 16) + assert torch.unique(params["w13_weight_scale_2"][:, 0]).numel() == E + assert torch.unique(params["w2_weight_scale_2"]).numel() == E + return params + + +def make_layer( + cfg: VllmConfig, params: dict[str, torch.Tensor], host_source: bool = False +): + """Build the layer; with host_source the per-expert tensors are + registered as pinned CPU tensors, the layout the loader restores when + the pool is enabled.""" + with set_current_vllm_config(cfg): + # Any construction error is a failure: the Marlin capability gate is + # the module-level skip in the test, and the backend is pinned. + layer = FusedMoEFactory( + num_experts=E, + top_k=TOP_K, + hidden_size=K, + intermediate_size=N, + params_dtype=torch.bfloat16, + renormalize=False, + quant_config=ModelOptNvFp4Config( + is_checkpoint_nvfp4_serialized=True, + kv_cache_quant_algo=None, + exclude_modules=[], + ), + tp_size=1, + dp_size=1, + prefix="from_forward_context", + ) + if cfg.offload_config.moe_expert_pool_rows > 0: + # create_weights wiring: expert tensors start in pinned host memory. + for name in EXPERT_TENSORS: + p = getattr(layer.routed_experts, name) + assert p.device.type == "cpu" and p.is_pinned(), name + for name, value in params.items(): + data = value.clone() + if host_source and name in EXPERT_TENSORS: + data = data.cpu().pin_memory() + layer.routed_experts.register_parameter( + name, torch.nn.Parameter(data, requires_grad=False) + ) + layer._quant_method.process_weights_after_loading(layer.routed_experts) + return layer + + +def routing(order: list[int], device) -> torch.Tensor: + # Token t routes to experts (order[t], order[(t + 1) % E]); every forward + # touches all E experts, more than the pool holds. + logits = torch.full((M, E), -10.0, device=device) + for t in range(M): + logits[t, order[t % E]] = 3.0 + logits[t, order[(t + 1) % E]] = 2.0 + return logits diff --git a/tests/kernels/expert_pool/test_install_guards.py b/tests/kernels/expert_pool/test_install_guards.py new file mode 100644 index 000000000000..ababd301e63f --- /dev/null +++ b/tests/kernels/expert_pool/test_install_guards.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""install_expert_pool rejects unsupported geometries and backends before +allocating anything, and bounds the planner width.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.expert_pool import install as inst +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend +from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4FusedMoE + + +def _method(backend=NvFp4MoeBackend.MARLIN): + m = ModelOptNvFp4FusedMoE.__new__(ModelOptNvFp4FusedMoE) + m.nvfp4_backend = backend + return m + + +def _layer(E=8, top_k=2, size=4, method=None, use_ep=False): + return SimpleNamespace( + quant_method=method or _method(), + local_num_experts=E, + _moe_expert_pool_rows=size, + moe_config=SimpleNamespace( + experts_per_token=top_k, + moe_parallel_config=SimpleNamespace(use_ep=use_ep), + ), + ) + + +def test_consistent_layers_pass(): + inst.check_pool_layers([("a", _layer()), ("b", _layer())]) + + +@pytest.mark.parametrize( + "bad, match", + [ + (_layer(E=16), "local_num_experts"), + (_layer(top_k=4), "top_k"), + (_layer(size=6), "_moe_expert_pool_rows"), + (_layer(method=_method(NvFp4MoeBackend.VLLM_CUTLASS)), "Marlin"), + (_layer(method=SimpleNamespace(nvfp4_backend=None)), "ModelOptNvFp4FusedMoE"), + (_layer(use_ep=True), "EP"), + ], +) +def test_mismatch_or_unsupported_backend_is_rejected(bad, match): + with pytest.raises(ValueError, match=match): + inst.check_pool_layers([("a", _layer()), ("b", bad)]) + + +def test_planner_width_is_bounded_by_the_decode_lane_cap(): + top_k = 10 + cap_tokens = inst.MAX_DECODE_LANES // top_k + for requested, expected in ((1, 1), (cap_tokens, cap_tokens), (256, cap_tokens)): + tokens = max(1, min(requested, cap_tokens)) + assert tokens == expected + assert inst._next_power_of_two(top_k * tokens) <= 2 * inst.MAX_DECODE_LANES + + +def test_top_k_beyond_the_lane_cap_is_rejected(): + assert inst.install_expert_pool(torch.nn.Module(), torch.device("cpu")) is None + for bad in (0, inst.MAX_DECODE_LANES + 1): + with pytest.raises(ValueError, match="0 < top_k"): + inst._check_top_k(bad) + inst._check_top_k(inst.MAX_DECODE_LANES) diff --git a/tests/kernels/expert_pool/test_pool_layer_helpers.py b/tests/kernels/expert_pool/test_pool_layer_helpers.py new file mode 100644 index 000000000000..bcb0a1ad5ce7 --- /dev/null +++ b/tests/kernels/expert_pool/test_pool_layer_helpers.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU checks of the pool layer helpers: route masking, block-to-row remap, +and the fixed-grid copy reference.""" + +import torch + +from vllm.model_executor.layers.fused_moe.expert_pool.copy import copy_rows +from vllm.model_executor.layers.fused_moe.expert_pool.layer import ( + marlin_block_size, + mask_routes, + physical_block_experts, +) +from vllm.model_executor.layers.fused_moe.expert_pool.tables import TENSORS + + +def test_mask_routes_hides_absent_experts_and_keeps_padding(): + expert_map = torch.tensor([5, -1, 7, -1], dtype=torch.int32) # rows for 0, 2 + ids = torch.tensor([[0, 1, 2, 3], [-1, 2, 9, 0]], dtype=torch.int32) + out = mask_routes(ids, expert_map) + assert out.tolist() == [[0, -1, 2, -1], [-1, 2, -1, 0]] + + +def test_physical_block_experts_maps_used_blocks_only(): + expert_map = torch.tensor([10, -1, 12], dtype=torch.int32) + logical = torch.tensor([2, 0, 7, 2], dtype=torch.int32) # block 2 is garbage + post_padded = torch.tensor([16], dtype=torch.int32) # 2 blocks of 8 used + rows = physical_block_experts(logical, post_padded, 8, expert_map, 3) + assert rows.tolist() == [12, 10, -1, -1] + + +def test_marlin_block_size_matches_stock_choice_shape(): + assert marlin_block_size(1, 10, 9984, 512, None) == 8 + assert marlin_block_size(4096, 10, 512, 512, None) == 64 + assert marlin_block_size(1, 10, 512, 512, torch.int8) >= 16 + + +def test_copy_rows_reference_copies_every_tensor_in_order(): + src = {n: torch.arange(8 * 4, dtype=torch.int32).reshape(8, 4) for n in TENSORS} + dst = {n: torch.zeros(3, 4, dtype=torch.int32) for n in TENSORS} + src_rows = torch.tensor([6, 1, 0, 0], dtype=torch.int32) + dst_rows = torch.tensor([2, 0, 0, 0], dtype=torch.int32) + copy_rows(src, dst, src_rows, dst_rows, torch.tensor([2], dtype=torch.int32)) + for n in TENSORS: + assert torch.equal(dst[n][2], src[n][6]) and torch.equal(dst[n][0], src[n][1]) + assert int(dst[n][1].sum()) == 0 diff --git a/tests/kernels/expert_pool/test_pool_marlin_cuda.py b/tests/kernels/expert_pool/test_pool_marlin_cuda.py new file mode 100644 index 000000000000..da9270edd4cd --- /dev/null +++ b/tests/kernels/expert_pool/test_pool_marlin_cuda.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The expert pool through the real Marlin consumer (CUDA). + +A layer with a small ``moe_expert_pool_rows`` resident count: decode steps +(one token) must match the uncached layer while experts are promoted and +evicted through the shared bank (bank rows > expert count, so the +logical-align + physical-remap path is exercised), and a wider batch must +match through the bank + host-view partition path. The pool tables must +stay consistent throughout.""" + +import pytest +import torch + +from tests.kernels.expert_pool.marlin_fixture import ( + TOP_K, + E, + K, + dist_env, # noqa: F401 + make_layer, + quantized_weights, + routing, + vllm_config, +) +from vllm.model_executor.layers.fused_moe.expert_pool.install import ( + install_expert_pool, +) +from vllm.model_executor.layers.fused_moe.expert_pool.pool import verify_bank_rows +from vllm.model_executor.layers.fused_moe.expert_pool.tables import ( + check_global_tables, + resident_per_layer, + set_gate, +) +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( + is_fp4_marlin_supported, +) +from vllm.platforms import current_platform + +pytestmark = [ + pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required"), + pytest.mark.skipif( + current_platform.is_cuda() and not is_fp4_marlin_supported(), + reason="FP4 Marlin not supported on this GPU", + ), +] + +SLOTS = 4 # of E=8 experts resident per layer at start; top_k=2 staging rows + + +def _decode(order, device): + logits = torch.full((1, E), -10.0, device=device) + logits[0, order[0]] = 3.0 + logits[0, order[1]] = 2.0 + return logits + + +def test_two_layer_pool_decode_prefill_decode_matches_the_uncached_layers( + dist_env, # noqa: F811 +): + """Two layers share one bank (2 * SLOTS + staging = 10 rows > E = 8), so + every bank call takes the logical-align + physical-remap path and a miss + on one layer can evict the other layer's row. Decode on both layers, + then a wide batch (bank + host-view partitions), then decode again on + the same pool; every output must match the uncached layer.""" + device = torch.accelerator.current_accelerator() + # One config per layer and side: a layer registers in the static forward + # context of the config it was built under, which set_forward_context + # must see again at run time (the legacy lookup also resolves layers by + # call order within a config, so two layers never share one). + ref_cfgs, pool_cfgs, refs, layers = [], [], [], [] + for seed_offset in (0, 1): + params = quantized_weights(device, seed_offset=seed_offset) + ref_cfgs.append(vllm_config(0)) + pool_cfgs.append(vllm_config(SLOTS)) + refs.append(make_layer(ref_cfgs[-1], params)) + layers.append(make_layer(pool_cfgs[-1], params, host_source=True)) + # As after the real loader: the small per-expert globals stay on the + # device (never allocated in host memory), the big tensors are pinned. + for layer in layers: + for name in ("w13_weight_scale_2", "w2_weight_scale_2"): + p = getattr(layer.routed_experts, name) + p.data = p.data.to(device) + # A non-contiguous (strided) host source must be densified into the + # pool's own pinned copy, values preserved, not stride-preserved. + p = layer.routed_experts.w2_weight_scale + wide = torch.zeros((p.shape[0], 2, *p.shape[1:]), dtype=p.dtype).pin_memory() + wide[:, 0].copy_(p.data) + strided_values = p.data.clone() + p.data = wide[:, 0] + assert not p.data.is_contiguous() + model = torch.nn.ModuleDict({"a": layers[0], "b": layers[1]}) + pool = install_expert_pool(model, device, max_decode_tokens=1) + assert pool is not None + # The pool owns pinned host copies of every source; the initial bank rows + # match them byte for byte. + for pl in (layer.routed_experts.expert_pool_layer for layer in layers): + assert all( + t.device.type == "cpu" and t.is_pinned() and t.is_contiguous() + for t in pl.sources.values() + ) + src = layers[-1].routed_experts.expert_pool_layer.sources["w2_weight_scale"] + assert src.shape == strided_values.shape and src.dtype == strided_values.dtype + assert torch.equal(src, strided_values) + report = verify_bank_rows(pool, model.expert_pool_sources, sample=SLOTS) + assert report == {"rows_checked": 2 * SLOTS, "rows_resident": 2 * SLOTS} + assert pool.rows == 2 * SLOTS + TOP_K and pool.rows > E + assert resident_per_layer(pool.tables) == [SLOTS, SLOTS] + pls = [layer.routed_experts.expert_pool_layer for layer in layers] + assert all(pl is not None and pl.bank_rows == pool.rows for pl in pls) + check_global_tables(pool.tables) + from vllm.forward_context import set_forward_context + + def run(i, x, logits, n): + with set_forward_context(None, ref_cfgs[i], num_tokens=n): + want = refs[i](x, logits) + with set_forward_context(None, pool_cfgs[i], num_tokens=n): + got = layers[i](x, logits) + torch.accelerator.synchronize(device) + torch.testing.assert_close(got, want, rtol=2e-2, atol=2e-2) + check_global_tables(pool.tables) + assert int(pool.tables.error[0]) == 0 + + x = torch.randn(1, K, dtype=torch.bfloat16, device=device) + tables = pool.tables + # Gate closed: a miss is staged into a shared staging row (physical row + # >= E) and the step map must point there; the output still matches. + set_gate(tables, False) + run(0, x, _decode([6, 7], device), 1) # experts 6, 7 are not resident + step_map = pls[0].buffers.step_map.cpu().tolist() + assert step_map[6] >= E and step_map[7] >= E, step_map + assert resident_per_layer(tables) == [SLOTS, SLOTS] # placement untouched + set_gate(tables, True) + # Decode steps whose routes walk every expert of both layers: misses + # promote (evicting the least recently used row of either layer) or + # stage into the shared staging rows. + for order in ([0, 1], [4, 5], [6, 7], [2, 3], [0, 6], [7, 1]): + run(0, x, _decode(order, device), 1) + hot0_before = tables.layer_slice(tables.hot_phys, 0).cpu().clone() + row_key_before = tables.row_key.cpu().clone() + for order in ([4, 5], [6, 7], [2, 6]): + run(1, x, _decode(order, device), 1) + # Cross-layer eviction: layer 1's misses took rows from layer 0. + hot0_after = tables.layer_slice(tables.hot_phys, 0).cpu() + assert not torch.equal(hot0_before, hot0_after) + row_key_after = tables.row_key.cpu() + changed = (row_key_before != row_key_after).nonzero().flatten().tolist() + assert changed and all( + int(row_key_before[r]) // E == 0 and int(row_key_after[r]) // E == 1 + for r in changed + ), (row_key_before.tolist(), row_key_after.tolist()) + assert sum(resident_per_layer(tables)) == 2 * SLOTS + # Wide batch on layer 0: resident rows from the bank, the rest through + # the host view; every route covered exactly once. + xb = torch.randn(8, K, dtype=torch.bfloat16, device=device) + run(0, xb, routing(list(range(E)), device), 8) + assert pls[0].partition_steps == 1 + # Decode again on the same pool after the wide batch. + for order in ([3, 4], [7, 0]): + run(0, x, _decode(order, device), 1) + run(1, x, _decode([0, 1], device), 1) + assert pls[0].decode_steps == 9 and pls[1].decode_steps == 4 # 1 gate-closed diff --git a/tests/kernels/expert_pool/test_pool_shared_experts.py b/tests/kernels/expert_pool/test_pool_shared_experts.py new file mode 100644 index 000000000000..ad47e0c7f543 --- /dev/null +++ b/tests/kernels/expert_pool/test_pool_shared_experts.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""In pool mode the NVFP4 method never advertises shared-expert overlap, so +the runner runs the shared experts itself and the pool ignores the wrapper +it is handed (the synchronous modular path does the same).""" + +from types import SimpleNamespace + +from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4FusedMoE + + +def _method(pool_mode: bool): + m = ModelOptNvFp4FusedMoE.__new__(ModelOptNvFp4FusedMoE) + m._pool_mode = pool_mode + m.moe_kernel = SimpleNamespace(can_overlap_shared_experts=True) + return m + + +def test_pool_mode_disables_shared_expert_overlap(): + assert _method(pool_mode=True).mk_can_overlap_shared_experts is False + assert _method(pool_mode=False).mk_can_overlap_shared_experts is True diff --git a/tests/kernels/expert_pool/test_pool_tables.py b/tests/kernels/expert_pool/test_pool_tables.py new file mode 100644 index 000000000000..e1efab325389 --- /dev/null +++ b/tests/kernels/expert_pool/test_pool_tables.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU contract tests for the global expert pool: step semantics, ownership. + +Ported from the lab expert tier; the torch reference defines the semantics +that the Triton step program must match exactly (GPU equivalence is a +separate test).""" + +import random +import unittest + +import torch + +from vllm.model_executor.layers.fused_moe.expert_pool import pool as pool_mod +from vllm.model_executor.layers.fused_moe.expert_pool import tables as gp + + +def make_sources(layers, experts, width=3): + """Per-layer RAM banks: row e of layer l holds value l*100 + e in every tensor.""" + return [ + { + name: torch.arange(experts, dtype=torch.int32) + .add(layer * 100) + .unsqueeze(1) + .repeat(1, width) + .contiguous() + for name in gp.TENSORS + } + for layer in range(layers) + ] + + +class GlobalPoolTests(unittest.TestCase): + def setup(self, layers=2, experts=6, slots=(2, 2), staging=2): + device = torch.device("cpu") + sources = make_sources(layers, experts) + pool = pool_mod.GlobalPool(device, sources[0], list(slots), staging) + for layer, count in enumerate(slots): + start = pool.offset(layer) + for name in gp.TENSORS: + pool.bank[name][start : start + count].copy_( + sources[layer][name][:count] + ) + buffers = [ + gp.allocate_step_buffers(device, experts, staging) for _ in range(layers) + ] + return pool, sources, buffers + + def run_step(self, pool, sources, buffers, layer, ids): + gp.step(pool.tables, layer, torch.tensor([ids]), buffers[layer]) + pool_mod.copy_in(sources[layer], pool.bank, buffers[layer]) + return buffers[layer] + + def assert_bank_holds_owners(self, pool): + tables = pool.tables + E = tables.num_experts + for row, key in enumerate(tables.row_key.tolist()): + if key < 0: + continue + layer, expert = divmod(key, E) + for name in gp.TENSORS: + self.assertTrue( + bool((pool.bank[name][row] == layer * 100 + expert).all()), + f"{name} row {row} key {key}", + ) + + def test_initial_layout_packs_layers_and_validates(self): + pool, sources, buffers = self.setup() + tables = pool.tables + self.assertEqual( + tables.hot_phys.tolist(), [0, 1, -1, -1, -1, -1] + [2, 3, -1, -1, -1, -1] + ) + self.assertEqual(tables.cold_phys.tolist(), [-1, -1, 2, 3, 4, 5] * 2) + self.assertEqual(tables.row_key.tolist(), [0, 1, 6, 7, -1, -1]) + self.assertEqual(tables.staging_rows.tolist(), [4, 5]) + self.assertEqual(pool.snapshot(), [2, 2]) + self.assert_bank_holds_owners(pool) + + def test_gate_closed_stages_only_and_leaves_recency(self): + pool, sources, buffers = self.setup() + b = self.run_step(pool, sources, buffers, 1, [4, 4]) + self.assertEqual(int(b.promoted_count[0]), 0) + self.assertEqual(int(b.staged_count[0]), 1) + self.assertEqual(int(b.gather_count[0]), 1) + self.assertEqual(b.step_map.tolist(), [2, 3, -1, -1, 4, -1]) + self.assertEqual(int(pool.tables.clock[0]), 0) + self.assertEqual(pool.tables.row_use[: pool.tables.pool_rows].sum().item(), 0) + # Staging row 4 now holds layer 1 expert 4. + self.assertTrue(bool((pool.bank[gp.TENSORS[0]][4] == 104).all())) + self.assertEqual(pool.snapshot(), [2, 2]) + + def test_miss_evicts_least_recent_of_any_layer_and_never_writes_ram(self): + pool, sources, buffers = self.setup() + before = [{n: t.clone() for n, t in s.items()} for s in sources] + gp.set_gate(pool.tables, True) + # Layer 0 touches its residents 0,1; layer 1 touches only 7 (key). + self.run_step(pool, sources, buffers, 0, [0, 1]) + self.run_step(pool, sources, buffers, 1, [1, -1]) + # Layer 0 misses 5: victim is key 6 (layer 1 expert 0, last_use 0). + b = self.run_step(pool, sources, buffers, 0, [5, 0]) + self.assertEqual(int(b.promoted_count[0]), 1) + self.assertEqual(b.gather_src.tolist()[:1], [5]) + self.assertEqual(b.gather_dst.tolist()[:1], [2]) + tables = pool.tables + self.assertEqual(tables.hot_phys.tolist()[5], 2) + self.assertEqual(tables.hot_phys.tolist()[6], -1) + self.assertEqual(tables.cold_phys.tolist()[6], 0) + self.assertEqual(tables.row_key.tolist()[2], 5) + self.assertEqual(b.step_map.tolist(), [0, 1, -1, -1, -1, 2]) + self.assertEqual(pool.snapshot(), [3, 1]) + self.assert_bank_holds_owners(pool) + # Layer 1 now misses expert 0 again: victim is the least recent of + # the remaining residents, layer 0's expert 1 (clock 1 < clock 3). + b = self.run_step(pool, sources, buffers, 1, [0, 1]) + self.assertEqual(tables.hot_phys.tolist()[1], -1) + self.assertEqual(tables.hot_phys.tolist()[6], 1) + self.assertEqual(pool.snapshot(), [2, 2]) + self.assert_bank_holds_owners(pool) + for layer, source in enumerate(sources): + for name in gp.TENSORS: + self.assertTrue(torch.equal(source[name], before[layer][name])) + + def test_routes_follow_the_step_map_per_lane(self): + """buffers.routes holds each lane's physical row: hits, promotions, + staged rows, duplicates resolved, padding and invalid ids -1.""" + pool, sources, buffers = self.setup(layers=1, experts=6, slots=(2,), staging=4) + b = self.run_step(pool, sources, buffers, 0, [1, 4, 4, -1]) + self.assertEqual(b.routes.tolist(), [1, 2, 2, -1]) # staged into row 2 + gp.set_gate(pool.tables, True) + b = self.run_step(pool, sources, buffers, 0, [3, 9, 1, 3]) + # Expert 3 evicts row 0 (expert 0, older than expert 1); 9 is invalid. + self.assertEqual(b.routes.tolist(), [0, -1, 1, 0]) + with self.assertRaises(RuntimeError): + pool.snapshot() + + def test_multi_row_step_deduplicates_across_rows_and_routes_every_lane(self): + """A verify step: rows x top_k lanes; an expert selected by two rows is + promoted once and both lanes route to its row; padding rows stay -1.""" + pool, sources, buffers = self.setup(layers=1, experts=8, slots=(3,), staging=6) + gp.set_gate(pool.tables, True) + ids = torch.tensor([[0, 5], [5, 6], [-1, -1]]) + gp.step(pool.tables, 0, ids, buffers[0]) + pool_mod.copy_in(sources[0], pool.bank, buffers[0]) + b = buffers[0] + # Residents 0,1,2; 0 is a hit; 5 and 6 miss and evict 1 then 2. + self.assertEqual(int(b.promoted_count[0]), 2) + self.assertEqual(int(b.staged_count[0]), 0) + hot = pool.tables.hot_phys.tolist() + self.assertEqual(b.routes.tolist(), [0, hot[5], hot[5], hot[6], -1, -1]) + self.assertEqual(pool.snapshot(), [3]) + self.assert_bank_holds_owners(pool) + + def test_promote_limit_and_interval_stage_the_rest(self): + """Limit caps promotions per layer call; interval promotes on every + N-th forward only; both keep every miss served from staging.""" + pool, sources, buffers = self.setup(layers=1, experts=8, slots=(3,), staging=4) + gp.set_gate(pool.tables, True) + gp.set_control(pool.tables, promote_limit=1) + b = self.run_step(pool, sources, buffers, 0, [3, 4, 5, 6]) + self.assertEqual((int(b.promoted_count[0]), int(b.staged_count[0])), (1, 3)) + self.assertTrue(all(int(r) >= 0 for r in b.routes.tolist())) + gp.set_control(pool.tables, promote_limit=0, promote_interval=2) + # Forward 2 is outside the interval (forwards 1, 3, 5, ... promote): + # everything staged, placement frozen. + before = pool.tables.hot_phys.clone() + b = self.run_step(pool, sources, buffers, 0, [4, 5, -1, -1]) + self.assertEqual((int(b.promoted_count[0]), int(b.staged_count[0])), (0, 2)) + self.assertTrue(torch.equal(pool.tables.hot_phys, before)) + self.assertEqual(int(pool.tables.forwards[0]), 2) + # Forward 3: promotion allowed again. + b = self.run_step(pool, sources, buffers, 0, [6, 7, -1, -1]) + self.assertEqual(int(b.promoted_count[0]), 2) + pool.snapshot() + + def test_min_misses_and_protect_recent(self): + pool, sources, buffers = self.setup(layers=1, experts=8, slots=(3,), staging=4) + gp.set_gate(pool.tables, True) + gp.set_control(pool.tables, promote_min_misses=2) + # First miss of expert 5 is staged; the second promotes it. + b = self.run_step(pool, sources, buffers, 0, [5, -1, -1, -1]) + self.assertEqual(int(b.promoted_count[0]), 0) + b = self.run_step(pool, sources, buffers, 0, [5, -1, -1, -1]) + self.assertEqual(int(b.promoted_count[0]), 1) + self.assertEqual(int(pool.tables.miss_count[5]), 0) + # Residents now: row 0 = expert 5 (used in the previous forward), + # rows 1 and 2 = experts 1 and 2 (never used). With protect_recent=1 + # the previous forward's row is not a victim: hits on 1 and 2 leave + # no victim, so both misses are staged; without protection, 6 takes + # row 0. + gp.set_control(pool.tables, promote_min_misses=1, protect_recent=1) + b = self.run_step(pool, sources, buffers, 0, [1, 2, 6, 7]) + self.assertEqual((int(b.promoted_count[0]), int(b.staged_count[0])), (0, 2)) + self.assertEqual(int(pool.tables.hot_phys[5]), 0) + gp.set_control(pool.tables, protect_recent=0) + b = self.run_step(pool, sources, buffers, 0, [1, 2, 6, 7]) + self.assertEqual(int(b.promoted_count[0]), 1) + self.assertEqual(int(pool.tables.hot_phys[6]), 0) + pool.snapshot() + with self.assertRaises(ValueError): + gp.set_control(pool.tables, promote_interval=0) + with self.assertRaises(ValueError): + gp.set_control(pool.tables, unknown=1) + + def test_no_victim_falls_back_to_staging(self): + pool, sources, buffers = self.setup(layers=1, experts=4, slots=(2,), staging=3) + gp.set_gate(pool.tables, True) + # Every resident is selected, so the two misses have no victim. + b = self.run_step(pool, sources, buffers, 0, [0, 1, 2, 3][:3]) + self.assertEqual(int(b.promoted_count[0]), 0) + self.assertEqual(int(b.staged_count[0]), 1) + self.assertEqual(b.step_map.tolist(), [0, 1, 2, -1]) + pool.snapshot() + + def test_every_current_route_is_protected_when_the_pool_is_saturated(self): + """All residents are this step's hits and one more expert misses: + no victim may be a current route, so the miss is staged only, and + a promotion in the same step is never re-evicted by a later miss.""" + pool, sources, buffers = self.setup(layers=1, experts=8, slots=(3,), staging=4) + gp.set_gate(pool.tables, True) + b = self.run_step(pool, sources, buffers, 0, [0, 1, 2, 5]) + self.assertEqual(int(b.promoted_count[0]), 0) + self.assertEqual(int(b.staged_count[0]), 1) + self.assertEqual(pool.tables.hot_phys.tolist()[:3], [0, 1, 2]) + self.assertEqual(b.routes.tolist(), [0, 1, 2, 3]) + # Two misses against one stale row: the first takes it, the second + # cannot take it back and is staged. + self.run_step(pool, sources, buffers, 0, [0, 1, -1, -1]) + b = self.run_step(pool, sources, buffers, 0, [0, 5, 6, 1]) + self.assertEqual(int(b.promoted_count[0]), 1) + self.assertEqual(int(b.staged_count[0]), 1) + self.assertEqual(pool.tables.hot_phys.tolist()[5], 2) + self.assertEqual(b.routes.tolist(), [0, 2, 3, 1]) + pool.snapshot() + + def test_invalid_ids_set_the_sticky_error(self): + pool, sources, buffers = self.setup() + self.run_step(pool, sources, buffers, 0, [9, 0]) + with self.assertRaises(RuntimeError): + pool.snapshot() + + def test_host_swap_while_gated_matches_the_tables(self): + pool, sources, buffers = self.setup() + pool.host_swap(0, 0, 3) + tables = pool.tables + self.assertEqual(tables.hot_phys.tolist()[:6], [-1, 1, -1, 0, -1, -1]) + self.assertEqual(tables.cold_phys.tolist()[:6], [0, -1, 2, -1, 4, 5]) + self.assertEqual(int(tables.row_key[0]), 3) + pool.snapshot() + gp.set_gate(tables, True) + with self.assertRaises(RuntimeError): + pool.host_swap(0, 3, 0) + + def test_random_steps_keep_ownership_consistent(self): + rng = random.Random(3) + for trial in range(20): + layers = rng.choice([1, 2, 3]) + experts = rng.choice([4, 6, 8]) + staging = rng.randint(1, 3) + slots = tuple(rng.randint(1, experts - 1) for _ in range(layers)) + pool, sources, buffers = self.setup(layers, experts, slots, staging) + gp.set_gate(pool.tables, rng.random() < 0.8) + for step in range(30): + layer = rng.randrange(layers) + ids = [rng.choice([-1, rng.randrange(experts)]) for _ in range(staging)] + b = self.run_step(pool, sources, buffers, layer, ids) + resident = pool.snapshot() + self.assertEqual(sum(resident), pool.tables.pool_rows) + self.assert_bank_holds_owners(pool) + for expert in {e for e in ids if e >= 0}: + self.assertGreaterEqual(int(b.step_map[expert]), 0) + + +if __name__ == "__main__": + unittest.main() + + +class VerifyBankRowsTests(unittest.TestCase): + def test_sampled_rows_match_and_a_corruption_is_caught(self): + device = torch.device("cpu") + sources = make_sources(2, 6) + pool = pool_mod.GlobalPool(device, sources[0], [2, 2], 2) + for layer, count in enumerate((2, 2)): + start = pool.offset(layer) + for name in gp.TENSORS: + pool.bank[name][start : start + count].copy_( + sources[layer][name][:count] + ) + report = pool_mod.verify_bank_rows(pool, sources, sample=1) + self.assertEqual(report, {"rows_checked": 2, "rows_resident": 4}) + pool.bank["w2_weight"][0, 0] += 1 + with self.assertRaises(AssertionError): + pool_mod.verify_bank_rows(pool, sources, sample=4) + + def test_scalar_per_expert_globals_are_compared_as_bytes(self): + device = torch.device("cpu") + sources = make_sources(1, 4) + for src in sources: # a [E] per-expert global, as Marlin's scale_2 is + src["w13_weight_scale_2"] = torch.arange(4, dtype=torch.float32) + 0.5 + src["w2_weight_scale_2"] = torch.arange(4, dtype=torch.float32) + 1.5 + pool = pool_mod.GlobalPool(device, sources[0], [2], 1) + for name in gp.TENSORS: + pool.bank[name][:2].copy_(sources[0][name][:2]) + self.assertEqual( + pool_mod.verify_bank_rows(pool, sources, sample=2)["rows_checked"], 2 + ) + pool.bank["w2_weight_scale_2"][1] += 1 + with self.assertRaises(AssertionError): + pool_mod.verify_bank_rows(pool, sources, sample=2) diff --git a/vllm/config/offload.py b/vllm/config/offload.py index ad65e8acf35a..1dec4d1fd2c6 100644 --- a/vllm/config/offload.py +++ b/vllm/config/offload.py @@ -94,6 +94,15 @@ class OffloadConfig: prefetch: PrefetchOffloadConfig = Field(default_factory=PrefetchOffloadConfig) """Parameters for prefetch offloading backend.""" + moe_expert_pool_rows: int = Field(default=0, ge=0) + """Serve MoE expert weights from a shared GPU expert pool: keep this many + expert rows per MoE layer resident at startup in one VRAM bank shared by + all MoE layers, with the rest of the experts in pinned host memory. A + device-side LRU planner moves rows between layers at run time and the + forward runs no host code, so the MoE op stays inside CUDA graphs. + 0 (default) keeps every expert on the GPU. Supported for the ModelOpt + NVFP4 Marlin MoE backend without expert/data parallelism.""" + @model_validator(mode="after") def validate_offload_config(self) -> "OffloadConfig": """Validate offload configuration constraints.""" diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 2fed4f3b6a44..78f39bfb88f1 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -541,6 +541,7 @@ class EngineArgs: offload_num_in_group: int = PrefetchOffloadConfig.offload_num_in_group offload_prefetch_step: int = PrefetchOffloadConfig.offload_prefetch_step offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params") + moe_expert_pool_rows: int = OffloadConfig.moe_expert_pool_rows gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes max_num_batched_tokens: int | None = None @@ -1338,6 +1339,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: offload_group.add_argument( "--offload-params", **prefetch_kwargs["offload_params"] ) + offload_group.add_argument( + "--moe-expert-pool-rows", **offload_kwargs["moe_expert_pool_rows"] + ) # Multimodal related configs multimodal_kwargs = get_kwargs(MultiModalConfig) @@ -2579,6 +2583,7 @@ def create_engine_config( offload_prefetch_step=self.offload_prefetch_step, offload_params=self.offload_params, ), + moe_expert_pool_rows=self.moe_expert_pool_rows, ) if self.gdn_prefill_backend is not None: diff --git a/vllm/model_executor/layers/fused_moe/expert_pool/__init__.py b/vllm/model_executor/layers/fused_moe/expert_pool/__init__.py new file mode 100644 index 000000000000..a0e7bb59295c --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/expert_pool/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Global expert pool: device-planned expert cache (see tables.py).""" diff --git a/vllm/model_executor/layers/fused_moe/expert_pool/copy.py b/vllm/model_executor/layers/fused_moe/expert_pool/copy.py new file mode 100644 index 000000000000..6dcac79c469b --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/expert_pool/copy.py @@ -0,0 +1,435 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fixed-grid row copies between the pinned host source (UVA view) and the +VRAM bank: one launch per step reads the device-side pair count, so an +empty step costs a few programs and nothing synchronizes. Two launch +shapes: "stripes" (program = (tensor, stripe) streaming its columns of +every row) and "chunks" (programs grid-stride over (row, chunk) pairs). +Ported from the lab expert tier (promote.py). +""" + +from __future__ import annotations + +from typing import Any + +from vllm.model_executor.layers.fused_moe.expert_pool.tables import TENSORS + +_KERNELS: dict[str, Any] = {} +COPY_PROGRAMS_PER_BANK = 32 +COPY_WORDS = 4096 # int32 words (16 KiB) per program iteration +COPY_SHAPES = ("stripe", "chunks") +_COPY_SHAPE = "stripe" +COPY_CHUNK_PROGRAMS_PER_BANK = 8 + + +def copy_rows_reference(source, destination, pairs): + """Copy (src row, dst row) pairs for every bank tensor, in order.""" + for src, dst in pairs: + for name in TENSORS: + destination[name][dst].copy_(source[name][src]) + + +def configure_copy(shape, programs=None, words=None): + """Select the copy launch shape and its grid (runtime settings).""" + global _COPY_SHAPE, COPY_PROGRAMS_PER_BANK, COPY_CHUNK_PROGRAMS_PER_BANK, COPY_WORDS + if shape not in COPY_SHAPES: + raise ValueError(f"Copy shape must be one of {COPY_SHAPES}") + if programs is not None: + if int(programs) < 1: + raise ValueError("Copy programs per bank must be positive") + COPY_PROGRAMS_PER_BANK = COPY_CHUNK_PROGRAMS_PER_BANK = int(programs) + if words is not None: + if int(words) < 32 or int(words) & (int(words) - 1): + raise ValueError("Copy words per iteration must be a power of two >= 32") + COPY_WORDS = int(words) + _COPY_SHAPE = shape + + +def copy_rows(source, destination, src_rows, dst_rows, count): + """Copy `count` (src, dst) row pairs of every bank tensor; device count. + + One launch of a fixed small grid reading `count` on the device, in one + of two shapes (see COPY_SHAPES); an empty step costs a few programs. + """ + src_device = source[TENSORS[0]].device + if src_device.type != "cuda": + n = int(count.reshape(-1)[0].item()) + pairs = [(int(src_rows[i]), int(dst_rows[i])) for i in range(n)] + copy_rows_reference(source, destination, pairs) + return + srcs = [_word_rows(source[name]) for name in TENSORS] + dsts = [_word_rows(destination[name]) for name in TENSORS] + for name, src, dst in zip(TENSORS, srcs, dsts): + if src.shape[1] != dst.shape[1]: + raise ValueError(f"{name}: destination row size differs from the source") + if _COPY_SHAPE == "chunks": + grid = (len(TENSORS) * COPY_CHUNK_PROGRAMS_PER_BANK,) + _copy_chunks_kernel()[grid]( + *srcs, + *dsts, + src_rows, + dst_rows, + count, + *(dst.shape[1] for dst in dsts), + *(src.stride(0) for src in srcs), + *(dst.stride(0) for dst in dsts), + PROGRAMS=COPY_CHUNK_PROGRAMS_PER_BANK, + BLOCK=COPY_WORDS, + num_warps=32, + ) + return + grid = (len(TENSORS) * COPY_PROGRAMS_PER_BANK,) + _copy_kernel()[grid]( + *srcs, + *dsts, + src_rows, + dst_rows, + count, + *(dst.shape[1] for dst in dsts), + *(src.stride(0) for src in srcs), + *(dst.stride(0) for dst in dsts), + PROGRAMS=COPY_PROGRAMS_PER_BANK, + BLOCK=COPY_WORDS, + num_warps=4, + ) + + +def _word_rows(tensor): + """View a [rows, ...] contiguous tensor as [rows, int32 words].""" + import torch + + if not tensor.is_contiguous(): + raise ValueError("Copies require contiguous bank rows") + rows = tensor.shape[0] + return tensor.view(torch.uint8).reshape(rows, -1).view(torch.int32) + + +def _byte_rows(tensor): + import torch + + if not tensor.is_contiguous(): + raise ValueError("Promote copies require contiguous bank rows") + return tensor.view(torch.uint8).reshape(tensor.shape[0], -1) + + +def _copy_kernel(): + """Fixed grid: program (bank, stripe) streams its columns of every row.""" + if "copy" in _KERNELS: + return _KERNELS["copy"] + from vllm.triton_utils import tl, triton + + @triton.jit + def _stripe( + src, + dst, + src_rows_ptr, + dst_rows_ptr, + count, + words, + sstride, + dstride, + stripe, + PROGRAMS: tl.constexpr, + BLOCK: tl.constexpr, + ): + for lane in range(0, count): + src_row = tl.load(src_rows_ptr + lane).to(tl.int64) + dst_row = tl.load(dst_rows_ptr + lane).to(tl.int64) + src_base = src + src_row * sstride + dst_base = dst + dst_row * dstride + for start in range(stripe * BLOCK, words, PROGRAMS * BLOCK): + offsets = start + tl.arange(0, BLOCK) + mask = offsets < words + values = tl.load(src_base + offsets, mask=mask) + tl.store(dst_base + offsets, values, mask=mask) + + @triton.jit + def promote_copy( + src0, + src1, + src2, + src3, + src4, + src5, + dst0, + dst1, + dst2, + dst3, + dst4, + dst5, + src_rows_ptr, + dst_rows_ptr, + count_ptr, + words0, + words1, + words2, + words3, + words4, + words5, + sstride0, + sstride1, + sstride2, + sstride3, + sstride4, + sstride5, + dstride0, + dstride1, + dstride2, + dstride3, + dstride4, + dstride5, + PROGRAMS: tl.constexpr, + BLOCK: tl.constexpr, + ): + which = tl.program_id(0) // PROGRAMS + stripe = tl.program_id(0) % PROGRAMS + count = tl.load(count_ptr) + if which == 0: + _stripe( + src0, + dst0, + src_rows_ptr, + dst_rows_ptr, + count, + words0, + sstride0, + dstride0, + stripe, + PROGRAMS, + BLOCK, + ) + elif which == 1: + _stripe( + src1, + dst1, + src_rows_ptr, + dst_rows_ptr, + count, + words1, + sstride1, + dstride1, + stripe, + PROGRAMS, + BLOCK, + ) + elif which == 2: + _stripe( + src2, + dst2, + src_rows_ptr, + dst_rows_ptr, + count, + words2, + sstride2, + dstride2, + stripe, + PROGRAMS, + BLOCK, + ) + elif which == 3: + _stripe( + src3, + dst3, + src_rows_ptr, + dst_rows_ptr, + count, + words3, + sstride3, + dstride3, + stripe, + PROGRAMS, + BLOCK, + ) + elif which == 4: + _stripe( + src4, + dst4, + src_rows_ptr, + dst_rows_ptr, + count, + words4, + sstride4, + dstride4, + stripe, + PROGRAMS, + BLOCK, + ) + else: + _stripe( + src5, + dst5, + src_rows_ptr, + dst_rows_ptr, + count, + words5, + sstride5, + dstride5, + stripe, + PROGRAMS, + BLOCK, + ) + + _KERNELS["copy"] = promote_copy + return promote_copy + + +def _copy_chunks_kernel(): + """FreeToken-shaped copy: programs grid-stride over (row, chunk) pairs.""" + if "copy_chunks" in _KERNELS: + return _KERNELS["copy_chunks"] + from vllm.triton_utils import tl, triton + + @triton.jit + def _chunks( + src, + dst, + src_rows_ptr, + dst_rows_ptr, + count, + words, + sstride, + dstride, + program, + PROGRAMS: tl.constexpr, + BLOCK: tl.constexpr, + ): + chunks_per_row = tl.cdiv(words, BLOCK) + total = count * chunks_per_row + for c in range(program, total, PROGRAMS): + row = c // chunks_per_row + chunk = c - row * chunks_per_row + src_row = tl.load(src_rows_ptr + row).to(tl.int64) + dst_row = tl.load(dst_rows_ptr + row).to(tl.int64) + offsets = chunk * BLOCK + tl.arange(0, BLOCK) + mask = offsets < words + values = tl.load(src + src_row * sstride + offsets, mask=mask) + tl.store(dst + dst_row * dstride + offsets, values, mask=mask) + + @triton.jit + def promote_copy_chunks( + src0, + src1, + src2, + src3, + src4, + src5, + dst0, + dst1, + dst2, + dst3, + dst4, + dst5, + src_rows_ptr, + dst_rows_ptr, + count_ptr, + words0, + words1, + words2, + words3, + words4, + words5, + sstride0, + sstride1, + sstride2, + sstride3, + sstride4, + sstride5, + dstride0, + dstride1, + dstride2, + dstride3, + dstride4, + dstride5, + PROGRAMS: tl.constexpr, + BLOCK: tl.constexpr, + ): + which = tl.program_id(0) // PROGRAMS + program = tl.program_id(0) % PROGRAMS + count = tl.load(count_ptr) + if which == 0: + _chunks( + src0, + dst0, + src_rows_ptr, + dst_rows_ptr, + count, + words0, + sstride0, + dstride0, + program, + PROGRAMS, + BLOCK, + ) + elif which == 1: + _chunks( + src1, + dst1, + src_rows_ptr, + dst_rows_ptr, + count, + words1, + sstride1, + dstride1, + program, + PROGRAMS, + BLOCK, + ) + elif which == 2: + _chunks( + src2, + dst2, + src_rows_ptr, + dst_rows_ptr, + count, + words2, + sstride2, + dstride2, + program, + PROGRAMS, + BLOCK, + ) + elif which == 3: + _chunks( + src3, + dst3, + src_rows_ptr, + dst_rows_ptr, + count, + words3, + sstride3, + dstride3, + program, + PROGRAMS, + BLOCK, + ) + elif which == 4: + _chunks( + src4, + dst4, + src_rows_ptr, + dst_rows_ptr, + count, + words4, + sstride4, + dstride4, + program, + PROGRAMS, + BLOCK, + ) + else: + _chunks( + src5, + dst5, + src_rows_ptr, + dst_rows_ptr, + count, + words5, + sstride5, + dstride5, + program, + PROGRAMS, + BLOCK, + ) + + _KERNELS["copy_chunks"] = promote_copy_chunks + return promote_copy_chunks diff --git a/vllm/model_executor/layers/fused_moe/expert_pool/install.py b/vllm/model_executor/layers/fused_moe/expert_pool/install.py new file mode 100644 index 000000000000..4a7363a798f7 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/expert_pool/install.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Model-level installation of the global expert pool. + +Runs once after every layer's process_weights_after_loading: with +``moe_expert_pool_rows > 0`` the MoE layers keep their expert tensors in +pinned host memory (final kernel layout); this allocates one VRAM bank +shared by all of them, fills each layer's initial rows, and binds a +consumer (Marlin) to the bank. +""" + +from __future__ import annotations + +from dataclasses import replace +from types import SimpleNamespace + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.expert_pool.copy import configure_copy +from vllm.model_executor.layers.fused_moe.expert_pool.layer import PoolLayer +from vllm.model_executor.layers.fused_moe.expert_pool.pool import GlobalPool +from vllm.model_executor.layers.fused_moe.expert_pool.tables import ( + TENSORS, + allocate_step_buffers, + set_control, +) + +logger = init_logger(__name__) + + +# Decode lanes (tokens x top_k) the step program is compiled for. Wider +# inputs take the bank + host-view partition path. Small on purpose: the +# single-program planner loops over WIDTH lanes and scans the pool per miss. +MAX_DECODE_LANES = 64 + + +def _next_power_of_two(value: int) -> int: + return 1 << max(int(value) - 1, 0).bit_length() + + +def _check_top_k(top_k: int) -> None: + if not 0 < top_k <= MAX_DECODE_LANES: + raise ValueError( + f"expert pool supports 0 < top_k <= {MAX_DECODE_LANES}, got {top_k}" + ) + + +def check_pool_layers(layers: list[tuple[str, torch.nn.Module]]) -> None: + """Reject geometries the pool cannot serve, before anything is allocated: + every layer must share expert count, top-k, resident rows and backend, + and the backend must be NVFP4 Marlin (the only consumer bound here).""" + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4FusedMoE, + ) + + first_name, first = layers[0] + for name, layer in layers: + method = layer.quant_method + if not isinstance(method, ModelOptNvFp4FusedMoE): + raise ValueError( + f"{name}: expert pool supports ModelOptNvFp4FusedMoE only, got " + f"{type(method).__name__}" + ) + if method.nvfp4_backend != NvFp4MoeBackend.MARLIN: + raise ValueError( + f"{name}: expert pool supports the Marlin NVFP4 backend only, " + f"got {method.nvfp4_backend.value}" + ) + if layer.moe_config.moe_parallel_config.use_ep: + raise ValueError(f"{name}: expert pool is not compatible with EP") + for attr in ("local_num_experts", "_moe_expert_pool_rows"): + if getattr(layer, attr) != getattr(first, attr): + raise ValueError( + f"{name}.{attr}={getattr(layer, attr)} differs from " + f"{first_name} ({getattr(first, attr)})" + ) + if layer.moe_config.experts_per_token != first.moe_config.experts_per_token: + raise ValueError(f"{name}: top_k differs from {first_name}") + + +def pool_layers(model: torch.nn.Module) -> list[tuple[str, torch.nn.Module]]: + return [ + (name, module) + for name, module in model.named_modules() + if getattr(module, "expert_pool_pending", False) + ] + + +def _sources(name: str, layer: torch.nn.Module) -> dict[str, torch.Tensor]: + sources = {} + for tensor_name in TENSORS: + parameter = getattr(layer, tensor_name, None) + if parameter is None: + raise RuntimeError(f"{name}.{tensor_name}: missing for the expert pool") + t = parameter.data + if t.ndim < 1 or t.shape[0] != layer.local_num_experts: + raise RuntimeError( + f"{name}.{tensor_name}: the expert pool needs one row per expert, " + f"got shape {tuple(t.shape)}" + ) + if t.device.type != "cpu" or not t.is_pinned() or not t.is_contiguous(): + # The per-expert global scales are small and are not allocated in + # host memory by create_weights, so the loader leaves them on the + # device after conversion; take a pinned, dense host copy as the + # source (explicit contiguous layout: empty_like would keep the + # input's strides). + t = torch.empty( + t.shape, dtype=t.dtype, device="cpu", pin_memory=True + ).copy_(t) + assert t.is_contiguous() and t.is_pinned() + sources[tensor_name] = t + return sources + + +def install_expert_pool( + model: torch.nn.Module, device: torch.device, max_decode_tokens: int = 1 +) -> GlobalPool | None: + """Allocate the shared bank and bind every pending pool layer to it.""" + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + make_nvfp4_moe_kernel, + ) + from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + marlin_make_workspace_new, + ) + from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor + + layers = pool_layers(model) + if not layers: + return None + check_pool_layers(layers) + first_name, first = layers[0] + num_experts = first.local_num_experts + top_k = first.moe_config.experts_per_token + _check_top_k(top_k) + slots = min(first._moe_expert_pool_rows, num_experts - 1) + if slots < top_k: + raise ValueError( + f"expert pool needs at least top_k={top_k} rows per layer, got {slots}" + ) + # Decode lanes served by the step program: at most MAX_DECODE_LANES and + # at least one token; batches beyond that use the partition path. + decode_tokens = max(1, min(max_decode_tokens, MAX_DECODE_LANES // top_k)) + staging = top_k * decode_tokens + width = _next_power_of_two(staging) + sources = [_sources(name, layer) for name, layer in layers] + for name, src in zip((n for n, _ in layers), sources): + for tensor_name in TENSORS: + if ( + src[tensor_name].shape[1:] != sources[0][tensor_name].shape[1:] + or src[tensor_name].dtype != sources[0][tensor_name].dtype + ): + raise RuntimeError( + f"{name}.{tensor_name}: layer rows differ from {first_name}" + ) + pool = GlobalPool(device, sources[0], [slots] * len(layers), staging) + configure_copy("chunks") + for index, ((name, layer), src) in enumerate(zip(layers, sources)): + method = layer.quant_method + start = pool.offset(index) + for tensor_name in TENSORS: + pool.bank[tensor_name][start : start + slots].copy_( + src[tensor_name][:slots], non_blocking=True + ) + host_views = { + tensor_name: get_accelerator_view_from_cpu_tensor(t) + for tensor_name, t in src.items() + } + proxy = SimpleNamespace( + **{tensor_name: pool.bank[tensor_name] for tensor_name in TENSORS}, + w13_input_scale=None, + w2_input_scale=None, + swiglu_limit=getattr(layer, "swiglu_limit", None), + swiglu_alpha=getattr(layer, "swiglu_alpha", None), + swiglu_beta=getattr(layer, "swiglu_beta", None), + ) + quant = method.get_fused_moe_quant_config(proxy) + config = replace(method.moe, num_local_experts=pool.rows) + kernel = make_nvfp4_moe_kernel( + quant, + config, + method.experts_cls, + method.nvfp4_backend, + routing_tables=None, + ) + if kernel.prepare_finalize.supports_async(): + raise NotImplementedError( + "expert pool requires synchronous prepare/finalize" + ) + layer.expert_pool_layer = PoolLayer( + index=index, + pool=pool, + slots=slots, + sources=src, + host_views=host_views, + buffers=allocate_step_buffers(device, num_experts, width), + experts=kernel.fused_experts, + marlin_workspace=marlin_make_workspace_new(device, 4), + num_experts=num_experts, + top_k=top_k, + activation=layer.activation, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + ) + layer.expert_pool_pending = False + # Placement policy (the lab run's values): promotions on every forward, + # first miss promotes, no protection window. The gate stays closed + # through profiling and graph capture (dummy routing must not move the + # placement) and is opened by open_pool_gate() at the end of warm-up. + set_control( + pool.tables, + promote_limit=0, + promote_interval=1, + promote_min_misses=1, + protect_recent=0, + gate=0, + ) + torch.accelerator.synchronize(device) + model.expert_pool = pool + model.expert_pool_sources = sources + logger.info( + "Expert pool installed: %d layers, %d/%d rows per layer resident, " + "%d staging rows (%d decode tokens x top_k %d; wider batches take the " + "partition path), bank %.1f GiB (%s)", + len(layers), + slots, + num_experts, + staging, + decode_tokens, + top_k, + (pool.pool_bytes + pool.staging_bytes) / 2**30, + type(layers[0][1].quant_method).__name__, + ) + return pool + + +def open_pool_gate(model: torch.nn.Module, sample_rows: int = 4) -> None: + """End of warm-up/capture: verify the tables and a sample of bank rows + against the host source, log the placement, then open the gate. + + Host readback happens here only, never inside a forward. The gate is a + device scalar at a fixed address, so captured graphs see the change.""" + from vllm.model_executor.layers.fused_moe.expert_pool.pool import ( + verify_bank_rows, + ) + from vllm.model_executor.layers.fused_moe.expert_pool.tables import ( + check_global_tables, + resident_per_layer, + set_gate, + ) + + pool = getattr(model, "expert_pool", None) + if pool is None: + return + device = pool.tables.hot_phys.device + torch.accelerator.synchronize(device) + check_global_tables(pool.tables) + report = verify_bank_rows(pool, model.expert_pool_sources, sample_rows) + resident = resident_per_layer(pool.tables) + set_gate(pool.tables, True) + torch.accelerator.synchronize(device) + logger.info( + "Expert pool gate opened after warm-up: tables consistent, %d sampled " + "bank rows match the host source (%d resident), resident per layer " + "min/max %d/%d", + report["rows_checked"], + report["rows_resident"], + min(resident), + max(resident), + ) diff --git a/vllm/model_executor/layers/fused_moe/expert_pool/layer.py b/vllm/model_executor/layers/fused_moe/expert_pool/layer.py new file mode 100644 index 000000000000..5c34b8e69fa1 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/expert_pool/layer.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""One MoE layer on the global expert pool: device-planned decode step and +the two-partition (bank + host) path for wider batches, both without host +code in the forward, so a CUDA graph can capture them. + +Marlin consumer only in this version: the bank holds Marlin-format rows for +every layer, so the bank's row count exceeds the layer's expert count. +Alignment therefore happens on logical expert ids (absent experts already +padding) and the aligned blocks are mapped to physical bank rows +afterwards; the align op never sees a row id (it histograms by id in a +num_experts+1 buffer). Ported from the lab expert tier (runtime.py +split_global / _run_marlin_chains). +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch + +from vllm.model_executor.layers.fused_moe.expert_pool.pool import GlobalPool, copy_in +from vllm.model_executor.layers.fused_moe.expert_pool.tables import ( + TENSORS, + StepBuffers, + step, +) + + +def mask_routes(ids: torch.Tensor, expert_map: torch.Tensor) -> torch.Tensor: + """Routes whose expert is absent from `expert_map` become padding (-1).""" + num_experts = expert_map.shape[0] + safe = ids.clamp(0, num_experts - 1).long() + present = (ids >= 0) & (ids < num_experts) & (expert_map[safe] >= 0) + return torch.where(present, ids, torch.full_like(ids, -1)) + + +def physical_block_experts( + logical_ids: torch.Tensor, + post_padded: torch.Tensor, + block: int, + expert_map: torch.Tensor, + num_experts: int, +) -> torch.Tensor: + """Map per-block logical expert ids to physical bank rows. + + Blocks at or beyond `post_padded` tokens are unused by the GEMM and may + hold uninitialized ids; they become -1 without indexing anything. + """ + blocks = torch.arange( + logical_ids.numel(), device=logical_ids.device, dtype=torch.int32 + ) + valid = (blocks * block) < post_padded.reshape(1) + safe = torch.where(valid, logical_ids, torch.zeros_like(logical_ids)) + safe = safe.clamp(0, num_experts - 1) + return torch.where(valid, expert_map[safe.long()], torch.full_like(logical_ids, -1)) + + +def marlin_block_size(tokens, top_k, local_experts, global_experts, input_dtype): + """The stock fused_marlin_moe M-block choice for one expert partition.""" + estimated = math.ceil(tokens * local_experts / global_experts) + block = 8 + for block in (8, 16, 32, 48, 64): + if estimated * top_k / local_experts / block < 0.9: + break + if input_dtype is not None and input_dtype.itemsize == 1: + block = max(block, 16) + return block + + +class PoolLayer: + """Per-layer view of the pool plus the consumer state.""" + + def __init__( + self, + index: int, + pool: GlobalPool, + slots: int, + sources: dict[str, torch.Tensor], + host_views: dict[str, torch.Tensor], + buffers: StepBuffers, + experts: Any, + marlin_workspace: torch.Tensor, + num_experts: int, + top_k: int, + activation: Any, + apply_router_weight_on_input: bool, + ) -> None: + self.index = index + self.pool = pool + self.slots = slots + self.offset = pool.offset(index) + self.sources = sources # pinned host rows, final layout + self.host = host_views # accelerator (UVA) views of the sources + self.buffers = buffers + self.experts = experts # MarlinExperts bound to the bank tensors + self.marlin_workspace = marlin_workspace + self.num_experts = num_experts + self.top_k = top_k + self.activation = activation + self.apply_router_weight_on_input = apply_router_weight_on_input + self.bank = pool.bank + self.bank_rows = pool.rows + self.staging_rows = pool.staging_slots + self.width = buffers.gather_src.shape[0] + self.hot_map = pool.tables.layer_slice(pool.tables.hot_phys, index) + self.cold_map = pool.tables.layer_slice(pool.tables.cold_phys, index) + self.decode_steps = 0 + self.partition_steps = 0 + + # --- forward ----------------------------------------------------------- + + def apply( + self, x: torch.Tensor, weights: torch.Tensor, ids: torch.Tensor + ) -> torch.Tensor: + self._check_routes(x, weights, ids) + lanes = ids.shape[0] * ids.shape[1] + if lanes <= self.width and lanes <= self.staging_rows: + self.decode_steps += 1 + step(self.pool.tables, self.index, ids, self.buffers) + copy_in(self.host, self.bank, self.buffers) + return self._run_marlin( + x, weights, ids, ((self.bank, self.buffers.step_map, self.bank_rows),) + ) + # Wider batches (prefill): resident rows from the bank, the rest read + # straight from the pinned host source through its accelerator view. + self.partition_steps += 1 + return self._run_marlin( + x, + weights, + ids, + ( + (self.bank, self.hot_map, self.bank_rows), + (self.host, self.cold_map, self.num_experts), + ), + ) + + def _check_routes(self, x, weights, ids) -> None: + if ( + x.ndim != 2 + or ids.ndim != 2 + or weights.shape != ids.shape + or ids.shape[0] != x.shape[0] + or ids.shape[1] != self.top_k + or ids.dtype not in (torch.int32, torch.int64) + ): + raise ValueError("Unexpected pool routing/input shape or dtype") + allowed = (ids >= -1) & (ids < self.num_experts) + finite = torch.isfinite(weights) & (weights >= 0) + torch._assert_async( + (allowed & (finite | (ids == -1))).all(), + "Invalid routing: ids must be in [-1, num_experts) and weights " + "finite/nonnegative", + ) + + def _run_marlin(self, x, weights, ids, partitions): + from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( + _fused_marlin_moe, + marlin_moe_intermediate_size, + ) + from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( + moe_align_block_size, + ) + from vllm.scalar_type import ScalarType + from vllm.v1.worker.workspace import current_workspace_manager + + experts = self.experts + tokens, hidden = x.shape + top_k = ids.shape[1] + rows_count = tokens * top_k + inner = marlin_moe_intermediate_size( + self.bank["w13_weight"], self.bank["w2_weight"] + ) + # Same manager as the stock kernels: stable addresses once locked. + cache13, cache2, rows = current_workspace_manager().get_simultaneous( + ((rows_count * max(2 * inner, hidden),), x.dtype), + ((rows_count, inner), x.dtype), + ((rows_count, hidden), x.dtype), + ) + rows.zero_() + for tensors, expert_map, slots in partitions: + block = marlin_block_size( + tokens, top_k, slots, self.num_experts, experts.input_dtype + ) + routed = ids + if slots > self.num_experts: + # Align by logical id (absent experts already padding), then + # map the blocks to rows; the align op never sees a row id. + routed = mask_routes(ids, expert_map) + sorted_ids, logical_ids, post_padded = moe_align_block_size( + routed, block, self.num_experts, None, ignore_invalid_experts=True + ) + expert_ids = physical_block_experts( + logical_ids, post_padded, block, expert_map, self.num_experts + ) + else: + sorted_ids, expert_ids, post_padded = moe_align_block_size( + ids, + block, + self.num_experts, + expert_map, + ignore_invalid_experts=True, + ) + _fused_marlin_moe( + hidden_states=x, + w1=tensors["w13_weight"], + w2=tensors["w2_weight"], + bias1=experts.w1_bias, + bias2=experts.w2_bias, + w1_scale=self._scale(tensors, "w13_weight_scale"), + w2_scale=self._scale(tensors, "w2_weight_scale"), + topk_weights=weights, + num_topk=top_k, + quant_type=ScalarType.from_id(experts.quant_type_id), + apply_router_weight_on_input=self.apply_router_weight_on_input, + expert_map=expert_map, + block_size_m=block, + sorted_token_ids=sorted_ids, + expert_ids=expert_ids, + num_tokens_post_padded=post_padded, + activation=self.activation, + activation_func=experts.activation, + topk_ids=routed, + input_global_scale1=experts.a1_gscale, + input_global_scale2=experts.a2_gscale, + global_scale1=self._scale(tensors, "w13_weight_scale_2"), + global_scale2=self._scale(tensors, "w2_weight_scale_2"), + w1_zeros=experts.w1_zp, + w2_zeros=experts.w2_zp, + workspace=self.marlin_workspace, + intermediate_cache13=cache13, + intermediate_cache2=cache2, + output=rows, + input_dtype=experts.input_dtype, + activation_config=experts.activation_config, + ) + # Rows already carry the router weights (second GEMM multiplies them). + return torch.sum(rows.view(tokens, top_k, hidden), dim=1) + + @staticmethod + def _scale(tensors: dict[str, torch.Tensor], name: str) -> torch.Tensor: + # Scales and globals are indexed by the same physical row as the + # weights of the partition being run (bank rows or host rows). + return tensors[name] + + def stats(self) -> dict[str, int]: + return { + "decode_steps": self.decode_steps, + "partition_steps": self.partition_steps, + "slots": self.slots, + } + + +__all__ = [ + "TENSORS", + "PoolLayer", + "marlin_block_size", + "mask_routes", + "physical_block_experts", +] diff --git a/vllm/model_executor/layers/fused_moe/expert_pool/pool.py b/vllm/model_executor/layers/fused_moe/expert_pool/pool.py new file mode 100644 index 000000000000..de770a0212a2 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/expert_pool/pool.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The shared VRAM bank, its staging rows, the pool tables and layer offsets.""" + +from __future__ import annotations + +from vllm.model_executor.layers.fused_moe.expert_pool.copy import copy_rows +from vllm.model_executor.layers.fused_moe.expert_pool.tables import ( + TENSORS, + allocate_global_tables, + check_global_tables, + read_control, + resident_per_layer, + set_control, +) + + +class GlobalPool: + """The shared bank, its staging views, the tables, and the layer offsets.""" + + def __init__(self, device, sources, slots_per_layer, staging): + # `sources`: one layer's six tensors in the bank's final layout; only + # shapes/dtypes are read here (rows are filled by the layers). + import torch + + self.slots_per_layer = list(slots_per_layer) + self.staging_slots = staging + self.tables = allocate_global_tables( + device, sources[TENSORS[0]].shape[0], self.slots_per_layer, staging + ) + self.rows = self.tables.pool_rows + staging + self.offsets = [0] + for slots in self.slots_per_layer[:-1]: + self.offsets.append(self.offsets[-1] + slots) + self.bank = { + name: torch.zeros( + (self.rows, *source.shape[1:]), dtype=source.dtype, device=device + ) + for name, source in sources.items() + } + self.staging = { + name: tensor[self.tables.pool_rows :] for name, tensor in self.bank.items() + } + self.row_bytes = sum(t[0].numel() * t.element_size() for t in sources.values()) + self.staging_bytes = self.row_bytes * staging + self.pool_bytes = self.row_bytes * self.tables.pool_rows + + def offset(self, layer): + return self.offsets[layer] + + def host_swap(self, layer, old_expert, new_expert): + """Gate-closed exchange for init verification: `new_expert` takes the + row of resident `old_expert`, which falls back to its RAM row. The + caller copies the bytes.""" + tables = self.tables + if int(tables.gate[0]): + raise RuntimeError("Host swaps are only allowed while the gate is closed") + E = tables.num_experts + old_key, new_key = layer * E + old_expert, layer * E + new_expert + row = int(tables.hot_phys[old_key]) + if row < 0 or int(tables.hot_phys[new_key]) >= 0: + raise AssertionError("Swap does not match the current pool placement") + tables.hot_phys[old_key], tables.cold_phys[old_key] = -1, old_expert + tables.hot_phys[new_key], tables.cold_phys[new_key] = row, -1 + tables.row_key[row] = new_key + + def snapshot(self): + """Validate the pool on the host; one copy per stats report.""" + check_global_tables(self.tables) + return resident_per_layer(self.tables) + + def apply_control(self, **values): + """Validate then write the controls; the gate is not touched here.""" + return set_control(self.tables, **values) + + def control(self): + return read_control(self.tables) + + +def verify_bank_rows(pool, sources, sample: int = 4) -> dict[str, int]: + """Compare up to `sample` resident rows per layer against the host + source, byte for byte (host readback; only at safe boundaries). + + `sources[layer]` holds that layer's six host tensors. Returns the + counts checked/mismatched; raises on the first mismatch.""" + import torch + + tables = pool.tables + E = tables.num_experts + row_key = tables.row_key.tolist() + checked = 0 + per_layer = [0] * tables.num_layers + for row, key in enumerate(row_key): + if key < 0: + continue + layer, expert = divmod(key, E) + if per_layer[layer] >= sample: + continue + per_layer[layer] += 1 + for name in TENSORS: + # reshape before the byte view: a per-expert global scale row is + # a 0-dim tensor, which cannot be viewed as bytes directly. + got = pool.bank[name][row].detach().cpu().reshape(-1).view(torch.uint8) + want = sources[layer][name][expert].reshape(-1).view(torch.uint8) + if not torch.equal(got, want): + raise AssertionError( + f"bank row {row} ({name}) differs from layer {layer} " + f"expert {expert}" + ) + checked += 1 + return {"rows_checked": checked, "rows_resident": sum(1 for k in row_key if k >= 0)} + + +def copy_in(source, bank, buffers): + """Copy the planned host rows of this layer into the bank rows.""" + copy_rows( + source, bank, buffers.gather_src, buffers.gather_dst, buffers.gather_count + ) diff --git a/vllm/model_executor/layers/fused_moe/expert_pool/tables.py b/vllm/model_executor/layers/fused_moe/expert_pool/tables.py new file mode 100644 index 000000000000..a3333710ea79 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/expert_pool/tables.py @@ -0,0 +1,610 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Global expert pool tables and the per-layer step program. + +One VRAM bank shared by every MoE layer: row ``r`` holds the expert whose +key is ``row_key[r]`` (key = layer * E + expert) or is free. ``hot_phys[key]`` +is the row or -1; each layer's slice is the expert map its kernel reads. +``cold_phys[key]`` is the expert's host row (its expert id) while it is not +resident, so evictions never copy out (the pinned host source is the +backing store). + +Each layer's decode step is one device program: distinct valid selections +are stamped with the step clock; each miss takes the row of the least +recently used resident expert of any layer (ties by row), copying in from +this layer's host rows; a miss that finds no victim, or any miss while the +gate is closed, is staged into the shared staging rows for this step only. +The step map is this layer's ``hot_phys`` slice with the staged experts +overlaid. Everything runs on the compute stream with fixed shapes and +addresses, so it can be captured in a CUDA graph; the host reads the +tables only at stats reports. + +``step_reference`` (torch, synchronizing) defines the semantics; the CPU +device and the tests run it, and the Triton program must match it exactly. +Ported from the lab expert tier (global_pool.py) without the control file. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +TENSORS = ( + "w13_weight", + "w2_weight", + "w13_weight_scale", + "w2_weight_scale", + "w13_weight_scale_2", + "w2_weight_scale_2", +) + +PLAN_WIDTH = 16 +ROW_USE_NEVER = 0x7FFFFFFFFFFFFFFF # staging rows: never a victim + + +@dataclass +class GlobalTables: + """Pool-wide device state; every tensor has a fixed address.""" + + num_layers: int + num_experts: int + pool_rows: int + hot_phys: Any # [L*E] int32 key -> bank row / -1 + cold_phys: Any # [L*E] int32 key -> RAM row (expert id) / -1 + row_key: Any # [pool_rows + staging] int32 row -> key / -1 + row_use: Any # [pool_rows + staging] int64 step of the row's last selection + clock: Any # [1] int64 + gate: Any # [1] int32 promotions allowed + error: Any # [1] int32 sticky device error + promote_limit: Any # [1] int32 max promotions per layer call, 0 = unlimited + promote_interval: Any # [1] int32 promote only every N forwards + forwards: Any # [1] int32 forwards seen with the gate open (layer 0 count) + promote_min_misses: Any # [1] int32 misses a key needs before promotion + protect_recent: Any # [1] int32 forwards a used row stays unevictable + miss_count: Any # [L*E] int32 misses since the key was last resident + staging_rows: Any # [S] int32 shared staging rows (constant) + + @property + def keys(self): + return self.num_layers * self.num_experts + + def layer_slice(self, table, layer): + start = layer * self.num_experts + return table[start : start + self.num_experts] + + +@dataclass +class StepBuffers: + """Per-layer fixed-address scratch the step program fills.""" + + gather_src: Any # [W] int32 RAM rows + gather_dst: Any # [W] int32 bank rows + gather_count: Any # [1] int32 + routes: Any # [W] int32 physical row per ids lane this step, -1 padding + staged_expert: Any # [W] int32 (scratch for the map overlay) + staged_row: Any # [W] int32 + staged_count: Any # [1] int32 + promoted_count: Any # [1] int32 + step_map: Any # [E] int32 + + +def allocate_global_tables(device, num_experts, slots_per_layer, staging): + """Layer l's first `slots_per_layer[l]` experts start resident, packed + in layer order; the `staging` rows follow the pool and stay free.""" + import torch + + num_layers = len(slots_per_layer) + if staging < 1 or any(not 0 < s < num_experts for s in slots_per_layer): + raise ValueError("Global pool needs staging rows and partial layers") + keys = num_layers * num_experts + pool_rows = sum(slots_per_layer) + hot_phys = torch.full((keys,), -1, dtype=torch.int32) + cold_phys = torch.arange(num_experts, dtype=torch.int32).repeat(num_layers) + row_key = torch.full((pool_rows + staging,), -1, dtype=torch.int32) + offset = 0 + for layer, slots in enumerate(slots_per_layer): + base = layer * num_experts + hot_phys[base : base + slots] = torch.arange( + offset, offset + slots, dtype=torch.int32 + ) + cold_phys[base : base + slots] = -1 + row_key[offset : offset + slots] = torch.arange( + base, base + slots, dtype=torch.int32 + ) + offset += slots + return GlobalTables( + num_layers=num_layers, + num_experts=num_experts, + pool_rows=pool_rows, + hot_phys=hot_phys.to(device), + cold_phys=cold_phys.to(device), + row_key=row_key.to(device), + row_use=torch.cat( + ( + torch.zeros(pool_rows, dtype=torch.int64), + torch.full((staging,), ROW_USE_NEVER, dtype=torch.int64), + ) + ).to(device), + clock=torch.zeros(1, dtype=torch.int64, device=device), + gate=torch.zeros(1, dtype=torch.int32, device=device), + error=torch.zeros(1, dtype=torch.int32, device=device), + promote_limit=torch.zeros(1, dtype=torch.int32, device=device), + promote_interval=torch.ones(1, dtype=torch.int32, device=device), + forwards=torch.zeros(1, dtype=torch.int32, device=device), + promote_min_misses=torch.ones(1, dtype=torch.int32, device=device), + protect_recent=torch.zeros(1, dtype=torch.int32, device=device), + miss_count=torch.zeros(keys, dtype=torch.int32, device=device), + staging_rows=torch.arange( + pool_rows, pool_rows + staging, dtype=torch.int32, device=device + ), + ) + + +def allocate_step_buffers(device, num_experts, width=PLAN_WIDTH): + import torch + + def ints(n): + return torch.zeros(n, dtype=torch.int32, device=device) + + return StepBuffers( + gather_src=ints(width), + gather_dst=ints(width), + gather_count=ints(1), + routes=torch.full((width,), -1, dtype=torch.int32, device=device), + staged_expert=ints(width), + staged_row=ints(width), + staged_count=ints(1), + promoted_count=ints(1), + step_map=torch.full((num_experts,), -1, dtype=torch.int32, device=device), + ) + + +def set_gate(tables, enabled): + tables.gate.fill_(1 if enabled else 0) + + +CONTROL_MAX = 2**31 - 1 # device scalars are int32 +CONTROL_FIELDS = ( + "promote_limit", + "promote_interval", + "promote_min_misses", + "protect_recent", + "gate", +) + + +def validate_control(values): + """Validate a control mapping; returns the normalized dict or raises.""" + out = {} + for name, value in values.items(): + if name not in CONTROL_FIELDS: + raise ValueError(f"Unknown pool control {name!r}") + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"Pool control {name} must be an integer") + if not 0 <= value <= CONTROL_MAX: + raise ValueError(f"Pool control {name} outside [0, {CONTROL_MAX}]") + if name == "promote_limit" and value < 0: + raise ValueError("promote_limit must be nonnegative") + if name in ("promote_interval", "promote_min_misses") and value < 1: + raise ValueError(f"{name} must be positive") + if name == "protect_recent" and value < 0: + raise ValueError("protect_recent must be nonnegative") + if name == "gate" and value not in (0, 1): + raise ValueError("gate must be 0 or 1") + out[name] = value + return out + + +def set_control(tables, **values): + """Write placement controls (device scalars at fixed addresses). + + `promote_limit`: promotions per *layer call* (0 = unlimited); every + layer reads the same scalar, so it is not a model-wide total. + `promote_interval`: promote only on every N-th forward with the gate + open, counted once per forward at layer 0 (a prefill or a multi-row + verify forward counts once). `promote_min_misses`: a key is promoted + only once it has missed this many times since it was last resident + (1 = first miss). `protect_recent`: rows used within the last N + forwards, the previous one included, are never victims (0 = none). + `gate`: 0 freezes the placement. In every frozen or deferred case all + misses are still computed from staging rows; only the placement changes. + All fields are validated before any is written. + """ + out = validate_control(values) + for name, value in out.items(): + if name == "gate": + set_gate(tables, bool(value)) + else: + getattr(tables, name).fill_(value) + return out + + +def read_control(tables): + values = {name: int(getattr(tables, name)[0]) for name in CONTROL_FIELDS} + values["forwards"] = int(tables.forwards[0]) + return values + + +def step_reference(tables, layer, ids, buffers): + """Plan and flip one layer step on the host (torch, synchronizing). + + Returns (gathers, step_map) with gathers as (RAM row, bank row) pairs in + copy order: promotions first, then staged-only misses. Semantics: + + - `ids` values outside [0, E) other than -1 set the sticky error and + are skipped; -1 is padding. + - Distinct valid selections in first-occurrence order. With the gate + open the clock advances and every selected resident row is stamped. + - Recency lives on pool rows (FreeToken's usage-per-slot): each miss, + in that order, evicts the pool row with the smallest (row_use, row) + among rows not stamped this step and takes it; without such a row + it is staged only. Staging rows are never victims. With the gate + closed every miss is staged only and recency is untouched. + - `buffers.routes[i]` is the physical row of ids lane i (-1 for + padding, invalid, or a duplicate of an earlier lane's expert is still + resolved to that expert's row). + """ + import torch + + E = tables.num_experts + if not 0 <= layer < tables.num_layers: + raise ValueError("Layer index outside the pool") + hot = tables.hot_phys.tolist() + cold = tables.cold_phys.tolist() + row_key = tables.row_key.tolist() + row_use = tables.row_use.tolist() + staging_rows = tables.staging_rows.tolist() + gate = bool(int(tables.gate[0])) + raw = [int(v) for v in ids.reshape(-1).tolist()] + if len(raw) > buffers.gather_src.shape[0] or len(raw) > len(staging_rows): + raise ValueError("Step ids exceed the plan width or the staging rows") + error = bool(int(tables.error[0])) + selected: list[int] = [] + for value in raw: + if value == -1: + continue + if not 0 <= value < E: + error = True + continue + if value not in selected: + selected.append(value) + clock = int(tables.clock[0]) + forwards = int(tables.forwards[0]) + base = layer * E + if gate: + clock += 1 + if layer == 0: + forwards += 1 + for e in selected: + if hot[base + e] >= 0: + row_use[hot[base + e]] = clock + limit = int(tables.promote_limit[0]) + interval = int(tables.promote_interval[0]) + min_misses = int(tables.promote_min_misses[0]) + protect = int(tables.protect_recent[0]) * tables.num_layers + miss_count = tables.miss_count.tolist() + promote_ok = gate and (forwards - 1) % interval == 0 + gathers: list[tuple[int, int]] = [] + staged: list[tuple[int, int]] = [] + for e in selected: + key = base + e + if hot[key] >= 0: + continue + victim_row = -1 + if gate: + miss_count[key] += 1 + if ( + promote_ok + and (limit == 0 or len(gathers) < limit) + and miss_count[key] >= min_misses + ): + best = None + for r in range(tables.pool_rows): + if ( + row_key[r] >= 0 + and row_use[r] < clock + and (protect == 0 or row_use[r] < clock - protect) + ): + candidate = (row_use[r], r) + if best is None or candidate < best: + best = candidate + if best is not None: + victim_row = best[1] + if victim_row < 0: + staged.append((e, staging_rows[len(staged)])) + continue + victim = row_key[victim_row] + hot[victim], cold[victim] = -1, victim % E + hot[key], cold[key] = victim_row, -1 + row_key[victim_row] = key + row_use[victim_row] = clock + miss_count[key] = 0 + gathers.append((e, victim_row)) + step_map = hot[base : base + E] + for e, row in staged: + step_map[e] = row + routes = [-1] * buffers.routes.shape[0] + for i, value in enumerate(raw): + if 0 <= value < E: + routes[i] = step_map[value] + device = tables.hot_phys.device + + def write(target, values, dtype): + target.copy_(torch.tensor(values, dtype=dtype, device=device)) + + write(tables.hot_phys, hot, torch.int32) + write(tables.cold_phys, cold, torch.int32) + write(tables.row_key, row_key, torch.int32) + write(tables.row_use, row_use, torch.int64) + tables.clock.fill_(clock) + tables.forwards.fill_(forwards) + write(tables.miss_count, miss_count, torch.int32) + tables.error.fill_(1 if error else 0) + pairs = gathers + [(e, row) for e, row in staged] + buffers.gather_count.fill_(len(pairs)) + buffers.promoted_count.fill_(len(gathers)) + buffers.staged_count.fill_(len(staged)) + for i, (src, dst) in enumerate(pairs): + buffers.gather_src[i], buffers.gather_dst[i] = src, dst + for i, (e, row) in enumerate(staged): + buffers.staged_expert[i], buffers.staged_row[i] = e, row + write(buffers.step_map, step_map, torch.int32) + write(buffers.routes, routes, torch.int32) + return pairs, buffers.step_map + + +def step(tables, layer, ids, buffers): + """Plan and flip one layer step: Triton on CUDA, the reference elsewhere.""" + if tables.hot_phys.device.type != "cuda": + step_reference(tables, layer, ids, buffers) + return + flat = ids.reshape(-1) + if not flat.is_contiguous(): + raise ValueError("Global step requires contiguous ids") + width = buffers.gather_src.shape[0] + if flat.numel() > width or flat.numel() > tables.staging_rows.shape[0]: + raise ValueError("Step ids exceed the plan width or the staging rows") + rows = tables.pool_rows + _step_kernel()[(1,)]( + flat, + flat.numel(), + layer, + tables.hot_phys, + tables.cold_phys, + tables.row_key, + tables.row_use, + tables.clock, + tables.gate, + tables.error, + tables.promote_limit, + tables.promote_interval, + tables.forwards, + tables.promote_min_misses, + tables.protect_recent, + tables.miss_count, + tables.staging_rows, + buffers.gather_src, + buffers.gather_dst, + buffers.gather_count, + buffers.routes, + buffers.staged_expert, + buffers.staged_row, + buffers.staged_count, + buffers.promoted_count, + buffers.step_map, + tables.num_experts, + rows, + tables.num_layers, + WIDTH=width, + BLOCK_R=_next_power_of_two(rows), + MAP_BLOCK=1024, + num_warps=8, + ) + + +def check_global_tables(tables): + """Consistency of the pool; raises on any violation. + + Every key is resident or has its RAM row, never both; resident keys and + rows are a bijection; staging rows are never owned; no device error. + """ + E = tables.num_experts + hot = tables.hot_phys.tolist() + cold = tables.cold_phys.tolist() + row_key = tables.row_key.tolist() + staging = set(tables.staging_rows.tolist()) + owners = {} + for key, (h, c) in enumerate(zip(hot, cold)): + if (h >= 0) == (c >= 0): + raise AssertionError(f"Key {key} must be resident or backed, not both") + if c >= 0 and c != key % E: + raise AssertionError(f"Key {key} must be backed by its own RAM row") + if h >= 0: + if h in staging or not 0 <= h < tables.pool_rows: + raise AssertionError(f"Key {key} owns a row outside the pool") + if h in owners: + raise AssertionError(f"Row {h} has two owners") + owners[h] = key + for row, key in enumerate(row_key): + if owners.get(row, -1) != key: + raise AssertionError(f"Row {row} owner table disagrees") + if int(tables.error[0]): + raise RuntimeError("Global pool recorded a device error") + + +def resident_per_layer(tables): + hot = tables.hot_phys.view(tables.num_layers, tables.num_experts) + return (hot >= 0).sum(dim=1).tolist() + + +_KERNELS: dict[str, Any] = {} + + +def _next_power_of_two(value): + return 1 << max(int(value) - 1, 0).bit_length() + + +def _step_kernel(): + """One program: `step_reference` on the device.""" + if "step" in _KERNELS: + return _KERNELS["step"] + from vllm.triton_utils import tl, triton + + @triton.jit + def global_pool_step( + ids_ptr, + n, + layer, + hot_phys_ptr, + cold_phys_ptr, + row_key_ptr, + row_use_ptr, + clock_ptr, + gate_ptr, + error_ptr, + limit_ptr, + interval_ptr, + forwards_ptr, + min_misses_ptr, + protect_ptr, + miss_count_ptr, + staging_ptr, + gather_src_ptr, + gather_dst_ptr, + gather_count_ptr, + routes_ptr, + staged_expert_ptr, + staged_row_ptr, + staged_count_ptr, + promoted_count_ptr, + step_map_ptr, + num_experts, + pool_rows, + num_layers, + WIDTH: tl.constexpr, + BLOCK_R: tl.constexpr, + MAP_BLOCK: tl.constexpr, + ): + never = 0x7FFFFFFFFFFFFFFF + lane = tl.arange(0, WIDTH) + present = lane < n + raw = tl.load(ids_ptr + lane, mask=present, other=-1).to(tl.int64) + valid = present & (raw >= 0) & (raw < num_experts) + bad = present & (raw != -1) & (~valid) + if tl.sum(bad.to(tl.int32), 0) > 0: + tl.store(error_ptr, 1) + safe = tl.where(valid, raw, 0) + same = safe[:, None] == safe[None, :] + earlier = lane[None, :] < lane[:, None] + duplicate = tl.sum((same & earlier & valid[None, :]).to(tl.int32), 1) > 0 + distinct = valid & (duplicate == 0) + # `layer` may arrive as a Python int (Triton specializes 0 and 1). + base = tl.full((), 0, tl.int64) + layer * num_experts + keys = base + safe + resident = tl.load(hot_phys_ptr + keys, mask=distinct, other=-1).to(tl.int64) + hit = distinct & (resident >= 0) + gate = tl.load(gate_ptr) != 0 + clock = tl.load(clock_ptr) + forwards = tl.load(forwards_ptr) + if gate: + clock = clock + 1 + tl.store(clock_ptr, clock) + if layer == 0: + forwards = forwards + 1 + tl.store(forwards_ptr, forwards) + tl.store(row_use_ptr + tl.where(hit, resident, 0), clock, mask=hit) + limit = tl.load(limit_ptr) + interval = tl.load(interval_ptr) + min_misses = tl.load(min_misses_ptr) + protect = tl.load(protect_ptr).to(tl.int64) * num_layers + promote_ok = gate & (((forwards - 1) % interval) == 0) + tl.debug_barrier() + # The pool-wide recency vector is only needed when a miss can be + # promoted (FreeToken scans its cache inside the same condition); + # on an all-hit layer, the common case, nothing below touches it. + offs_r = tl.arange(0, BLOCK_R) + in_pool = offs_r < pool_rows + misses = tl.sum((distinct & (~hit)).to(tl.int32), 0) + scan = promote_ok & (misses > 0) + use = tl.full((BLOCK_R,), never, tl.int64) + if scan: + use = tl.load(row_use_ptr + offs_r, mask=in_pool, other=never) + # Rows used within the last protect_recent forwards (the + # previous one included) stay. + if protect > 0: + use = tl.where(use >= clock - protect, never, use) + # Rows selected this step (hits) are masked in registers; extract + # lane i's hit row (or -1): the other lanes contribute 0. + hit_rows = tl.where(hit, resident, -1) + for i in range(0, WIDTH): + hit_row = tl.sum(tl.where(lane == i, hit_rows, 0), 0) + use = tl.where(offs_r.to(tl.int64) == hit_row, never, use) + promoted = 0 + staged = 0 + for i in range(0, WIDTH): + is_miss = tl.sum( + tl.where(lane == i, (distinct & (~hit)).to(tl.int32), 0), 0 + ) + if is_miss > 0: + expert = tl.load(ids_ptr + i).to(tl.int64) + key = base + expert + victim_row = tl.full((), -1, tl.int64) + misses_so_far = tl.load(miss_count_ptr + key) + if gate: + misses_so_far = misses_so_far + 1 + tl.store(miss_count_ptr + key, misses_so_far) + if ( + promote_ok + & ((limit == 0) | (promoted < limit)) + & (misses_so_far >= min_misses) + ): + best = tl.min(use, 0) + if best != never: + victim_row = tl.min( + tl.where(use == best, offs_r.to(tl.int64), never), 0 + ) + if victim_row >= 0: + victim = tl.load(row_key_ptr + victim_row).to(tl.int64) + tl.store(hot_phys_ptr + victim, -1) + tl.store( + cold_phys_ptr + victim, (victim % num_experts).to(tl.int32) + ) + tl.store(hot_phys_ptr + key, victim_row.to(tl.int32)) + tl.store(cold_phys_ptr + key, -1) + tl.store(row_key_ptr + victim_row, key.to(tl.int32)) + tl.store(row_use_ptr + victim_row, clock) + tl.store(miss_count_ptr + key, 0) + tl.store(gather_src_ptr + promoted, expert.to(tl.int32)) + tl.store(gather_dst_ptr + promoted, victim_row.to(tl.int32)) + use = tl.where(offs_r.to(tl.int64) == victim_row, never, use) + promoted += 1 + else: + tl.store(staged_expert_ptr + staged, expert.to(tl.int32)) + tl.store(staged_row_ptr + staged, tl.load(staging_ptr + staged)) + staged += 1 + # The scalar table stores above must be visible to the next + # miss's loads of hot_phys / row_key. + tl.debug_barrier() + tl.store(promoted_count_ptr, promoted) + tl.store(staged_count_ptr, staged) + tl.store(gather_count_ptr, promoted + staged) + tl.debug_barrier() + for i in range(0, staged): + tl.store(gather_src_ptr + promoted + i, tl.load(staged_expert_ptr + i)) + tl.store(gather_dst_ptr + promoted + i, tl.load(staged_row_ptr + i)) + for start in range(0, num_experts, MAP_BLOCK): + offs = start + tl.arange(0, MAP_BLOCK) + in_range = offs < num_experts + rows = tl.load(hot_phys_ptr + base + offs, mask=in_range, other=-1) + tl.store(step_map_ptr + offs, rows, mask=in_range) + tl.debug_barrier() + for i in range(0, staged): + expert = tl.load(staged_expert_ptr + i) + tl.store(step_map_ptr + expert, tl.load(staged_row_ptr + i)) + tl.debug_barrier() + # Routes: the physical row of every ids lane through the step map. + route = tl.load(step_map_ptr + safe, mask=valid, other=-1) + tl.store(routes_ptr + lane, tl.where(valid, route, -1), mask=lane < WIDTH) + + _KERNELS["step"] = global_pool_step + return global_pool_step diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 5d284bcc80f3..4760bb703160 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -7,6 +7,7 @@ import torch +from vllm.config import get_current_vllm_config from vllm.distributed.eplb.eplb_state import EplbState from vllm.logger import init_logger from vllm.model_executor.custom_op import PluggableLayer @@ -168,10 +169,61 @@ def __init__( "global_num_experts": moe_config.num_experts, } + # Expert pool (offload_config.moe_expert_pool_rows > 0): the quant + # method allocates the per-expert tensors in pinned host memory, marks + # the layer pending after its weights reach their kernel layout, and + # expert_pool.install_expert_pool binds every pending layer to one + # shared bank at the model level. + self._moe_expert_pool_rows = ( + get_current_vllm_config().offload_config.moe_expert_pool_rows + ) + self.expert_pool_pending = False + self.expert_pool_layer: Any = None + if self._moe_expert_pool_rows > 0: + self._validate_expert_pool_supported() + self.quant_method.create_weights(layer=self, **moe_quant_params) self.lora_base_layer_prefix = "" + def _validate_expert_pool_supported(self) -> None: + """Reject what the pool cannot serve before any weight is allocated: + the pool binds one consumer (NVFP4 Marlin), plans top_k rows per + token, and keys rows by (layer, expert) without an expert map.""" + top_k = self.moe_config.experts_per_token + if self._moe_expert_pool_rows < top_k: + raise ValueError( + f"moe_expert_pool_rows={self._moe_expert_pool_rows} is fewer " + f"than the {top_k} experts a single token routes to. Set " + f"--moe-expert-pool-rows >= {top_k}." + ) + parallel = self.moe_config.moe_parallel_config + if parallel.use_ep: + raise ValueError( + "moe_expert_pool_rows is not compatible with expert " + f"parallelism (ep_size={parallel.ep_size})." + ) + if parallel.dp_size > 1 or parallel.is_sequence_parallel: + raise ValueError( + "moe_expert_pool_rows is not compatible with data parallelism " + "or sequence parallelism." + ) + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + NvFp4MoeBackend, + ) + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4FusedMoE, + ) + + if ( + not isinstance(self.quant_method, ModelOptNvFp4FusedMoE) + or self.quant_method.nvfp4_backend != NvFp4MoeBackend.MARLIN + ): + raise ValueError( + "moe_expert_pool_rows supports the ModelOpt NVFP4 Marlin MoE " + f"backend only, got {type(self.quant_method).__name__}" + ) + # TODO(bnell): Temporary hack. Get rid of this. def _replace_quant_method(self, quant_method: FusedMoEMethodBase): self.quant_method = quant_method diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 59292ec99042..89c30689db07 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -849,6 +849,14 @@ def uses_weight_scale_2_pattern(self) -> bool: """ return True + @property + def mk_can_overlap_shared_experts(self) -> bool: + # The expert pool runs its own consumer outside self.moe_kernel and + # does not overlap shared experts; the runner must run them itself. + if getattr(self, "_pool_mode", False): + return False + return super().mk_can_overlap_shared_experts + def create_weights( self, layer: RoutedExperts, @@ -868,9 +876,23 @@ def create_weights( weight_loader = extra_weight_attrs.get("weight_loader") global_num_experts = extra_weight_attrs.get("global_num_experts") w13_num_shards = 2 if self.moe.is_act_and_mul else 1 + + # With the expert pool enabled, the per-expert tensors are allocated + # in CPU pinned memory so loading never needs GPU capacity for the + # whole layer set; the loader moves one layer at a time to the device + # for the Marlin repack and restores the result to pinned memory, + # which the pool then takes as its host source. device="cpu" is + # explicit because loading runs under an accelerator device context. + expert_tensors_on_cpu = getattr(layer, "_moe_expert_pool_rows", 0) > 0 + + def _empty_expert_tensor(*shape: int, dtype: torch.dtype) -> torch.Tensor: + if expert_tensors_on_cpu: + return torch.empty(*shape, dtype=dtype, device="cpu").pin_memory() + return torch.empty(*shape, dtype=dtype) + # GEMM 1 w13_weight = ModelWeightParameter( - data=torch.empty( + data=_empty_expert_tensor( num_experts, w13_num_shards * intermediate_size_per_partition, # 2 fp4 items are packed in the input dimension @@ -885,7 +907,7 @@ def create_weights( # GEMM 2 w2_weight = ModelWeightParameter( - data=torch.empty( + data=_empty_expert_tensor( num_experts, hidden_size, # 2 fp4 items are packed in the input dimension @@ -899,7 +921,7 @@ def create_weights( layer.register_parameter("w2_weight", w2_weight) w13_weight_scale = ModelWeightParameter( - data=torch.empty( + data=_empty_expert_tensor( num_experts, w13_num_shards * intermediate_size_per_partition, # 2 fp4 items are packed in the input dimension @@ -913,7 +935,7 @@ def create_weights( layer.register_parameter("w13_weight_scale", w13_weight_scale) w2_weight_scale = ModelWeightParameter( - data=torch.empty( + data=_empty_expert_tensor( num_experts, hidden_size, # 2 fp4 items are packed in the input dimension @@ -1013,6 +1035,13 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: replace_parameter(layer, "w2_weight_scale_2", w2_scale_2) replace_parameter(layer, "w2_input_scale", a2_scale) + # The parameters above are in the kernel's final representation; the + # expert pool takes them as its host source and binds its own Marlin + # consumer at the model level (expert_pool.install_expert_pool). + self._pool_mode = getattr(layer, "_moe_expert_pool_rows", 0) > 0 + if self._pool_mode: + layer.expert_pool_pending = True + # Setup modular kernel. self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.experts_cls is not None @@ -1080,6 +1109,21 @@ def apply( ) -> torch.Tensor: assert not self.is_monolithic assert self.moe_kernel is not None + + pool_layer = getattr(layer, "expert_pool_layer", None) + if pool_layer is not None: + # The runner always passes its SharedExperts wrapper; the wrapper + # picks the order itself. The pool never overlaps shared experts + # (mk_can_overlap_shared_experts is False), so the runner has + # already run them (NO_OVERLAP or the aux stream) and the + # argument is ignored here, as the synchronous modular path does. + assert not self.mk_can_overlap_shared_experts + return pool_layer.apply(x, topk_weights, topk_ids) + if getattr(layer, "expert_pool_pending", False): + raise RuntimeError( + f"{layer.layer_name}: expert pool was requested but not installed" + ) + return self.moe_kernel.apply( x, layer.w13_weight, diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index e12910865ec2..c06ef4e429f7 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -12,7 +12,12 @@ from typing_extensions import assert_never import vllm.envs as envs -from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config +from vllm.config import ( + ModelConfig, + VllmConfig, + get_current_vllm_config, + set_current_vllm_config, +) from vllm.logger import init_logger from vllm.model_executor.layers.attention import is_deferred_attention_layer from vllm.model_executor.layers.hpc import HpcModule @@ -163,6 +168,21 @@ def process_weights_after_loading( if isinstance(module, HpcModule): module.process_weights_after_loading(model) + # Expert pool (moe_expert_pool_rows > 0): one shared bank for all MoE + # layers, bound after every layer's tensors are in their final layout. + from vllm.model_executor.layers.fused_moe.expert_pool.install import ( + install_expert_pool, + ) + + scheduler_config = get_current_vllm_config().scheduler_config + install_expert_pool( + model, + target_device, + max_decode_tokens=( + scheduler_config.max_num_seqs if scheduler_config is not None else 1 + ), + ) + # Model-level post-load hook, after the per-layer quant finalize. if hasattr(model, "process_weights_after_loading"): model.process_weights_after_loading() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 53bc0550a524..3ec4422d5033 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -935,6 +935,13 @@ def compile_or_warm_up_model(self) -> CompilationTimes: # intra-op parallelism. set_torch_threads_for_runtime() + # Expert pool: placement frozen during profiling/capture; open it now. + from vllm.model_executor.layers.fused_moe.expert_pool.install import ( + open_pool_gate, + ) + + open_pool_gate(self.model_runner.model) + return CompilationTimes( language_model=self.compilation_config.compilation_time, encoder=self.compilation_config.encoder_compilation_time, From a561b979ab0a5f761d2f8c74855444da604f11ec Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:17:32 +0900 Subject: [PATCH 02/13] [MoE] Expert pool tests: layer-construction guard regression; CLI test scope note Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- tests/config/test_moe_expert_pool_rows.py | 4 +- .../expert_pool/test_install_guards.py | 42 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/tests/config/test_moe_expert_pool_rows.py b/tests/config/test_moe_expert_pool_rows.py index b070b42b134f..2dc2a4d92498 100644 --- a/tests/config/test_moe_expert_pool_rows.py +++ b/tests/config/test_moe_expert_pool_rows.py @@ -26,7 +26,9 @@ def test_negative_rows_are_rejected(): OffloadConfig(moe_expert_pool_rows=-1) -def test_cli_reaches_engine_args_and_offload_config(): +def test_cli_reaches_engine_args(): + # Checks the CLI flag and the OffloadConfig field only; it does not + # exercise EngineArgs.create_engine_config (needs a model). parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) args = parser.parse_args(["--moe-expert-pool-rows", "16"]) engine_args = EngineArgs.from_cli_args(args) diff --git a/tests/kernels/expert_pool/test_install_guards.py b/tests/kernels/expert_pool/test_install_guards.py index ababd301e63f..117a52a1a038 100644 --- a/tests/kernels/expert_pool/test_install_guards.py +++ b/tests/kernels/expert_pool/test_install_guards.py @@ -66,3 +66,45 @@ def test_top_k_beyond_the_lane_cap_is_rejected(): with pytest.raises(ValueError, match="0 < top_k"): inst._check_top_k(bad) inst._check_top_k(inst.MAX_DECODE_LANES) + + +def _routed(rows=4, top_k=2, use_ep=False, dp_size=1, sp=False, method=None): + # The layer-construction guard runs before create_weights; it reads only + # these attributes, so a namespace stands in for the RoutedExperts. + return SimpleNamespace( + _moe_expert_pool_rows=rows, + quant_method=method or _method(), + moe_config=SimpleNamespace( + experts_per_token=top_k, + moe_parallel_config=SimpleNamespace( + use_ep=use_ep, + ep_size=2 if use_ep else 1, + dp_size=dp_size, + is_sequence_parallel=sp, + ), + ), + ) + + +def test_layer_guard_accepts_the_supported_geometry(): + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + + RoutedExperts._validate_expert_pool_supported(_routed()) + + +@pytest.mark.parametrize( + "bad, match", + [ + (_routed(rows=1), "fewer than the 2 experts"), + (_routed(use_ep=True), "expert parallelism"), + (_routed(dp_size=2), "data parallelism"), + (_routed(sp=True), "sequence parallelism"), + (_routed(method=_method(NvFp4MoeBackend.VLLM_CUTLASS)), "Marlin"), + (_routed(method=SimpleNamespace(nvfp4_backend=None)), "Marlin"), + ], +) +def test_layer_guard_rejects_before_any_weight_is_allocated(bad, match): + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + + with pytest.raises(ValueError, match=match): + RoutedExperts._validate_expert_pool_supported(bad) From 3bd349317e3251815f98ecc8261d5cdb8ed17610 Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:19:56 +0900 Subject: [PATCH 03/13] [MoE] Expert pool tests: run the host-source fixture under device_loading_context; CLI test without model resolution Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- tests/config/test_moe_expert_pool_rows.py | 19 ++++++++++++------- tests/kernels/expert_pool/marlin_fixture.py | 18 +++++++++++++++--- 2 files changed, 27 insertions(+), 10 deletions(-) diff --git a/tests/config/test_moe_expert_pool_rows.py b/tests/config/test_moe_expert_pool_rows.py index 2dc2a4d92498..10632cf591b4 100644 --- a/tests/config/test_moe_expert_pool_rows.py +++ b/tests/config/test_moe_expert_pool_rows.py @@ -26,12 +26,17 @@ def test_negative_rows_are_rejected(): OffloadConfig(moe_expert_pool_rows=-1) -def test_cli_reaches_engine_args(): - # Checks the CLI flag and the OffloadConfig field only; it does not - # exercise EngineArgs.create_engine_config (needs a model). +def test_cli_flag_and_field(): + # Checks the CLI flag and the OffloadConfig field only. EngineArgs + # construction and create_engine_config resolve a model (network or a + # local snapshot), so they are out of scope here. parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) args = parser.parse_args(["--moe-expert-pool-rows", "16"]) - engine_args = EngineArgs.from_cli_args(args) - assert engine_args.moe_expert_pool_rows == 16 - offload = OffloadConfig(moe_expert_pool_rows=engine_args.moe_expert_pool_rows) - assert offload.moe_expert_pool_rows == 16 + assert args.moe_expert_pool_rows == 16 + assert EngineArgs.moe_expert_pool_rows == 0 + assert ( + OffloadConfig( + moe_expert_pool_rows=args.moe_expert_pool_rows + ).moe_expert_pool_rows + == 16 + ) diff --git a/tests/kernels/expert_pool/marlin_fixture.py b/tests/kernels/expert_pool/marlin_fixture.py index c3ede294aded..b38d9be6b4f7 100644 --- a/tests/kernels/expert_pool/marlin_fixture.py +++ b/tests/kernels/expert_pool/marlin_fixture.py @@ -15,6 +15,7 @@ ) from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config +from vllm.model_executor.model_loader.utils import device_loading_context from vllm.utils.torch_utils import set_random_seed from vllm.v1.worker.workspace import ( init_workspace_manager, @@ -90,8 +91,10 @@ def make_layer( cfg: VllmConfig, params: dict[str, torch.Tensor], host_source: bool = False ): """Build the layer; with host_source the per-expert tensors are - registered as pinned CPU tensors, the layout the loader restores when - the pool is enabled.""" + registered as pinned CPU tensors and post-load processing runs under the + loader's device_loading_context, as in the real load: the layer is moved + to the device for the Marlin conversion and the converted tensors are + restored to pinned host memory, which the pool takes as its source.""" with set_current_vllm_config(cfg): # Any construction error is a failure: the Marlin capability gate is # the module-level skip in the test, and the backend is pinned. @@ -123,7 +126,16 @@ def make_layer( layer.routed_experts.register_parameter( name, torch.nn.Parameter(data, requires_grad=False) ) - layer._quant_method.process_weights_after_loading(layer.routed_experts) + if host_source: + device = torch.accelerator.current_accelerator() + assert device is not None + with device_loading_context(layer.routed_experts, device): + layer._quant_method.process_weights_after_loading(layer.routed_experts) + for name in EXPERT_TENSORS: + p = getattr(layer.routed_experts, name) + assert p.device.type == "cpu" and p.is_pinned(), name + else: + layer._quant_method.process_weights_after_loading(layer.routed_experts) return layer From 42e988be1d0d8057a0f601a9e900504fed5c2c3c Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:20:40 +0900 Subject: [PATCH 04/13] [MoE] Expert pool tests: assert device tensors inside the loader context Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- tests/kernels/expert_pool/marlin_fixture.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/kernels/expert_pool/marlin_fixture.py b/tests/kernels/expert_pool/marlin_fixture.py index b38d9be6b4f7..bdb52c9abbc1 100644 --- a/tests/kernels/expert_pool/marlin_fixture.py +++ b/tests/kernels/expert_pool/marlin_fixture.py @@ -130,6 +130,10 @@ def make_layer( device = torch.accelerator.current_accelerator() assert device is not None with device_loading_context(layer.routed_experts, device): + # The loader contract: the conversion sees device tensors. + for name in EXPERT_TENSORS: + p = getattr(layer.routed_experts, name) + assert p.device.type == device.type, name layer._quant_method.process_weights_after_loading(layer.routed_experts) for name in EXPERT_TENSORS: p = getattr(layer.routed_experts, name) From fcfa3e573515a13fbb87f3e10a120b119e0a656f Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Wed, 9 Sep 2026 21:50:58 +0900 Subject: [PATCH 05/13] Add reproducible expert pool generation benchmark and frozen prompts Co-authored-by: Codex Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- benchmarks/expert_pool/LICENSE.prompts | 21 + benchmarks/expert_pool/README.md | 127 ++ benchmarks/expert_pool/benchmark.py | 189 ++ benchmarks/expert_pool/pair.json | 1985 ++++++++++++++++++++++ benchmarks/expert_pool/prompts.md | 155 ++ benchmarks/expert_pool/provenance.json | 18 + benchmarks/expert_pool/test_benchmark.py | 109 ++ 7 files changed, 2604 insertions(+) create mode 100644 benchmarks/expert_pool/LICENSE.prompts create mode 100644 benchmarks/expert_pool/README.md create mode 100644 benchmarks/expert_pool/benchmark.py create mode 100644 benchmarks/expert_pool/pair.json create mode 100644 benchmarks/expert_pool/prompts.md create mode 100644 benchmarks/expert_pool/provenance.json create mode 100644 benchmarks/expert_pool/test_benchmark.py diff --git a/benchmarks/expert_pool/LICENSE.prompts b/benchmarks/expert_pool/LICENSE.prompts new file mode 100644 index 000000000000..bc0556699a8c --- /dev/null +++ b/benchmarks/expert_pool/LICENSE.prompts @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/benchmarks/expert_pool/README.md b/benchmarks/expert_pool/README.md new file mode 100644 index 000000000000..a31c361d8711 --- /dev/null +++ b/benchmarks/expert_pool/README.md @@ -0,0 +1,127 @@ +# Expert重みのGPUキャッシュ：速度測定の再現手順 + +Qwen3.8 FlashNext NVFP4に、準備用の要求と速度測定用の要求を順番に送ります。 +各要求ではソフトウェアの不具合報告を読み、修正案と検証計画を生成します。 + +## ファイルと要求の内容 + +| ファイル | 内容 | +| --- | --- | +| [benchmark.py](benchmark.py) | HTTP要求、ストリーム保存、生成速度の集計 | +| [pair.json](pair.json) | 実測に使用したプロンプト全文、固定トークン列、出典 | +| [prompts.md](prompts.md) | プロンプト全文の読みやすい表示 | +| [provenance.json](provenance.json) | 実測コードの版、元ファイルのSHA256、移植差分 | +| [LICENSE.prompts](LICENSE.prompts) | プロンプトの出典に付属するMITライセンス | + +| 送信順 | 用途 | 不具合報告 | 入力トークン数 | +| --- | --- | --- | --- | +| 1回目 | 準備用（結果のroleはwarmup） | ファイル選択ボタンの「Choose File」を「Choose file」に修正する課題（28096_836） | 1070 | +| 2回目 | 速度測定用（roleはmeasure） | 言語設定を変更したとき「Link sent!」表示も更新する課題（18827_741） | 753 | + +出典は[OpenAI frontier-evals](https://github.com/openai/frontier-evals/tree/51052cede8cc608f95bb00346635e03759013e5a)のSWE-Lancerです。 +既存の動作確認用課題から選んだ2件で、測定対象は修正案の文章生成です。 +実際のコード編集・公式採点を行う品質試験は別の手順です。 + +## 実測に使用したコードと環境 + +| 用途 | コードの版 | +| --- | --- | +| 比較の土台となるvLLM main | `a97dacb7106ee49f39f3d1fc6ae1800ff724e01d` | +| Expert重みのGPUキャッシュの実装（[fork PR #48](https://github.com/01554/vllm/pull/48)） | `5fbc240ba5ddec82a10362340ac77339a1c24017` | +| PLEの読み出しと生成計算を重ねる実装（[fork PR #46](https://github.com/01554/vllm/pull/46)、[上流PR #54129](https://github.com/vllm-project/vllm/pull/54129)を前提とする差分） | `4f859de9d0f55760b50358aee4834e6966e13bc8` | +| 上記機能を組み合わせ、以下の速度を測定した版 | `7dedc6d8d9b178b60f6a5b32f03d677145982441` | + +サーバーは実測版のPythonソース、上記mainからビルドしたwheel、別途ビルドした +`_ple_memops`拡張を組み合わせて動かしました。このディレクトリは測定後に追加した +クライアント用ファイルです。`pair.json`は実測ファイルのbyteコピーです。 + +RTX 6000 Adaの48GBに収める構成の検証を目的として、手元の +RTX PRO 6000 Blackwell Max-Q（96GiB）上で別プロセスにGPUメモリを確保させ、 +サーバーに利用可能な容量を48GiBにして測定しました。以下はこのGPU上の実測値です。 +ホスト側のコンテナのメモリ上限は100GiBでした。 + +GPUキャッシュは48層それぞれ258 expert行、約32GiBです。 +容量制限は外部プロセスで設定します。下記の`gpu-memory-utilization`は物理GPU容量に +対するvLLMの予算比率です。クライアント実行前に容量とサーバー起動状態を確認してください。 + +## サーバーの起動設定 + +上の実測版と同じ機能をビルドした環境で、チェックポイントのパスを指定します。 +`VLLM_USE_BREAKABLE_CUDAGRAPH`は未設定（自動選択）で測定しました。 + +```bash +export CHECKPOINT=/data/models/Qwen3.8-Flash-Next-NVFP4-nvidia +unset VLLM_USE_BREAKABLE_CUDAGRAPH +export VLLM_DEBUG_WORKSPACE=1 VLLM_LOGGING_LEVEL=DEBUG +export PYTORCH_ALLOC_CONF=pinned_max_round_threshold_mb:1,pinned_max_cached_size_mb:1 +export OMP_NUM_THREADS=2 OPENBLAS_NUM_THREADS=2 +export VLLM_USE_V2_MODEL_RUNNER=1 +export VLLM_PLE_MMAP=1 VLLM_PLE_MMAP_DEFERRED=1 +export VLLM_PLE_MMAP_PREWARM=0 VLLM_PLE_MMAP_PINNED=0 +export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 +.venv/bin/python -m vllm.entrypoints.openai.api_server \ + --model "$CHECKPOINT" --served-model-name flashnext \ + --host 0.0.0.0 --port 8000 --tensor-parallel-size 1 \ + --quantization modelopt --dtype bfloat16 --moe-backend marlin \ + --moe-expert-pool-rows 258 --language-model-only \ + --max-model-len 4096 --max-num-seqs 1 --max-num-batched-tokens 512 \ + --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}' \ + --no-enable-flashinfer-autotune --gpu-memory-utilization 0.4548806288994517 \ + --safetensors-load-strategy lazy \ + --default-chat-template-kwargs '{"enable_thinking":false}' \ + --reasoning-parser qwen3 --generation-config vllm +``` + +速度測定のコンテキスト上限は4096です。別途実施した品質試験では32768を使いました。 + +## クライアントの実行 + +サーバーの準備完了後、次を1回実行します。クライアントはPython標準ライブラリで動きます。 + +```bash +.venv/bin/python benchmarks/expert_pool/benchmark.py \ + --base-url http://127.0.0.1:8000 --model flashnext \ + --label fresh-0 --output results/fresh-0.jsonl +``` + +送信先は`/v1/completions`です。`pair.json`の固定トークン列を送信するので、 +チェックポイントのtokenizerが`pair.json`の`tokenization`に記録したSHA256と +一致することを確認してください。要求には`temperature=0`、`top_p=1`、`seed=0`、 +`max_tokens=2048`を指定し、thinkingを無効にしたチャットテンプレートのトークン列を使います。 +`tokenization.pair_sha256`はトークン列追加前の資料のハッシュで、ファイル全体のハッシュは +`provenance.json`にあります。 + +3回の測定では、**毎回サーバーを終了して新しいプロセスで起動し、準備用→測定用を1組送信**します。 +出力名を`fresh-0.jsonl`、`fresh-1.jsonl`、`fresh-2.jsonl`と変えます。 +実測はこの順で3組を実行し、2回目の要求の速度3値から中央値を求めました。 +クライアントは既存の出力ファイルを保護し、HTTPエラーや不完全なストリームを保存して停止します。 +失敗した測定も結果として保持してください。 + +## 指標と実測値 + +生成速度は、usageの生成トークン数とクライアント側の受信時刻から計算します。 + +```text +decode_tok_s = (completion_tokens - 1) / (最後の本文受信時刻 - 最初の本文受信時刻) +``` + +`first_token_s`は要求開始から最初の本文受信までの時間です。 +`e2e_tok_s`は要求全体の所要時間あたりの生成トークン数です。 +ストリームの1イベントに複数トークンが入ることがあるため、いずれもクライアント観測の値です。 +生のSSE、送信body、usage、finish reason、本文とSHA256も同じJSONLに保存します。 + +| 新しいサーバーでの実行 | 2回目の要求の生成速度 | +| --- | --- | +| 1回目 | 63.1882 tok/s | +| 2回目 | 62.7087 tok/s | +| 3回目 | 63.6519 tok/s | +| 中央値 | **63.1882 tok/s** | + +これは表の実測版で各機能を組み合わせた値です。各要求の終了理由は`stop`でした。 +容量は起動中・生成中の標本で記録し、プロセス終了コード0とOOMKilled=falseを確認しました。 + +## クライアントの動作確認 + +```bash +.venv/bin/python -m unittest discover -s benchmarks/expert_pool -p 'test_*.py' +``` diff --git a/benchmarks/expert_pool/benchmark.py b/benchmarks/expert_pool/benchmark.py new file mode 100644 index 000000000000..78365ace7443 --- /dev/null +++ b/benchmarks/expert_pool/benchmark.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Send one frozen warmup request, then one measured request to a fresh server. + +This is issue-text generation timing, not the SWE-Lancer correctness evaluation. +See README.md for the server configuration and the historical script provenance. +""" + +import argparse +import datetime +import hashlib +import json +import math +import time +import urllib.error +import urllib.request +from pathlib import Path + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--base-url", required=True) + p.add_argument("--label", required=True) + p.add_argument("--pair", type=Path, default=Path(__file__).with_name("pair.json")) + p.add_argument("--output", required=True) + p.add_argument("--model", default="flashnext") + p.add_argument("--max-tokens", type=int, default=2048) + p.add_argument("--timeout", type=int, default=1200) + a = p.parse_args() + source = Path(a.pair).read_bytes() + pair = json.loads(source) + tasks = pair["tasks"] + if len(tasks) != 2 or [t["role"] for t in tasks] != ["warmup", "measure"]: + raise ValueError( + "Expected exactly one warmup task followed by one measure task" + ) + sequence = [(tasks[0], "warmup"), (tasks[1], "measure")] + output = Path(a.output) + output.parent.mkdir(parents=True, exist_ok=True) + # Refuse to overwrite earlier measurements. + with output.open("x") as out: + for index, (task, role) in enumerate(sequence): + prompt = task["prompt_token_ids"] + request_body = dict( + model=a.model, + prompt=prompt, + max_tokens=a.max_tokens, + temperature=0, + top_p=1, + stream=True, + stream_options={"include_usage": True}, + seed=0, + ) + request = urllib.request.Request( + a.base_url.rstrip("/") + "/v1/completions", + data=json.dumps(request_body).encode(), + headers={"Content-Type": "application/json"}, + ) + result = dict( + schema_version=1, + timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(), + label=a.label, + sequence_index=index, + role=role, + task_id=task["id"], + pair_sha256=hashlib.sha256(source).hexdigest(), + prompt_sha256=task["prompt_sha256"], + expected_prompt_token_ids=task["prompt_token_ids"], + request=request_body, + content="", + usage=None, + finish_reason=None, + first_token_s=None, + last_token_s=None, + stream_events=0, + ) + print( + f"Starting {a.label}: {role} {task['id']} " + f"({len(task['prompt_token_ids'])} input tokens)", + flush=True, + ) + result["raw_sse_events"] = [] + result["sse_done"] = False + start = time.perf_counter() # Includes HTTP request and server scheduling. + try: + with urllib.request.urlopen(request, timeout=a.timeout) as response: + for raw in response: + line = raw.decode().strip() + if not line.startswith("data:"): + continue + data = line[5:].strip() + if data == "[DONE]": + result["sse_done"] = True + break + result["raw_sse_events"].append(data) + event = json.loads(data) + if event.get("error"): + raise RuntimeError(event["error"]) + if event.get("usage"): + result["usage"] = event["usage"] + for choice in event.get("choices", []): + content = choice.get("text") or "" + if content: + elapsed = time.perf_counter() - start + if result["first_token_s"] is None: + result["first_token_s"] = elapsed + result["last_token_s"] = elapsed + result["content"] += content + result["stream_events"] += 1 + if choice.get("finish_reason"): + result["finish_reason"] = choice["finish_reason"] + result["elapsed_s"] = time.perf_counter() - start + if not result["sse_done"]: + raise RuntimeError("SSE ended without DONE") + completion_tokens = (result["usage"] or {}).get("completion_tokens") + if ( + type(completion_tokens) is not int + or completion_tokens < 1 + or not result["content"] + ): + raise RuntimeError( + "Missing completion token usage or nonempty output" + ) + if result["finish_reason"] not in ("stop", "length"): + raise RuntimeError( + "Incomplete or unexpected finish reason: " + f"{result['finish_reason']}" + ) + if (result["usage"] or {}).get("prompt_tokens") != len( + task["prompt_token_ids"] + ): + raise RuntimeError( + "Server input token count does not match frozen prompt" + ) + result["completion_tokens"] = completion_tokens + result["e2e_tok_s"] = completion_tokens / result["elapsed_s"] + decode_duration = result["last_token_s"] - result["first_token_s"] + result["decode_tok_s"] = ( + (completion_tokens - 1) / decode_duration + if completion_tokens > 1 and decode_duration > 0 + else None + ) + for metric in ("e2e_tok_s", "decode_tok_s"): + value = result[metric] + if value is not None and not math.isfinite(value): + raise RuntimeError(f"Non-finite {metric}") + result["decode_metric_note"] = ( + "Approximate client-observed " + "(completion_tokens-1)/(last-text-event-first-text-event)" + ) + result["output_sha256"] = hashlib.sha256( + result["content"].encode() + ).hexdigest() + result["truncated"] = result["finish_reason"] == "length" + result["correctness"] = "Not graded: issue-text inference workload only" + except Exception as exc: + result["elapsed_s"] = time.perf_counter() - start + result["error"] = repr(exc) + if isinstance(exc, urllib.error.HTTPError): + result["error_body"] = exc.read().decode(errors="replace") + out.write(json.dumps(result, ensure_ascii=False) + "\n") + out.flush() + print( + json.dumps( + { + k: result.get(k) + for k in ( + "label", + "role", + "task_id", + "completion_tokens", + "first_token_s", + "elapsed_s", + "decode_tok_s", + "e2e_tok_s", + "finish_reason", + "error", + ) + }, + ensure_ascii=False, + ), + flush=True, + ) + if result.get("error"): + raise RuntimeError(result["error"]) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/expert_pool/pair.json b/benchmarks/expert_pool/pair.json new file mode 100644 index 000000000000..d99d19404c22 --- /dev/null +++ b/benchmarks/expert_pool/pair.json @@ -0,0 +1,1985 @@ +{ + "schema_version": 1, + "benchmark": "SWE-Lancer IC SWE Diamond — issue-text inference workload", + "source": { + "repository": "https://github.com/openai/frontier-evals", + "revision": "51052cede8cc608f95bb00346635e03759013e5a", + "dataset_sha256": "5c3a6d4570b49be0d9fced98f5b32487420b16f25c98d6658830e31fa03f049a", + "license": "MIT", + "license_path": "sources/LICENSE.md" + }, + "selection": { + "rule": "First two tasks, in existing sanity3 order, from the user's pre-existing local subset evaluation. Fixed before running either new backend.", + "source": "https://github.com/01554/swelancer-local-subset-evals/blob/999e25b72f16e68ca09159ec79bc7b48e0c5f141/SELECTION.md", + "limitation": "The historical sanity3 was selected from previously solved, small-input tasks. This pair is not random or representative of benchmark difficulty." + }, + "protocol": { + "temperature": 0, + "max_tokens": 2048, + "stream": true, + "order": [ + "28096_836", + "18827_741" + ], + "sequence": "Same loaded backend: first task warmup, second task measure. Restart the comparison sequence for the other backend.", + "repeat_label": "warmed_repeat", + "scope": "Single response per issue for inference timing; no repository tools, edits, execution, or official score.", + "input_transformation": "Exact CSV title and description, expanding only literal backslash-n to newlines, joined by two newlines. Other Unicode and all links retained. Same short system message added to both tasks.", + "metrics": [ + "actual prompt tokens", + "actual generated tokens", + "TTFT", + "total wall seconds", + "decode tokens per second", + "finish_reason", + "cache state" + ], + "token_limit_note": "Check exact tokenizer counts against backend context size before requesting 2048 output tokens. The limit is a cap, not a forced generation length." + }, + "tasks": [ + { + "id": "28096_836", + "role": "warmup", + "variant": "ic_swe", + "split": "diamond", + "title": "[HOLD for payment 2023-10-10] [$500] Scan - Second word of 'Choose File' button in scan request money should not be capital", + "description": "If you haven’t already, check out our [contributing guidelines](https://github.com/Expensify/ReactNativeChat/blob/main/contributingGuides/CONTRIBUTING.md) for onboarding and email contributors@expensify.com to request to join our Slack channel!\n___\n\n## Action Performed:\n1. Open the app\n2. Click on plus and click Request money\n3. Select Scan and observe that 'Choose File' button has both capital first letter for both words\n\n## Expected Result:\nApp should keep 'Choose File' button text as first word with capital and other words with first letter small as we do throughout the app like 'New workspace' and even in Spanish version of 'Choose File'\n\n## Actual Result:\nApp displays 'Choose File' button text with first letter capital for both the words in English\n\n## Workaround:\nUnknown\n\n## Platforms:\n\nWhich of our officially supported platforms is this issue occurring on?\n- [x] Android / native\n- [x] Android / Chrome\n- [x] iOS / native\n- [x] iOS / Safari\n- [x] MacOS / Chrome / Safari\n- [x] MacOS / Desktop\n\n**Version Number:** 1.3.73.0\n**Reproducible in staging?:** y\n**Reproducible in production?:** y\n**If this was caught during regression testing, add the test name, ID and link from TestRail:**\n**Email or phone of affected tester (no customers):**\n**Logs:** https://stackoverflow.com/c/expensify/questions/4856\n**Notes/Photos/Videos:** Any additional supporting documentation\n\n![Choose file capital mac chrome](https://github.com/Expensify/App/assets/93399543/28409289-c4c0-42c6-a2c0-002c7603a77b)\n![image (51)](https://github.com/Expensify/App/assets/93399543/aa9d538a-06ea-4e8c-b47b-590a07c95231)\n![image (50)](https://github.com/Expensify/App/assets/93399543/96dda57b-285c-43bb-a7c1-1638df6d3188)\n\n![Screenshot_20230927_165942_Chrome](https://github.com/Expensify/App/assets/93399543/91c7ce4e-de3e-4a3c-ad74-f282e1acfa4c)\n\n![Choose file capital mac ios (1)](https://github.com/Expensify/App/assets/93399543/59b04315-c4ad-42de-93f9-f16d07ebecd5)\n\n**Expensify/Expensify Issue URL:**\n**Issue reported by:** @ahmedGaber93\n**Slack conversation:** https://expensify.slack.com/archives/C049HHMV9SM/p1691345773866069\n\n[View all open jobs on GitHub](https://github.com/Expensify/App/issues?q=is%3Aopen+is%3Aissue+label%3A%22Help+Wanted%22)\n\n
Upwork Automation - Do Not Edit\n
    \n
  • Upwork Job URL: https://www.upwork.com/jobs/~01723d17e27c5f12de
    • \n
  • Upwork Job ID: 1706012807304642560
    • \n
  • Last Price Increase: 2023-10-01
    • \n
  • Automatic offers:
    • \n
      \n
    • ahmedGaber93 | Contributor | 26971588
      • \n
    • ahmedGaber93 | Reporter | 26971589
      • \n
    \n
\n
", + "price_usd": 500.0, + "prompt": "[HOLD for payment 2023-10-10] [$500] Scan - Second word of 'Choose File' button in scan request money should not be capital\n\nIf you haven’t already, check out our [contributing guidelines](https://github.com/Expensify/ReactNativeChat/blob/main/contributingGuides/CONTRIBUTING.md) for onboarding and email contributors@expensify.com to request to join our Slack channel!\n___\n\n## Action Performed:\n1. Open the app\n2. Click on plus and click Request money\n3. Select Scan and observe that 'Choose File' button has both capital first letter for both words\n\n## Expected Result:\nApp should keep 'Choose File' button text as first word with capital and other words with first letter small as we do throughout the app like 'New workspace' and even in Spanish version of 'Choose File'\n\n## Actual Result:\nApp displays 'Choose File' button text with first letter capital for both the words in English\n\n## Workaround:\nUnknown\n\n## Platforms:\n\nWhich of our officially supported platforms is this issue occurring on?\n- [x] Android / native\n- [x] Android / Chrome\n- [x] iOS / native\n- [x] iOS / Safari\n- [x] MacOS / Chrome / Safari\n- [x] MacOS / Desktop\n\n**Version Number:** 1.3.73.0\n**Reproducible in staging?:** y\n**Reproducible in production?:** y\n**If this was caught during regression testing, add the test name, ID and link from TestRail:**\n**Email or phone of affected tester (no customers):**\n**Logs:** https://stackoverflow.com/c/expensify/questions/4856\n**Notes/Photos/Videos:** Any additional supporting documentation\n\n![Choose file capital mac chrome](https://github.com/Expensify/App/assets/93399543/28409289-c4c0-42c6-a2c0-002c7603a77b)\n![image (51)](https://github.com/Expensify/App/assets/93399543/aa9d538a-06ea-4e8c-b47b-590a07c95231)\n![image (50)](https://github.com/Expensify/App/assets/93399543/96dda57b-285c-43bb-a7c1-1638df6d3188)\n\n![Screenshot_20230927_165942_Chrome](https://github.com/Expensify/App/assets/93399543/91c7ce4e-de3e-4a3c-ad74-f282e1acfa4c)\n\n![Choose file capital mac ios (1)](https://github.com/Expensify/App/assets/93399543/59b04315-c4ad-42de-93f9-f16d07ebecd5)\n\n**Expensify/Expensify Issue URL:**\n**Issue reported by:** @ahmedGaber93\n**Slack conversation:** https://expensify.slack.com/archives/C049HHMV9SM/p1691345773866069\n\n[View all open jobs on GitHub](https://github.com/Expensify/App/issues?q=is%3Aopen+is%3Aissue+label%3A%22Help+Wanted%22)\n\n
Upwork Automation - Do Not Edit\n
    \n
  • Upwork Job URL: https://www.upwork.com/jobs/~01723d17e27c5f12de
    • \n
  • Upwork Job ID: 1706012807304642560
    • \n
  • Last Price Increase: 2023-10-01
    • \n
  • Automatic offers:
    • \n
      \n
    • ahmedGaber93 | Contributor | 26971588
      • \n
    • ahmedGaber93 | Reporter | 26971589
      • \n
    \n
\n
", + "messages": [ + { + "role": "system", + "content": "You are a software engineer. From the supplied issue text, propose a concrete fix and a verification plan. State any assumptions about unavailable repository code." + }, + { + "role": "user", + "content": "[HOLD for payment 2023-10-10] [$500] Scan - Second word of 'Choose File' button in scan request money should not be capital\n\nIf you haven’t already, check out our [contributing guidelines](https://github.com/Expensify/ReactNativeChat/blob/main/contributingGuides/CONTRIBUTING.md) for onboarding and email contributors@expensify.com to request to join our Slack channel!\n___\n\n## Action Performed:\n1. Open the app\n2. Click on plus and click Request money\n3. Select Scan and observe that 'Choose File' button has both capital first letter for both words\n\n## Expected Result:\nApp should keep 'Choose File' button text as first word with capital and other words with first letter small as we do throughout the app like 'New workspace' and even in Spanish version of 'Choose File'\n\n## Actual Result:\nApp displays 'Choose File' button text with first letter capital for both the words in English\n\n## Workaround:\nUnknown\n\n## Platforms:\n\nWhich of our officially supported platforms is this issue occurring on?\n- [x] Android / native\n- [x] Android / Chrome\n- [x] iOS / native\n- [x] iOS / Safari\n- [x] MacOS / Chrome / Safari\n- [x] MacOS / Desktop\n\n**Version Number:** 1.3.73.0\n**Reproducible in staging?:** y\n**Reproducible in production?:** y\n**If this was caught during regression testing, add the test name, ID and link from TestRail:**\n**Email or phone of affected tester (no customers):**\n**Logs:** https://stackoverflow.com/c/expensify/questions/4856\n**Notes/Photos/Videos:** Any additional supporting documentation\n\n![Choose file capital mac chrome](https://github.com/Expensify/App/assets/93399543/28409289-c4c0-42c6-a2c0-002c7603a77b)\n![image (51)](https://github.com/Expensify/App/assets/93399543/aa9d538a-06ea-4e8c-b47b-590a07c95231)\n![image (50)](https://github.com/Expensify/App/assets/93399543/96dda57b-285c-43bb-a7c1-1638df6d3188)\n\n![Screenshot_20230927_165942_Chrome](https://github.com/Expensify/App/assets/93399543/91c7ce4e-de3e-4a3c-ad74-f282e1acfa4c)\n\n![Choose file capital mac ios (1)](https://github.com/Expensify/App/assets/93399543/59b04315-c4ad-42de-93f9-f16d07ebecd5)\n\n**Expensify/Expensify Issue URL:**\n**Issue reported by:** @ahmedGaber93\n**Slack conversation:** https://expensify.slack.com/archives/C049HHMV9SM/p1691345773866069\n\n[View all open jobs on GitHub](https://github.com/Expensify/App/issues?q=is%3Aopen+is%3Aissue+label%3A%22Help+Wanted%22)\n\n
Upwork Automation - Do Not Edit\n
    \n
  • Upwork Job URL: https://www.upwork.com/jobs/~01723d17e27c5f12de
    • \n
  • Upwork Job ID: 1706012807304642560
    • \n
  • Last Price Increase: 2023-10-01
    • \n
  • Automatic offers:
    • \n
      \n
    • ahmedGaber93 | Contributor | 26971588
      • \n
    • ahmedGaber93 | Reporter | 26971589
      • \n
    \n
\n
" + } + ], + "prompt_sha256": "8f88b14928ebe98d09c47434ace645bee0fb98bc5d3fb52192091cd06ad7f0bc", + "messages_sha256": "38c1a6e6c5f139e815d74aa5dffa8d6cb30eca9d65d26c939567600e200cdd1d", + "prompt_characters": 2843, + "prompt_utf8_bytes": 2845, + "source": { + "repository": "https://github.com/openai/frontier-evals", + "revision": "51052cede8cc608f95bb00346635e03759013e5a", + "dataset_url": "https://raw.githubusercontent.com/openai/frontier-evals/51052cede8cc608f95bb00346635e03759013e5a/project/swelancer/all_swelancer_tasks.csv", + "dataset_sha256": "5c3a6d4570b49be0d9fced98f5b32487420b16f25c98d6658830e31fa03f049a", + "row_key": { + "question_id": "28096_836" + }, + "original_row_path": "tasks/28096_836/source_row.json", + "original_row_sha256": "1b61147435c4f8cf019f5a3496f48d3696393a4fee24dceeac49c216dae7c494", + "public_issue_url": "https://github.com/Expensify/App/issues/28096", + "official_issue_assets_url": "https://github.com/openai/frontier-evals/tree/51052cede8cc608f95bb00346635e03759013e5a/project/swelancer/issues/28096_836" + }, + "repository_context": { + "required_for_official_evaluation": true, + "included_in_speed_messages": false, + "repository": "https://github.com/Expensify/App", + "cwd_in_official_container": "/app/expensify", + "base_commit": "2b791c9f3053c1682ddcb50ab036deb3e55a7542", + "setup": "Official setup checks out commit_id, then applies the task's bug_reintroduce.patch (these two revert_command files are empty). Base commit alone is not the broken task state.", + "image_tag_from_official_readme": "swelancer/swelancer_x86_28096_836:releasev1", + "image_digest_verified": false, + "media": "Original attachment URLs retained as text; images/video not downloaded or supplied. Written requirements and reproduction steps are present." + }, + "official_prompt_path": "tasks/28096_836/official_prompt.txt", + "official_messages_path": "tasks/28096_836/official_messages.json", + "rendered_prompt": "<|im_start|>system\nYou are a software engineer. From the supplied issue text, propose a concrete fix and a verification plan. State any assumptions about unavailable repository code.<|im_end|>\n<|im_start|>user\n[HOLD for payment 2023-10-10] [$500] Scan - Second word of 'Choose File' button in scan request money should not be capital\n\nIf you haven’t already, check out our [contributing guidelines](https://github.com/Expensify/ReactNativeChat/blob/main/contributingGuides/CONTRIBUTING.md) for onboarding and email contributors@expensify.com to request to join our Slack channel!\n___\n\n## Action Performed:\n1. Open the app\n2. Click on plus and click Request money\n3. Select Scan and observe that 'Choose File' button has both capital first letter for both words\n\n## Expected Result:\nApp should keep 'Choose File' button text as first word with capital and other words with first letter small as we do throughout the app like 'New workspace' and even in Spanish version of 'Choose File'\n\n## Actual Result:\nApp displays 'Choose File' button text with first letter capital for both the words in English\n\n## Workaround:\nUnknown\n\n## Platforms:\n\nWhich of our officially supported platforms is this issue occurring on?\n- [x] Android / native\n- [x] Android / Chrome\n- [x] iOS / native\n- [x] iOS / Safari\n- [x] MacOS / Chrome / Safari\n- [x] MacOS / Desktop\n\n**Version Number:** 1.3.73.0\n**Reproducible in staging?:** y\n**Reproducible in production?:** y\n**If this was caught during regression testing, add the test name, ID and link from TestRail:**\n**Email or phone of affected tester (no customers):**\n**Logs:** https://stackoverflow.com/c/expensify/questions/4856\n**Notes/Photos/Videos:** Any additional supporting documentation\n\n![Choose file capital mac chrome](https://github.com/Expensify/App/assets/93399543/28409289-c4c0-42c6-a2c0-002c7603a77b)\n![image (51)](https://github.com/Expensify/App/assets/93399543/aa9d538a-06ea-4e8c-b47b-590a07c95231)\n![image (50)](https://github.com/Expensify/App/assets/93399543/96dda57b-285c-43bb-a7c1-1638df6d3188)\n\n![Screenshot_20230927_165942_Chrome](https://github.com/Expensify/App/assets/93399543/91c7ce4e-de3e-4a3c-ad74-f282e1acfa4c)\n\n![Choose file capital mac ios (1)](https://github.com/Expensify/App/assets/93399543/59b04315-c4ad-42de-93f9-f16d07ebecd5)\n\n**Expensify/Expensify Issue URL:**\n**Issue reported by:** @ahmedGaber93\n**Slack conversation:** https://expensify.slack.com/archives/C049HHMV9SM/p1691345773866069\n\n[View all open jobs on GitHub](https://github.com/Expensify/App/issues?q=is%3Aopen+is%3Aissue+label%3A%22Help+Wanted%22)\n\n
Upwork Automation - Do Not Edit\n
    \n
  • Upwork Job URL: https://www.upwork.com/jobs/~01723d17e27c5f12de
    • \n
  • Upwork Job ID: 1706012807304642560
    • \n
  • Last Price Increase: 2023-10-01
    • \n
  • Automatic offers:
    • \n
      \n
    • ahmedGaber93 | Contributor | 26971588
      • \n
    • ahmedGaber93 | Reporter | 26971589
      • \n
    \n
\n
<|im_end|>\n<|im_start|>assistant\n\n\n\n\n", + "prompt_token_ids": [ + 248045, + 8678, + 198, + 2523, + 513, + 264, + 3061, + 22839, + 13, + 5369, + 279, + 16713, + 4125, + 1414, + 11, + 28647, + 264, + 13769, + 4884, + 321, + 264, + 22188, + 3019, + 13, + 3130, + 866, + 30800, + 883, + 33829, + 12182, + 1970, + 13, + 248046, + 198, + 248045, + 846, + 198, + 77615, + 7615, + 364, + 7903, + 220, + 17, + 15, + 17, + 18, + 12, + 16, + 15, + 12, + 16, + 15, + 60, + 23503, + 20, + 15, + 15, + 60, + 30090, + 471, + 10130, + 3299, + 314, + 359, + 23298, + 2793, + 6, + 3037, + 303, + 8307, + 1622, + 3117, + 1220, + 524, + 381, + 6511, + 271, + 2592, + 488, + 8719, + 1357, + 2582, + 11, + 1716, + 680, + 1004, + 498, + 45537, + 10292, + 16981, + 9251, + 2349, + 1074, + 5039, + 877, + 14, + 7780, + 704, + 1386, + 14, + 14370, + 20166, + 15213, + 33673, + 14905, + 30211, + 1285, + 10292, + 16196, + 287, + 14, + 5609, + 2301, + 15483, + 1658, + 20668, + 8, + 364, + 383, + 36325, + 321, + 2469, + 19726, + 31, + 4431, + 704, + 1386, + 877, + 310, + 1622, + 310, + 4973, + 1004, + 55327, + 5323, + 0, + 198, + 5785, + 271, + 550, + 5411, + 3495, + 9853, + 25, + 198, + 16, + 13, + 5097, + 279, + 878, + 198, + 17, + 13, + 8916, + 383, + 5346, + 321, + 4066, + 5952, + 3117, + 198, + 18, + 13, + 8167, + 30090, + 321, + 22270, + 421, + 359, + 23298, + 2793, + 6, + 3037, + 682, + 2107, + 6511, + 1118, + 6321, + 364, + 2107, + 4105, + 271, + 550, + 30003, + 5536, + 25, + 198, + 2095, + 1220, + 2426, + 359, + 23298, + 2793, + 6, + 3037, + 1414, + 430, + 1118, + 3299, + 440, + 6511, + 321, + 975, + 4105, + 440, + 1118, + 6321, + 2526, + 430, + 567, + 635, + 6600, + 279, + 878, + 1040, + 359, + 3446, + 26622, + 6, + 321, + 1442, + 303, + 14712, + 2243, + 314, + 359, + 23298, + 2793, + 6, + 271, + 550, + 32254, + 5536, + 25, + 198, + 2095, + 18126, + 359, + 23298, + 2793, + 6, + 3037, + 1414, + 440, + 1118, + 6321, + 6511, + 364, + 2107, + 279, + 4105, + 303, + 6163, + 271, + 550, + 5374, + 18864, + 25, + 198, + 13394, + 271, + 550, + 91826, + 25, + 198, + 6164, + 12, + 198, + 3840, + 974, + 866, + 14960, + 421, + 513, + 11164, + 539, + 411, + 4125, + 198, + 312, + 397, + 198, + 22365, + 314, + 1004, + 18003, + 7021, + 14960, + 369, + 411, + 4125, + 29853, + 383, + 30, + 198, + 12, + 498, + 87, + 60, + 8253, + 593, + 9575, + 198, + 12, + 498, + 87, + 60, + 8253, + 593, + 16717, + 198, + 12, + 498, + 87, + 60, + 15575, + 593, + 9575, + 198, + 12, + 498, + 87, + 60, + 15575, + 593, + 27819, + 198, + 12, + 498, + 87, + 60, + 86631, + 593, + 16717, + 593, + 27819, + 198, + 12, + 498, + 87, + 60, + 86631, + 593, + 34130, + 271, + 332, + 5460, + 5447, + 64700, + 220, + 16, + 13, + 18, + 13, + 22, + 18, + 13, + 15, + 198, + 332, + 674, + 751, + 75600, + 303, + 46170, + 4666, + 332, + 374, + 198, + 332, + 674, + 751, + 75600, + 303, + 5492, + 4666, + 332, + 374, + 198, + 332, + 2592, + 411, + 557, + 10255, + 2261, + 29551, + 7262, + 11, + 884, + 279, + 1228, + 803, + 11, + 2937, + 321, + 2569, + 494, + 3284, + 90805, + 64700, + 198, + 332, + 4628, + 466, + 4392, + 314, + 11164, + 35881, + 318, + 2083, + 6112, + 188912, + 198, + 332, + 49334, + 64700, + 3577, + 1074, + 40645, + 877, + 2805, + 63250, + 704, + 1386, + 41888, + 14, + 19, + 23, + 20, + 21, + 198, + 332, + 21003, + 14, + 30729, + 26355, + 5023, + 64700, + 5586, + 4945, + 12250, + 9417, + 271, + 20077, + 23298, + 999, + 6511, + 8707, + 25584, + 9251, + 2349, + 1074, + 5039, + 877, + 14, + 7780, + 704, + 1386, + 41331, + 20785, + 14, + 24, + 18, + 18, + 24, + 24, + 20, + 19, + 18, + 14, + 17, + 23, + 19, + 15, + 24, + 17, + 23, + 24, + 1723, + 19, + 66, + 15, + 12, + 19, + 17, + 66, + 21, + 7174, + 17, + 66, + 15, + 12, + 15, + 15, + 17, + 66, + 22, + 21, + 15, + 18, + 64, + 22, + 22, + 65, + 8, + 198, + 20077, + 1742, + 318, + 20, + 16, + 7025, + 7, + 2349, + 1074, + 5039, + 877, + 14, + 7780, + 704, + 1386, + 41331, + 20785, + 14, + 24, + 18, + 18, + 24, + 24, + 20, + 19, + 18, + 14, + 5137, + 24, + 67, + 20, + 18, + 23, + 64, + 12, + 15, + 21, + 12150, + 12, + 19, + 68, + 23, + 66, + 1402, + 19, + 22, + 65, + 12, + 20, + 24, + 15, + 64, + 15, + 22, + 66, + 24, + 20, + 17, + 18, + 16, + 8, + 198, + 20077, + 1742, + 318, + 20, + 15, + 7025, + 7, + 2349, + 1074, + 5039, + 877, + 14, + 7780, + 704, + 1386, + 41331, + 20785, + 14, + 24, + 18, + 18, + 24, + 24, + 20, + 19, + 18, + 14, + 24, + 21, + 69061, + 20, + 22, + 65, + 12, + 17, + 23, + 20, + 66, + 12, + 19, + 18, + 5876, + 7174, + 22, + 66, + 16, + 12, + 16, + 21, + 18, + 23, + 2846, + 21, + 67, + 18, + 16, + 23, + 23, + 8, + 271, + 20077, + 60416, + 62, + 17, + 15, + 17, + 18, + 15, + 24, + 17, + 22, + 62, + 16, + 21, + 20, + 24, + 19, + 17, + 26693, + 6439, + 9251, + 2349, + 1074, + 5039, + 877, + 14, + 7780, + 704, + 1386, + 41331, + 20785, + 14, + 24, + 18, + 18, + 24, + 24, + 20, + 19, + 18, + 14, + 24, + 16, + 66, + 22, + 341, + 19, + 68, + 6596, + 18, + 68, + 12, + 19, + 64, + 18, + 66, + 24927, + 22, + 19, + 2149, + 17, + 23, + 17, + 68, + 16, + 565, + 3510, + 19, + 66, + 8, + 271, + 20077, + 23298, + 999, + 6511, + 8707, + 26639, + 318, + 16, + 7025, + 7, + 2349, + 1074, + 5039, + 877, + 14, + 7780, + 704, + 1386, + 41331, + 20785, + 14, + 24, + 18, + 18, + 24, + 24, + 20, + 19, + 18, + 14, + 20, + 24, + 65, + 15, + 19, + 18, + 16, + 20, + 1723, + 19, + 327, + 12, + 19, + 17, + 442, + 12, + 24, + 18, + 69, + 24, + 2149, + 16, + 21, + 67, + 15, + 22, + 2967, + 35975, + 20, + 8, + 271, + 332, + 7780, + 704, + 1386, + 14, + 7780, + 704, + 1386, + 24425, + 5375, + 64700, + 198, + 332, + 40619, + 4800, + 539, + 64700, + 554, + 1413, + 1993, + 38, + 41716, + 24, + 18, + 198, + 332, + 7207, + 463, + 10125, + 64700, + 3577, + 1074, + 4431, + 704, + 1386, + 24305, + 463, + 877, + 57949, + 1821, + 10968, + 15, + 19, + 24, + 22456, + 64375, + 24, + 9221, + 4181, + 16, + 21, + 24, + 16, + 18, + 19, + 20, + 22, + 22, + 18, + 23, + 21, + 21, + 15, + 21, + 24, + 271, + 58, + 825, + 660, + 1724, + 6672, + 383, + 31038, + 9251, + 2349, + 1074, + 5039, + 877, + 14, + 7780, + 704, + 1386, + 41331, + 37463, + 42332, + 89124, + 4, + 18, + 32, + 2428, + 10, + 284, + 4, + 18, + 32, + 10835, + 10, + 1448, + 4, + 18, + 32, + 4, + 17, + 17, + 12325, + 94506, + 7327, + 4, + 17, + 17, + 8, + 271, + 27, + 14441, + 1721, + 1648, + 29, + 2248, + 1715, + 51973, + 471, + 3054, + 2717, + 8081, + 510, + 1648, + 29, + 198, + 262, + 361, + 355, + 29, + 198, + 285, + 361, + 719, + 29, + 2248, + 1715, + 11669, + 5375, + 25, + 3577, + 1074, + 2068, + 17126, + 1715, + 877, + 4291, + 5308, + 51990, + 15, + 16, + 22, + 17, + 18, + 67, + 16, + 22, + 68, + 17, + 22, + 66, + 20, + 69, + 16, + 17, + 442, + 510, + 719, + 29, + 198, + 285, + 361, + 719, + 29, + 2248, + 1715, + 11669, + 2937, + 25, + 220, + 16, + 22, + 15, + 21, + 15, + 16, + 17, + 23, + 15, + 22, + 18, + 15, + 19, + 21, + 19, + 17, + 20, + 21, + 15, + 510, + 719, + 29, + 198, + 285, + 361, + 719, + 92081, + 8223, + 38321, + 25, + 220, + 17, + 15, + 17, + 18, + 12, + 16, + 15, + 12, + 15, + 16, + 510, + 719, + 29, + 198, + 8546, + 29, + 60676, + 5891, + 25, + 672, + 719, + 29, + 198, + 19045, + 29, + 198, + 8546, + 29, + 1413, + 1993, + 38, + 41716, + 24, + 18, + 735, + 62534, + 735, + 220, + 17, + 21, + 24, + 22, + 16, + 20, + 23, + 23, + 510, + 719, + 29, + 198, + 8546, + 29, + 1413, + 1993, + 38, + 41716, + 24, + 18, + 735, + 45412, + 735, + 220, + 17, + 21, + 24, + 22, + 16, + 20, + 23, + 24, + 510, + 719, + 29, + 198, + 510, + 355, + 1419, + 719, + 29, + 198, + 262, + 672, + 355, + 29, + 198, + 510, + 14441, + 29, + 248046, + 198, + 248045, + 74455, + 198, + 248068, + 271, + 248069, + 271 + ], + "prompt_tokens": 1070 + }, + { + "id": "18827_741", + "role": "measure", + "variant": "ic_swe", + "split": "diamond", + "title": "[HOLD for payment 2023-06-21] [$1000] 'Link sent!' message is not dynamically updated when changing language preference", + "description": "If you haven’t already, check out our [contributing guidelines](https://github.com/Expensify/ReactNativeChat/blob/main/contributingGuides/CONTRIBUTING.md) for onboarding and email contributors@expensify.com to request to join our Slack channel!\n___\n\n## Action Performed:\n\n1. Add a secondary login account, and leave it unverified\n2. Open new tab\n3. Navigate to sign in page\n4. Enter the unverified email address\n5. Click on Unlink button\n6. Verify that 'Link sent!' message is displayed\n7. Scroll down, and select Spanish as a language preference\n\n## Expected Result: \n'Link sent!' message should be dynamically updated to Spanish\n## Actual Result:\n 'Link sent!' message is not dynamically updated to Spanish\n\n## Workaround:\nCan the user still use Expensify without this being fixed? Have you informed them of the workaround?\n\n## Platforms:\n\nWhich of our officially supported platforms is this issue occurring on?\n- [x] Android / native\n- [x] Android / Chrome\n- [x] iOS / native\n- [x] iOS / Safari\n- [x] MacOS / Chrome / Safari\n- [x] MacOS / Desktop\n\n**Version Number:** 1.3.13.1\n**Reproducible in staging?:** y\n**Reproducible in production?:** y\n**If this was caught during regression testing, add the test name, ID and link from TestRail:**\n**Email or phone of affected tester (no customers):**\n**Logs:** https://stackoverflow.com/c/expensify/questions/4856\n**Notes/Photos/Videos:** Any additional supporting documentation\n\nhttps://github.com/Expensify/App/assets/43996225/2f37361c-60f6-4de9-8b6c-a1e93a3b3b99\n\n\n**Expensify/Expensify Issue URL:**\n**Issue reported by:** @natnael-guchima\n**Slack conversation:** https://expensify.slack.com/archives/C049HHMV9SM/p1683649359591549\n\n[View all open jobs on GitHub](https://github.com/Expensify/App/issues?q=is%3Aopen+is%3Aissue+label%3A%22Help+Wanted%22)\n\n
Upwork Automation - Do Not Edit\n
    \n
  • Upwork Job URL: https://www.upwork.com/jobs/~0198d8720bb9bfb004
    • \n
  • Upwork Job ID: 1659222928961388544
    • \n
  • Last Price Increase: 2023-05-18
    • \n
\n
", + "price_usd": 1000.0, + "prompt": "[HOLD for payment 2023-06-21] [$1000] 'Link sent!' message is not dynamically updated when changing language preference\n\nIf you haven’t already, check out our [contributing guidelines](https://github.com/Expensify/ReactNativeChat/blob/main/contributingGuides/CONTRIBUTING.md) for onboarding and email contributors@expensify.com to request to join our Slack channel!\n___\n\n## Action Performed:\n\n1. Add a secondary login account, and leave it unverified\n2. Open new tab\n3. Navigate to sign in page\n4. Enter the unverified email address\n5. Click on Unlink button\n6. Verify that 'Link sent!' message is displayed\n7. Scroll down, and select Spanish as a language preference\n\n## Expected Result: \n'Link sent!' message should be dynamically updated to Spanish\n## Actual Result:\n 'Link sent!' message is not dynamically updated to Spanish\n\n## Workaround:\nCan the user still use Expensify without this being fixed? Have you informed them of the workaround?\n\n## Platforms:\n\nWhich of our officially supported platforms is this issue occurring on?\n- [x] Android / native\n- [x] Android / Chrome\n- [x] iOS / native\n- [x] iOS / Safari\n- [x] MacOS / Chrome / Safari\n- [x] MacOS / Desktop\n\n**Version Number:** 1.3.13.1\n**Reproducible in staging?:** y\n**Reproducible in production?:** y\n**If this was caught during regression testing, add the test name, ID and link from TestRail:**\n**Email or phone of affected tester (no customers):**\n**Logs:** https://stackoverflow.com/c/expensify/questions/4856\n**Notes/Photos/Videos:** Any additional supporting documentation\n\nhttps://github.com/Expensify/App/assets/43996225/2f37361c-60f6-4de9-8b6c-a1e93a3b3b99\n\n\n**Expensify/Expensify Issue URL:**\n**Issue reported by:** @natnael-guchima\n**Slack conversation:** https://expensify.slack.com/archives/C049HHMV9SM/p1683649359591549\n\n[View all open jobs on GitHub](https://github.com/Expensify/App/issues?q=is%3Aopen+is%3Aissue+label%3A%22Help+Wanted%22)\n\n
Upwork Automation - Do Not Edit\n
    \n
  • Upwork Job URL: https://www.upwork.com/jobs/~0198d8720bb9bfb004
    • \n
  • Upwork Job ID: 1659222928961388544
    • \n
  • Last Price Increase: 2023-05-18
    • \n
\n
", + "messages": [ + { + "role": "system", + "content": "You are a software engineer. From the supplied issue text, propose a concrete fix and a verification plan. State any assumptions about unavailable repository code." + }, + { + "role": "user", + "content": "[HOLD for payment 2023-06-21] [$1000] 'Link sent!' message is not dynamically updated when changing language preference\n\nIf you haven’t already, check out our [contributing guidelines](https://github.com/Expensify/ReactNativeChat/blob/main/contributingGuides/CONTRIBUTING.md) for onboarding and email contributors@expensify.com to request to join our Slack channel!\n___\n\n## Action Performed:\n\n1. Add a secondary login account, and leave it unverified\n2. Open new tab\n3. Navigate to sign in page\n4. Enter the unverified email address\n5. Click on Unlink button\n6. Verify that 'Link sent!' message is displayed\n7. Scroll down, and select Spanish as a language preference\n\n## Expected Result: \n'Link sent!' message should be dynamically updated to Spanish\n## Actual Result:\n 'Link sent!' message is not dynamically updated to Spanish\n\n## Workaround:\nCan the user still use Expensify without this being fixed? Have you informed them of the workaround?\n\n## Platforms:\n\nWhich of our officially supported platforms is this issue occurring on?\n- [x] Android / native\n- [x] Android / Chrome\n- [x] iOS / native\n- [x] iOS / Safari\n- [x] MacOS / Chrome / Safari\n- [x] MacOS / Desktop\n\n**Version Number:** 1.3.13.1\n**Reproducible in staging?:** y\n**Reproducible in production?:** y\n**If this was caught during regression testing, add the test name, ID and link from TestRail:**\n**Email or phone of affected tester (no customers):**\n**Logs:** https://stackoverflow.com/c/expensify/questions/4856\n**Notes/Photos/Videos:** Any additional supporting documentation\n\nhttps://github.com/Expensify/App/assets/43996225/2f37361c-60f6-4de9-8b6c-a1e93a3b3b99\n\n\n**Expensify/Expensify Issue URL:**\n**Issue reported by:** @natnael-guchima\n**Slack conversation:** https://expensify.slack.com/archives/C049HHMV9SM/p1683649359591549\n\n[View all open jobs on GitHub](https://github.com/Expensify/App/issues?q=is%3Aopen+is%3Aissue+label%3A%22Help+Wanted%22)\n\n
Upwork Automation - Do Not Edit\n
    \n
  • Upwork Job URL: https://www.upwork.com/jobs/~0198d8720bb9bfb004
    • \n
  • Upwork Job ID: 1659222928961388544
    • \n
  • Last Price Increase: 2023-05-18
    • \n
\n
" + } + ], + "prompt_sha256": "60e8e374df1d98ff7fa8bd6d6c0043c8d360d27b6c5aeb0290d58a3bae8e9a1e", + "messages_sha256": "5c2f324902890c133c6d974e52edf71b2a019714c2ed4c83e8474f2c92aacd65", + "prompt_characters": 2262, + "prompt_utf8_bytes": 2264, + "source": { + "repository": "https://github.com/openai/frontier-evals", + "revision": "51052cede8cc608f95bb00346635e03759013e5a", + "dataset_url": "https://raw.githubusercontent.com/openai/frontier-evals/51052cede8cc608f95bb00346635e03759013e5a/project/swelancer/all_swelancer_tasks.csv", + "dataset_sha256": "5c3a6d4570b49be0d9fced98f5b32487420b16f25c98d6658830e31fa03f049a", + "row_key": { + "question_id": "18827_741" + }, + "original_row_path": "tasks/18827_741/source_row.json", + "original_row_sha256": "e52dfe84235c47bd9672e128530cae61d562c80a0ef8f13a2f5d19ee0b08d76f", + "public_issue_url": "https://github.com/Expensify/App/issues/18827", + "official_issue_assets_url": "https://github.com/openai/frontier-evals/tree/51052cede8cc608f95bb00346635e03759013e5a/project/swelancer/issues/18827_741" + }, + "repository_context": { + "required_for_official_evaluation": true, + "included_in_speed_messages": false, + "repository": "https://github.com/Expensify/App", + "cwd_in_official_container": "/app/expensify", + "base_commit": "2b791c9f3053c1682ddcb50ab036deb3e55a7542", + "setup": "Official setup checks out commit_id, then applies the task's bug_reintroduce.patch (these two revert_command files are empty). Base commit alone is not the broken task state.", + "image_tag_from_official_readme": "swelancer/swelancer_x86_18827_741:releasev1", + "image_digest_verified": false, + "media": "Original attachment URLs retained as text; images/video not downloaded or supplied. Written requirements and reproduction steps are present." + }, + "official_prompt_path": "tasks/18827_741/official_prompt.txt", + "official_messages_path": "tasks/18827_741/official_messages.json", + "rendered_prompt": "<|im_start|>system\nYou are a software engineer. From the supplied issue text, propose a concrete fix and a verification plan. State any assumptions about unavailable repository code.<|im_end|>\n<|im_start|>user\n[HOLD for payment 2023-06-21] [$1000] 'Link sent!' message is not dynamically updated when changing language preference\n\nIf you haven’t already, check out our [contributing guidelines](https://github.com/Expensify/ReactNativeChat/blob/main/contributingGuides/CONTRIBUTING.md) for onboarding and email contributors@expensify.com to request to join our Slack channel!\n___\n\n## Action Performed:\n\n1. Add a secondary login account, and leave it unverified\n2. Open new tab\n3. Navigate to sign in page\n4. Enter the unverified email address\n5. Click on Unlink button\n6. Verify that 'Link sent!' message is displayed\n7. Scroll down, and select Spanish as a language preference\n\n## Expected Result: \n'Link sent!' message should be dynamically updated to Spanish\n## Actual Result:\n 'Link sent!' message is not dynamically updated to Spanish\n\n## Workaround:\nCan the user still use Expensify without this being fixed? Have you informed them of the workaround?\n\n## Platforms:\n\nWhich of our officially supported platforms is this issue occurring on?\n- [x] Android / native\n- [x] Android / Chrome\n- [x] iOS / native\n- [x] iOS / Safari\n- [x] MacOS / Chrome / Safari\n- [x] MacOS / Desktop\n\n**Version Number:** 1.3.13.1\n**Reproducible in staging?:** y\n**Reproducible in production?:** y\n**If this was caught during regression testing, add the test name, ID and link from TestRail:**\n**Email or phone of affected tester (no customers):**\n**Logs:** https://stackoverflow.com/c/expensify/questions/4856\n**Notes/Photos/Videos:** Any additional supporting documentation\n\nhttps://github.com/Expensify/App/assets/43996225/2f37361c-60f6-4de9-8b6c-a1e93a3b3b99\n\n\n**Expensify/Expensify Issue URL:**\n**Issue reported by:** @natnael-guchima\n**Slack conversation:** https://expensify.slack.com/archives/C049HHMV9SM/p1683649359591549\n\n[View all open jobs on GitHub](https://github.com/Expensify/App/issues?q=is%3Aopen+is%3Aissue+label%3A%22Help+Wanted%22)\n\n
Upwork Automation - Do Not Edit\n
    \n
  • Upwork Job URL: https://www.upwork.com/jobs/~0198d8720bb9bfb004
    • \n
  • Upwork Job ID: 1659222928961388544
    • \n
  • Last Price Increase: 2023-05-18
    • \n
\n
<|im_end|>\n<|im_start|>assistant\n\n\n\n\n", + "prompt_token_ids": [ + 248045, + 8678, + 198, + 2523, + 513, + 264, + 3061, + 22839, + 13, + 5369, + 279, + 16713, + 4125, + 1414, + 11, + 28647, + 264, + 13769, + 4884, + 321, + 264, + 22188, + 3019, + 13, + 3130, + 866, + 30800, + 883, + 33829, + 12182, + 1970, + 13, + 248046, + 198, + 248045, + 846, + 198, + 77615, + 7615, + 364, + 7903, + 220, + 17, + 15, + 17, + 18, + 12, + 15, + 21, + 12, + 17, + 16, + 60, + 23503, + 16, + 15, + 15, + 15, + 60, + 220, + 359, + 3806, + 3106, + 30351, + 1876, + 369, + 524, + 40624, + 5860, + 948, + 9722, + 3992, + 21257, + 271, + 2592, + 488, + 8719, + 1357, + 2582, + 11, + 1716, + 680, + 1004, + 498, + 45537, + 10292, + 16981, + 9251, + 2349, + 1074, + 5039, + 877, + 14, + 7780, + 704, + 1386, + 14, + 14370, + 20166, + 15213, + 33673, + 14905, + 30211, + 1285, + 10292, + 16196, + 287, + 14, + 5609, + 2301, + 15483, + 1658, + 20668, + 8, + 364, + 383, + 36325, + 321, + 2469, + 19726, + 31, + 4431, + 704, + 1386, + 877, + 310, + 1622, + 310, + 4973, + 1004, + 55327, + 5323, + 0, + 198, + 5785, + 271, + 550, + 5411, + 3495, + 9853, + 25, + 271, + 16, + 13, + 2604, + 264, + 13838, + 5677, + 2605, + 11, + 321, + 5106, + 424, + 632, + 20392, + 198, + 17, + 13, + 5097, + 491, + 5474, + 198, + 18, + 13, + 78920, + 310, + 1777, + 303, + 2081, + 198, + 19, + 13, + 10925, + 279, + 632, + 20392, + 2469, + 2534, + 198, + 20, + 13, + 8916, + 383, + 1188, + 2012, + 3037, + 198, + 21, + 13, + 24624, + 421, + 359, + 3806, + 3106, + 30351, + 1876, + 369, + 12234, + 198, + 22, + 13, + 21694, + 1441, + 11, + 321, + 3186, + 14712, + 430, + 264, + 3992, + 21257, + 271, + 550, + 30003, + 5536, + 25, + 695, + 6, + 3806, + 3106, + 30351, + 1876, + 1220, + 381, + 40624, + 5860, + 310, + 14712, + 198, + 550, + 32254, + 5536, + 25, + 198, + 359, + 3806, + 3106, + 30351, + 1876, + 369, + 524, + 40624, + 5860, + 310, + 14712, + 271, + 550, + 5374, + 18864, + 25, + 198, + 6503, + 279, + 1156, + 1990, + 958, + 7541, + 704, + 1386, + 1973, + 411, + 1602, + 8097, + 30, + 11893, + 488, + 15515, + 1070, + 314, + 279, + 57035, + 30, + 271, + 550, + 91826, + 25, + 198, + 6164, + 12, + 198, + 3840, + 974, + 866, + 14960, + 421, + 513, + 11164, + 539, + 411, + 4125, + 198, + 312, + 397, + 198, + 22365, + 314, + 1004, + 18003, + 7021, + 14960, + 369, + 411, + 4125, + 29853, + 383, + 30, + 198, + 12, + 498, + 87, + 60, + 8253, + 593, + 9575, + 198, + 12, + 498, + 87, + 60, + 8253, + 593, + 16717, + 198, + 12, + 498, + 87, + 60, + 15575, + 593, + 9575, + 198, + 12, + 498, + 87, + 60, + 15575, + 593, + 27819, + 198, + 12, + 498, + 87, + 60, + 86631, + 593, + 16717, + 593, + 27819, + 198, + 12, + 498, + 87, + 60, + 86631, + 593, + 34130, + 271, + 332, + 5460, + 5447, + 64700, + 220, + 16, + 13, + 18, + 13, + 16, + 18, + 13, + 16, + 198, + 332, + 674, + 751, + 75600, + 303, + 46170, + 4666, + 332, + 374, + 198, + 332, + 674, + 751, + 75600, + 303, + 5492, + 4666, + 332, + 374, + 198, + 332, + 2592, + 411, + 557, + 10255, + 2261, + 29551, + 7262, + 11, + 884, + 279, + 1228, + 803, + 11, + 2937, + 321, + 2569, + 494, + 3284, + 90805, + 64700, + 198, + 332, + 4628, + 466, + 4392, + 314, + 11164, + 35881, + 318, + 2083, + 6112, + 188912, + 198, + 332, + 49334, + 64700, + 3577, + 1074, + 40645, + 877, + 2805, + 63250, + 704, + 1386, + 41888, + 14, + 19, + 23, + 20, + 21, + 198, + 332, + 21003, + 14, + 30729, + 26355, + 5023, + 64700, + 5586, + 4945, + 12250, + 9417, + 271, + 2349, + 1074, + 5039, + 877, + 14, + 7780, + 704, + 1386, + 41331, + 20785, + 14, + 19, + 18, + 24, + 24, + 21, + 17, + 17, + 20, + 14, + 17, + 69, + 18, + 22, + 18, + 21, + 16, + 66, + 12, + 21, + 15, + 69, + 21, + 12, + 19, + 442, + 24, + 12, + 23, + 65, + 21, + 66, + 7174, + 16, + 68, + 24, + 18, + 64, + 18, + 65, + 18, + 65, + 24, + 24, + 1358, + 332, + 7780, + 704, + 1386, + 14, + 7780, + 704, + 1386, + 24425, + 5375, + 64700, + 198, + 332, + 40619, + 4800, + 539, + 64700, + 554, + 32194, + 3267, + 300, + 2294, + 1339, + 7287, + 198, + 332, + 7207, + 463, + 10125, + 64700, + 3577, + 1074, + 4431, + 704, + 1386, + 24305, + 463, + 877, + 57949, + 1821, + 10968, + 15, + 19, + 24, + 22456, + 64375, + 24, + 9221, + 4181, + 16, + 21, + 23, + 18, + 21, + 19, + 24, + 18, + 20, + 24, + 20, + 24, + 16, + 20, + 19, + 24, + 271, + 58, + 825, + 660, + 1724, + 6672, + 383, + 31038, + 9251, + 2349, + 1074, + 5039, + 877, + 14, + 7780, + 704, + 1386, + 41331, + 37463, + 42332, + 89124, + 4, + 18, + 32, + 2428, + 10, + 284, + 4, + 18, + 32, + 10835, + 10, + 1448, + 4, + 18, + 32, + 4, + 17, + 17, + 12325, + 94506, + 7327, + 4, + 17, + 17, + 8, + 271, + 27, + 14441, + 1721, + 1648, + 29, + 2248, + 1715, + 51973, + 471, + 3054, + 2717, + 8081, + 510, + 1648, + 29, + 198, + 262, + 361, + 355, + 29, + 198, + 285, + 361, + 719, + 29, + 2248, + 1715, + 11669, + 5375, + 25, + 3577, + 1074, + 2068, + 17126, + 1715, + 877, + 4291, + 5308, + 51990, + 15, + 16, + 24, + 23, + 67, + 23, + 22, + 17, + 15, + 5876, + 24, + 65, + 10481, + 15, + 15, + 19, + 510, + 719, + 29, + 198, + 285, + 361, + 719, + 29, + 2248, + 1715, + 11669, + 2937, + 25, + 220, + 16, + 21, + 20, + 24, + 17, + 17, + 17, + 24, + 17, + 23, + 24, + 21, + 16, + 18, + 23, + 23, + 20, + 19, + 19, + 510, + 719, + 29, + 198, + 285, + 361, + 719, + 92081, + 8223, + 38321, + 25, + 220, + 17, + 15, + 17, + 18, + 12, + 15, + 20, + 12, + 16, + 23, + 510, + 719, + 29, + 198, + 262, + 672, + 355, + 29, + 198, + 510, + 14441, + 29, + 248046, + 198, + 248045, + 74455, + 198, + 248068, + 271, + 248069, + 271 + ], + "prompt_tokens": 753 + } + ], + "tokenization": { + "model": "/model", + "pair_sha256": "95b45ff592a7055975d3cb35fdffb6934e4271c7bf03f51385956c259485b2ca", + "enable_thinking": false, + "add_generation_prompt": true, + "max_model_len": 4096, + "max_tokens": 2048, + "tokenizer_files_sha256": { + "tokenizer.json": "0997f410c57a1f4e53b09e4be8f4a172d90edd9564368fb0847030937229b9f3", + "tokenizer_config.json": "b11349aafa7cdc6a320767cf7ceb29ed82f7eda5d65e8e0819e76f0ce947bf27", + "chat_template.jinja": "c3cf9e34abf4f9e36c2d72165aa9c132d3e2a725b6c2586aaa3a8af9d7a81041" + } + } +} diff --git a/benchmarks/expert_pool/prompts.md b/benchmarks/expert_pool/prompts.md new file mode 100644 index 000000000000..c99f5cde5e2e --- /dev/null +++ b/benchmarks/expert_pool/prompts.md @@ -0,0 +1,155 @@ +# 速度測定に使用したプロンプト全文 + +`pair.json`の`rendered_prompt`をそのまま表示しています。実際の送信は同ファイルの固定トークン列を使います。 + +## 1回目：28096_836 + +```text +<|im_start|>system +You are a software engineer. From the supplied issue text, propose a concrete fix and a verification plan. State any assumptions about unavailable repository code.<|im_end|> +<|im_start|>user +[HOLD for payment 2023-10-10] [$500] Scan - Second word of 'Choose File' button in scan request money should not be capital + +If you haven’t already, check out our [contributing guidelines](https://github.com/Expensify/ReactNativeChat/blob/main/contributingGuides/CONTRIBUTING.md) for onboarding and email contributors@expensify.com to request to join our Slack channel! +___ + +## Action Performed: +1. Open the app +2. Click on plus and click Request money +3. Select Scan and observe that 'Choose File' button has both capital first letter for both words + +## Expected Result: +App should keep 'Choose File' button text as first word with capital and other words with first letter small as we do throughout the app like 'New workspace' and even in Spanish version of 'Choose File' + +## Actual Result: +App displays 'Choose File' button text with first letter capital for both the words in English + +## Workaround: +Unknown + +## Platforms: + +Which of our officially supported platforms is this issue occurring on? +- [x] Android / native +- [x] Android / Chrome +- [x] iOS / native +- [x] iOS / Safari +- [x] MacOS / Chrome / Safari +- [x] MacOS / Desktop + +**Version Number:** 1.3.73.0 +**Reproducible in staging?:** y +**Reproducible in production?:** y +**If this was caught during regression testing, add the test name, ID and link from TestRail:** +**Email or phone of affected tester (no customers):** +**Logs:** https://stackoverflow.com/c/expensify/questions/4856 +**Notes/Photos/Videos:** Any additional supporting documentation + +![Choose file capital mac chrome](https://github.com/Expensify/App/assets/93399543/28409289-c4c0-42c6-a2c0-002c7603a77b) +![image (51)](https://github.com/Expensify/App/assets/93399543/aa9d538a-06ea-4e8c-b47b-590a07c95231) +![image (50)](https://github.com/Expensify/App/assets/93399543/96dda57b-285c-43bb-a7c1-1638df6d3188) + +![Screenshot_20230927_165942_Chrome](https://github.com/Expensify/App/assets/93399543/91c7ce4e-de3e-4a3c-ad74-f282e1acfa4c) + +![Choose file capital mac ios (1)](https://github.com/Expensify/App/assets/93399543/59b04315-c4ad-42de-93f9-f16d07ebecd5) + +**Expensify/Expensify Issue URL:** +**Issue reported by:** @ahmedGaber93 +**Slack conversation:** https://expensify.slack.com/archives/C049HHMV9SM/p1691345773866069 + +[View all open jobs on GitHub](https://github.com/Expensify/App/issues?q=is%3Aopen+is%3Aissue+label%3A%22Help+Wanted%22) + +
Upwork Automation - Do Not Edit +
    +
  • Upwork Job URL: https://www.upwork.com/jobs/~01723d17e27c5f12de
    • +
  • Upwork Job ID: 1706012807304642560
    • +
  • Last Price Increase: 2023-10-01
    • +
  • Automatic offers:
    • +
      +
    • ahmedGaber93 | Contributor | 26971588
      • +
    • ahmedGaber93 | Reporter | 26971589
      • +
    +
+
<|im_end|> +<|im_start|>assistant + + + + + +``` + +## 2回目：18827_741 + +```text +<|im_start|>system +You are a software engineer. From the supplied issue text, propose a concrete fix and a verification plan. State any assumptions about unavailable repository code.<|im_end|> +<|im_start|>user +[HOLD for payment 2023-06-21] [$1000] 'Link sent!' message is not dynamically updated when changing language preference + +If you haven’t already, check out our [contributing guidelines](https://github.com/Expensify/ReactNativeChat/blob/main/contributingGuides/CONTRIBUTING.md) for onboarding and email contributors@expensify.com to request to join our Slack channel! +___ + +## Action Performed: + +1. Add a secondary login account, and leave it unverified +2. Open new tab +3. Navigate to sign in page +4. Enter the unverified email address +5. Click on Unlink button +6. Verify that 'Link sent!' message is displayed +7. Scroll down, and select Spanish as a language preference + +## Expected Result: +'Link sent!' message should be dynamically updated to Spanish +## Actual Result: + 'Link sent!' message is not dynamically updated to Spanish + +## Workaround: +Can the user still use Expensify without this being fixed? Have you informed them of the workaround? + +## Platforms: + +Which of our officially supported platforms is this issue occurring on? +- [x] Android / native +- [x] Android / Chrome +- [x] iOS / native +- [x] iOS / Safari +- [x] MacOS / Chrome / Safari +- [x] MacOS / Desktop + +**Version Number:** 1.3.13.1 +**Reproducible in staging?:** y +**Reproducible in production?:** y +**If this was caught during regression testing, add the test name, ID and link from TestRail:** +**Email or phone of affected tester (no customers):** +**Logs:** https://stackoverflow.com/c/expensify/questions/4856 +**Notes/Photos/Videos:** Any additional supporting documentation + +https://github.com/Expensify/App/assets/43996225/2f37361c-60f6-4de9-8b6c-a1e93a3b3b99 + + +**Expensify/Expensify Issue URL:** +**Issue reported by:** @natnael-guchima +**Slack conversation:** https://expensify.slack.com/archives/C049HHMV9SM/p1683649359591549 + +[View all open jobs on GitHub](https://github.com/Expensify/App/issues?q=is%3Aopen+is%3Aissue+label%3A%22Help+Wanted%22) + +
Upwork Automation - Do Not Edit +
    +
  • Upwork Job URL: https://www.upwork.com/jobs/~0198d8720bb9bfb004
    • +
  • Upwork Job ID: 1659222928961388544
    • +
  • Last Price Increase: 2023-05-18
    • +
+
<|im_end|> +<|im_start|>assistant + + + + + +``` diff --git a/benchmarks/expert_pool/provenance.json b/benchmarks/expert_pool/provenance.json new file mode 100644 index 000000000000..392e9f0bb972 --- /dev/null +++ b/benchmarks/expert_pool/provenance.json @@ -0,0 +1,18 @@ +{ + "measured_integration_commit": "7dedc6d8d9b178b60f6a5b32f03d677145982441", + "pool_pr": 48, + "pool_runtime_commit": "5fbc240ba5ddec82a10362340ac77339a1c24017", + "historical_files_sha256": { + "benchmark.py": "2a66a48731eb8a4c6930ed1f6c3becd7b2397a6585b6d6e4fffd196dc3a8a1e8", + "pair.json": "3d4e64ecde61dd943d431a90b6b2bf6c32e89e7f1a3901cc7fb9b41a0856926a", + "run.sh": "5cf9c93ab6c1ecaec8cbad8d7f30f7ec23ebefe4662cb6cf894d165f7803fb39" + }, + "changes_from_historical_runner": [ + "Removed unused FreeToken text-prompt option and cache-verification hooks.", + "Removed optional same-server measured repeats; always sends warmup then measure once.", + "Removed client-host nvidia-smi and proc snapshots; capacity monitoring remains external.", + "Defaults to the adjacent unchanged pair.json and removes the internal engine label.", + "Preserves vLLM request fields, SSE parsing, decode formula, error persistence, and exclusive output creation.", + "Rejects malformed completion counts and non-finite computed rates; valid-response arithmetic is unchanged." + ] +} diff --git a/benchmarks/expert_pool/test_benchmark.py b/benchmarks/expert_pool/test_benchmark.py new file mode 100644 index 000000000000..dfe6cfbc2548 --- /dev/null +++ b/benchmarks/expert_pool/test_benchmark.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Check the client wire contract and failure evidence without a GPU.""" + +import contextlib +import io +import json +import tempfile +import unittest +import urllib.error +from pathlib import Path +from unittest.mock import patch + +import benchmark + + +class BenchmarkTest(unittest.TestCase): + def run_client(self, mode="ok"): + pair = json.loads(Path(benchmark.__file__).with_name("pair.json").read_text()) + captured = [] + + def respond(request, **kwargs): + body = json.loads(request.data) + captured.append(body) + if mode == "http": + raise urllib.error.HTTPError( + request.full_url, 400, "Bad Request", {}, io.BytesIO(b"bad body") + ) + events = [ + {"choices": [{"text": "hello", "finish_reason": None}]}, + {"choices": [{"text": " world", "finish_reason": "stop"}]}, + { + "choices": [], + "usage": { + "prompt_tokens": len(body["prompt"]), + "completion_tokens": 3, + }, + }, + ] + if mode == "error": + events = [{"error": {"message": "stream error"}}] + wire = "".join("data: " + json.dumps(e) + "\n\n" for e in events) + if mode != "cut": + wire += "data: [DONE]\n\n" + return io.BytesIO(wire.encode()) + + with tempfile.TemporaryDirectory() as tmp: + output = Path(tmp) / "result.jsonl" + argv = [ + "benchmark.py", + "--base-url", + "http://localhost:8000", + "--label", + "test", + "--output", + str(output), + ] + with ( + patch("sys.argv", argv), + patch("urllib.request.urlopen", side_effect=respond), + contextlib.redirect_stdout(io.StringIO()), + ): + if mode == "ok": + benchmark.main() + with self.assertRaises(FileExistsError): + benchmark.main() + else: + with self.assertRaises(RuntimeError): + benchmark.main() + rows = [json.loads(line) for line in output.read_text().splitlines()] + return pair, captured, rows + + def test_frozen_requests_usage_only_event_and_measurement(self): + pair, bodies, rows = self.run_client() + self.assertEqual(len(bodies), 2) + for task, body, row in zip(pair["tasks"], bodies, rows): + self.assertEqual( + body, + { + "model": "flashnext", + "prompt": task["prompt_token_ids"], + "max_tokens": 2048, + "temperature": 0, + "top_p": 1, + "stream": True, + "stream_options": {"include_usage": True}, + "seed": 0, + }, + ) + self.assertEqual(row["content"], "hello world") + self.assertTrue(row["sse_done"]) + self.assertEqual(row["role"], task["role"]) + self.assertEqual( + row["decode_tok_s"], 2 / (row["last_token_s"] - row["first_token_s"]) + ) + + def test_failure_is_saved_and_second_request_is_not_sent(self): + for mode in ("http", "error", "cut"): + with self.subTest(mode=mode): + _, bodies, rows = self.run_client(mode) + self.assertEqual(len(bodies), 1) + self.assertEqual(len(rows), 1) + self.assertIn("error", rows[0]) + if mode == "http": + self.assertEqual(rows[0]["error_body"], "bad body") + + +if __name__ == "__main__": + unittest.main() From e00923f74fd6b894cd556a3853a73b5d32becedf Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Thu, 10 Sep 2026 10:23:57 +0900 Subject: [PATCH 06/13] [MoE] Expert pool: docs page, CI registration, English benchmark README Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- .buildkite/test_areas/kernels.yaml | 2 + benchmarks/expert_pool/README.md | 153 ++++++++++++++++------------- benchmarks/expert_pool/prompts.md | 8 +- docs/features/moe_expert_pool.md | 54 ++++++++++ 4 files changed, 144 insertions(+), 73 deletions(-) create mode 100644 docs/features/moe_expert_pool.md diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 2b43938060e9..b539c0338937 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -231,6 +231,7 @@ steps: - csrc/quantization/cutlass_w8a8/moe/ - csrc/moe/ - tests/kernels/moe + - tests/kernels/expert_pool - vllm/model_executor/layers/fused_moe/ - vllm/distributed/device_communicators/ - vllm/envs.py @@ -238,6 +239,7 @@ steps: commands: - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s kernels/expert_pool --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT parallelism: 5 mirror: amd: diff --git a/benchmarks/expert_pool/README.md b/benchmarks/expert_pool/README.md index a31c361d8711..678de38b5be6 100644 --- a/benchmarks/expert_pool/README.md +++ b/benchmarks/expert_pool/README.md @@ -1,53 +1,61 @@ -# Expert重みのGPUキャッシュ：速度測定の再現手順 +# Expert pool: reproducing the generation-speed measurement -Qwen3.8 FlashNext NVFP4に、準備用の要求と速度測定用の要求を順番に送ります。 -各要求ではソフトウェアの不具合報告を読み、修正案と検証計画を生成します。 +Sends a warmup request and then a measured request to Qwen3.8-Flash-Next +NVFP4. Each request is a software bug report; the model writes a fix +proposal and a verification plan. -## ファイルと要求の内容 +## Files and requests -| ファイル | 内容 | +| file | content | | --- | --- | -| [benchmark.py](benchmark.py) | HTTP要求、ストリーム保存、生成速度の集計 | -| [pair.json](pair.json) | 実測に使用したプロンプト全文、固定トークン列、出典 | -| [prompts.md](prompts.md) | プロンプト全文の読みやすい表示 | -| [provenance.json](provenance.json) | 実測コードの版、元ファイルのSHA256、移植差分 | -| [LICENSE.prompts](LICENSE.prompts) | プロンプトの出典に付属するMITライセンス | +| [benchmark.py](benchmark.py) | HTTP requests, stream capture, decode-speed computation | +| [pair.json](pair.json) | the two prompts as measured: full text, frozen token ids, provenance | +| [prompts.md](prompts.md) | the prompt texts, readable | +| [provenance.json](provenance.json) | code versions measured, SHA256 of the original files, changes from the historical runner | +| [LICENSE.prompts](LICENSE.prompts) | MIT license of the prompt source | -| 送信順 | 用途 | 不具合報告 | 入力トークン数 | +| order | role | bug report | input tokens | | --- | --- | --- | --- | -| 1回目 | 準備用（結果のroleはwarmup） | ファイル選択ボタンの「Choose File」を「Choose file」に修正する課題（28096_836） | 1070 | -| 2回目 | 速度測定用（roleはmeasure） | 言語設定を変更したとき「Link sent!」表示も更新する課題（18827_741） | 753 | +| 1st | warmup (`role` = `warmup`) | rename the "Choose File" button to "Choose file" (task 28096_836) | 1070 | +| 2nd | measured (`role` = `measure`) | update the "Link sent!" message when the language setting changes (task 18827_741) | 753 | -出典は[OpenAI frontier-evals](https://github.com/openai/frontier-evals/tree/51052cede8cc608f95bb00346635e03759013e5a)のSWE-Lancerです。 -既存の動作確認用課題から選んだ2件で、測定対象は修正案の文章生成です。 -実際のコード編集・公式採点を行う品質試験は別の手順です。 +The prompts come from SWE-Lancer in +[OpenAI frontier-evals](https://github.com/openai/frontier-evals/tree/51052cede8cc608f95bb00346635e03759013e5a), +two of its existing sanity tasks. This measures fix-proposal text +generation only; the code-editing and official-grading quality checks are a +separate procedure. -## 実測に使用したコードと環境 +## Code and environment measured -| 用途 | コードの版 | +| purpose | version | | --- | --- | -| 比較の土台となるvLLM main | `a97dacb7106ee49f39f3d1fc6ae1800ff724e01d` | -| Expert重みのGPUキャッシュの実装（[fork PR #48](https://github.com/01554/vllm/pull/48)） | `5fbc240ba5ddec82a10362340ac77339a1c24017` | -| PLEの読み出しと生成計算を重ねる実装（[fork PR #46](https://github.com/01554/vllm/pull/46)、[上流PR #54129](https://github.com/vllm-project/vllm/pull/54129)を前提とする差分） | `4f859de9d0f55760b50358aee4834e6966e13bc8` | -| 上記機能を組み合わせ、以下の速度を測定した版 | `7dedc6d8d9b178b60f6a5b32f03d677145982441` | - -サーバーは実測版のPythonソース、上記mainからビルドしたwheel、別途ビルドした -`_ple_memops`拡張を組み合わせて動かしました。このディレクトリは測定後に追加した -クライアント用ファイルです。`pair.json`は実測ファイルのbyteコピーです。 - -RTX 6000 Adaの48GBに収める構成の検証を目的として、手元の -RTX PRO 6000 Blackwell Max-Q（96GiB）上で別プロセスにGPUメモリを確保させ、 -サーバーに利用可能な容量を48GiBにして測定しました。以下はこのGPU上の実測値です。 -ホスト側のコンテナのメモリ上限は100GiBでした。 - -GPUキャッシュは48層それぞれ258 expert行、約32GiBです。 -容量制限は外部プロセスで設定します。下記の`gpu-memory-utilization`は物理GPU容量に -対するvLLMの予算比率です。クライアント実行前に容量とサーバー起動状態を確認してください。 - -## サーバーの起動設定 - -上の実測版と同じ機能をビルドした環境で、チェックポイントのパスを指定します。 -`VLLM_USE_BREAKABLE_CUDAGRAPH`は未設定（自動選択）で測定しました。 +| vLLM main used as the base | `a97dacb7106ee49f39f3d1fc6ae1800ff724e01d` | +| expert pool implementation (this PR, at the time of measurement) | `5fbc240ba5ddec82a10362340ac77339a1c24017` | +| deferred PLE rows ([01554/vllm#46](https://github.com/01554/vllm/pull/46), a diff on top of upstream PR [#54129](https://github.com/vllm-project/vllm/pull/54129)) | `4f859de9d0f55760b50358aee4834e6966e13bc8` | +| the combination of the above that produced the numbers below | `7dedc6d8d9b178b60f6a5b32f03d677145982441` | + +The server ran the Python sources of the measured combination over a wheel +built from the base commit above, plus a separately built `_ple_memops` +extension. This directory holds client-side files added after the +measurement; `pair.json` is a byte copy of the file used. + +The goal was a configuration that fits an RTX 6000 Ada (48 GB). Measurements +were taken on an RTX PRO 6000 Blackwell Max-Q (96 GiB) with a separate +process holding GPU memory so that the server had 48 GiB available; the +container's host memory limit was 100 GiB. The numbers below are from that GPU. + +The expert pool holds 258 expert rows per layer for 48 layers, about +32 GiB. The memory limit is set by the external process; the +`gpu-memory-utilization` below is vLLM's budget as a fraction of the +physical GPU. Check the available memory and that the server is up before +running the client. + +## Server launch + +Same features as the measured combination, checkpoint path as needed. +`VLLM_USE_BREAKABLE_CUDAGRAPH` was unset (automatic selection). +`VLLM_DEBUG_WORKSPACE`, `VLLM_LOGGING_LEVEL`, `PYTORCH_ALLOC_CONF` and the +thread counts are the values used during measurement, not requirements. ```bash export CHECKPOINT=/data/models/Qwen3.8-Flash-Next-NVFP4-nvidia @@ -72,11 +80,13 @@ export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 --reasoning-parser qwen3 --generation-config vllm ``` -速度測定のコンテキスト上限は4096です。別途実施した品質試験では32768を使いました。 +The speed measurement uses a context limit of 4096. The separate quality +checks used 32768. -## クライアントの実行 +## Client -サーバーの準備完了後、次を1回実行します。クライアントはPython標準ライブラリで動きます。 +Run once after the server is ready. The client needs only the Python +standard library. ```bash .venv/bin/python benchmarks/expert_pool/benchmark.py \ @@ -84,43 +94,48 @@ export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 --label fresh-0 --output results/fresh-0.jsonl ``` -送信先は`/v1/completions`です。`pair.json`の固定トークン列を送信するので、 -チェックポイントのtokenizerが`pair.json`の`tokenization`に記録したSHA256と -一致することを確認してください。要求には`temperature=0`、`top_p=1`、`seed=0`、 -`max_tokens=2048`を指定し、thinkingを無効にしたチャットテンプレートのトークン列を使います。 -`tokenization.pair_sha256`はトークン列追加前の資料のハッシュで、ファイル全体のハッシュは -`provenance.json`にあります。 +Requests go to `/v1/completions` with the frozen token ids from +`pair.json`, so confirm that the checkpoint's tokenizer matches the SHA256 +recorded under `tokenization` in `pair.json`. Each request uses +`temperature=0`, `top_p=1`, `seed=0`, `max_tokens=2048` and the +chat-template token sequence with thinking disabled. +`tokenization.pair_sha256` is the hash of the material before the token ids +were added; the hash of the whole file is in `provenance.json`. -3回の測定では、**毎回サーバーを終了して新しいプロセスで起動し、準備用→測定用を1組送信**します。 -出力名を`fresh-0.jsonl`、`fresh-1.jsonl`、`fresh-2.jsonl`と変えます。 -実測はこの順で3組を実行し、2回目の要求の速度3値から中央値を求めました。 -クライアントは既存の出力ファイルを保護し、HTTPエラーや不完全なストリームを保存して停止します。 -失敗した測定も結果として保持してください。 +For the three-run measurement, **stop the server and start a new process +before each run, then send warmup -> measure once**, writing to +`fresh-0.jsonl`, `fresh-1.jsonl`, `fresh-2.jsonl`. The reported value is the +median of the three measured-request speeds. The client refuses to +overwrite an existing output file, and saves HTTP errors and incomplete +streams before stopping; keep failed runs as results too. -## 指標と実測値 +## Metric and measured values -生成速度は、usageの生成トークン数とクライアント側の受信時刻から計算します。 +Decode speed is computed from the usage completion-token count and the +client-side receive times: ```text -decode_tok_s = (completion_tokens - 1) / (最後の本文受信時刻 - 最初の本文受信時刻) +decode_tok_s = (completion_tokens - 1) / (last text event time - first text event time) ``` -`first_token_s`は要求開始から最初の本文受信までの時間です。 -`e2e_tok_s`は要求全体の所要時間あたりの生成トークン数です。 -ストリームの1イベントに複数トークンが入ることがあるため、いずれもクライアント観測の値です。 -生のSSE、送信body、usage、finish reason、本文とSHA256も同じJSONLに保存します。 +`first_token_s` is the time from request start to the first text event; +`e2e_tok_s` is completion tokens over the whole request time. A stream +event may carry more than one token, so both are client-observed values. +The raw SSE events, request body, usage, finish reason, text and its SHA256 +are stored in the same JSONL. -| 新しいサーバーでの実行 | 2回目の要求の生成速度 | +| fresh server run | measured-request decode speed | | --- | --- | -| 1回目 | 63.1882 tok/s | -| 2回目 | 62.7087 tok/s | -| 3回目 | 63.6519 tok/s | -| 中央値 | **63.1882 tok/s** | +| 1 | 63.1882 tok/s | +| 2 | 62.7087 tok/s | +| 3 | 63.6519 tok/s | +| median | **63.1882 tok/s** | -これは表の実測版で各機能を組み合わせた値です。各要求の終了理由は`stop`でした。 -容量は起動中・生成中の標本で記録し、プロセス終了コード0とOOMKilled=falseを確認しました。 +Measured on the combination listed above. Every request finished with +`stop`. Memory was sampled during startup and generation; the process +exited 0 with OOMKilled=false. -## クライアントの動作確認 +## Client self-check ```bash .venv/bin/python -m unittest discover -s benchmarks/expert_pool -p 'test_*.py' diff --git a/benchmarks/expert_pool/prompts.md b/benchmarks/expert_pool/prompts.md index c99f5cde5e2e..d976ef310dfe 100644 --- a/benchmarks/expert_pool/prompts.md +++ b/benchmarks/expert_pool/prompts.md @@ -1,8 +1,8 @@ -# 速度測定に使用したプロンプト全文 +# Prompts used for the speed measurement -`pair.json`の`rendered_prompt`をそのまま表示しています。実際の送信は同ファイルの固定トークン列を使います。 +The `rendered_prompt` field of `pair.json`, shown as is. The client sends the frozen token ids from the same file. -## 1回目：28096_836 +## 1st (warmup): 28096_836 ```text <|im_start|>system @@ -81,7 +81,7 @@ Which of our officially supported platforms is this issue occurring on? ``` -## 2回目：18827_741 +## 2nd (measured): 18827_741 ```text <|im_start|>system diff --git a/docs/features/moe_expert_pool.md b/docs/features/moe_expert_pool.md new file mode 100644 index 000000000000..2a411d6eadba --- /dev/null +++ b/docs/features/moe_expert_pool.md @@ -0,0 +1,54 @@ +# MoE Expert Pool + +An opt-in way to serve a MoE model whose expert weights do not fit in GPU +memory: the expert weights stay in pinned host memory, and one GPU bank +shared by all MoE layers holds a subset of expert rows that is managed on +the device. + +```bash +vllm serve --moe-expert-pool-rows 258 +``` + +`--moe-expert-pool-rows N` (config field `OffloadConfig.moe_expert_pool_rows`, +default `0` = off) sets how many expert rows per layer are resident at +startup. The bank size is `N x number of MoE layers x row bytes`; rows can +move between layers at run time. + +## How it works + +- Loading: the NVFP4 expert tensors are allocated in pinned host memory + instead of the GPU. The loader still moves one layer at a time to the GPU + for the Marlin conversion and restores the converted tensors to pinned + memory, which becomes the pool's source. +- Installation: after every layer has been processed, the pool allocates the + bank, fills each layer's initial rows and binds a Marlin consumer to the + bank. The placement is frozen (gate closed) during profiling and CUDA + graph capture and opened at the end of warm-up. +- Decode: a device-side step program looks up each routed expert in the + bank, promotes misses (device LRU over all layers), copies the needed rows + from the host source with fixed-grid kernels, and runs Marlin on the bank + with logical expert alignment and a physical-row remap. No host-side + routing readback or cache planning sits between routing and the GEMM, so + the MoE layers stay inside the captured CUDA graph (`FULL_DECODE_ONLY`). +- Wider batches (prefill): resident experts are read from the bank and the + rest directly from the pinned host source through its accelerator view. + +## Requirements and limits + +- ModelOpt NVFP4 checkpoints with the Marlin MoE backend + (`--moe-backend marlin`). Other quantization methods or NVFP4 backends + are rejected when the layer is built. +- No expert, data or sequence parallelism. +- `N` must be at least `top_k` and every MoE layer must have the same expert + count and top-k. +- Prefill of long inputs reads the non-resident experts from host memory on + every chunk; expect prefill to be bounded by host-to-device bandwidth. + +## Example + +Qwen3.8-Flash-Next NVFP4 on a 48 GiB GPU budget with 258 rows per layer +(about a 32 GiB bank): decode about 63 tok/s at 4096 context with +`--compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}'`, versus +about 7 tok/s with `--offload-backend uva --cpu-offload-gb 40`. A +reproducible client and the exact launch flags are in +`benchmarks/expert_pool/`. From feba2ca997d89d20d023ec5744dd9bccb2345a61 Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Thu, 10 Sep 2026 10:27:12 +0900 Subject: [PATCH 07/13] [MoE] Expert pool docs: bank sizing with clamp and staging, loading detail, measurement provenance, prefill wording; AMD mirror test dependency Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- .buildkite/test_areas/kernels.yaml | 1 + docs/features/moe_expert_pool.md | 39 ++++++++++++++++++++---------- 2 files changed, 27 insertions(+), 13 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index b539c0338937..6a376088815f 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -251,6 +251,7 @@ steps: - csrc/quantization/cutlass_w8a8/moe/ - csrc/moe/ - tests/kernels/moe + - tests/kernels/expert_pool - vllm/model_executor/layers/fused_moe/ - vllm/distributed/device_communicators/ - vllm/envs.py diff --git a/docs/features/moe_expert_pool.md b/docs/features/moe_expert_pool.md index 2a411d6eadba..b8bbf9f003ef 100644 --- a/docs/features/moe_expert_pool.md +++ b/docs/features/moe_expert_pool.md @@ -11,15 +11,21 @@ vllm serve --moe-expert-pool-rows 258 `--moe-expert-pool-rows N` (config field `OffloadConfig.moe_expert_pool_rows`, default `0` = off) sets how many expert rows per layer are resident at -startup. The bank size is `N x number of MoE layers x row bytes`; rows can -move between layers at run time. +startup; the effective count is `min(N, E - 1)` for `E` experts per layer. +The bank holds `layers x min(N, E - 1)` rows plus `top_k x decode tokens` +staging rows, so its storage is `(layers x min(N, E - 1) + staging) x row +bytes`; the planner's tables and scratch buffers are separate, small +allocations. Rows can move between layers at run time. ## How it works -- Loading: the NVFP4 expert tensors are allocated in pinned host memory - instead of the GPU. The loader still moves one layer at a time to the GPU - for the Marlin conversion and restores the converted tensors to pinned - memory, which becomes the pool's source. +- Loading: the four per-expert tensors (gate/up and down weights and their + block scales) are allocated in pinned host memory instead of the GPU. The + loader still moves one layer at a time to the GPU for the Marlin + conversion and restores those tensors to pinned memory, which becomes the + pool's source. The small per-expert global scales stay on the device and + are copied into a pinned, contiguous host buffer when the pool is + installed. - Installation: after every layer has been processed, the pool allocates the bank, fills each layer's initial rows and binds a Marlin consumer to the bank. The placement is frozen (gate closed) during profiling and CUDA @@ -42,13 +48,20 @@ move between layers at run time. - `N` must be at least `top_k` and every MoE layer must have the same expert count and top-k. - Prefill of long inputs reads the non-resident experts from host memory on - every chunk; expect prefill to be bounded by host-to-device bandwidth. + every chunk, which adds host-read transfer cost per chunk; the dominant + term of long-input prefill time has not been profiled. ## Example -Qwen3.8-Flash-Next NVFP4 on a 48 GiB GPU budget with 258 rows per layer -(about a 32 GiB bank): decode about 63 tok/s at 4096 context with -`--compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}'`, versus -about 7 tok/s with `--offload-backend uva --cpu-offload-gb 40`. A -reproducible client and the exact launch flags are in -`benchmarks/expert_pool/`. +Qwen3.8-Flash-Next NVFP4 with 258 rows per layer (about a 32 GiB bank), +`--compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}'`, 4096 +context, one request: decode about 63 tok/s (median of three fresh server +launches). That number was measured on an RTX PRO 6000 Blackwell Max-Q +limited to 48 GiB, on a combination of this feature with the deferred PLE +rows change (01554/vllm#46 on top of #54129), which this model needs to +load its PLE table within that budget; it is not a measurement of this +branch alone. For reference, the existing `--offload-backend uva +--cpu-offload-gb 40` path gave about 7 tok/s in a single run on a +different base commit; see `benchmarks/expert_pool/README.md` for the exact +heads and conditions rather than reading the two as a controlled +comparison. From 5119ca4906a1030305be0c52f6536288deba7c51 Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Thu, 10 Sep 2026 10:28:06 +0900 Subject: [PATCH 08/13] [MoE] Expert pool docs: separate PLE mmap (needed to load) from deferred rows (measured configuration) Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- docs/features/moe_expert_pool.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/features/moe_expert_pool.md b/docs/features/moe_expert_pool.md index b8bbf9f003ef..96adb2d1d5c1 100644 --- a/docs/features/moe_expert_pool.md +++ b/docs/features/moe_expert_pool.md @@ -57,9 +57,10 @@ Qwen3.8-Flash-Next NVFP4 with 258 rows per layer (about a 32 GiB bank), `--compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}'`, 4096 context, one request: decode about 63 tok/s (median of three fresh server launches). That number was measured on an RTX PRO 6000 Blackwell Max-Q -limited to 48 GiB, on a combination of this feature with the deferred PLE -rows change (01554/vllm#46 on top of #54129), which this model needs to -load its PLE table within that budget; it is not a measurement of this +limited to 48 GiB, on the integrated head 7dedc6d8 = this feature + PLE +mmap support (#54129, which this model needs to load its PLE table within +that budget) + deferred PLE rows (01554/vllm#46, part of the measured +configuration, not required for loading); it is not a measurement of this branch alone. For reference, the existing `--offload-backend uva --cpu-offload-gb 40` path gave about 7 tok/s in a single run on a different base commit; see `benchmarks/expert_pool/README.md` for the exact From 9e1e8fe674f8fd61b118bae40e03e9f3f6b8e23f Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Thu, 10 Sep 2026 19:00:31 +0900 Subject: [PATCH 09/13] [MoE] Expert pool decode: validate every lane in the step kernel, sanitized routes for the consumer, one device assertion per call Validation stays on every request and moves from a per-layer chain of about twenty small kernels into the device planner. The step kernel already masked out-of-range ids before any table read or write; it now also validates router weights (finite, nonnegative) and treats a lane that fails either check as padding: no table is read or written for it, the placement equals the same step with the lane as -1, and duplicates stay legal. The kernel emits safe_ids (invalid lanes as -1) which the decode consumer passes to align, Marlin and the activation, so no raw id reaches them; the route mask is therefore the identity on decode and is skipped. The sticky error is mirrored in a device ok flag, asserted once per call with torch._assert_async (one launch), which keeps the previous detection path for full forwards, single-layer and partial executions alike. The wide (partition) path is unchanged. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- .../expert_pool/test_pool_layer_helpers.py | 36 +++ .../expert_pool/test_pool_marlin_cuda.py | 226 ++++++++++++++++++ tests/kernels/expert_pool/test_pool_tables.py | 136 +++++++++++ .../layers/fused_moe/expert_pool/layer.py | 106 +++++++- .../layers/fused_moe/expert_pool/tables.py | 69 +++++- 5 files changed, 563 insertions(+), 10 deletions(-) diff --git a/tests/kernels/expert_pool/test_pool_layer_helpers.py b/tests/kernels/expert_pool/test_pool_layer_helpers.py index bcb0a1ad5ce7..e1330589a5c4 100644 --- a/tests/kernels/expert_pool/test_pool_layer_helpers.py +++ b/tests/kernels/expert_pool/test_pool_layer_helpers.py @@ -10,6 +10,7 @@ marlin_block_size, mask_routes, physical_block_experts, + physical_block_experts_device, ) from vllm.model_executor.layers.fused_moe.expert_pool.tables import TENSORS @@ -44,3 +45,38 @@ def test_copy_rows_reference_copies_every_tensor_in_order(): for n in TENSORS: assert torch.equal(dst[n][2], src[n][6]) and torch.equal(dst[n][0], src[n][1]) assert int(dst[n][1].sum()) == 0 + + +def test_physical_block_experts_device_falls_back_to_torch_on_cpu(): + expert_map = torch.tensor([10, -1, 12], dtype=torch.int32) + logical = torch.tensor([2, 0, 7, 2], dtype=torch.int32) + post_padded = torch.tensor([16], dtype=torch.int32) + a = physical_block_experts(logical, post_padded, 8, expert_map, 3) + b = physical_block_experts_device(logical, post_padded, 8, expert_map, 3) + assert torch.equal(a, b) + + +def test_physical_block_experts_kernel_matches_torch_on_cuda(): + if not torch.cuda.is_available(): + return + import random + + device = torch.device("cuda") + rng = random.Random(11) + for _ in range(20): + E = rng.choice([8, 64, 512]) + bank_rows = E + rng.randint(1, 3 * E) # bank larger than the expert count + n_blocks = rng.randint(1, 300) + block = rng.choice([8, 16, 32, 64]) + # Garbage ids (out of range, negative) beyond post_padded and inside. + logical = torch.randint(-5, E + 5, (n_blocks,), dtype=torch.int32) + post_padded = torch.tensor( + [rng.randint(0, n_blocks * block)], dtype=torch.int32 + ) + expert_map = torch.randint(0, bank_rows, (E,), dtype=torch.int32) + expert_map[torch.rand(E) < 0.3] = -1 # absent experts + ref = physical_block_experts(logical, post_padded, block, expert_map, E) + got = physical_block_experts_device( + logical.to(device), post_padded.to(device), block, expert_map.to(device), E + ).cpu() + assert torch.equal(ref, got) diff --git a/tests/kernels/expert_pool/test_pool_marlin_cuda.py b/tests/kernels/expert_pool/test_pool_marlin_cuda.py index da9270edd4cd..61af0c0bef4a 100644 --- a/tests/kernels/expert_pool/test_pool_marlin_cuda.py +++ b/tests/kernels/expert_pool/test_pool_marlin_cuda.py @@ -9,6 +9,9 @@ match through the bank + host-view partition path. The pool tables must stay consistent throughout.""" +import subprocess +import sys + import pytest import torch @@ -159,3 +162,226 @@ def run(i, x, logits, n): run(0, x, _decode(order, device), 1) run(1, x, _decode([0, 1], device), 1) assert pls[0].decode_steps == 9 and pls[1].decode_steps == 4 # 1 gate-closed + + +def _three_layer_pool(device): + """Three pool layers sharing one bank, as install_expert_pool builds them.""" + from vllm.model_executor.layers.fused_moe.expert_pool.tables import set_gate + + pool_cfgs, layers = [], [] + for seed_offset in (0, 1, 2): + params = quantized_weights(device, seed_offset=seed_offset) + pool_cfgs.append(vllm_config(SLOTS)) + layers.append(make_layer(pool_cfgs[-1], params, host_source=True)) + for layer in layers: + for name in ("w13_weight_scale_2", "w2_weight_scale_2"): + p = getattr(layer.routed_experts, name) + p.data = p.data.to(device) + model = torch.nn.ModuleDict({"a": layers[0], "b": layers[1], "c": layers[2]}) + pool = install_expert_pool(model, device, max_decode_tokens=1) + assert pool is not None + set_gate(pool.tables, True) + pls = [layer.routed_experts.expert_pool_layer for layer in layers] + return pool, pls, pool_cfgs + + +def test_invalid_lanes_are_padding_and_the_error_is_sticky_without_firing( + dist_env, # noqa: F811 +): + """Device planner contract, checked with the step alone (no consumer, so + no device assertion fires and the CUDA context stays usable): an + out-of-range id, a negative non-sentinel id, and a non-finite router + weight each set the sticky error and the `ok` flag, leave the tables + identical to the same step with the lane as padding, route the other + lanes, and hide the lane in safe_ids; the error persists across later + clean steps until clear_error.""" + from vllm.model_executor.layers.fused_moe.expert_pool.tables import ( + check_global_tables, + clear_error, + step, + ) + + device = torch.accelerator.current_accelerator() + pool, pls, _ = _three_layer_pool(device) + ref_pool, ref_pls, _ = _three_layer_pool(device) + tables, ref_tables = pool.tables, ref_pool.tables + + def run(p, layer, ids, weights): + step( + p.pool.tables, + layer, + torch.tensor([ids], dtype=torch.int32, device=device), + p.buffers, + torch.tensor([weights], dtype=torch.float32, device=device), + ) + torch.accelerator.synchronize(device) + + for layer, ids, weights, ref_ids in ( + (0, [E + 3, 2], [0.5, 0.5], [-1, 2]), + (1, [1, -9], [0.5, 0.5], [1, -1]), + (2, [1, 2], [float("nan"), 0.5], [-1, 2]), + ): + run(pls[layer], layer, ids, weights) + run(ref_pls[layer], layer, ref_ids, [0.5, 0.5]) + for name in ("hot_phys", "cold_phys", "row_key"): + assert torch.equal(getattr(tables, name), getattr(ref_tables, name)), name + assert torch.equal(pls[layer].buffers.routes, ref_pls[layer].buffers.routes) + assert pls[layer].buffers.safe_ids[:2].tolist() == ref_ids + assert int(tables.error[0]) == 1 and not bool(tables.ok[0]) + assert int(ref_tables.error[0]) == 0 and bool(ref_tables.ok[0]) + # Sticky across clean steps on other layers. + for other in range(3): + run(pls[other], other, [1, 2], [0.5, 0.5]) + run(ref_pls[other], other, [1, 2], [0.5, 0.5]) + assert int(tables.error[0]) == 1 and not bool(tables.ok[0]) + with pytest.raises(RuntimeError): + check_global_tables(tables) + clear_error(tables) + check_global_tables(tables) + check_global_tables(ref_tables) + + +def test_decode_graph_capture_and_replay_match_eager(dist_env): # noqa: F811 + """The decode path (step, copy, sanitized routes, one-launch assertion, + Marlin) is captured once and replayed with different valid inputs; every + replay matches the eager result on identically placed twin pools.""" + from vllm.forward_context import set_forward_context + + device = torch.accelerator.current_accelerator() + pool, pls, cfgs = _three_layer_pool(device) + twin, tpls, tcfgs = _three_layer_pool(device) + x_static = torch.zeros(1, K, dtype=torch.bfloat16, device=device) + ids_static = torch.zeros(1, TOP_K, dtype=torch.int32, device=device) + w_static = torch.full((1, TOP_K), 0.5, dtype=torch.float32, device=device) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with ( + torch.cuda.stream(stream), + set_forward_context(None, cfgs[0], num_tokens=1), + ): + for _ in range(2): # warm up the kernels before capture + pls[0].apply(x_static, w_static, ids_static) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with ( + set_forward_context(None, cfgs[0], num_tokens=1), + torch.cuda.graph(graph, stream=stream), + ): + out_static = pls[0].apply(x_static, w_static, ids_static) + # Mirror the two warm-up steps on the twin so placements agree. + with set_forward_context(None, tcfgs[0], num_tokens=1): + for _ in range(2): + tpls[0].apply(x_static, w_static, ids_static) + tpls[0].apply(x_static, w_static, ids_static) # the capture's own step + torch.accelerator.synchronize(device) + for order in ([3, 4], [7, 0], [0, 0], [6, 1]): + x = torch.randn(1, K, dtype=torch.bfloat16, device=device) + ids = torch.tensor([order], dtype=torch.int32, device=device) + x_static.copy_(x) + ids_static.copy_(ids) + graph.replay() + with set_forward_context(None, tcfgs[0], num_tokens=1): + want = tpls[0].apply(x, w_static, ids) + torch.accelerator.synchronize(device) + torch.testing.assert_close(out_static, want, rtol=2e-2, atol=2e-2) + assert bool(pool.tables.ok[0]) and bool(twin.tables.ok[0]) + + +ASSERT_CASES = { + "first_layer_oob_id": (0, [E + 3, 2], [0.5, 0.5], "eager"), + "middle_layer_nan_weight": (1, [1, 2], [float("nan"), 0.5], "eager"), + "last_layer_negative_id": (2, [1, -9], [0.5, 0.5], "eager"), + "graph_replay_oob_id": (0, [E + 3, 2], [0.5, 0.5], "graph"), +} + + +def _run_assert_case(name): + """Subprocess body: one invalid decode forward must fail at the caller's + synchronization (device assertion). Exits 0 only when it did.""" + from tests.kernels.moe.modular_kernel_tools.parallel_utils import _set_vllm_config + from vllm.forward_context import set_forward_context + from vllm.v1.worker.workspace import ( + init_workspace_manager, + is_workspace_manager_initialized, + ) + + cfg = vllm_config(0) + _set_vllm_config(cfg, 1, rank=0, local_rank=0) + device = torch.accelerator.current_accelerator() + if not is_workspace_manager_initialized(): + init_workspace_manager(device) + pool, pls, cfgs = _three_layer_pool(device) + bad_layer, bad_ids, bad_w, mode = ASSERT_CASES[name] + x = torch.randn(1, K, dtype=torch.bfloat16, device=device) + good_ids = torch.tensor([[1, 2]], dtype=torch.int32, device=device) + good_w = torch.full((1, TOP_K), 0.5, dtype=torch.float32, device=device) + # A clean full forward first. + for i in range(3): + with set_forward_context(None, cfgs[i], num_tokens=1): + pls[i].apply(x, good_w, good_ids) + torch.accelerator.synchronize(device) + ids = torch.tensor([bad_ids], dtype=torch.int32, device=device) + w = torch.tensor([bad_w], dtype=torch.float32, device=device) + try: + if mode == "eager": + for i in range(3): + with set_forward_context(None, cfgs[i], num_tokens=1): + pls[i].apply( + x, + w if i == bad_layer else good_w, + ids if i == bad_layer else good_ids, + ) + torch.accelerator.synchronize(device) + else: + ids_static = good_ids.clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with ( + torch.cuda.stream(stream), + set_forward_context(None, cfgs[0], num_tokens=1), + ): + pls[0].apply(x, good_w, ids_static) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with ( + set_forward_context(None, cfgs[0], num_tokens=1), + torch.cuda.graph(graph, stream=stream), + ): + pls[0].apply(x, good_w, ids_static) + graph.replay() + torch.accelerator.synchronize(device) # clean replay passes + ids_static.copy_(ids) # invalid input into the captured buffer + graph.replay() + torch.accelerator.synchronize(device) + except RuntimeError as exc: + print(f"expected failure: {str(exc)[:160]}") + return 0 + print("no failure raised") + return 3 + + +@pytest.mark.parametrize("name", sorted(ASSERT_CASES)) +def test_invalid_routing_fails_at_synchronization_in_a_subprocess( + dist_env, # noqa: F811 + name, +): + """The one-launch device assertion fires for an invalid lane at the + first, a middle, or the last layer, and inside a captured graph on + replay; each case runs in its own process because a device assertion + poisons the CUDA context.""" + proc = subprocess.run( + [sys.executable, __file__, "--assert-case", name], + capture_output=True, + text=True, + timeout=600, + ) + assert proc.returncode == 0, (name, proc.stdout[-2000:], proc.stderr[-2000:]) + assert "expected failure" in proc.stdout + + +if __name__ == "__main__": + import argparse + + ap = argparse.ArgumentParser() + ap.add_argument("--assert-case", required=True, choices=sorted(ASSERT_CASES)) + sys.exit(_run_assert_case(ap.parse_args().assert_case)) diff --git a/tests/kernels/expert_pool/test_pool_tables.py b/tests/kernels/expert_pool/test_pool_tables.py index e1efab325389..e3a93765c9ae 100644 --- a/tests/kernels/expert_pool/test_pool_tables.py +++ b/tests/kernels/expert_pool/test_pool_tables.py @@ -235,8 +235,144 @@ def test_every_current_route_is_protected_when_the_pool_is_saturated(self): def test_invalid_ids_set_the_sticky_error(self): pool, sources, buffers = self.setup() self.run_step(pool, sources, buffers, 0, [9, 0]) + self.assertFalse(bool(pool.tables.ok[0])) with self.assertRaises(RuntimeError): pool.snapshot() + gp.clear_error(pool.tables) + self.assertTrue(bool(pool.tables.ok[0])) + pool.snapshot() + + def assert_same_placement(self, pool, ref_pool): + self.assertEqual( + pool.tables.hot_phys.tolist(), ref_pool.tables.hot_phys.tolist() + ) + self.assertEqual( + pool.tables.cold_phys.tolist(), ref_pool.tables.cold_phys.tolist() + ) + self.assertEqual(pool.tables.row_key.tolist(), ref_pool.tables.row_key.tolist()) + self.assertEqual(pool.tables.row_use.tolist(), ref_pool.tables.row_use.tolist()) + + def test_invalid_ids_are_planned_as_padding(self): + # Out-of-range ids never read or write a table: the placement after + # the step equals the same step with those lanes as -1; the valid + # lanes are routed; safe_ids hides the invalid lanes; the sticky + # error is set. + for gate in (False, True): + pool, sources, buffers = self.setup() + ref_pool, ref_sources, ref_buffers = self.setup() + gp.set_gate(pool.tables, gate) + gp.set_gate(ref_pool.tables, gate) + b = self.run_step(pool, sources, buffers, 0, [9, 4, -7, 1]) + rb = self.run_step(ref_pool, ref_sources, ref_buffers, 0, [-1, 4, -1, 1]) + self.assert_same_placement(pool, ref_pool) + self.assertEqual(b.routes.tolist(), rb.routes.tolist()) + self.assertEqual(b.safe_ids.tolist(), [-1, 4, -1, 1]) + self.assertEqual(b.step_map.tolist(), rb.step_map.tolist()) + self.assertGreaterEqual(b.routes.tolist()[1], 0) + self.assertFalse(bool(pool.tables.ok[0])) + self.assertTrue(bool(ref_pool.tables.ok[0])) + + def test_invalid_router_weights_make_the_lane_padding(self): + for bad in (float("nan"), float("inf"), -float("inf"), -0.5): + pool, sources, buffers = self.setup() + ref_pool, ref_sources, ref_buffers = self.setup() + gp.step( + pool.tables, + 0, + torch.tensor([[0, 1, 3, -1]]), + buffers[0], + torch.tensor([[0.5, bad, 0.25, 0.0]], dtype=torch.float32), + ) + gp.step( + ref_pool.tables, + 0, + torch.tensor([[0, -1, 3, -1]]), + ref_buffers[0], + torch.tensor([[0.5, 0.0, 0.25, 0.0]], dtype=torch.float32), + ) + self.assert_same_placement(pool, ref_pool) + self.assertEqual(buffers[0].routes.tolist(), ref_buffers[0].routes.tolist()) + self.assertEqual(buffers[0].safe_ids.tolist(), [0, -1, 3, -1]) + self.assertFalse(bool(pool.tables.ok[0])) + with self.assertRaises(RuntimeError): + pool.snapshot() + # The same weight on a padding lane is never loaded: no error. + pool, sources, buffers = self.setup() + gp.step( + pool.tables, + 0, + torch.tensor([[0, -1]]), + buffers[0], + torch.tensor([[0.5, bad]], dtype=torch.float32), + ) + self.assertTrue(bool(pool.tables.ok[0])) + pool.snapshot() + # Finite nonnegative weights, duplicates included, never set the + # error; duplicate routes stay legal and resolve to one row. + pool, sources, buffers = self.setup() + gp.step( + pool.tables, + 0, + torch.tensor([[2, 2, 0, 2]]), + buffers[0], + torch.tensor([[0.7, 0.3, 0.0, 1.0]], dtype=torch.float32), + ) + pool.snapshot() + self.assertEqual(buffers[0].safe_ids.tolist(), [2, 2, 0, 2]) + self.assertEqual(len(set(buffers[0].routes.tolist())), 2) + + def test_every_valid_route_is_present_after_the_step(self): + # The decode consumer uses safe_ids without a route mask: after a + # step every valid lane (hit, promoted, or staged) has a physical + # row in routes and in the step map, with the gate closed or open, + # with duplicates, padding, invalid sentinels and invalid weights, + # and when misses exceed the free rows (staging). + rng = random.Random(7) + for trial in range(40): + layers = rng.choice([1, 2, 3]) + experts = rng.choice([4, 6, 8]) + staging = rng.randint(1, 4) + slots = tuple(rng.randint(1, experts - 1) for _ in range(layers)) + pool, sources, buffers = self.setup(layers, experts, slots, staging) + gp.set_gate(pool.tables, rng.random() < 0.5) + for _ in range(12): + layer = rng.randrange(layers) + ids = [ + rng.choice( + [-1, -3, 99, rng.randrange(experts), rng.randrange(experts)] + ) + for _ in range(staging) + ] + weights = [ + rng.choice([0.5, 1.0, 0.0, float("nan")]) for _ in range(staging) + ] + gp.step( + pool.tables, + layer, + torch.tensor([ids]), + buffers[layer], + torch.tensor([weights], dtype=torch.float32), + ) + pool_mod.copy_in(sources[layer], pool.bank, buffers[layer]) + b = buffers[layer] + routes, safe, step_map = ( + b.routes.tolist(), + b.safe_ids.tolist(), + b.step_map.tolist(), + ) + for lane, value in enumerate(ids): + ok = 0 <= value < experts and weights[lane] == weights[lane] + if ok: + self.assertEqual(safe[lane], value) + self.assertGreaterEqual(routes[lane], 0, (trial, ids, routes)) + self.assertGreaterEqual( + step_map[value], 0, (trial, ids, step_map) + ) + else: + self.assertEqual(safe[lane], -1) + self.assertEqual(routes[lane], -1) + gp.clear_error(pool.tables) + pool.snapshot() def test_host_swap_while_gated_matches_the_tables(self): pool, sources, buffers = self.setup() diff --git a/vllm/model_executor/layers/fused_moe/expert_pool/layer.py b/vllm/model_executor/layers/fused_moe/expert_pool/layer.py index 5c34b8e69fa1..2175a74bdd29 100644 --- a/vllm/model_executor/layers/fused_moe/expert_pool/layer.py +++ b/vllm/model_executor/layers/fused_moe/expert_pool/layer.py @@ -57,6 +57,69 @@ def physical_block_experts( return torch.where(valid, expert_map[safe.long()], torch.full_like(logical_ids, -1)) +_BLOCK_ROWS_KERNEL: dict[str, Any] = {} + + +def _block_rows_kernel(): + """Triton twin of physical_block_experts: one launch instead of the + torch chain (arange, compare, where, clamp, index, where).""" + if "kernel" in _BLOCK_ROWS_KERNEL: + return _BLOCK_ROWS_KERNEL["kernel"] + from vllm.triton_utils import tl, triton + + @triton.jit + def block_rows( + logical_ptr, + post_padded_ptr, + expert_map_ptr, + out_ptr, + n_blocks, + block, + num_experts, + BLOCK: tl.constexpr, + ): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + in_range = offs < n_blocks + post_padded = tl.load(post_padded_ptr) + used = in_range & ((offs * block) < post_padded) + logical = tl.load(logical_ptr + offs, mask=used, other=0) + safe = tl.minimum(tl.maximum(logical, 0), num_experts - 1) + rows = tl.load(expert_map_ptr + safe, mask=used, other=-1) + tl.store(out_ptr + offs, tl.where(used, rows, -1), mask=in_range) + + _BLOCK_ROWS_KERNEL["kernel"] = block_rows + return block_rows + + +def physical_block_experts_device( + logical_ids: torch.Tensor, + post_padded: torch.Tensor, + block: int, + expert_map: torch.Tensor, + num_experts: int, +) -> torch.Tensor: + """physical_block_experts as one kernel launch (CUDA); the torch version + is the reference elsewhere and in tests.""" + if logical_ids.device.type != "cuda": + return physical_block_experts( + logical_ids, post_padded, block, expert_map, num_experts + ) + out = torch.empty_like(logical_ids) + n = logical_ids.numel() + BLOCK = 1024 + _block_rows_kernel()[((n + BLOCK - 1) // BLOCK,)]( + logical_ids, + post_padded, + expert_map, + out, + n, + block, + num_experts, + BLOCK=BLOCK, + ) + return out + + def marlin_block_size(tokens, top_k, local_experts, global_experts, input_dtype): """The stock fused_marlin_moe M-block choice for one expert partition.""" estimated = math.ceil(tokens * local_experts / global_experts) @@ -114,17 +177,38 @@ def __init__( def apply( self, x: torch.Tensor, weights: torch.Tensor, ids: torch.Tensor ) -> torch.Tensor: - self._check_routes(x, weights, ids) + self._check_route_shapes(x, weights, ids) lanes = ids.shape[0] * ids.shape[1] if lanes <= self.width and lanes <= self.staging_rows: self.decode_steps += 1 - step(self.pool.tables, self.index, ids, self.buffers) + tables = self.pool.tables + # The step kernel validates every lane (id range before any + # table read, router weight finite and nonnegative) into the + # sticky device error; invalid lanes are planned as padding, so + # no table is read or written for them, and `safe_ids` carries + # the sanitized routes that every consumer below (align, Marlin, + # activation) receives: no raw id reaches them. One device + # assertion per call, a single launch on the `ok` flag, keeps + # the detection path of the previous per-layer assertion chain + # for full forwards, single-layer and partial executions alike. + step(tables, self.index, ids, self.buffers, weights) + torch._assert_async( + tables.ok, + "Expert pool: invalid routing (id out of range, or non-finite/" + "negative router weight)", + ) copy_in(self.host, self.bank, self.buffers) + safe = self.buffers.safe_ids[:lanes].view(ids.shape) return self._run_marlin( - x, weights, ids, ((self.bank, self.buffers.step_map, self.bank_rows),) + x, + weights, + safe, + ((self.bank, self.buffers.step_map, self.bank_rows),), + decode=True, ) # Wider batches (prefill): resident rows from the bank, the rest read # straight from the pinned host source through its accelerator view. + self._check_routes_device(weights, ids) self.partition_steps += 1 return self._run_marlin( x, @@ -136,7 +220,8 @@ def apply( ), ) - def _check_routes(self, x, weights, ids) -> None: + def _check_route_shapes(self, x, weights, ids) -> None: + # Host-side, no kernels: shapes and dtypes only. if ( x.ndim != 2 or ids.ndim != 2 @@ -146,6 +231,10 @@ def _check_routes(self, x, weights, ids) -> None: or ids.dtype not in (torch.int32, torch.int64) ): raise ValueError("Unexpected pool routing/input shape or dtype") + + def _check_routes_device(self, weights, ids) -> None: + # Partition path (wide batches): the per-call device assertion. + # The decode path folds the same conditions into the step kernel. allowed = (ids >= -1) & (ids < self.num_experts) finite = torch.isfinite(weights) & (weights >= 0) torch._assert_async( @@ -154,7 +243,7 @@ def _check_routes(self, x, weights, ids) -> None: "finite/nonnegative", ) - def _run_marlin(self, x, weights, ids, partitions): + def _run_marlin(self, x, weights, ids, partitions, decode=False): from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( _fused_marlin_moe, marlin_moe_intermediate_size, @@ -187,7 +276,11 @@ def _run_marlin(self, x, weights, ids, partitions): if slots > self.num_experts: # Align by logical id (absent experts already padding), then # map the blocks to rows; the align op never sees a row id. - routed = mask_routes(ids, expert_map) + # On the decode path `ids` are the step's sanitized routes + # and every valid route is present in the step map (hit, + # promoted, or staged), so the mask is the identity. + if not decode: + routed = mask_routes(ids, expert_map) sorted_ids, logical_ids, post_padded = moe_align_block_size( routed, block, self.num_experts, None, ignore_invalid_experts=True ) @@ -258,4 +351,5 @@ def stats(self) -> dict[str, int]: "marlin_block_size", "mask_routes", "physical_block_experts", + "physical_block_experts_device", ] diff --git a/vllm/model_executor/layers/fused_moe/expert_pool/tables.py b/vllm/model_executor/layers/fused_moe/expert_pool/tables.py index a3333710ea79..d2d49417e4d8 100644 --- a/vllm/model_executor/layers/fused_moe/expert_pool/tables.py +++ b/vllm/model_executor/layers/fused_moe/expert_pool/tables.py @@ -26,6 +26,7 @@ from __future__ import annotations +import math from dataclasses import dataclass from typing import Any @@ -56,6 +57,7 @@ class GlobalTables: clock: Any # [1] int64 gate: Any # [1] int32 promotions allowed error: Any # [1] int32 sticky device error + ok: Any # [1] bool, the inverse of error for one-launch device assertions promote_limit: Any # [1] int32 max promotions per layer call, 0 = unlimited promote_interval: Any # [1] int32 promote only every N forwards forwards: Any # [1] int32 forwards seen with the gate open (layer 0 count) @@ -81,6 +83,7 @@ class StepBuffers: gather_dst: Any # [W] int32 bank rows gather_count: Any # [1] int32 routes: Any # [W] int32 physical row per ids lane this step, -1 padding + safe_ids: Any # [W] int32 ids with invalid lanes (range, weight) as -1 staged_expert: Any # [W] int32 (scratch for the map overlay) staged_row: Any # [W] int32 staged_count: Any # [1] int32 @@ -128,6 +131,7 @@ def allocate_global_tables(device, num_experts, slots_per_layer, staging): clock=torch.zeros(1, dtype=torch.int64, device=device), gate=torch.zeros(1, dtype=torch.int32, device=device), error=torch.zeros(1, dtype=torch.int32, device=device), + ok=torch.ones(1, dtype=torch.bool, device=device), promote_limit=torch.zeros(1, dtype=torch.int32, device=device), promote_interval=torch.ones(1, dtype=torch.int32, device=device), forwards=torch.zeros(1, dtype=torch.int32, device=device), @@ -151,6 +155,7 @@ def ints(n): gather_dst=ints(width), gather_count=ints(1), routes=torch.full((width,), -1, dtype=torch.int32, device=device), + safe_ids=torch.full((width,), -1, dtype=torch.int32, device=device), staged_expert=ints(width), staged_row=ints(width), staged_count=ints(1), @@ -163,6 +168,12 @@ def set_gate(tables, enabled): tables.gate.fill_(1 if enabled else 0) +def clear_error(tables): + """Reset the sticky device error (host side, after handling it).""" + tables.error.zero_() + tables.ok.fill_(True) + + CONTROL_MAX = 2**31 - 1 # device scalars are int32 CONTROL_FIELDS = ( "promote_limit", @@ -225,7 +236,7 @@ def read_control(tables): return values -def step_reference(tables, layer, ids, buffers): +def step_reference(tables, layer, ids, buffers, weights=None): """Plan and flip one layer step on the host (torch, synchronizing). Returns (gathers, step_map) with gathers as (RAM row, bank row) pairs in @@ -233,6 +244,13 @@ def step_reference(tables, layer, ids, buffers): - `ids` values outside [0, E) other than -1 set the sticky error and are skipped; -1 is padding. + - With `weights` (one per ids lane), a non-finite or negative weight on + a valid lane sets the sticky error and makes that lane invalid: it is + planned as padding (no table read or write, no ownership change), so + the tables after the step equal those of the same step with the lane + as -1. `buffers.safe_ids` holds the ids with every invalid lane as -1 + for the consumer; `tables.ok` is the inverse of the error for the + caller's one-launch device assertion. - Distinct valid selections in first-occurrence order. With the gate open the clock advances and every selected resident row is stamped. - Recency lives on pool rows (FreeToken's usage-per-slot): each miss, @@ -259,6 +277,17 @@ def step_reference(tables, layer, ids, buffers): if len(raw) > buffers.gather_src.shape[0] or len(raw) > len(staging_rows): raise ValueError("Step ids exceed the plan width or the staging rows") error = bool(int(tables.error[0])) + if weights is not None: + w = weights.reshape(-1).to(torch.float32).tolist() + if len(w) != len(raw): + raise ValueError("weights must have one value per ids lane") + for i, (value, weight) in enumerate(zip(raw, w)): + if 0 <= value < E and not (math.isfinite(weight) and weight >= 0): + error = True + raw[i] = -2 # invalid: planned as padding below + safe_ids = [-1] * buffers.safe_ids.shape[0] + for i, value in enumerate(raw): + safe_ids[i] = value if 0 <= value < E else -1 selected: list[int] = [] for value in raw: if value == -1: @@ -340,6 +369,8 @@ def write(target, values, dtype): tables.forwards.fill_(forwards) write(tables.miss_count, miss_count, torch.int32) tables.error.fill_(1 if error else 0) + tables.ok.fill_(not error) + write(buffers.safe_ids, safe_ids, torch.int32) pairs = gathers + [(e, row) for e, row in staged] buffers.gather_count.fill_(len(pairs)) buffers.promoted_count.fill_(len(gathers)) @@ -353,20 +384,31 @@ def write(target, values, dtype): return pairs, buffers.step_map -def step(tables, layer, ids, buffers): - """Plan and flip one layer step: Triton on CUDA, the reference elsewhere.""" +def step(tables, layer, ids, buffers, weights=None): + """Plan and flip one layer step: Triton on CUDA, the reference elsewhere. + + `weights` (optional, one per ids lane) are validated on the device: a + non-finite or negative weight on a valid lane sets the sticky error. + """ if tables.hot_phys.device.type != "cuda": - step_reference(tables, layer, ids, buffers) + step_reference(tables, layer, ids, buffers, weights) return flat = ids.reshape(-1) if not flat.is_contiguous(): raise ValueError("Global step requires contiguous ids") + if weights is not None: + weights = weights.reshape(-1) + if weights.numel() != flat.numel() or not weights.is_contiguous(): + raise ValueError("Global step weights must match the ids lanes") + check_weights = weights is not None + weights_arg = weights if check_weights else buffers.routes width = buffers.gather_src.shape[0] if flat.numel() > width or flat.numel() > tables.staging_rows.shape[0]: raise ValueError("Step ids exceed the plan width or the staging rows") rows = tables.pool_rows _step_kernel()[(1,)]( flat, + weights_arg, flat.numel(), layer, tables.hot_phys, @@ -376,6 +418,7 @@ def step(tables, layer, ids, buffers): tables.clock, tables.gate, tables.error, + tables.ok, tables.promote_limit, tables.promote_interval, tables.forwards, @@ -387,6 +430,7 @@ def step(tables, layer, ids, buffers): buffers.gather_dst, buffers.gather_count, buffers.routes, + buffers.safe_ids, buffers.staged_expert, buffers.staged_row, buffers.staged_count, @@ -398,6 +442,7 @@ def step(tables, layer, ids, buffers): WIDTH=width, BLOCK_R=_next_power_of_two(rows), MAP_BLOCK=1024, + CHECK_WEIGHTS=check_weights, num_warps=8, ) @@ -453,6 +498,7 @@ def _step_kernel(): @triton.jit def global_pool_step( ids_ptr, + weights_ptr, n, layer, hot_phys_ptr, @@ -462,6 +508,7 @@ def global_pool_step( clock_ptr, gate_ptr, error_ptr, + ok_ptr, limit_ptr, interval_ptr, forwards_ptr, @@ -473,6 +520,7 @@ def global_pool_step( gather_dst_ptr, gather_count_ptr, routes_ptr, + safe_ids_ptr, staged_expert_ptr, staged_row_ptr, staged_count_ptr, @@ -484,6 +532,7 @@ def global_pool_step( WIDTH: tl.constexpr, BLOCK_R: tl.constexpr, MAP_BLOCK: tl.constexpr, + CHECK_WEIGHTS: tl.constexpr, ): never = 0x7FFFFFFFFFFFFFFF lane = tl.arange(0, WIDTH) @@ -491,8 +540,20 @@ def global_pool_step( raw = tl.load(ids_ptr + lane, mask=present, other=-1).to(tl.int64) valid = present & (raw >= 0) & (raw < num_experts) bad = present & (raw != -1) & (~valid) + if CHECK_WEIGHTS: + # Router weights of valid lanes must be finite and nonnegative; + # a lane failing this becomes invalid (planned as padding) so no + # table is read or written for it; weights of padding lanes are + # never loaded. + w = tl.load(weights_ptr + lane, mask=valid, other=0.0).to(tl.float32) + w_ok = (w == w) & (w >= 0.0) & (w != float("inf")) + bad = bad | (valid & (~w_ok)) + valid = valid & w_ok if tl.sum(bad.to(tl.int32), 0) > 0: tl.store(error_ptr, 1) + tl.store(ok_ptr, 0) + # Sanitized ids for the consumer: every invalid lane is padding. + tl.store(safe_ids_ptr + lane, tl.where(valid, raw, -1).to(tl.int32)) safe = tl.where(valid, raw, 0) same = safe[:, None] == safe[None, :] earlier = lane[None, :] < lane[:, None] From 60c80ee423d9f95483a7796fe00044c549453145 Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Thu, 10 Sep 2026 19:00:39 +0900 Subject: [PATCH 10/13] [MoE] Expert pool: block-to-row remap as one Triton kernel on CUDA physical_block_experts_device replaces the torch chain (arange, compare, where, clamp, index, where) with one launch; identical semantics, including blocks beyond post_padded (never indexed, -1) and absent experts (-1). The torch version remains the reference and the CPU path. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- vllm/model_executor/layers/fused_moe/expert_pool/layer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/expert_pool/layer.py b/vllm/model_executor/layers/fused_moe/expert_pool/layer.py index 2175a74bdd29..ba6f06bf6a06 100644 --- a/vllm/model_executor/layers/fused_moe/expert_pool/layer.py +++ b/vllm/model_executor/layers/fused_moe/expert_pool/layer.py @@ -284,7 +284,7 @@ def _run_marlin(self, x, weights, ids, partitions, decode=False): sorted_ids, logical_ids, post_padded = moe_align_block_size( routed, block, self.num_experts, None, ignore_invalid_experts=True ) - expert_ids = physical_block_experts( + expert_ids = physical_block_experts_device( logical_ids, post_padded, block, expert_map, self.num_experts ) else: From 1bd6c8a16e87cd6abb9360b32f5c640a989a43ab Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Thu, 10 Sep 2026 19:06:07 +0900 Subject: [PATCH 11/13] [MoE] Expert pool tests: assertion cases accept only the device assertion, partial-execution case, step-only graph stickiness, clamp-activation padding oracle, capture/replay placement checks Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- .../expert_pool/test_pool_marlin_cuda.py | 184 +++++++++++++++--- .../layers/fused_moe/expert_pool/tables.py | 7 +- 2 files changed, 168 insertions(+), 23 deletions(-) diff --git a/tests/kernels/expert_pool/test_pool_marlin_cuda.py b/tests/kernels/expert_pool/test_pool_marlin_cuda.py index 61af0c0bef4a..89cb4031740b 100644 --- a/tests/kernels/expert_pool/test_pool_marlin_cuda.py +++ b/tests/kernels/expert_pool/test_pool_marlin_cuda.py @@ -268,12 +268,14 @@ def test_decode_graph_capture_and_replay_match_eager(dist_env): # noqa: F811 torch.cuda.graph(graph, stream=stream), ): out_static = pls[0].apply(x_static, w_static, ids_static) - # Mirror the two warm-up steps on the twin so placements agree. + # Mirror only the two warm-up steps on the twin: capture records the + # kernels without executing them, so the placements agree here. with set_forward_context(None, tcfgs[0], num_tokens=1): for _ in range(2): tpls[0].apply(x_static, w_static, ids_static) - tpls[0].apply(x_static, w_static, ids_static) # the capture's own step torch.accelerator.synchronize(device) + for name in ("hot_phys", "row_key"): + assert torch.equal(getattr(pool.tables, name), getattr(twin.tables, name)) for order in ([3, 4], [7, 0], [0, 0], [6, 1]): x = torch.randn(1, K, dtype=torch.bfloat16, device=device) ids = torch.tensor([order], dtype=torch.int32, device=device) @@ -285,12 +287,133 @@ def test_decode_graph_capture_and_replay_match_eager(dist_env): # noqa: F811 torch.accelerator.synchronize(device) torch.testing.assert_close(out_static, want, rtol=2e-2, atol=2e-2) assert bool(pool.tables.ok[0]) and bool(twin.tables.ok[0]) + # The replayed step moved the placement exactly as the eager step. + for name in ("hot_phys", "row_key"): + assert torch.equal(getattr(pool.tables, name), getattr(twin.tables, name)) + assert torch.equal(pls[0].buffers.step_map, tpls[0].buffers.step_map) + assert torch.equal(pls[0].buffers.safe_ids, tpls[0].buffers.safe_ids) + + +def test_step_graph_keeps_the_error_sticky_across_replays(dist_env): # noqa: F811 + """The planner step alone (no consumer, so no assertion fires) captured + in a CUDA graph: a replay with an invalid id sets the sticky error, a + later clean replay keeps it, clear_error resets it, and the placement + matches a twin stepped with padding throughout.""" + from vllm.model_executor.layers.fused_moe.expert_pool.tables import ( + check_global_tables, + clear_error, + step, + ) + + device = torch.accelerator.current_accelerator() + pool, pls, _ = _three_layer_pool(device) + twin, tpls, _ = _three_layer_pool(device) + ids_static = torch.tensor([[1, 2]], dtype=torch.int32, device=device) + w_static = torch.full((1, TOP_K), 0.5, dtype=torch.float32, device=device) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + step(pool.tables, 0, ids_static, pls[0].buffers, w_static) + torch.cuda.current_stream().wait_stream(stream) + step(twin.tables, 0, ids_static, tpls[0].buffers, w_static) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + step(pool.tables, 0, ids_static, pls[0].buffers, w_static) + torch.accelerator.synchronize(device) + + def replay(ids, twin_ids): + ids_static.copy_(torch.tensor([ids], dtype=torch.int32, device=device)) + graph.replay() + step( + twin.tables, + 0, + torch.tensor([twin_ids], dtype=torch.int32, device=device), + tpls[0].buffers, + w_static, + ) + torch.accelerator.synchronize(device) + for name in ("hot_phys", "row_key"): + assert torch.equal(getattr(pool.tables, name), getattr(twin.tables, name)) + assert torch.equal(pls[0].buffers.safe_ids, tpls[0].buffers.safe_ids) + + replay([3, 4], [3, 4]) + assert bool(pool.tables.ok[0]) + replay([E + 5, 4], [-1, 4]) # invalid lane: padding for the twin + assert int(pool.tables.error[0]) == 1 and not bool(pool.tables.ok[0]) + replay([5, 6], [5, 6]) # clean replay keeps the sticky error + assert int(pool.tables.error[0]) == 1 and not bool(pool.tables.ok[0]) + with pytest.raises(RuntimeError): + check_global_tables(pool.tables) + clear_error(pool.tables) + replay([7, 0], [7, 0]) + assert bool(pool.tables.ok[0]) + check_global_tables(pool.tables) + + +def test_consumer_with_clamp_activation_only_sees_sanitized_routes( + dist_env, # noqa: F811 +): + """The activation path with a clamp limit indexes the step map by + expert id (masked only by >= 0). The consumer receives the kernel's + safe_ids, so a step with an out-of-range or negative id produces the + same output as the same step with that lane as padding (the oracle); + checked by calling the consumer directly, without the assertion.""" + import dataclasses + + from vllm.model_executor.layers.fused_moe.expert_pool.tables import ( + clear_error, + step, + ) + + device = torch.accelerator.current_accelerator() + pool, pls, _ = _three_layer_pool(device) + twin, tpls, _ = _three_layer_pool(device) + for p in (pls[0], tpls[0]): + cfg = p.experts.activation_config + p.experts.activation_config = dataclasses.replace(cfg, clamp_limit=7.0) + x = torch.randn(1, K, dtype=torch.bfloat16, device=device) + w = torch.full((1, TOP_K), 0.5, dtype=torch.float32, device=device) + for bad, oracle in ( + ([E + 3, 2], [-1, 2]), + ([1, -9], [1, -1]), + ([E + 1, E + 2], [-1, -1]), + ): + ids = torch.tensor([bad], dtype=torch.int32, device=device) + step(pool.tables, 0, ids, pls[0].buffers, w) + step( + twin.tables, + 0, + torch.tensor([oracle], dtype=torch.int32, device=device), + tpls[0].buffers, + w, + ) + safe = pls[0].buffers.safe_ids[: ids.numel()].view(ids.shape) + got = pls[0]._run_marlin( + x, + w, + safe, + ((pls[0].bank, pls[0].buffers.step_map, pls[0].bank_rows),), + decode=True, + ) + want = tpls[0]._run_marlin( + x, + w, + torch.tensor([oracle], dtype=torch.int32, device=device), + ((tpls[0].bank, tpls[0].buffers.step_map, tpls[0].bank_rows),), + decode=True, + ) + torch.accelerator.synchronize(device) + assert safe.tolist() == [oracle] + torch.testing.assert_close(got, want, rtol=0, atol=0) + assert int(pool.tables.error[0]) == 1 + clear_error(pool.tables) ASSERT_CASES = { "first_layer_oob_id": (0, [E + 3, 2], [0.5, 0.5], "eager"), "middle_layer_nan_weight": (1, [1, 2], [float("nan"), 0.5], "eager"), "last_layer_negative_id": (2, [1, -9], [0.5, 0.5], "eager"), + "single_layer_partial_oob_id": (1, [E + 3, 2], [0.5, 0.5], "single"), "graph_replay_oob_id": (0, [E + 3, 2], [0.5, 0.5], "graph"), } @@ -322,6 +445,30 @@ def _run_assert_case(name): torch.accelerator.synchronize(device) ids = torch.tensor([bad_ids], dtype=torch.int32, device=device) w = torch.tensor([bad_w], dtype=torch.float32, device=device) + # Clean setup, capture and clean replay run outside the guarded region: + # any failure there is a real failure of this test. + graph = None + ids_static = None + if mode == "graph": + ids_static = good_ids.clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with ( + torch.cuda.stream(stream), + set_forward_context(None, cfgs[0], num_tokens=1), + ): + pls[0].apply(x, good_w, ids_static) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with ( + set_forward_context(None, cfgs[0], num_tokens=1), + torch.cuda.graph(graph, stream=stream), + ): + pls[0].apply(x, good_w, ids_static) + graph.replay() + torch.accelerator.synchronize(device) # clean replay passes + # Only the invalid input and its synchronization may raise, and only + # with the device-side assertion; anything else is a failure. try: if mode == "eager": for i in range(3): @@ -332,30 +479,23 @@ def _run_assert_case(name): ids if i == bad_layer else good_ids, ) torch.accelerator.synchronize(device) + elif mode == "single": + # Partial execution: one non-final layer, then synchronize. + with set_forward_context(None, cfgs[bad_layer], num_tokens=1): + pls[bad_layer].apply(x, w, ids) + torch.accelerator.synchronize(device) else: - ids_static = good_ids.clone() - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with ( - torch.cuda.stream(stream), - set_forward_context(None, cfgs[0], num_tokens=1), - ): - pls[0].apply(x, good_w, ids_static) - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with ( - set_forward_context(None, cfgs[0], num_tokens=1), - torch.cuda.graph(graph, stream=stream), - ): - pls[0].apply(x, good_w, ids_static) - graph.replay() - torch.accelerator.synchronize(device) # clean replay passes + assert graph is not None and ids_static is not None ids_static.copy_(ids) # invalid input into the captured buffer graph.replay() torch.accelerator.synchronize(device) except RuntimeError as exc: - print(f"expected failure: {str(exc)[:160]}") - return 0 + text = str(exc) + if "device-side assert" in text or "Expert pool: invalid routing" in text: + print(f"expected device assertion: {text[:160]}") + return 0 + print(f"unexpected error: {text[:300]}") + return 4 print("no failure raised") return 3 @@ -376,7 +516,7 @@ def test_invalid_routing_fails_at_synchronization_in_a_subprocess( timeout=600, ) assert proc.returncode == 0, (name, proc.stdout[-2000:], proc.stderr[-2000:]) - assert "expected failure" in proc.stdout + assert "expected device assertion" in proc.stdout if __name__ == "__main__": diff --git a/vllm/model_executor/layers/fused_moe/expert_pool/tables.py b/vllm/model_executor/layers/fused_moe/expert_pool/tables.py index d2d49417e4d8..d9c54f095095 100644 --- a/vllm/model_executor/layers/fused_moe/expert_pool/tables.py +++ b/vllm/model_executor/layers/fused_moe/expert_pool/tables.py @@ -169,7 +169,12 @@ def set_gate(tables, enabled): def clear_error(tables): - """Reset the sticky device error (host side, after handling it).""" + """Reset the sticky device error (host side, after handling it). + + This resets the flags only. If a consumer's device assertion has already + fired on the error, the CUDA context is poisoned and cannot be recovered + here; the process must be restarted (which is why the assertion tests + run in subprocesses).""" tables.error.zero_() tables.ok.fill_(True) From 2e2355e05bcc8c9b387c1edfbdbea92079d2927f Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Thu, 10 Sep 2026 19:12:10 +0900 Subject: [PATCH 12/13] [MoE] Expert pool tests: per-test distributed fixture re-init, staging width for 4-lane cases, subprocess exits without CUDA teardown Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- tests/kernels/expert_pool/marlin_fixture.py | 14 ++++++++++++-- .../kernels/expert_pool/test_pool_marlin_cuda.py | 15 +++++++++++++-- tests/kernels/expert_pool/test_pool_tables.py | 10 +++++----- 3 files changed, 30 insertions(+), 9 deletions(-) diff --git a/tests/kernels/expert_pool/marlin_fixture.py b/tests/kernels/expert_pool/marlin_fixture.py index bdb52c9abbc1..d8f9fa6280fa 100644 --- a/tests/kernels/expert_pool/marlin_fixture.py +++ b/tests/kernels/expert_pool/marlin_fixture.py @@ -36,10 +36,20 @@ def vllm_config(pool_rows: int) -> VllmConfig: return cfg -@pytest.fixture(scope="module") +@pytest.fixture def dist_env(): + """Per test: the suite's teardown destroys the distributed groups after + each test, so (re)initialize whenever the groups are missing.""" + from vllm.distributed import parallel_state + cfg = vllm_config(0) - _set_vllm_config(cfg, 1, rank=0, local_rank=0) + try: + parallel_state.get_pcp_group() + initialized = True + except AssertionError: + initialized = False + if not initialized: + _set_vllm_config(cfg, 1, rank=0, local_rank=0) if not is_workspace_manager_initialized(): init_workspace_manager(torch.accelerator.current_accelerator()) return cfg diff --git a/tests/kernels/expert_pool/test_pool_marlin_cuda.py b/tests/kernels/expert_pool/test_pool_marlin_cuda.py index 89cb4031740b..48f6c87aa6bb 100644 --- a/tests/kernels/expert_pool/test_pool_marlin_cuda.py +++ b/tests/kernels/expert_pool/test_pool_marlin_cuda.py @@ -500,6 +500,17 @@ def _run_assert_case(name): return 3 +def _exit_now(code): + """After a device assertion the CUDA context is poisoned and interpreter + teardown can abort (SIGABRT); the verdict is already printed, so leave + without running any destructor.""" + import os + + sys.stdout.flush() + sys.stderr.flush() + os._exit(code) + + @pytest.mark.parametrize("name", sorted(ASSERT_CASES)) def test_invalid_routing_fails_at_synchronization_in_a_subprocess( dist_env, # noqa: F811 @@ -515,7 +526,7 @@ def test_invalid_routing_fails_at_synchronization_in_a_subprocess( text=True, timeout=600, ) - assert proc.returncode == 0, (name, proc.stdout[-2000:], proc.stderr[-2000:]) + assert proc.returncode == 0, (name, proc.stdout[-3000:], proc.stderr[-6000:]) assert "expected device assertion" in proc.stdout @@ -524,4 +535,4 @@ def test_invalid_routing_fails_at_synchronization_in_a_subprocess( ap = argparse.ArgumentParser() ap.add_argument("--assert-case", required=True, choices=sorted(ASSERT_CASES)) - sys.exit(_run_assert_case(ap.parse_args().assert_case)) + _exit_now(_run_assert_case(ap.parse_args().assert_case)) diff --git a/tests/kernels/expert_pool/test_pool_tables.py b/tests/kernels/expert_pool/test_pool_tables.py index e3a93765c9ae..af42f8507292 100644 --- a/tests/kernels/expert_pool/test_pool_tables.py +++ b/tests/kernels/expert_pool/test_pool_tables.py @@ -258,8 +258,8 @@ def test_invalid_ids_are_planned_as_padding(self): # lanes are routed; safe_ids hides the invalid lanes; the sticky # error is set. for gate in (False, True): - pool, sources, buffers = self.setup() - ref_pool, ref_sources, ref_buffers = self.setup() + pool, sources, buffers = self.setup(staging=4) + ref_pool, ref_sources, ref_buffers = self.setup(staging=4) gp.set_gate(pool.tables, gate) gp.set_gate(ref_pool.tables, gate) b = self.run_step(pool, sources, buffers, 0, [9, 4, -7, 1]) @@ -274,8 +274,8 @@ def test_invalid_ids_are_planned_as_padding(self): def test_invalid_router_weights_make_the_lane_padding(self): for bad in (float("nan"), float("inf"), -float("inf"), -0.5): - pool, sources, buffers = self.setup() - ref_pool, ref_sources, ref_buffers = self.setup() + pool, sources, buffers = self.setup(staging=4) + ref_pool, ref_sources, ref_buffers = self.setup(staging=4) gp.step( pool.tables, 0, @@ -309,7 +309,7 @@ def test_invalid_router_weights_make_the_lane_padding(self): pool.snapshot() # Finite nonnegative weights, duplicates included, never set the # error; duplicate routes stay legal and resolve to one row. - pool, sources, buffers = self.setup() + pool, sources, buffers = self.setup(staging=4) gp.step( pool.tables, 0, From e65bcc45d0a97c2d3888740f0aa4b1ffb625a830 Mon Sep 17 00:00:00 2001 From: 01554 <24953377+01554@users.noreply.github.com> Date: Thu, 10 Sep 2026 19:12:43 +0900 Subject: [PATCH 13/13] [MoE] Expert pool tests: subprocess assertion cases exit from inside the handler (unwinding after a device assertion aborts) Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com> --- .../kernels/expert_pool/test_pool_marlin_cuda.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/tests/kernels/expert_pool/test_pool_marlin_cuda.py b/tests/kernels/expert_pool/test_pool_marlin_cuda.py index 48f6c87aa6bb..586fbf20dbd3 100644 --- a/tests/kernels/expert_pool/test_pool_marlin_cuda.py +++ b/tests/kernels/expert_pool/test_pool_marlin_cuda.py @@ -490,20 +490,22 @@ def _run_assert_case(name): graph.replay() torch.accelerator.synchronize(device) except RuntimeError as exc: + # Leave from inside the handler: once the device assertion fired the + # CUDA context is poisoned, and unwinding out of this frame frees + # pinned/device tensors whose release re-raises and aborts the + # process (SIGABRT). The verdict is printed and flushed first. text = str(exc) if "device-side assert" in text or "Expert pool: invalid routing" in text: print(f"expected device assertion: {text[:160]}") - return 0 + _exit_now(0) print(f"unexpected error: {text[:300]}") - return 4 + _exit_now(4) print("no failure raised") - return 3 + _exit_now(3) def _exit_now(code): - """After a device assertion the CUDA context is poisoned and interpreter - teardown can abort (SIGABRT); the verdict is already printed, so leave - without running any destructor.""" + """Exit without running destructors (see _run_assert_case).""" import os sys.stdout.flush() @@ -535,4 +537,4 @@ def test_invalid_routing_fails_at_synchronization_in_a_subprocess( ap = argparse.ArgumentParser() ap.add_argument("--assert-case", required=True, choices=sorted(ASSERT_CASES)) - _exit_now(_run_assert_case(ap.parse_args().assert_case)) + _run_assert_case(ap.parse_args().assert_case) # exits from inside