diff --git a/benchmarks/kernels/benchmark_moe_permute_unpermute.py b/benchmarks/kernels/benchmark_moe_permute_unpermute.py index 6867274ebe28..f5aa21fb4bd5 100644 --- a/benchmarks/kernels/benchmark_moe_permute_unpermute.py +++ b/benchmarks/kernels/benchmark_moe_permute_unpermute.py @@ -56,13 +56,9 @@ def benchmark_permute( qhidden_states, input_gating, topk, False ) scratch = MoEPermuteScratch( - max_num_tokens=num_tokens, - topk=topk, num_experts=num_experts, num_local_experts=num_experts, device=qhidden_states.device, - hidden_size=hidden_size, - hidden_dtype=qhidden_states.dtype, ) def prepare(i: int): @@ -135,13 +131,9 @@ def benchmark_unpermute( qhidden_states, input_gating, topk, False ) scratch = MoEPermuteScratch( - max_num_tokens=num_tokens, - topk=topk, num_experts=num_experts, num_local_experts=num_experts, device=qhidden_states.device, - hidden_size=hidden_size, - hidden_dtype=qhidden_states.dtype, ) def prepare(): diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 8d6deac4f025..0f786deb5498 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -1437,6 +1437,7 @@ def test_humming_global_valid_shape_m( def test_humming_permute_scratch_is_shared_by_config( monkeypatch: pytest.MonkeyPatch, ): + """Layers share one scratch whose key is independent of batch shape.""" from types import SimpleNamespace from unittest.mock import Mock @@ -1447,34 +1448,19 @@ def test_humming_permute_scratch_is_shared_by_config( 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]) + scratch = Mock() + scratch_type = Mock(side_effect=[scratch]) monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True) 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(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 - - assert scratch_type.call_count == 2 - first_call, second_call = scratch_type.call_args_list - assert first_call.kwargs["max_num_tokens"] == 1024 - assert first_call.kwargs["topk"] == 6 - assert second_call.kwargs["max_num_tokens"] == 6144 - assert second_call.kwargs["topk"] == 1 - + assert humming.HummingExpertsBase._get_permute_scratch(experts) is scratch + assert humming.HummingExpertsBase._get_permute_scratch(other_layer) is scratch 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) + assert humming.HummingExpertsBase._get_permute_scratch(experts) is scratch + assert scratch_type.call_count == 1 def test_humming_delegates_to_instance_activation(): @@ -1531,7 +1517,7 @@ def test_humming_grouped_apply_forwards_valid_prefix( num_experts=2, estimate_local_valid_shape_m=lambda _: 6, prepare_buffers=lambda *_: buffers, - _get_permute_scratch=lambda _, *, indices_only=False: object(), + _get_permute_scratch=lambda: object(), process_input=Mock( side_effect=lambda _, **kwargs: ( kwargs["inputs"], diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 76f4b57cf07f..0ab5760c3590 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -169,13 +169,9 @@ def test_moe_permute_unpermute( 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, ) ( @@ -247,13 +243,9 @@ def test_moe_permute_reuses_scratch_buffers( _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False) scratch_config = dict( - max_num_tokens=n_token, - topk=topk, num_experts=n_expert, num_local_experts=n_expert, device=hidden_states.device, - hidden_size=n_hidden, - hidden_dtype=hidden_states.dtype, ) scratch = get_moe_permute_scratch(**scratch_config) @@ -300,15 +292,10 @@ def test_moe_permute_reuses_scratch_buffers( permuted_idx_2, ) = second - torch.testing.assert_close(permuted_hidden_states_1, permuted_hidden_states_2) torch.testing.assert_close(expert_first_token_offset_1, expert_first_token_offset_2) torch.testing.assert_close(inv_permuted_idx_1, inv_permuted_idx_2) torch.testing.assert_close(permuted_idx_1, permuted_idx_2) - assert ( - permuted_hidden_states_1.untyped_storage().data_ptr() - == permuted_hidden_states_2.untyped_storage().data_ptr() - ) assert ( expert_first_token_offset_1.untyped_storage().data_ptr() == expert_first_token_offset_2.untyped_storage().data_ptr() @@ -330,14 +317,11 @@ def test_moe_permute_scratch_reused_across_graph_sizes(workspace_init) -> None: 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, ) + scratch.reserve(128 * 6) current_workspace_manager().lock() runs = [] for n_token in (128, 1, 33): @@ -372,13 +356,9 @@ def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> Non 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)): @@ -396,7 +376,7 @@ def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> Non 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 + assert len({run[-1].permuted_idx.data_ptr() for run in runs}) == 4 manager.lock() for _ in range(3): for i, (stream, graph, hidden, topk_ids, _, scratch) in enumerate(runs): @@ -423,8 +403,6 @@ def test_moe_permute_scratch_without_manager(monkeypatch) -> None: monkeypatch.setattr(workspace, "_manager", None) config = dict( - max_num_tokens=4, - topk=2, num_experts=4, num_local_experts=4, device=torch.device("cuda"), @@ -432,6 +410,8 @@ def test_moe_permute_scratch_without_manager(monkeypatch) -> None: first = get_moe_permute_scratch(**config) second = get_moe_permute_scratch(**config) + first.reserve(8) + second.reserve(8) assert ( first.token_expert_indices.data_ptr() != second.token_expert_indices.data_ptr() ) @@ -450,13 +430,9 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: topk_ids = torch.tensor([[0], [-1], [1], [4], [2]], device="cuda") expert_map = torch.tensor([0, 1, -1, -1], dtype=torch.int32, device="cuda") scratch = MoEPermuteScratch( - max_num_tokens=5, - topk=1, num_experts=4, num_local_experts=2, device=hidden_states.device, - hidden_size=16, - hidden_dtype=hidden_states.dtype, ) permuted, _, expert_offsets, inverse, _ = moe_permute( @@ -498,3 +474,39 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: expected_offsets = torch.tensor([0, 5, 5], dtype=torch.int64, device="cuda") torch.testing.assert_close(indices.flatten(), expected_inverse) torch.testing.assert_close(expert_offsets, expected_offsets) + + +def test_moe_permute_scratch_grows_until_workspace_lock(workspace_init) -> None: + """Capacity follows the inputs seen in warmup, whatever their topk or size.""" + 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( + num_experts=16, num_local_experts=16, device=device + ) + + def run(n_token: int, topk: int, dtype: torch.dtype) -> None: + hidden = torch.randn((n_token, 64), device=device).to(dtype) + topk_ids = torch.randint(16, (n_token, topk), device=device) + actual = moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + expected = moe_permute(hidden, None, topk_ids, 16) + for a, e in zip(actual, expected): + if a is not None: + torch.testing.assert_close(a.float(), e.float()) + + run(8, 6, torch.bfloat16) + run(64, 6, torch.bfloat16) + run(384, 1, torch.float8_e4m3fn) + assert scratch.max_expanded_rows == 384 + + current_workspace_manager().lock() + run(64, 6, torch.bfloat16) + with pytest.raises(AssertionError, match="Scratch growth is not allowed"): + run(65, 6, torch.bfloat16) + + # Experts rebuilt after the lock (e.g. a weight reload) keep the capacity. + rebuilt = get_moe_permute_scratch( + num_experts=16, num_local_experts=16, device=device + ) + assert rebuilt.max_expanded_rows == 384 diff --git a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py index a03651c0a93a..6f9a62f282ce 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py @@ -20,6 +20,7 @@ ) from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( MoEPermuteScratch, + get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, moe_unpermute, @@ -304,7 +305,6 @@ def __init__( self.ab_strides2 = ab_strides2 self.c_strides1 = c_strides1 self.c_strides2 = ab_strides1_c_strides2 - self._permute_scratch: MoEPermuteScratch | None = None @staticmethod def _supports_current_device() -> bool: @@ -335,15 +335,13 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceDelegate() def _get_permute_scratch(self) -> MoEPermuteScratch | None: - if self._permute_scratch is None and moe_permute_unpermute_supported(): - self._permute_scratch = MoEPermuteScratch( - max_num_tokens=self.moe_config.max_num_tokens, - topk=self.moe_config.experts_per_token, - num_experts=self.moe_config.num_experts, - num_local_experts=self.moe_config.num_local_experts, - device=torch.device(self.moe_config.device), - ) - return self._permute_scratch + if not moe_permute_unpermute_supported(): + return None + return get_moe_permute_scratch( + num_experts=self.moe_config.num_experts, + num_local_experts=self.moe_config.num_local_experts, + device=torch.device(self.moe_config.device), + ) def apply( self, @@ -1296,7 +1294,6 @@ def __init__( self.s_strides2[:, 0] = k self.group_size = group_size - self._permute_scratch: MoEPermuteScratch | None = None @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -1357,15 +1354,13 @@ def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: return self.out_dtype if self.out_dtype is not None else act_dtype def _get_permute_scratch(self) -> MoEPermuteScratch | None: - if self._permute_scratch is None and moe_permute_unpermute_supported(): - self._permute_scratch = MoEPermuteScratch( - max_num_tokens=self.moe_config.max_num_tokens, - topk=self.moe_config.experts_per_token, - num_experts=self.moe_config.num_experts, - num_local_experts=self.moe_config.num_local_experts, - device=torch.device(self.moe_config.device), - ) - return self._permute_scratch + if not moe_permute_unpermute_supported(): + return None + return get_moe_permute_scratch( + num_experts=self.moe_config.num_experts, + num_local_experts=self.moe_config.num_local_experts, + device=torch.device(self.moe_config.device), + ) def workspace_shapes( self, 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 a6d44c3141f7..66e03229a9e1 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 @@ -258,25 +258,14 @@ def humming_forward( **kwargs, ) - def _get_permute_scratch( - self, topk: int, indices_only: bool = False - ) -> MoEPermuteScratch | None: + def _get_permute_scratch(self) -> MoEPermuteScratch | None: if not moe_permute_unpermute_supported(): return None - 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=None if indices_only else self.moe_config.hidden_dim, - hidden_dtype=None if indices_only else self.moe_config.in_dtype, ) def get_global_valid_shape_m(self, topk_ids: torch.Tensor): @@ -911,14 +900,18 @@ def apply( ) buffers["output"] = output + scratch = self._get_permute_scratch() if a1q_scale is None: - scratch = self._get_permute_scratch(topk_ids.size(1), indices_only=True) assert scratch is not None scatter_metadata = moe_prepare_scatter(topk_ids, expert_map, scratch) expert_offsets, scatter_idx = scatter_metadata inv_perm = scatter_idx.flatten() num_valid_tokens = expert_offsets[-1:] if expert_map is not None else None else: + # Prequantized inputs skip w13 quantization, so its buffer is free + # to hold the permuted hidden states. + permuted_hidden_states = buffers["quanted_gate_up_input"] + assert permuted_hidden_states.dtype == hidden_states.dtype hidden_states, a1q_scale, expert_offsets, inv_perm, _ = moe_permute( hidden_states=hidden_states, a1q_scale=a1q_scale, @@ -926,7 +919,8 @@ def apply( n_expert=global_num_experts, n_local_expert=self.num_experts, expert_map=expert_map, - scratch=self._get_permute_scratch(topk_ids.size(1)), + permuted_hidden_states=permuted_hidden_states, + scratch=scratch, ) scatter_idx = None num_valid_tokens = None 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 d83152fc696a..b4bf3c89163d 100644 --- a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py +++ b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py @@ -13,90 +13,70 @@ @dataclass class MoEPermuteScratch: - # Reused metadata buffers for repeated grouped-MoE permutes. - max_num_tokens: int - topk: int + """Reused metadata buffers for repeated grouped-MoE permutes. + + Capacity follows the largest input seen, like the shared workspace: it grows + during warmup and is fixed once the workspace manager is locked. + """ + num_experts: int num_local_experts: int device: torch.device - hidden_size: int | None = None - hidden_dtype: torch.dtype | None = None + max_expanded_rows: int = field(init=False, default=0) token_expert_indices: torch.Tensor = field(init=False) expert_first_token_offset: torch.Tensor = field(init=False) permuted_idx: torch.Tensor = field(init=False) inv_permuted_idx: torch.Tensor = field(init=False) - permuted_hidden_states: torch.Tensor | None = field(init=False, default=None) sort_workspace: torch.Tensor = field(init=False) permuted_experts_id: torch.Tensor = field(init=False) sorted_row_idx: torch.Tensor = field(init=False) topk_ids_int32: torch.Tensor = field(init=False) topk_ids_for_sort: torch.Tensor = field(init=False) - max_expanded_rows: int = field(init=False) def __post_init__(self) -> None: - assert self.max_num_tokens > 0 - assert self.topk > 0 assert self.num_experts > 0 assert self.num_local_experts > 0 - if self.hidden_size is None: - assert self.hidden_dtype is None - else: - assert self.hidden_dtype is not None - - self.max_expanded_rows = self.max_num_tokens * self.topk - self.token_expert_indices = torch.arange( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) self.expert_first_token_offset = torch.empty( self.num_local_experts + 1, dtype=torch.int64, device=self.device ) - self.permuted_idx = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - self.inv_permuted_idx = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - if self.hidden_size is not None: - hidden_numel = self.max_expanded_rows * self.hidden_size - self.permuted_hidden_states = torch.empty( - hidden_numel, dtype=self.hidden_dtype, device=self.device + # torch.device("cuda") in config, after initialized, + # will be changed to cuda:{index}, so we need to refresh here. + self.device = self.expert_first_token_offset.device + + def reserve(self, expanded_rows: int) -> None: + """Ensure capacity for ``expanded_rows`` (num_tokens * topk) rows.""" + if expanded_rows <= self.max_expanded_rows: + return + if ( + is_workspace_manager_initialized() + and current_workspace_manager().is_locked() + ): + raise AssertionError( + f"Workspace is locked but MoE permute scratch requires " + f"{expanded_rows} rows, current capacity is " + f"{self.max_expanded_rows} rows. Scratch growth is not allowed " + "after locking." ) - self.permuted_experts_id = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - self.sorted_row_idx = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - self.topk_ids_int32 = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - self.topk_ids_for_sort = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device + + def empty_rows() -> torch.Tensor: + return torch.empty(expanded_rows, dtype=torch.int32, device=self.device) + + self.token_expert_indices = torch.arange( + expanded_rows, dtype=torch.int32, device=self.device ) + self.permuted_idx = empty_rows() + self.inv_permuted_idx = empty_rows() + self.permuted_experts_id = empty_rows() + self.sorted_row_idx = empty_rows() + self.topk_ids_int32 = empty_rows() + self.topk_ids_for_sort = empty_rows() sorter_size = torch.ops._moe_C.moe_permute_sort_workspace_size( - self.max_expanded_rows, self.num_experts + expanded_rows, self.num_experts ) self.sort_workspace = torch.empty( sorter_size, dtype=torch.int8, device=self.device ) - # torch.device("cuda") in config, after initialized, - # will be changed to cuda:{index}, so we need to refresh here. - self.device = self.token_expert_indices.device - - def validate(self, hidden_states: torch.Tensor, topk_ids: torch.Tensor) -> None: - n_token, n_hidden = hidden_states.shape - assert hidden_states.device == self.device - assert topk_ids.device == self.device - assert n_token <= self.max_num_tokens - assert topk_ids.size(1) == self.topk - assert topk_ids.size(0) == n_token - if self.hidden_size is not None: - assert n_hidden == self.hidden_size - assert hidden_states.dtype == self.hidden_dtype - assert self.permuted_hidden_states is not None - - def token_expert_indices_view(self, n_token: int) -> torch.Tensor: - return self.token_expert_indices[: n_token * self.topk].view(n_token, self.topk) + self.max_expanded_rows = expanded_rows def prepare_topk_ids(self, topk_ids: torch.Tensor) -> torch.Tensor: if topk_ids.dtype == torch.int32: @@ -114,17 +94,16 @@ def moe_prepare_scatter( ) -> tuple[torch.Tensor, torch.Tensor]: """Generate expert offsets and shared scatter/unpermute destination indices.""" assert topk_ids.device == scratch.device - assert topk_ids.size(0) <= scratch.max_num_tokens - assert topk_ids.size(1) == scratch.topk n_token, topk = topk_ids.shape expanded_rows = topk_ids.numel() + scratch.reserve(expanded_rows) inverse = scratch.inv_permuted_idx[:expanded_rows].view(n_token, topk) if expanded_rows == 0: scratch.expert_first_token_offset.zero_() return scratch.expert_first_token_offset, inverse torch.ops._moe_C.moe_prepare_scatter( scratch.prepare_topk_ids(topk_ids), - scratch.token_expert_indices_view(n_token), + scratch.token_expert_indices[:expanded_rows].view(n_token, topk), expert_map, scratch.num_experts, scratch.num_local_experts, @@ -140,13 +119,9 @@ def moe_prepare_scatter( 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. @@ -155,28 +130,15 @@ def get_moe_permute_scratch( 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, - ) + key = (MoEPermuteScratch, num_experts, num_local_experts, device) return current_workspace_manager().get_persistent_resource(key, create_scratch) @@ -204,9 +166,8 @@ def moe_permute( parallel shard. permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor. If None, the output tensor will be created in this function. - scratch (Optional[MoEPermuteScratch]): Optional preallocated scratch - buffers. Validated against hidden_states and topk_ids when given, - otherwise the buffers are allocated in this function. + scratch (Optional[MoEPermuteScratch]): Optional reusable metadata + buffers. Otherwise they are allocated in this function. Returns: - permuted_hidden_states (torch.Tensor): permuted activation. @@ -227,20 +188,11 @@ def moe_permute( if n_local_expert == -1: n_local_expert = n_expert if permuted_hidden_states is None: - if scratch is None: - permuted_hidden_states = torch.empty( - (permuted_row_size, n_hidden), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - else: - scratch.validate(hidden_states, topk_ids) - hidden_numel = permuted_row_size * n_hidden - scratch_hidden_states = scratch.permuted_hidden_states - assert scratch_hidden_states is not None - permuted_hidden_states = scratch_hidden_states[:hidden_numel].view( - permuted_row_size, n_hidden - ) + permuted_hidden_states = torch.empty( + (permuted_row_size, n_hidden), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) assert permuted_hidden_states.size() == (permuted_row_size, n_hidden), ( f"Expected permuted hidden states to be {(permuted_row_size, n_hidden)}" f" but got {permuted_hidden_states.size()}" @@ -278,10 +230,15 @@ def moe_permute( permuted_idx, ) else: - scratch.validate(hidden_states, topk_ids) + assert hidden_states.device == scratch.device + assert topk_ids.device == scratch.device + assert topk_ids.size(0) == n_token assert n_expert == scratch.num_experts assert n_local_expert == scratch.num_local_experts - token_expert_indices = scratch.token_expert_indices_view(n_token) + scratch.reserve(permuted_row_size) + token_expert_indices = scratch.token_expert_indices[:permuted_row_size].view( + n_token, topk + ) expert_first_token_offset = scratch.expert_first_token_offset permuted_idx = scratch.permuted_idx[:permuted_row_size] permuted_idx.fill_(permuted_row_size) diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index e69e55bd8216..a05534281c11 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -139,7 +139,8 @@ def get_persistent_resource(self, key: Hashable, factory: Callable[[], T]) -> T: 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 + resize or replace tensor storage held by a cached resource once the + manager is locked. Uses of the same resource must not overlap; include the CUDA stream in the key when streams can execute concurrently.