iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Core][Kernel] Share persistent workspaces for Marlin and Humming (#57421)

Browse Source 
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: ErinYin <liang.yin@daocloud.io>
This commit is contained in:
Misha Goin
2026-09-20 17:28:33 +08:00
committed by GitHub
co-authored by ErinYin
parent 0e110f696d
commit e5fce7b56b
21 changed files with 598 additions and 211 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/test_moe.py
+17 -3
View File
@@ -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
tests/kernels/moe/test_moe_permute_unpermute.py
+163 -12
View File
@@ -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.")
tests/kernels/quantization/test_marlin_gemm.py
+58
View File
@@ -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
tests/kernels/quantization/test_marlin_tile_padding.py
+7 -6
View File
@@ -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(
tests/model_executor/model_loader/test_reload.py
+44 -23
View File
@@ -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):
tests/v1/worker/test_workspace.py
+63
View File
@@ -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)
vllm/model_executor/kernels/linear/mixed_precision/marlin.py
+1 -7
View File
@@ -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],
vllm/model_executor/kernels/linear/mxfp4/marlin.py
+1 -1
View File
@@ -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,
vllm/model_executor/kernels/linear/mxfp8/marlin.py
+1 -1
View File
@@ -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,
vllm/model_executor/kernels/linear/nvfp4/marlin.py
+1 -1
View File
@@ -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,
vllm/model_executor/kernels/linear/scaled_mm/marlin.py
+1 -1
View File
@@ -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,
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+15 -19
View File
@@ -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()
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+2 -2
View File
@@ -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(
vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
+47
View File
@@ -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,
vllm/model_executor/layers/quantization/auto_awq.py
-4
View File
@@ -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,
vllm/model_executor/layers/quantization/auto_gptq.py
-8
View File
@@ -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:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
-14
View File
@@ -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
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+91 -24
View File
@@ -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,
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+3 -23
View File
@@ -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
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+8 -41
View File
@@ -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
vllm/v1/worker/workspace.py
+75 -21
View File
@@ -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