A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 09:45 UTC
[Bug] Fix torch device issue for MOE permute (#44005)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
May 29, 2026 at 14:55 UTC
739096a028214991cff4af23a558795905ad35a1
1 file changed
+3
vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
+3
@@ -74,6 +74,9 @@ class MoEPermuteScratch:
74
self.sort_workspace = torch.empty(
75
sorter_size, dtype=torch.int8, device=self.device
76
)
77
+ # torch.device("cuda") in config, after initialized,
78
+ # will be changed to cuda:{index}, so we need to refresh here.
79
+ self.device = self.token_expert_indices.device
80
81
def validate(self, hidden_states: torch.Tensor, topk_ids: torch.Tensor) -> None:
82
n_token, n_hidden = hidden_states.shape