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 14, 05:30 UTC
[Bugfix] Fix platform-specific routing in CustomOp implementations (#24444)
Signed-off-by: Konrad Zawora <kzawora@habana.ai>
Konrad Zawora committed
Sep 11, 2025 at 19:15 UTC
4aa23892d6dfdc3ec6de7ff74e982b833adf6c52
8 files changed
+53
-30
vllm/model_executor/layers/activation.py
+4
-1
@@ -454,7 +454,7 @@ class XIELU(CustomOp):
454
)
455
return result.view(original_shape)
456
457
- def forward(self, input: torch.Tensor) -> torch.Tensor:
457
+ def forward_native(self, input: torch.Tensor) -> torch.Tensor:
458
if self._xielu_cuda_obj is not None and input.is_cuda:
459
if not torch._dynamo.is_compiling():
460
return self._xielu_cuda_fn(input)
@@ -464,6 +464,9 @@ class XIELU(CustomOp):
464
)
465
return self._xielu_python(input)
466
467
+ def forward_cuda(self, input: torch.Tensor) -> torch.Tensor:
468
+ return self.forward_native(input)
469
+
470
471
class ScaledActivation(nn.Module):
472
"""An activation function with post-scale parameters.
vllm/model_executor/layers/fused_moe/layer.py
+8
-1
@@ -1593,7 +1593,7 @@ class FusedMoE(CustomOp):
1593
else:
1594
return tensor_model_parallel_all_reduce(final_hidden_states)
1595
1596
- def forward(
1596
+ def forward_native(
1597
self,
1598
hidden_states: torch.Tensor,
1599
router_logits: torch.Tensor,
@@ -1627,6 +1627,13 @@ class FusedMoE(CustomOp):
1627
return (shared_output[..., :og_hidden_states],
1628
fused_output[..., :og_hidden_states])
1629
1630
+ def forward_cuda(
1631
+ self,
1632
+ hidden_states: torch.Tensor,
1633
+ router_logits: torch.Tensor,
1634
+ ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
1635
+ return self.forward_native(hidden_states, router_logits)
1636
+
1637
def forward_impl_chunked(
1638
self,
1639
full_hidden_states: torch.Tensor,
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+10
-1
@@ -88,7 +88,7 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
88
cache = torch.cat((cos, sin), dim=-1)
89
return cache
90
91
- def forward(
91
+ def forward_native(
92
self,
93
positions: torch.Tensor,
94
query: torch.Tensor,
@@ -129,3 +129,12 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
129
query = query_rot
130
key = key_rot
131
return query, key
132
+
133
+ def forward_cuda(
134
+ self,
135
+ positions: torch.Tensor,
136
+ query: torch.Tensor,
137
+ key: Optional[torch.Tensor] = None,
138
+ offsets: Optional[torch.Tensor] = None,
139
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
140
+ return self.forward_native(positions, query, key, offsets)
vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py
+10
-1
@@ -111,7 +111,7 @@ class DualChunkRotaryEmbedding(CustomOp):
111
device=self.device)
112
return q_cache, qc_cache, k_cache, qc_no_clamp_cache, q_inter_cache
113
114
- def forward(
114
+ def forward_native(
115
self,
116
positions: torch.Tensor,
117
query: torch.Tensor,
@@ -161,6 +161,15 @@ class DualChunkRotaryEmbedding(CustomOp):
161
dim=-1)
162
return query, key
163
164
+ def forward_cuda(
165
+ self,
166
+ positions: torch.Tensor,
167
+ query: torch.Tensor,
168
+ key: torch.Tensor,
169
+ offsets: Optional[torch.Tensor] = None,
170
+ ) -> tuple[torch.Tensor, torch.Tensor]:
171
+ return self.forward_native(positions, query, key, offsets)
172
+
173
def _apply_rotary_embedding(self, cos_sin, hidden_rot, hidden_pass):
174
cos, sin = cos_sin.chunk(2, dim=-1)
175
if self.is_neox_style:
vllm/model_executor/layers/rotary_embedding/ernie45_vl_rope.py
+9
-1
@@ -12,7 +12,7 @@ from .mrope import MRotaryEmbedding
12
class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):
13
"""3D rotary positional embedding. 3D is t:time h:height w:width"""
14
15
- def forward(
15
+ def forward_native( # type: ignore[override]
16
self,
17
positions: torch.Tensor,
18
query: torch.Tensor,
@@ -70,3 +70,11 @@ class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):
70
self.is_neox_style)
71
key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
72
return query, key
73
+
74
+ def forward_cuda( # type: ignore[override]
75
+ self,
76
+ positions: torch.Tensor,
77
+ query: torch.Tensor,
78
+ key: Optional[torch.Tensor] = None,
79
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
80
+ return self.forward_native(positions, query, key)
\ No newline at end of file
vllm/model_executor/layers/rotary_embedding/llama4_vision_rope.py
+8
-1
@@ -53,7 +53,7 @@ class Llama4VisionRotaryEmbedding(RotaryEmbedding):
53
torch.stack([torch.cos(freqs), torch.sin(freqs)], dim=-1))
54
return cache
55
56
- def forward(
56
+ def forward_native( # type: ignore[override]
57
self,
58
query: torch.Tensor,
59
key: Optional[torch.Tensor] = None,
@@ -72,3 +72,10 @@ class Llama4VisionRotaryEmbedding(RotaryEmbedding):
72
query_out = torch.view_as_real(query_ * freqs_ci).flatten(3)
73
key_out = torch.view_as_real(key_ * freqs_ci).flatten(3)
74
return query_out.type_as(query), key_out.type_as(key)
75
+
76
+ def forward_cuda( # type: ignore[override]
77
+ self,
78
+ query: torch.Tensor,
79
+ key: Optional[torch.Tensor] = None,
80
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
81
+ return self.forward_native(query, key)
vllm/model_executor/layers/rotary_embedding/mrope.py
-23
@@ -8,7 +8,6 @@ import numpy as np
8
import torch
9
from transformers import PretrainedConfig
10
11
-from vllm.platforms import current_platform
11
from vllm.triton_utils import tl, triton
12
13
from .base import RotaryEmbedding
@@ -202,28 +201,6 @@ class MRotaryEmbedding(RotaryEmbedding):
201
if self.mrope_section:
202
assert sum(self.mrope_section) == rotary_dim // 2
203
205
- self.use_triton = current_platform.is_cuda_alike()
206
-
207
- def forward(
208
- self,
209
- positions: torch.Tensor,
210
- query: torch.Tensor,
211
- key: Optional[torch.Tensor] = None,
212
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
213
- """MRope forward.
214
-
215
- Args:
216
- positions:
217
- [num_tokens,] (text only) or
218
- [3, num_tokens] (T/H/W positions with multimodal inputs)
219
- query: [num_tokens, num_heads * head_size]
220
- key: [num_tokens, num_kv_heads * head_size]
221
- """
222
- if self.use_triton:
223
- return self.forward_cuda(positions, query, key)
224
- else:
225
- return self.forward_native(positions, query, key)
226
-
204
def forward_native(
205
self,
206
positions: torch.Tensor,
vllm/model_executor/layers/vocab_parallel_embedding.py
+4
-1
@@ -399,7 +399,7 @@ class VocabParallelEmbedding(CustomOp):
399
param[:loaded_weight.shape[0]].data.copy_(loaded_weight)
400
param[loaded_weight.shape[0]:].data.fill_(0)
401
402
- def forward(self, input_):
402
+ def forward_native(self, input_):
403
if self.tp_size > 1:
404
# Build the mask.
405
masked_input, input_mask = get_masked_input_and_mask(
@@ -420,6 +420,9 @@ class VocabParallelEmbedding(CustomOp):
420
output = tensor_model_parallel_all_reduce(output_parallel)
421
return output
422
423
+ def forward_cuda(self, input_):
424
+ return self.forward_native(input_)
425
+
426
def extra_repr(self) -> str:
427
s = f"num_embeddings={self.num_embeddings_per_partition}"
428
s += f", embedding_dim={self.embedding_dim}"