From e5fce7b56b07cdf992febd0862bfd599d6dca3fd Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Sun, 20 Sep 2026 05:28:33 -0400 Subject: [PATCH] [Core][Kernel] Share persistent workspaces for Marlin and Humming (#57421) Signed-off-by: mgoin Co-authored-by: ErinYin --- tests/kernels/moe/test_moe.py | 20 +- .../kernels/moe/test_moe_permute_unpermute.py | 175 ++++++++++++++++-- .../kernels/quantization/test_marlin_gemm.py | 58 ++++++ .../quantization/test_marlin_tile_padding.py | 13 +- .../model_loader/test_reload.py | 67 ++++--- tests/v1/worker/test_workspace.py | 63 +++++++ .../kernels/linear/mixed_precision/marlin.py | 8 +- .../kernels/linear/mxfp4/marlin.py | 2 +- .../kernels/linear/mxfp8/marlin.py | 2 +- .../kernels/linear/nvfp4/marlin.py | 2 +- .../kernels/linear/scaled_mm/marlin.py | 2 +- .../fused_moe/experts/fused_humming_moe.py | 34 ++-- .../layers/fused_moe/experts/marlin_moe.py | 4 +- .../layers/fused_moe/moe_permute_unpermute.py | 47 +++++ .../layers/quantization/auto_awq.py | 4 - .../layers/quantization/auto_gptq.py | 8 - .../compressed_tensors_moe_wna16.py | 14 -- .../layers/quantization/utils/marlin_utils.py | 115 +++++++++--- .../quantization/utils/marlin_utils_fp4.py | 26 +-- .../quantization/utils/marlin_utils_fp8.py | 49 +---- vllm/v1/worker/workspace.py | 96 +++++++--- 21 files changed, 598 insertions(+), 211 deletions(-) diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 8b4efaf913..1aca8593ed 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -1384,24 +1384,33 @@ def test_humming_global_valid_shape_m( assert result == expected -def test_humming_permute_scratch_is_keyed_by_runtime_topk( +def test_humming_permute_scratch_is_shared_by_config( monkeypatch: pytest.MonkeyPatch, ): from types import SimpleNamespace from unittest.mock import Mock import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming + import vllm.model_executor.layers.fused_moe.moe_permute_unpermute as permute + import vllm.v1.worker.workspace as workspace + + manager = workspace.WorkspaceManager(torch.device("cpu")) + monkeypatch.setattr(workspace, "_manager", manager) scratch_topk6 = Mock() scratch_topk1 = Mock() scratch_type = Mock(side_effect=[scratch_topk6, scratch_topk1]) monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True) - monkeypatch.setattr(humming, "MoEPermuteScratch", scratch_type) + monkeypatch.setattr(permute, "MoEPermuteScratch", scratch_type) moe_config = make_dummy_moe_config(max_num_tokens=512, experts_per_token=6) moe_config.moe_parallel_config.dp_size = 2 - experts = SimpleNamespace(_permute_scratch={}, moe_config=moe_config) + experts = SimpleNamespace(moe_config=moe_config) + other_layer = SimpleNamespace(moe_config=moe_config) assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 + assert ( + humming.HummingExpertsBase._get_permute_scratch(other_layer, 6) is scratch_topk6 + ) assert humming.HummingExpertsBase._get_permute_scratch(experts, 1) is scratch_topk1 assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 @@ -1412,6 +1421,11 @@ def test_humming_permute_scratch_is_keyed_by_runtime_topk( assert second_call.kwargs["max_num_tokens"] == 6144 assert second_call.kwargs["topk"] == 1 + manager.lock() + assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 + with pytest.raises(AssertionError, match="was not allocated during warmup"): + humming.HummingExpertsBase._get_permute_scratch(experts, 2) + def test_humming_delegates_to_instance_activation(): from types import SimpleNamespace diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 76d1e57f58..52fc8963ec 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -15,12 +15,14 @@ from vllm.model_executor.layers.fused_moe.expert_map_manager import ( ) from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( MoEPermuteScratch, + get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, moe_unpermute, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +from vllm.v1.worker.workspace import current_workspace_manager NUM_EXPERTS = [16, 64, 256] TOP_KS = [2, 6, 8] @@ -120,6 +122,7 @@ def torch_unpermute( @pytest.mark.parametrize("topk", TOP_KS) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("ep_size", EP_SIZE) +@pytest.mark.parametrize("use_scratch", [False, True]) def test_moe_permute_unpermute( n_token: int, n_hidden: int, @@ -127,6 +130,8 @@ def test_moe_permute_unpermute( n_expert: int, ep_size: int, dtype: torch.dtype, + use_scratch: bool, + workspace_init, ): if not moe_permute_unpermute_supported(): pytest.skip("moe_permute_unpermute is not supported on this platform.") @@ -160,6 +165,18 @@ def test_moe_permute_unpermute( expert_map=expert_map, ) + scratch = None + if use_scratch: + scratch = get_moe_permute_scratch( + max_num_tokens=n_token, + topk=topk, + num_experts=n_expert, + num_local_experts=n_local_expert, + device=hidden_states.device, + hidden_size=n_hidden, + hidden_dtype=dtype, + ) + ( permuted_hidden_states, _, @@ -173,6 +190,7 @@ def test_moe_permute_unpermute( n_expert=n_expert, n_local_expert=n_local_expert, expert_map=expert_map, + scratch=scratch, ) # check expert_first_token_offset @@ -212,20 +230,22 @@ def test_moe_permute_unpermute( torch.testing.assert_close(result4, gold4, atol=2e-2, rtol=0) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype): +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("n_token", [1, 33, 128]) +@pytest.mark.parametrize("topk", [1, 6, 8]) +def test_moe_permute_reuses_scratch_buffers( + dtype: torch.dtype, n_token: int, topk: int, workspace_init +): if not moe_permute_unpermute_supported(): pytest.skip("moe_permute_unpermute is not supported on this platform.") - n_token = 64 n_hidden = 2048 n_expert = 16 - topk = 4 hidden_states = torch.randn((n_token, n_hidden), device="cuda").to(dtype) gating_output = torch.randn((n_token, n_expert), device="cuda").to(dtype) _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False) - scratch = MoEPermuteScratch( + scratch_config = dict( max_num_tokens=n_token, topk=topk, num_experts=n_expert, @@ -234,6 +254,7 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype): hidden_size=n_hidden, hidden_dtype=hidden_states.dtype, ) + scratch = get_moe_permute_scratch(**scratch_config) first = moe_permute( hidden_states=hidden_states, @@ -242,13 +263,26 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype): n_expert=n_expert, scratch=scratch, ) - second = moe_permute( - hidden_states=hidden_states, - a1q_scale=None, - topk_ids=topk_ids, - n_expert=n_expert, - scratch=scratch, - ) + current_workspace_manager().lock() + assert get_moe_permute_scratch(**scratch_config) is scratch + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + second = moe_permute( + hidden_states=hidden_states, + a1q_scale=None, + topk_ids=topk_ids, + n_expert=n_expert, + scratch=get_moe_permute_scratch(**scratch_config), + ) + + for _ in range(2): + hidden_states.add_(1) + topk_ids.copy_(topk_ids.roll(1, 0)) + graph.replay() + expected = moe_permute(hidden_states, None, topk_ids, n_expert) + for actual, reference in zip(second, expected): + if actual is not None: + torch.testing.assert_close(actual, reference) ( permuted_hidden_states_1, @@ -288,6 +322,123 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype): ) +def test_moe_permute_scratch_reused_across_graph_sizes(workspace_init) -> None: + """Switching graph sizes must not expose metadata left by a larger batch.""" + if not moe_permute_unpermute_supported(): + pytest.skip("moe_permute_unpermute is not supported on this platform.") + + device = torch.device("cuda:0") + scratch = get_moe_permute_scratch( + max_num_tokens=128, + topk=6, + num_experts=16, + num_local_experts=16, + device=device, + hidden_size=2048, + hidden_dtype=torch.bfloat16, + ) + current_workspace_manager().lock() + runs = [] + for n_token in (128, 1, 33): + hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device) + topk_ids = torch.randint(16, (n_token, 6), device=device) + moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + runs.append((graph, hidden, topk_ids, output)) + + for graph, hidden, topk_ids, output in runs + runs[::-1]: + hidden.add_(1) + topk_ids.copy_((topk_ids + 1) % 16) + graph.replay() + expected = moe_permute(hidden, None, topk_ids, 16) + for actual, reference in zip(output, expected): + if actual is not None: + torch.testing.assert_close(actual, reference) + + +def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> None: + """Concurrent graph replays in different ubatches/lanes cannot share scratch.""" + if not moe_permute_unpermute_supported(): + pytest.skip("moe_permute_unpermute is not supported on this platform.") + + from vllm.v1.worker import workspace + + device = torch.device("cuda:0") + manager = workspace.WorkspaceManager(device, num_ubatches=2, num_lanes=2) + monkeypatch.setattr(workspace, "_manager", manager) + ubatch = 0 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: ubatch) + scratch_config = dict( + max_num_tokens=64, + topk=6, + num_experts=16, + num_local_experts=16, + device=device, + hidden_size=2048, + hidden_dtype=torch.bfloat16, + ) + runs = [] + for i, n_token in enumerate((1, 7, 33, 64)): + ubatch = i // 2 + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream): + hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device) + topk_ids = torch.randint(16, (n_token, 6), device=device) + scratch = get_moe_permute_scratch(**scratch_config) + moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + runs.append((stream, graph, hidden, topk_ids, output, scratch)) + torch.cuda.current_stream().wait_stream(stream) + + assert len({run[-1].permuted_hidden_states.data_ptr() for run in runs}) == 4 + manager.lock() + for _ in range(3): + for i, (stream, graph, hidden, topk_ids, _, scratch) in enumerate(runs): + ubatch = i // 2 + with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream): + assert get_moe_permute_scratch(**scratch_config) is scratch + hidden.add_(1) + topk_ids.copy_((topk_ids + 1) % 16) + graph.replay() + for stream, _, hidden, topk_ids, output, _ in runs: + torch.cuda.current_stream().wait_stream(stream) + expected = moe_permute(hidden, None, topk_ids, 16) + for actual, reference in zip(output, expected): + if actual is not None: + torch.testing.assert_close(actual, reference) + + +def test_moe_permute_scratch_without_manager(monkeypatch) -> None: + """Standalone calls get independent scratch with initialized row indices.""" + if not moe_permute_unpermute_supported(): + pytest.skip("moe_permute_unpermute is not supported on this platform.") + + from vllm.v1.worker import workspace + + monkeypatch.setattr(workspace, "_manager", None) + config = dict( + max_num_tokens=4, + topk=2, + num_experts=4, + num_local_experts=4, + device=torch.device("cuda"), + ) + first = get_moe_permute_scratch(**config) + second = get_moe_permute_scratch(**config) + + assert ( + first.token_expert_indices.data_ptr() != second.token_expert_indices.data_ptr() + ) + torch.testing.assert_close( + first.token_expert_indices, torch.arange(8, dtype=torch.int32, device="cuda") + ) + + def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: if not moe_permute_unpermute_supported(): pytest.skip("moe_permute_unpermute is not supported on this platform.") diff --git a/tests/kernels/quantization/test_marlin_gemm.py b/tests/kernels/quantization/test_marlin_gemm.py index 4c290dada6..29910390bb 100644 --- a/tests/kernels/quantization/test_marlin_gemm.py +++ b/tests/kernels/quantization/test_marlin_gemm.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import ( per_token_quant_int8, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + get_marlin_workspace, marlin_make_empty, marlin_make_workspace_new, marlin_permute_bias, @@ -43,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.platforms import current_platform from vllm.scalar_type import scalar_types +from vllm.v1.worker.workspace import current_workspace_manager if current_platform.is_rocm(): pytest.skip( @@ -477,6 +479,62 @@ def test_marlin_gemm( assert max_diff < 0.04 +def test_marlin_persistent_locks_cuda_graph_streams(workspace_init): + """Compiled GEMMs select stream-local locks at runtime and survive replay.""" + size_m, size_k, size_n = 32, 1024, 2048 + inputs = [rand_data((size_m, size_k)) for _ in range(2)] + weight = rand_data((size_k, size_n)) + quant_type = scalar_types.uint4b8 + w_ref, qweight, scales = marlin_quantize(weight, quant_type, 128) + zeros = marlin_make_empty(scales.device) + streams = [torch.cuda.Stream() for _ in inputs] + graphs = [torch.cuda.CUDAGraph() for _ in inputs] + locks = [] + outputs = [] + + @torch.compile(backend="eager", fullgraph=True) + def run(x): + return torch.ops.vllm.marlin_gemm( + x, + qweight, + None, + scales, + None, + None, + zeros, + None, + quant_type.id, + size_m, + size_n, + size_k, + use_atomic_add=False, + use_fp32_reduce=True, + ) + + for stream, graph, x in zip(streams, graphs, inputs): + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + locks.append(get_marlin_workspace(x.device)) + run(x) + with torch.cuda.graph(graph, stream=stream): + outputs.append(run(x)) + torch.cuda.current_stream().wait_stream(stream) + + assert locks[0].data_ptr() != locks[1].data_ptr() + current_workspace_manager().lock() + for _ in range(3): + for stream, graph, x, lock in zip(streams, graphs, inputs, locks): + with torch.cuda.stream(stream): + assert get_marlin_workspace(x.device) is lock + x.add_(0.1) + graph.replay() + for stream in streams: + torch.cuda.current_stream().wait_stream(stream) + for x, output, lock in zip(inputs, outputs, locks): + assert compute_max_diff(output, x @ w_ref) < 0.04 + assert torch.count_nonzero(lock) == 0 + + def test_marlin_gemm_subset_input(): quant_type = scalar_types.uint4b8 group_size = 128 diff --git a/tests/kernels/quantization/test_marlin_tile_padding.py b/tests/kernels/quantization/test_marlin_tile_padding.py index 027ffbf10c..d224b0a2b2 100644 --- a/tests/kernels/quantization/test_marlin_tile_padding.py +++ b/tests/kernels/quantization/test_marlin_tile_padding.py @@ -14,6 +14,7 @@ from vllm import _custom_ops as ops from vllm.model_executor.layers.quantization.utils.marlin_utils import ( GPTQ_MARLIN_TILE, apply_gptq_marlin_linear, + get_marlin_workspace, marlin_make_empty, marlin_make_workspace_new, marlin_moe_padded_intermediate, @@ -225,7 +226,7 @@ def test_fp8_marlin_padded_round_trip(shape, use_bias): input=x, weight=layer.weight, weight_scale=layer.weight_scale, - workspace=layer.workspace, + workspace=get_marlin_workspace(layer.weight.device), size_n=size_n, size_k=size_k, bias=layer.bias if use_bias else None, @@ -288,7 +289,7 @@ def test_nvfp4_marlin_padded_round_trip(shape): weight=layer.weight, weight_scale=layer.weight_scale, weight_global_scale=layer.weight_global_scale, - workspace=layer.workspace, + workspace=get_marlin_workspace(layer.weight.device), size_n=size_n, size_k=size_k, ) @@ -392,7 +393,7 @@ def test_fp8_block_marlin_padded_round_trip(shape): input=x, weight=layer.weight, weight_scale=layer.weight_scale_inv, - workspace=layer.workspace, + workspace=get_marlin_workspace(layer.weight.device), size_n=size_n, size_k=size_k, bias=None, @@ -440,7 +441,7 @@ def test_mxfp8_marlin_padded_round_trip(shape): input=x, weight=layer.weight, weight_scale=layer.weight_scale, - workspace=layer.workspace, + workspace=get_marlin_workspace(layer.weight.device), size_n=size_n, size_k=size_k, ) @@ -756,7 +757,7 @@ def test_fp8_marlin_moe_padded_round_trip(shape, quant): topk_ids, quant_type_id=scalar_types.float8_e4m3fn.id, global_num_experts=e, - workspace=layer.workspace, + workspace=get_marlin_workspace(device), ) with set_current_vllm_config(VllmConfig()): ref = torch_experts( @@ -833,7 +834,7 @@ def test_mxfp8_marlin_moe_padded_round_trip(shape): topk_ids, quant_type_id=scalar_types.float8_e4m3fn.id, global_num_experts=e, - workspace=layer.workspace, + workspace=get_marlin_workspace(device), ) with set_current_vllm_config(VllmConfig()): ref = torch_experts( diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index fde8b86cda..155957160f 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -330,8 +330,8 @@ def _load_marlin_checkpoint_format_weights(layer): ) -def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_init): - """Marlin must reuse the workspace storage captured by CUDA graphs.""" +def test_marlin_post_load_does_not_own_workspace(monkeypatch, dist_init): + """Weight reload must not create layer-owned Marlin lock storage.""" _stub_marlin_ops(monkeypatch) kernel = _make_marlin_kernel() @@ -339,19 +339,17 @@ def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_i _load_marlin_checkpoint_format_weights(layer) kernel.process_weights_after_loading(layer) - workspace_ptr = kernel.workspace.data_ptr() + assert not hasattr(kernel, "workspace") _load_marlin_checkpoint_format_weights(layer) kernel.process_weights_after_loading(layer) - assert kernel.workspace.data_ptr() == workspace_ptr - assert torch.all(kernel.workspace == 0) + assert not hasattr(kernel, "workspace") @pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"]) -def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant): - """The Marlin fallback prepare_* functions rerun on weight reload and must - reuse the workspace storage whose address captured CUDA graphs hold.""" +def test_marlin_prepare_layer_does_not_own_workspace(monkeypatch, variant): + """Weight preparation must not attach runtime workspace to model layers.""" from vllm import _custom_ops as ops from vllm.model_executor.layers.quantization.utils import ( marlin_utils, @@ -417,34 +415,57 @@ def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant): load_checkpoint_format_weights() prepare(layer) - workspace_ptr = layer.workspace.data_ptr() + assert not hasattr(layer, "workspace") # Reload: fresh checkpoint-format tensors, prepare runs again load_checkpoint_format_weights() prepare(layer) - assert layer.workspace.data_ptr() == workspace_ptr - assert torch.all(layer.workspace == 0) + assert not hasattr(layer, "workspace") -def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch): - """An incompatible existing workspace means the address captured by CUDA - graphs is already unusable; allocating a replacement would hide that.""" +def test_marlin_workspace_uses_persistent_workspace_manager(monkeypatch): + """Calls reuse initialized locks; independent streams get separate storage.""" from vllm.model_executor.layers.quantization.utils import marlin_utils + from vllm.utils import torch_utils + from vllm.v1.worker import workspace as workspace_module monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) device = torch.device("cpu") + manager = workspace_module.WorkspaceManager(device) + monkeypatch.setattr(workspace_module, "_manager", manager) + stream = "main" + monkeypatch.setattr(torch_utils, "current_stream", lambda: stream) - workspace = marlin_utils.marlin_make_workspace_new(device) - reused = marlin_utils.marlin_make_workspace_new(device, existing=workspace) - assert reused is workspace + workspace = marlin_utils.get_marlin_workspace(device) + assert workspace.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,) + assert workspace.dtype == torch.int32 + assert torch.count_nonzero(workspace) == 0 - with pytest.raises(ValueError, match="incompatible"): - marlin_utils.marlin_make_workspace_new(device, 4, existing=workspace) - with pytest.raises(ValueError, match="incompatible"): - marlin_utils.marlin_make_workspace_new( - device, existing=workspace.to(torch.int64) - ) + workspace.fill_(1) + assert marlin_utils.get_marlin_workspace(device) is workspace + assert torch.all(workspace == 1) + + stream = "aux" + aux_workspace = marlin_utils.get_marlin_workspace(device) + assert aux_workspace.data_ptr() != workspace.data_ptr() + assert torch.count_nonzero(aux_workspace) == 0 + + manager.lock() + assert marlin_utils.get_marlin_workspace(device) is aux_workspace + + +def test_marlin_workspace_without_manager(monkeypatch): + from vllm.model_executor.layers.quantization.utils import marlin_utils + from vllm.v1.worker import workspace as workspace_module + + monkeypatch.setattr(workspace_module, "_manager", None) + monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) + first = marlin_utils.get_marlin_workspace(torch.device("cpu")) + second = marlin_utils.get_marlin_workspace(torch.device("cpu")) + assert first.data_ptr() != second.data_ptr() + assert first.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,) + assert torch.count_nonzero(first) == 0 def test_model_cleanup(dist_init, default_vllm_config): diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index 7824c8000d..3d03f25e0e 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import cast +from unittest.mock import Mock import pytest import torch @@ -124,3 +125,65 @@ def test_workspace_lane_validation(monkeypatch) -> None: with pytest.raises(ValueError, match="at least one"): workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0) + + +def test_persistent_resources_are_lazy_and_isolated(monkeypatch) -> None: + """Only the requesting slot allocates; each ubatch/lane owns its resource.""" + active_ubatch = 0 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + factory = Mock(side_effect=lambda: {"buffer": torch.zeros(8)}) + resources = [] + for ubatch in range(2): + active_ubatch = ubatch + for lane in range(2): + with workspace.use_workspace_lane(lane): + resource = manager.get_persistent_resource("scratch", factory) + resources.append(resource) + resource["buffer"].fill_(len(resources)) + assert manager.get_persistent_resource("scratch", factory) is resource + assert factory.call_count == len(resources) + + assert len({r["buffer"].data_ptr() for r in resources}) == 4 + for i, resource in enumerate(resources, 1): + assert torch.all(resource["buffer"] == i) + + +def test_persistent_resource_lock_applies_to_each_slot(monkeypatch) -> None: + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=2) + factory = Mock(side_effect=object) + first = manager.get_persistent_resource("scratch", factory) + manager.lock() + assert manager.get_persistent_resource("scratch", factory) is first + with pytest.raises(AssertionError, match="was not allocated during warmup"): + manager.get_persistent_resource("new", factory) + with workspace.use_workspace_lane(1): + with pytest.raises(AssertionError, match="was not allocated during warmup"): + manager.get_persistent_resource("scratch", factory) + assert factory.call_count == 1 + manager.unlock() + assert manager.get_persistent_resource("scratch", factory) is not first + assert manager.get_persistent_resource("scratch", factory) is first + + +def test_persistent_tensor_preserves_contents_and_rejects_changes() -> None: + """Initialization happens once and transient scratch cannot overwrite locks.""" + manager = workspace.WorkspaceManager(torch.device("cpu")) + first = manager.get_persistent("locks", (8,), torch.int32, zero_init=True) + assert torch.count_nonzero(first) == 0 + first.fill_(1) + (scratch,) = manager.get_simultaneous(((1024,), torch.uint8)) + scratch.zero_() + assert manager.get_persistent("locks", (8,), torch.int32, zero_init=True) is first + assert torch.all(first == 1) + with pytest.raises(ValueError, match="requested shape"): + manager.get_persistent("locks", (16,), torch.int32) + with pytest.raises(ValueError, match="requested shape"): + manager.get_persistent("locks", (8,), torch.int64) + manager.lock() + assert manager.get_persistent("locks", (8,), torch.int32) is first + with pytest.raises(AssertionError, match="was not allocated during warmup"): + manager.get_persistent("new", (8,), torch.int32) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py index 387c33c110..60687fc860 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py @@ -10,7 +10,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( apply_gptq_marlin_linear, marlin_act_int8_process_scales, marlin_make_empty, - marlin_make_workspace_new, marlin_pad_dim, marlin_pad_qweight, marlin_pad_scales, @@ -91,11 +90,6 @@ class MarlinLinearKernel(MPLinearKernel): size_k, size_n = c.partition_weight_shape padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size) - # Allocate marlin workspace, reusing existing storage on reload. - self.workspace = marlin_make_workspace_new( - device, existing=getattr(self, "workspace", None) - ) - # Default name since marlin requires empty parameter for zp, # TODO: remove this requirement from marlin (allow optional tensors) if self.w_zp_name is None: @@ -197,7 +191,7 @@ class MarlinLinearKernel(MPLinearKernel): weight=w_q, weight_scale=w_s, weight_zp=w_zp, # type: ignore - workspace=self.workspace, + workspace=None, wtype=c.weight_type, input_size_per_partition=c.partition_weight_shape[0], output_size_per_partition=c.partition_weight_shape[1], diff --git a/vllm/model_executor/kernels/linear/mxfp4/marlin.py b/vllm/model_executor/kernels/linear/mxfp4/marlin.py index 30ff0d58d4..108f835e64 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/marlin.py +++ b/vllm/model_executor/kernels/linear/mxfp4/marlin.py @@ -58,7 +58,7 @@ class MarlinMxFp4LinearKernel(MxFp4LinearKernel): weight=layer.weight, weight_scale=layer.weight_scale, weight_global_scale=None, - workspace=layer.workspace, + workspace=None, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, bias=bias, diff --git a/vllm/model_executor/kernels/linear/mxfp8/marlin.py b/vllm/model_executor/kernels/linear/mxfp8/marlin.py index bec54cd942..6ebdbd6474 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/marlin.py +++ b/vllm/model_executor/kernels/linear/mxfp8/marlin.py @@ -46,7 +46,7 @@ class MarlinMxfp8LinearKernel(Mxfp8LinearKernel): input=x, weight=layer.weight, weight_scale=layer.weight_scale, - workspace=layer.workspace, + workspace=None, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, bias=bias, diff --git a/vllm/model_executor/kernels/linear/nvfp4/marlin.py b/vllm/model_executor/kernels/linear/nvfp4/marlin.py index a05d6823c8..1ab5336876 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/marlin.py +++ b/vllm/model_executor/kernels/linear/nvfp4/marlin.py @@ -50,7 +50,7 @@ class MarlinNvFp4LinearKernel(NvFp4LinearKernel): weight=layer.weight, weight_scale=layer.weight_scale, weight_global_scale=layer.weight_global_scale, - workspace=layer.workspace, + workspace=None, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, bias=bias, diff --git a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py index 3ca89d3cea..1df6ad5541 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py @@ -93,7 +93,7 @@ class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): input=x, weight=layer.weight, weight_scale=weight_scale, - workspace=layer.workspace, + workspace=None, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, input_dtype=self.marlin_input_dtype, diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 52b6d4e753..18a2fb5cc6 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -27,6 +27,7 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( from vllm.model_executor.layers.fused_moe.moe_fused_mul_sum import moe_fused_mul_sum from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( MoEPermuteScratch, + get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, moe_unpermute, @@ -176,7 +177,6 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): max_num_tokens=max_num_tokens, num_dispatchers=num_dispatchers, ) - self._permute_scratch: dict[int, MoEPermuteScratch] = {} def init_humming_moe(self): from vllm.utils.humming import get_heuristics_config @@ -259,24 +259,20 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): if not moe_permute_unpermute_supported(): return None - scratch = self._permute_scratch.get(topk) - if scratch is None: - max_expanded_rows = ( - self.moe_config.max_num_tokens - * self.moe_config.dp_size - * self.moe_config.experts_per_token - ) - scratch = MoEPermuteScratch( - max_num_tokens=math.ceil(max_expanded_rows / topk), - topk=topk, - num_experts=self.moe_config.num_experts, - num_local_experts=self.moe_config.num_local_experts, - device=torch.device(self.moe_config.device), - hidden_size=self.moe_config.hidden_dim, - hidden_dtype=self.moe_config.in_dtype, - ) - self._permute_scratch[topk] = scratch - return scratch + max_expanded_rows = ( + self.moe_config.max_num_tokens + * self.moe_config.dp_size + * self.moe_config.experts_per_token + ) + return get_moe_permute_scratch( + max_num_tokens=math.ceil(max_expanded_rows / topk), + topk=topk, + num_experts=self.moe_config.num_experts, + num_local_experts=self.moe_config.num_local_experts, + device=torch.device(self.moe_config.device), + hidden_size=self.moe_config.hidden_dim, + hidden_dtype=self.moe_config.in_dtype, + ) def get_global_valid_shape_m(self, topk_ids: torch.Tensor): ctx = get_forward_context() diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index a6a70ca952..cfd78de28d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -35,7 +35,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( from vllm.model_executor.layers.fused_moe.utils import _resize_cache from vllm.model_executor.layers.quantization.utils.marlin_utils import ( get_marlin_input_dtype, - marlin_make_workspace_new, + get_marlin_workspace, marlin_moe_intermediate_size, marlin_quant_input, ) @@ -95,7 +95,7 @@ def _fused_marlin_moe( N = marlin_moe_intermediate_size(w1, w2) w13_num_shards = 2 if activation.is_gated else 1 if workspace is None: - workspace = marlin_make_workspace_new(hidden_states.device, 4) + workspace = get_marlin_workspace(hidden_states.device) if intermediate_cache13 is None: intermediate_cache13 = torch.empty( diff --git a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py index 7a53a960e8..0aa9a38071 100644 --- a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py +++ b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py @@ -5,6 +5,11 @@ from dataclasses import dataclass, field import torch +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) + @dataclass class MoEPermuteScratch: @@ -102,6 +107,48 @@ class MoEPermuteScratch: return topk_ids_int32 +def get_moe_permute_scratch( + *, + max_num_tokens: int, + topk: int, + num_experts: int, + num_local_experts: int, + device: torch.device, + hidden_size: int | None = None, + hidden_dtype: torch.dtype | None = None, +) -> MoEPermuteScratch: + """Share scratch across sequential layers in the current ubatch and lane. + + Without a workspace manager, allocate new scratch for each call. + """ + + def create_scratch() -> MoEPermuteScratch: + return MoEPermuteScratch( + max_num_tokens=max_num_tokens, + topk=topk, + num_experts=num_experts, + num_local_experts=num_local_experts, + device=device, + hidden_size=hidden_size, + hidden_dtype=hidden_dtype, + ) + + if not is_workspace_manager_initialized(): + return create_scratch() + + key = ( + MoEPermuteScratch, + max_num_tokens, + topk, + num_experts, + num_local_experts, + device, + hidden_size, + hidden_dtype, + ) + return current_workspace_manager().get_persistent_resource(key, create_scratch) + + def moe_permute( hidden_states: torch.Tensor, a1q_scale: torch.Tensor | None, diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 71b858e919..92132f3ffb 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -48,7 +48,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( check_marlin_supports_layer, check_moe_marlin_supports_layer, get_marlin_input_dtype, - marlin_make_workspace_new, verify_marlin_supported, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -642,9 +641,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): layer.register_parameter("w2_qzeros", w2_qzeros) set_weight_attrs(w2_qzeros, extra_weight_attrs) - device = layer.w13_qweight.device - layer.workspace = marlin_make_workspace_new(device, 4) - def process_weights_after_loading(self, layer: RoutedExperts) -> None: converted = convert_to_wna16_moe_kernel_format( backend=self.wna16_moe_backend, diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index e4478e897d..ffbc6fba22 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -16,7 +16,6 @@ from vllm.model_executor.kernels.linear import ( ) from vllm.model_executor.layers.fused_moe import ( FusedMoEConfig, - FusedMoEExpertsModular, FusedMoEMethodBase, FusedMoEQuantConfig, FusedMoeWeightScaleSupported, @@ -46,7 +45,6 @@ from vllm.model_executor.layers.quantization.utils.gptq_utils import ( from vllm.model_executor.layers.quantization.utils.marlin_utils import ( check_moe_marlin_supports_layer, get_marlin_input_dtype, - marlin_make_workspace_new, marlin_repeat_scales_on_all_ranks, verify_marlin_supported, ) @@ -601,12 +599,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): layer.register_parameter("w2_bias", w2_bias) set_weight_attrs(w2_bias, extra_weight_attrs) - if self.experts_cls is not None and issubclass( - self.experts_cls, FusedMoEExpertsModular - ): - device = layer.w13_qweight.device - layer.workspace = marlin_make_workspace_new(device, 4) - def process_weights_after_loading(self, layer: RoutedExperts) -> None: def replace_or_register(name: str, val: torch.Tensor | None): if val is None: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index bcbb687e3b..feddb77cff 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -11,7 +11,6 @@ from compressed_tensors.quantization import ( from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( - FusedMoEExpertsModular, RoutedExperts, SharedExperts, ) @@ -36,7 +35,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compress from vllm.model_executor.layers.quantization.utils.marlin_utils import ( check_moe_marlin_supports_config, get_marlin_input_dtype, - marlin_make_workspace_new, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, @@ -466,18 +464,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): torch.nn.Parameter(w2_input_global_scale, requires_grad=False), ) - # Marlin workspace — only needed for Marlin-family backends, not emulation. - if ( - self.experts_cls is not None - and issubclass(self.experts_cls, FusedMoEExpertsModular) - and self.wna16_backend != WNA16MoEBackend.EMULATION - ): - layer.workspace = marlin_make_workspace_new( - layer.w13_weight_packed.device, - 4, - existing=getattr(layer, "workspace", None), - ) - # Alias packed weights to w13_weight/w2_weight for the modular kernel interface layer.w13_weight = layer.w13_weight_packed layer.w2_weight = layer.w2_weight_packed diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index 82201dff3b..1ee52a1780 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -8,7 +8,6 @@ import numpy import torch import vllm.envs as envs -from vllm import _custom_ops as ops from vllm.distributed.utils import verify_group_size_divides_partition from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import RoutedExperts @@ -23,6 +22,7 @@ from vllm.platforms import current_platform from vllm.scalar_type import ScalarType, scalar_types from vllm.utils.math_utils import round_up from vllm.utils.platform_utils import num_compute_units +from vllm.utils.torch_utils import direct_register_custom_op from .quant_utils import pack_cols, unpack_cols @@ -32,6 +32,7 @@ GPTQ_MARLIN_TILE = 16 GPTQ_MARLIN_MIN_THREAD_N = 64 GPTQ_MARLIN_MIN_THREAD_K = 128 GPTQ_MARLIN_MAX_PARALLEL = 16 +MARLIN_MAX_BLOCKS_PER_SM = 4 MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] @@ -407,29 +408,96 @@ def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tenso def marlin_make_workspace_new( device: torch.device, max_blocks_per_sm: int = 1, - existing: torch.Tensor | None = None, ) -> torch.Tensor: # In the new marlin kernel, we use the num of threadblocks as workspace # size. The num of threadblocks is sms_count * max_blocks_per_sm. sms = num_compute_units(device.index) - size = sms * max_blocks_per_sm - # On weight reload, reuse the existing storage so the workspace address - # captured by CUDA graphs stays valid. - if existing is not None: - if ( - existing.device != device - or existing.dtype != torch.int - or existing.numel() != size - ): - raise ValueError( - f"Existing Marlin workspace is incompatible " - f"(device={existing.device}, dtype={existing.dtype}, " - f"numel={existing.numel()}; expected device={device}, " - f"dtype={torch.int}, numel={size}). Reload must reuse the " - f"workspace storage captured by CUDA graphs." - ) - return existing.zero_() - return torch.zeros(size, dtype=torch.int, device=device, requires_grad=False) + return torch.zeros(sms * max_blocks_per_sm, dtype=torch.int, device=device) + + +def get_marlin_workspace(device: torch.device) -> torch.Tensor: + """Get zero-initialized locks for the current stream, ubatch, and lane. + + Marlin resets its locks after each call, so ordered calls can reuse them. + """ + from vllm.utils.torch_utils import current_stream + from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, + ) + + if not is_workspace_manager_initialized(): + return marlin_make_workspace_new(device, MARLIN_MAX_BLOCKS_PER_SM) + + return current_workspace_manager().get_persistent( + ("marlin_locks", current_stream()), + (num_compute_units(device.index) * MARLIN_MAX_BLOCKS_PER_SM,), + torch.int, + zero_init=True, + ) + + +def _marlin_gemm( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_zeros: torch.Tensor | None, + workspace: torch.Tensor | None, + b_q_type_id: int, + size_m: int, + size_n: int, + size_k: int, + use_atomic_add: bool = False, + use_fp32_reduce: bool = False, + is_zp_float: bool = False, +) -> torch.Tensor: + if workspace is None: + workspace = get_marlin_workspace(a.device) + return torch.ops._C.marlin_gemm( + a, + None, + b_q_weight, + b_bias, + b_scales, + a_scales, + global_scale, + b_zeros, + workspace, + b_q_type_id, + size_m, + size_n, + size_k, + use_atomic_add, + use_fp32_reduce, + is_zp_float, + ) + + +def _marlin_gemm_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_zeros: torch.Tensor | None, + workspace: torch.Tensor | None, + b_q_type_id: int, + size_m: int, + size_n: int, + size_k: int, + use_atomic_add: bool = False, + use_fp32_reduce: bool = False, + is_zp_float: bool = False, +) -> torch.Tensor: + dtype = a.dtype if a.dtype in (torch.float16, torch.bfloat16) else b_scales.dtype + return torch.empty((size_m, size_n), device=a.device, dtype=dtype) + + +direct_register_custom_op("marlin_gemm", _marlin_gemm, fake_impl=_marlin_gemm_fake) def marlin_repeat_scales_on_all_ranks(group_size: int, is_row_parallel: bool) -> bool: @@ -686,7 +754,7 @@ def apply_gptq_marlin_linear( weight: torch.Tensor, weight_scale: torch.Tensor, weight_zp: torch.Tensor, - workspace: torch.Tensor, + workspace: torch.Tensor | None, wtype: ScalarType, output_size_per_partition: int, input_size_per_partition: int, @@ -723,9 +791,8 @@ def apply_gptq_marlin_linear( reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype) - output = ops.marlin_gemm( + output = torch.ops.vllm.marlin_gemm( reshaped_x, - None, weight, bias, weight_scale, @@ -733,7 +800,7 @@ def apply_gptq_marlin_linear( None, weight_zp, workspace, - wtype, + wtype.id, size_m=reshaped_x.shape[0], size_n=padded_n, size_k=padded_k, diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index a12d8b87e9..7fb5621f2f 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -10,7 +10,6 @@ from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts from vllm.model_executor.layers.quantization.utils.marlin_utils import ( USE_FP32_REDUCE_DEFAULT, get_marlin_input_dtype, - marlin_make_workspace_new, marlin_pad_dim, marlin_pad_qweight, marlin_pad_scales, @@ -161,7 +160,7 @@ def apply_fp4_marlin_linear( weight: torch.Tensor, weight_scale: torch.Tensor, weight_global_scale: torch.Tensor | None, - workspace: torch.Tensor, + workspace: torch.Tensor | None, size_n: int, size_k: int, bias: torch.Tensor | None = None, @@ -196,9 +195,8 @@ def apply_fp4_marlin_linear( inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn) - output = ops.marlin_gemm( + output = torch.ops.vllm.marlin_gemm( a=inputs, - c=None, b_q_weight=weight, b_bias=bias, b_scales=weight_scale, @@ -206,7 +204,7 @@ def apply_fp4_marlin_linear( global_scale=weight_global_scale, b_zeros=None, workspace=workspace, - b_q_type=scalar_types.float4_e2m1f, + b_q_type_id=scalar_types.float4_e2m1f.id, size_m=reshaped_x.size(0), size_n=padded_n, size_k=padded_k, @@ -237,13 +235,6 @@ def prepare_fp4_layer_for_marlin( assert layer.weight.shape == (part_size_n, part_size_k // 2) - device = layer.weight.device - - # WORKSPACE - layer.workspace = marlin_make_workspace_new( - device, existing=getattr(layer, "workspace", None) - ) - # WEIGHT # Repack weights to marlin format qweight = layer.weight.view(torch.int32).T.contiguous() @@ -389,15 +380,9 @@ def prepare_nvfp4_moe_layer_for_marlin( return x return torch.nn.functional.pad(x, (0, (padded_N - N) // packing)) - device = w13.device param_dtype = layer.params_dtype is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 - # WORKSPACE - layer.workspace = marlin_make_workspace_new( - device, 4, existing=getattr(layer, "workspace", None) - ) - # WEIGHT # Repack weights to marlin format def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: @@ -481,12 +466,7 @@ def prepare_moe_fp4_layer_for_marlin( k = layer.moe_config.hidden_dim n = layer.moe_config.intermediate_size_per_partition - # WORKSPACE - device = layer.w13_weight.device param_dtype = layer.params_dtype - layer.workspace = marlin_make_workspace_new( - device, 4, existing=getattr(layer, "workspace", None) - ) is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 # WEIGHT diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py index 4a5f5ee59a..257c5bf445 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py @@ -9,7 +9,6 @@ from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.marlin_utils import ( USE_FP32_REDUCE_DEFAULT, get_marlin_input_dtype, - marlin_make_workspace_new, marlin_moe_padded_intermediate, marlin_pad_dim, marlin_pad_qweight, @@ -50,7 +49,7 @@ def apply_fp8_marlin_linear( input: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, - workspace: torch.Tensor, + workspace: torch.Tensor | None, size_n: int, size_k: int, bias: torch.Tensor | None, @@ -80,9 +79,8 @@ def apply_fp8_marlin_linear( # inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn) raise RuntimeError("Marlin W8A8 is not supported.") - output = ops.marlin_gemm( + output = torch.ops.vllm.marlin_gemm( a=inputs, - c=None, b_q_weight=weight, b_bias=bias, b_scales=weight_scale, @@ -90,7 +88,7 @@ def apply_fp8_marlin_linear( global_scale=None, b_zeros=None, workspace=workspace, - b_q_type=scalar_types.float8_e4m3fn, + b_q_type_id=scalar_types.float8_e4m3fn.id, size_m=reshaped_x.size(0), size_n=padded_n, size_k=padded_k, @@ -127,13 +125,6 @@ def prepare_fp8_layer_for_marlin( else: assert layer.weight.shape == (part_size_n, part_size_k) - device = layer.weight.device - - # WORKSPACE - layer.workspace = marlin_make_workspace_new( - device, existing=getattr(layer, "workspace", None) - ) - # WEIGHT # Repack weights to marlin format qweight = pack_fp8_to_int32(layer.weight, size_k_first) @@ -241,12 +232,7 @@ def prepare_fp8_moe_layer_for_marlin( w13_weight_scale: torch.Tensor, w2_weight_scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Shuffle weights and scales into marlin format. - - Note that this function has the side effect of adding a `workspace` - attribute to the layer. This `workspace` does not need to be - registered as a Parameter as it is not used during weight reloading. - """ + """Shuffle weights and scales into marlin format.""" logger.warning_once( "Your GPU does not have native support for FP8 computation but " "FP8 quantization is being used. Weight-only FP8 compression will " @@ -272,12 +258,6 @@ def prepare_fp8_moe_layer_for_marlin( w13_weight = _moe_pad_shard_rows(w13_weight, n, padded_n) w2_weight = _moe_pad_last(w2_weight, n, padded_n) - # WORKSPACE - device = layer.w13_weight.device - layer.workspace = marlin_make_workspace_new( - device, 4, existing=getattr(layer, "workspace", None) - ) - # WEIGHT # Repack weights to marlin format def repack_weight(name: str, weight: torch.Tensor) -> torch.Tensor: @@ -396,7 +376,7 @@ def apply_mxfp8_marlin_linear( input: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, - workspace: torch.Tensor, + workspace: torch.Tensor | None, size_n: int, size_k: int, bias: torch.Tensor | None = None, @@ -416,9 +396,8 @@ def apply_mxfp8_marlin_linear( dtype=input.dtype, ) - output = ops.marlin_gemm( + output = torch.ops.vllm.marlin_gemm( a=reshaped_x, - c=None, b_q_weight=weight, b_bias=bias, b_scales=weight_scale, @@ -426,7 +405,7 @@ def apply_mxfp8_marlin_linear( global_scale=None, b_zeros=None, workspace=workspace, - b_q_type=scalar_types.float8_e4m3fn, + b_q_type_id=scalar_types.float8_e4m3fn.id, size_m=reshaped_x.size(0), size_n=padded_n, size_k=padded_k, @@ -451,13 +430,6 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None: group_size = 32 # MX standard block size padded_n, padded_k = marlin_padded_nk(part_size_n, part_size_k, group_size) - device = layer.weight.device - - # WORKSPACE - layer.workspace = marlin_make_workspace_new( - device, existing=getattr(layer, "workspace", None) - ) - # WEIGHT - repack FP8 weights to Marlin format qweight = pack_fp8_to_int32(layer.weight, size_k_first=False) qweight = qweight.T.contiguous() @@ -512,7 +484,7 @@ def prepare_mxfp8_moe_layer_for_marlin( """Repack MXFP8 MoE weights and scales into Marlin kernel format. Args: - layer: MoE layer (used to read params_dtype and attach workspace). + layer: MoE layer. w13: [E, 2*N, K] float8_e4m3fn weights. w2: [E, K, N] float8_e4m3fn weights. w13_scale: [E, 2*N, K//32] uint8 e8m0 scales. @@ -537,13 +509,8 @@ def prepare_mxfp8_moe_layer_for_marlin( n = padded_n w13_n = w13.shape[1] - device = w13.device param_dtype = torch.get_default_dtype() - layer.workspace = marlin_make_workspace_new( - device, 4, existing=getattr(layer, "workspace", None) - ) - def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: if "w13" in name: size_n, size_k = w13_n, k diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index cddb760acd..a86110f2f8 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -3,11 +3,12 @@ import inspect import os -from collections.abc import Iterator +from collections.abc import Callable, Hashable, Iterator from contextlib import contextmanager from contextvars import ContextVar from itertools import accumulate from math import prod +from typing import TypeVar, cast import torch @@ -17,6 +18,7 @@ from vllm.utils.math_utils import round_up from vllm.v1.worker.ubatching import dbo_current_ubatch_id logger = init_logger(__name__) +T = TypeVar("T") def _compute_bytes(shape: tuple[int, ...], dtype: torch.dtype) -> int: @@ -47,8 +49,8 @@ def use_workspace_lane(lane: int) -> Iterator[None]: class WorkspaceManager: """Manager for workspace allocation. - Manages one workspace buffer per active ``(ubatch, lane)`` slot. - Can be locked to prevent further growth during execution. + Manages shared scratch and keyed persistent resources per ``(ubatch, lane)``. + Can be locked to prevent growth or new resources during execution. """ def __init__( @@ -66,6 +68,9 @@ class WorkspaceManager: self._current_workspaces: list[torch.Tensor | None] = [None] * ( self._num_ubatches * self._num_lanes ) + self._persistent_resources: list[dict[Hashable, object]] = [ + {} for _ in self._current_workspaces + ] self._locked: bool = False @staticmethod @@ -76,10 +81,10 @@ class WorkspaceManager: return workspace.numel() * workspace.element_size() def lock(self) -> None: - """Lock the workspace to prevent further growth. + """Lock the workspace to prevent growth or new persistent resources. - After locking, any attempt to allocate a larger workspace will raise - an assertion error. This ensures workspace size is fixed during execution. + Larger scratch requests and new persistent resources raise an assertion + error. Existing resources retain their storage throughout execution. """ self._locked = True if envs.VLLM_DEBUG_WORKSPACE: @@ -113,6 +118,63 @@ class WorkspaceManager: """Check if workspace is locked.""" return self._locked + def _get_workspace_id(self) -> int: + lane = _workspace_lane.get() + if lane >= self._num_lanes: + raise RuntimeError( + f"Workspace lane {lane} is not configured; manager has " + f"{self._num_lanes} lane(s)." + ) + return dbo_current_ubatch_id() * self._num_lanes + lane + + def get_persistent_resource(self, key: Hashable, factory: Callable[[], T]) -> T: + """Return the resource cached for ``key`` in the current ubatch and lane. + + The first request calls ``factory`` and stores its result for the lifetime + of this manager. Later requests return the same object without calling + the factory again. Each ubatch/lane pair has a separate cache. Once the + manager is locked, requesting a key missing from that cache raises an + AssertionError. + + Keys must identify the resource type and configuration. Callers must not + resize or replace tensor storage held by a cached resource. Uses of the + same resource must not overlap; include the CUDA stream in the key when + streams can execute concurrently. + + With torch.compile, call this inside a runtime custom op so resource + selection happens during execution. + """ + resources = self._persistent_resources[self._get_workspace_id()] + if key not in resources: + if self._locked: + raise AssertionError( + f"Workspace is locked but {key!r} was not allocated during " + "warmup for the current ubatch and lane." + ) + resources[key] = factory() + return cast(T, resources[key]) + + def get_persistent( + self, + key: Hashable, + shape: tuple[int, ...], + dtype: torch.dtype, + *, + zero_init: bool = False, + ) -> torch.Tensor: + """Get a fixed-shape tensor, optionally zeroed on first allocation only.""" + factory = torch.zeros if zero_init else torch.empty + workspace = self.get_persistent_resource( + (torch.Tensor, key), + lambda: factory(shape, dtype=dtype, device=self._device), + ) + if workspace.shape != shape or workspace.dtype != dtype: + raise ValueError( + f"Persistent workspace {key!r} has shape={workspace.shape}, " + f"dtype={workspace.dtype}; requested shape={shape}, dtype={dtype}." + ) + return workspace + def get_simultaneous( self, *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype] ) -> list[torch.Tensor]: @@ -151,14 +213,7 @@ class WorkspaceManager: The current workspace tensor. """ - ubatch_id = dbo_current_ubatch_id() - lane = _workspace_lane.get() - if lane >= self._num_lanes: - raise RuntimeError( - f"Workspace lane {lane} is not configured; manager has " - f"{self._num_lanes} lane(s)." - ) - workspace_id = ubatch_id * self._num_lanes + lane + workspace_id = self._get_workspace_id() current_workspace = self._current_workspaces[workspace_id] current_size = self._workspace_size_bytes(current_workspace) @@ -218,8 +273,8 @@ class WorkspaceManager: get_caller_info(), current_size / _MB, required_bytes / _MB, - ubatch_id, - lane, + workspace_id // self._num_lanes, + workspace_id % self._num_lanes, ) return current_workspace @@ -277,12 +332,11 @@ def init_workspace_manager( def lock_workspace() -> None: - """Lock the workspace to prevent further growth. + """Lock the workspace to prevent growth or new persistent resources. - After calling this function, any attempt to allocate a workspace larger - than the current size will raise an AssertionError. This ensures that - workspace size is fixed during execution and prevents unexpected memory - allocations in the hot path. + Larger scratch requests and new persistent resources raise an AssertionError. + All required resources must be initialized in each execution slot during + warmup, so their storage remains fixed during execution. Example: # During initialization