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

[Kernel] Remove unused fake implementation (#55535)

Browse Source 
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
This commit is contained in:
Jee Jee Li
2026-09-06 17:54:32 -07:00
committed by GitHub
parent 4df80187b5
commit 199cb9b964
39 changed files with 7 additions and 956 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/_aiter_ops.py
-97
View File
@@ -387,18 +387,6 @@ def _rocm_aiter_topk_softmax_impl(
)
def _rocm_aiter_topk_softmax_fake(
topk_weights: torch.Tensor,
topk_indices: torch.Tensor,
token_expert_indices: torch.Tensor,
gating_output: torch.Tensor,
renormalize: bool,
num_shared_experts: int = 0,
shared_expert_scoring_func: str = "",
) -> None:
pass
def _rocm_aiter_topk_sigmoid_impl(
topk_weights: torch.Tensor,
topk_indices: torch.Tensor,
@@ -409,14 +397,6 @@ def _rocm_aiter_topk_sigmoid_impl(
topk_sigmoid(topk_weights, topk_indices, gating_output)
def _rocm_aiter_topk_sigmoid_fake(
topk_weights: torch.Tensor,
topk_indices: torch.Tensor,
gating_output: torch.Tensor,
) -> None:
pass
def _rocm_aiter_biased_grouped_topk_impl(
gating_output: torch.Tensor,
correction_bias: torch.Tensor,
@@ -441,19 +421,6 @@ def _rocm_aiter_biased_grouped_topk_impl(
)
def _rocm_aiter_biased_grouped_topk_fake(
gating_output: torch.Tensor,
correction_bias: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
num_expert_group: int,
topk_group: int,
need_renorm: bool,
routed_scaling_factor: float = 1.0, # mul to topk_weights
) -> None:
pass
def _rocm_aiter_grouped_topk_impl(
gating_output: torch.Tensor,
topk_weights: torch.Tensor,
@@ -479,19 +446,6 @@ def _rocm_aiter_grouped_topk_impl(
)
def _rocm_aiter_grouped_topk_fake(
gating_output: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
num_expert_group: int,
topk_group: int,
need_renorm: bool,
scoring_func: str = "softmax",
routed_scaling_factor: float = 1.0, # mul to topk_weights
) -> None:
pass
def _rocm_aiter_fused_topk_impl(
x: torch.Tensor,
router_logits: torch.Tensor,
@@ -618,29 +572,6 @@ def _rocm_aiter_mla_decode_fwd_impl(
)
def _rocm_aiter_mla_decode_fwd_fake(
q: torch.Tensor,
kv_buffer: torch.Tensor,
o: torch.Tensor,
qo_indptr: torch.Tensor,
max_seqlen_qo: int,
kv_indptr: torch.Tensor | None = None,
kv_indices: torch.Tensor | None = None,
kv_last_page_lens: torch.Tensor | None = None,
sm_scale: float = 1.0,
logit_cap: float = 0.0,
q_scale: torch.Tensor | None = None,
kv_scale: torch.Tensor | None = None,
work_meta_data: torch.Tensor | None = None,
work_indptr: torch.Tensor | None = None,
work_info_set: torch.Tensor | None = None,
reduce_indptr: torch.Tensor | None = None,
reduce_final_map: torch.Tensor | None = None,
reduce_partial_map: torch.Tensor | None = None,
) -> None:
pass
def _rocm_aiter_w8a8_gemm_impl(
A: torch.Tensor,
B: torch.Tensor,
@@ -1104,15 +1035,6 @@ def _rocm_aiter_per_tensor_quant_impl(
static_per_tensor_quant(out, x, scale)
def _rocm_aiter_per_tensor_quant_fake(
out: torch.Tensor,
x: torch.Tensor,
scale: torch.Tensor,
is_dynamic: bool,
) -> None:
pass
def _rocm_aiter_per_token_quant_impl(
x: torch.Tensor, quant_dtype: torch.dtype, scale: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -1597,18 +1519,6 @@ def _triton_rotary_embedding_impl(
key = key.view(key_shape)
def _triton_rotary_embedding_fake(
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor,
head_size: int,
cos_sin_cache: torch.Tensor,
is_neox_style: bool,
offsets: torch.Tensor | None = None,
) -> None:
return
def _rocm_aiter_fp8_attn_impl(
q: torch.Tensor,
k: torch.Tensor,
@@ -2132,7 +2042,6 @@ class rocm_aiter_ops:
op_name="rocm_aiter_topk_softmax",
op_func=_rocm_aiter_topk_softmax_impl,
mutates_args=["topk_weights", "topk_indices", "token_expert_indices"],
fake_impl=_rocm_aiter_topk_softmax_fake,
dispatch_key=current_platform.dispatch_key,
)
@@ -2140,7 +2049,6 @@ class rocm_aiter_ops:
op_name="rocm_aiter_topk_sigmoid",
op_func=_rocm_aiter_topk_sigmoid_impl,
mutates_args=["topk_weights", "topk_indices"],
fake_impl=_rocm_aiter_topk_sigmoid_fake,
dispatch_key=current_platform.dispatch_key,
)
@@ -2148,7 +2056,6 @@ class rocm_aiter_ops:
op_name="rocm_aiter_biased_grouped_topk",
op_func=_rocm_aiter_biased_grouped_topk_impl,
mutates_args=["topk_weights", "topk_ids"],
fake_impl=_rocm_aiter_biased_grouped_topk_fake,
dispatch_key=current_platform.dispatch_key,
)
@@ -2156,7 +2063,6 @@ class rocm_aiter_ops:
op_name="rocm_aiter_grouped_topk",
op_func=_rocm_aiter_grouped_topk_impl,
mutates_args=["topk_weights", "topk_ids"],
fake_impl=_rocm_aiter_grouped_topk_fake,
dispatch_key=current_platform.dispatch_key,
)
@@ -2172,7 +2078,6 @@ class rocm_aiter_ops:
op_name="rocm_aiter_mla_decode_fwd",
op_func=_rocm_aiter_mla_decode_fwd_impl,
mutates_args=["o"],
fake_impl=_rocm_aiter_mla_decode_fwd_fake,
)
direct_register_custom_op(
@@ -2266,7 +2171,6 @@ class rocm_aiter_ops:
op_name="rocm_aiter_per_tensor_quant",
op_func=_rocm_aiter_per_tensor_quant_impl,
mutates_args=["out", "scale"],
fake_impl=_rocm_aiter_per_tensor_quant_fake,
dispatch_key=current_platform.dispatch_key,
)
@@ -2306,7 +2210,6 @@ class rocm_aiter_ops:
op_name="rocm_aiter_triton_rotary_embedding",
op_func=_triton_rotary_embedding_impl,
mutates_args=["query", "key"], # These tensors are modified in-place
fake_impl=_triton_rotary_embedding_fake,
)
direct_register_custom_op(
vllm/_custom_ops.py
-56
View File
@@ -100,17 +100,6 @@ if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "scaled_fp4_quant"):
m = input.numel() // n
return create_fp4_output_tensors(m, n, input.device, is_sf_swizzled_layout)
@register_fake("_C::scaled_fp4_quant.out")
def _scaled_fp4_quant_out_fake(
input: torch.Tensor,
input_scale: torch.Tensor,
is_sf_swizzled_layout: bool,
*,
output: torch.Tensor,
output_scale: torch.Tensor,
) -> None:
return None
# page attention ops
def paged_attention_rocm(
@@ -740,27 +729,6 @@ def moe_gptq_gemm_rdna3(
)
if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3"):
@register_fake("_rocm_C::moe_gptq_gemm_rdna3")
def _moe_gptq_gemm_rdna3_fake(
a: torch.Tensor,
c: torch.Tensor,
b_q_weight: torch.Tensor,
b_scales: torch.Tensor,
b_qzeros: torch.Tensor,
topk_weights: torch.Tensor,
sorted_token_ids: torch.Tensor,
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor,
top_k: int,
block_size_m: int,
mul_topk_weight: bool,
output_topk: int = 0,
) -> None:
return
if hasattr(torch.ops._C, "allspark_w8a16_gemm"):
@register_fake("_C::allspark_w8a16_gemm")
@@ -2437,17 +2405,6 @@ def fp32_router_gemm(
return output
if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"):
@register_fake("_C::fp32_router_gemm")
def fp32_router_gemm_fake(
output: torch.Tensor,
mat_a: torch.Tensor,
mat_b: torch.Tensor,
) -> None:
return
def topk_softmax(
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
@@ -4363,19 +4320,6 @@ def safeFusedQuantizeNv(
return
if hasattr(torch.ops._qutlass_C, "fusedQuantizeNv"):
@register_fake("vllm::safeFusedQuantizeNv")
def _fake_fused_quantize_nv(
a: torch.Tensor,
b: torch.Tensor,
xh_e2m1: torch.Tensor,
xh_e4m3: torch.Tensor,
global_scale: torch.Tensor,
) -> None:
return
def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor:
"""
Perform Hadamard transforms using [Hadacore](https://arxiv.org/abs/2412.08832)
vllm/_xpu_ops.py
-60
View File
@@ -125,15 +125,6 @@ def _gemma_rms_norm_impl(
torch.ops._C.gemma_rms_norm(out, input, weight, epsilon)
def _gemma_rms_norm_fake(
out: torch.Tensor,
input: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
) -> None:
return None
def _fused_add_gemma_rms_norm_impl(
input: torch.Tensor,
residual: torch.Tensor,
@@ -143,15 +134,6 @@ def _fused_add_gemma_rms_norm_impl(
torch.ops._C.fused_add_gemma_rms_norm(input, residual, weight, epsilon)
def _fused_add_gemma_rms_norm_fake(
input: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
) -> None:
return None
def _gdn_attention_core_xpu_impl(
core_attn_out: torch.Tensor,
z: torch.Tensor,
@@ -240,16 +222,6 @@ def _gdn_attention_core_xpu_impl(
)
def _gdn_attention_core_xpu_fake(
core_attn_out: torch.Tensor,
z: torch.Tensor,
projected_states_qkvz: torch.Tensor,
projected_states_ba: torch.Tensor,
layer_name: str,
) -> None:
return
def _xpu_ops_deepseek_scaling_rope_impl(
positions: torch.Tensor,
query: torch.Tensor,
@@ -471,19 +443,6 @@ def _xpu_deepseek_fused_indexer_q_rope_fp8_impl(
)
def _xpu_deepseek_fused_indexer_q_rope_fp8_fake(
index_q: torch.Tensor,
positions: torch.Tensor,
index_q_cos_sin_cache: torch.Tensor,
index_weights: torch.Tensor,
index_weights_softmax_scale: float,
index_weights_head_scale: float,
index_q_fp8: torch.Tensor,
index_weights_out: torch.Tensor,
) -> None:
return
def _xpu_deepseek_fused_indexer_q_rope_mxfp4_impl(
index_q: torch.Tensor,
positions: torch.Tensor,
@@ -528,20 +487,6 @@ def _xpu_deepseek_fused_indexer_q_rope_mxfp4_impl(
)
def _xpu_deepseek_fused_indexer_q_rope_mxfp4_fake(
index_q: torch.Tensor,
positions: torch.Tensor,
index_q_cos_sin_cache: torch.Tensor,
index_weights: torch.Tensor,
index_weights_softmax_scale: float,
index_weights_head_scale: float,
index_q_packed: torch.Tensor,
index_q_scale: torch.Tensor,
index_weights_out: torch.Tensor,
) -> None:
return
def _xpu_mxfp8_quantize_impl(
x: torch.Tensor, dtype: torch.dtype | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -1262,14 +1207,12 @@ class xpu_ops:
op_name="xpu_gemma_rms_norm",
op_func=_gemma_rms_norm_impl,
mutates_args=["out"],
fake_impl=_gemma_rms_norm_fake,
)
direct_register_custom_op(
op_name="xpu_fused_add_gemma_rms_norm",
op_func=_fused_add_gemma_rms_norm_impl,
mutates_args=["input", "residual"],
fake_impl=_fused_add_gemma_rms_norm_fake,
)
direct_register_custom_op(
@@ -1314,7 +1257,6 @@ class xpu_ops:
op_name="gdn_attention_core_xpu",
op_func=eager_break_during_capture(_gdn_attention_core_xpu_impl),
mutates_args=["core_attn_out", "z"],
fake_impl=_gdn_attention_core_xpu_fake,
)
direct_register_custom_op(
@@ -1327,7 +1269,6 @@ class xpu_ops:
op_name="xpu_deepseek_fused_indexer_q_rope_fp8",
op_func=_xpu_deepseek_fused_indexer_q_rope_fp8_impl,
mutates_args=["index_q_fp8", "index_weights_out"],
fake_impl=_xpu_deepseek_fused_indexer_q_rope_fp8_fake,
)
direct_register_custom_op(
@@ -1338,7 +1279,6 @@ class xpu_ops:
"index_q_scale",
"index_weights_out",
],
fake_impl=_xpu_deepseek_fused_indexer_q_rope_mxfp4_fake,
)
_OPS_REGISTERED = True
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
-19
View File
@@ -293,24 +293,6 @@ if flashinfer_comm is not None:
or num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,
)
def call_trtllm_fused_allreduce_norm_fake(
allreduce_in: torch.Tensor,
residual: torch.Tensor,
rms_gamma: torch.Tensor,
rms_eps: float,
world_size: int,
launch_with_pdl: bool,
fp32_acc: bool,
max_token_num: int,
pattern_code: int,
norm_out: torch.Tensor | None = None,
quant_out: torch.Tensor | None = None,
scale_out: torch.Tensor | None = None,
scale_factor: torch.Tensor | None = None,
weight_bias: float = 0.0,
) -> None:
pass
direct_register_custom_op(
op_name="flashinfer_trtllm_fused_allreduce_norm",
op_func=call_trtllm_fused_allreduce_norm,
@@ -321,7 +303,6 @@ if flashinfer_comm is not None:
"quant_out",
"scale_out",
],
fake_impl=call_trtllm_fused_allreduce_norm_fake,
)
flashinfer_trtllm_fused_allreduce_norm = (
torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default
vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py
-10
View File
@@ -100,15 +100,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |
return result
def fake_impl(
result: torch.Tensor, # [num_tokens, hidden_size]
input: torch.Tensor, # [num_tokens, hidden_size]
scale: torch.Tensor, # [num_tokens, 1]
scale_ub: torch.Tensor | None = None, # scalar tensor
) -> None:
return
def baseline(
result: torch.Tensor, # [num_tokens, hidden_size]
input: torch.Tensor, # [num_tokens, hidden_size]
@@ -135,7 +126,6 @@ def baseline(
mutates_args=["result", "scale"],
config_picker=pick_config,
input_generator=generate_inputs,
fake_impl=fake_impl,
helion_settings=helion.Settings(
autotune_baseline_fn=baseline,
ignore_warnings=[helion.exc.TensorOperationInWrapper],
vllm/kernels/helion/ops/fused_qk_norm_rope.py
-18
View File
@@ -159,23 +159,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |
return result
def fake_impl(
qkv: torch.Tensor, # [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]
num_heads_q: int,
num_heads_k: int,
num_heads_v: int,
head_dim: int,
eps: float,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
cos_sin_cache: torch.Tensor, # [max_position, rotary_dim]
is_neox: bool,
position_ids: torch.Tensor, # [num_tokens],
forced_token_heads_per_warp: int = -1, # dummy
) -> None:
return
def baseline(
qkv: torch.Tensor, # [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]
num_heads_q: int,
@@ -216,7 +199,6 @@ def baseline(
mutates_args=["qkv"],
config_picker=pick_config,
input_generator=generate_inputs,
fake_impl=fake_impl,
helion_settings=helion.Settings(
autotune_baseline_fn=baseline,
autotune_baseline_atol=5e-2,
vllm/kernels/helion/ops/per_token_group_fp8_quant.py
-18
View File
@@ -135,23 +135,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |
return result
def fake_impl(
input: torch.Tensor, # [num_tokens, hidden_size]
output_q: torch.Tensor, # [num_tokens, hidden_size]
output_s: torch.Tensor, # [num_tokens, groups_per_row]
group_size: int,
eps: float,
fp8_min: float,
fp8_max: float,
scale_ue8m0: bool,
# Unused dummy args
# Kept for consistency with existing kernel interface
dummy_is_scale_transposed: bool = False,
dummy_is_tma_aligned: bool = False,
) -> None:
return
def baseline(
input: torch.Tensor, # [num_tokens, hidden_size]
output_q: torch.Tensor, # [num_tokens, hidden_size]
@@ -181,7 +164,6 @@ def baseline(
mutates_args=["output_q", "output_s"],
config_picker=pick_config,
input_generator=generate_inputs,
fake_impl=fake_impl,
helion_settings=helion.Settings(
autotune_baseline_fn=baseline,
),
vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py
-13
View File
@@ -119,18 +119,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |
return result
def fake_impl(
result: torch.Tensor, # [num_tokens, hidden_size]
input: torch.Tensor, # [num_tokens, hidden_size]
weight: torch.Tensor, # [hidden_size]
scale: torch.Tensor, # [num_tokens, 1]
epsilon: float,
scale_ub: torch.Tensor | None = None, # []
residual: torch.Tensor | None = None, # [num_tokens, hidden_size]
) -> None:
return
def baseline(
result: torch.Tensor, # [num_tokens, hidden_size]
input: torch.Tensor, # [num_tokens, hidden_size]
@@ -179,7 +167,6 @@ def baseline(
mutates_args=["result", "scale", "residual"],
config_picker=pick_config,
input_generator=generate_inputs,
fake_impl=fake_impl,
helion_settings=helion.Settings(
autotune_baseline_fn=baseline,
ignore_warnings=[helion.exc.TensorOperationInWrapper],
vllm/kernels/helion/ops/rms_norm_per_block_quant.py
-15
View File
@@ -153,20 +153,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |
return result
def fake_impl(
result: torch.Tensor, # [num_tokens, hidden_size]
input: torch.Tensor, # [num_tokens, hidden_size]
weight: torch.Tensor, # [hidden_size]
scale: torch.Tensor, # [num_tokens, groups_per_row]
epsilon: float,
scale_ub: torch.Tensor | None, # []
residual: torch.Tensor | None, # [num_tokens, hidden_size]
group_size: int,
is_scale_transposed: bool, # dummy
) -> None:
return
def baseline(
result: torch.Tensor, # [num_tokens, hidden_size]
input: torch.Tensor, # [num_tokens, hidden_size]
@@ -219,7 +205,6 @@ def baseline(
mutates_args=["result", "scale", "residual"],
config_picker=pick_config,
input_generator=generate_inputs,
fake_impl=fake_impl,
helion_settings=helion.Settings(
autotune_baseline_fn=baseline,
ignore_warnings=[helion.exc.TensorOperationInWrapper],
vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py
-12
View File
@@ -139,17 +139,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |
return result
def fake_impl(
out: torch.Tensor, # [num_tokens, intermediate_size]
input: torch.Tensor, # [num_tokens, 2 * intermediate_size]
scales: torch.Tensor, # [num_tokens, groups_per_row]
group_size: int,
scale_ub: torch.Tensor | None = None, # scalar tensor
is_scale_transposed: bool = False,
) -> None:
return
def baseline(
out: torch.Tensor, # [num_tokens, intermediate_size]
input: torch.Tensor, # [num_tokens, 2 * intermediate_size]
@@ -193,7 +182,6 @@ def baseline(
mutates_args=["out", "scales"],
config_picker=pick_config,
input_generator=generate_inputs,
fake_impl=fake_impl,
helion_settings=helion.Settings(
autotune_baseline_fn=baseline,
ignore_warnings=[helion.exc.TensorOperationInWrapper],
vllm/kernels/helion/register.py
+7 -2
View File
@@ -38,6 +38,7 @@ Key Classes
from __future__ import annotations
import inspect
from collections.abc import Callable
from typing import Any
@@ -258,7 +259,7 @@ class HelionKernelWrapper:
self,
raw_kernel_func: Callable,
op_name: str,
fake_impl: Callable,
fake_impl: Callable | None,
config_picker: ConfigPicker,
mutates_args: list[str] | None = None,
helion_settings: helion.Settings | None = None,
@@ -445,8 +446,12 @@ def register_kernel(
f"Use a different op_name or check for duplicate registrations."
)
# A void-returning kernel that mutates its args needs no fake impl:
# PyTorch auto-generates a trivial one for mutable, no-return schemas.
returns_none = inspect.signature(kernel_func).return_annotation is None
final_fake_impl = fake_impl
if final_fake_impl is None:
if final_fake_impl is None and not (mutates_args and returns_none):
final_fake_impl = infer_fake_impl(kernel_func, helion_settings)
logger.debug(
"Auto-generated fake_impl for Helion kernel '%s'",
vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py
-133
View File
@@ -869,156 +869,23 @@ def _fused_moe_lora_fp8(
)
def _fused_moe_lora_fp8_fake(
output: torch.Tensor,
qcurr_hidden_states: torch.Tensor,
lora_a_stacked: list[torch.Tensor],
lora_b_stacked: list[torch.Tensor],
topk_weights: torch.Tensor,
sorted_token_ids: torch.Tensor | None,
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor | None,
token_lora_mapping: torch.Tensor,
max_lora_rank: int,
top_k_num: int,
lora_ids: torch.Tensor,
num_active_loras: int,
adapter_enabled: torch.Tensor,
shrink_block_size_m: int,
shrink_block_size_n: int,
shrink_block_size_k: int,
shrink_group_size_m: int,
shrink_num_warps: int,
shrink_num_stages: int,
shrink_split_k: int,
expand_block_size_m: int,
expand_block_size_n: int,
expand_block_size_k: int,
expand_group_size_m: int,
expand_num_warps: int,
expand_num_stages: int,
expand_split_k: int,
lora_a_scale_stacked: list[torch.Tensor],
lora_b_scale_stacked: list[torch.Tensor],
mul_routed_weight: bool = False,
fully_sharded: bool = False,
offset: int = 0,
shrink_act_scale: torch.Tensor | None = None,
expand_act_scale: torch.Tensor | None = None,
use_fp8_w8a8: bool = False,
use_int8_w8a8: bool = False,
use_int8_w8a16: bool = False,
per_channel_quant: bool = False,
block_shape: List[int] | None = None, # noqa: UP006, UP007
) -> None:
return
def _fused_moe_lora_shrink_fp8_fake(
a_intermediate_cache1: torch.Tensor,
qcurr_hidden_states: torch.Tensor,
lora_a_stacked: list[torch.Tensor],
topk_weights: torch.Tensor,
sorted_token_ids: torch.Tensor | None,
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor | None,
token_lora_mapping: torch.Tensor,
top_k_num: int,
lora_ids: torch.Tensor,
adapter_enabled: torch.Tensor,
device: torch.device,
N: int,
M: int,
EM: int,
K: int,
num_tokens: int,
num_experts: int,
num_slices: int,
block_size_m: int,
block_size_n: int,
block_size_k: int,
group_size_m: int,
num_warps: int,
num_stages: int,
split_k: int,
num_active_loras: int,
lora_a_scale_stacked: list[torch.Tensor],
mul_routed_weight: bool = False,
use_gdc: bool = False,
act_scale: torch.Tensor | None = None,
use_fp8_w8a8: bool = False,
use_int8_w8a8: bool = False,
use_int8_w8a16: bool = False,
per_channel_quant: bool = False,
block_shape: List[int] | None = None, # noqa: UP006, UP007
) -> None:
return
def _fused_moe_lora_expand_fp8_fake(
output: torch.Tensor,
a_intermediate_cache1: torch.Tensor,
lora_b_stacked: list[torch.Tensor],
topk_weights: torch.Tensor,
sorted_token_ids: torch.Tensor | None,
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor | None,
token_lora_mapping: torch.Tensor,
top_k_num: int,
lora_ids: torch.Tensor,
adapter_enabled: torch.Tensor,
device: torch.device,
N: int,
M: int,
EM: int,
K: int,
num_tokens: int,
num_experts: int,
num_slices: int,
max_lora_rank: int,
w1_output_dim_size: int,
block_size_m: int,
block_size_n: int,
block_size_k: int,
group_size_m: int,
num_warps: int,
num_stages: int,
split_k: int,
num_active_loras: int,
act_scale: torch.Tensor,
lora_b_scale_stacked: list[torch.Tensor],
mul_routed_weight: bool = False,
offset: int = 0,
use_fp8_w8a8: bool = False,
use_int8_w8a8: bool = False,
use_int8_w8a16: bool = False,
per_channel_quant: bool = False,
block_shape: List[int] | None = None, # noqa: UP006, UP007
use_gdc: bool = False,
) -> None:
return
try:
direct_register_custom_op(
op_name="fused_moe_lora_fp8",
op_func=_fused_moe_lora_fp8,
mutates_args=["output"],
fake_impl=_fused_moe_lora_fp8_fake,
)
direct_register_custom_op(
op_name="fused_moe_lora_shrink_fp8",
op_func=_fused_moe_lora_shrink_fp8,
mutates_args=["a_intermediate_cache1"],
fake_impl=_fused_moe_lora_shrink_fp8_fake,
)
direct_register_custom_op(
op_name="fused_moe_lora_expand_fp8",
op_func=_fused_moe_lora_expand_fp8,
mutates_args=["output"],
fake_impl=_fused_moe_lora_expand_fp8_fake,
)
fused_moe_lora_fp8 = torch.ops.vllm.fused_moe_lora_fp8
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
-113
View File
@@ -1662,136 +1662,23 @@ def _fused_moe_lora(
)
def _fused_moe_lora_fake(
output: torch.Tensor,
qcurr_hidden_states: torch.Tensor,
lora_a_stacked: list[torch.Tensor],
lora_b_stacked: list[torch.Tensor],
topk_weights: torch.Tensor,
sorted_token_ids: torch.Tensor | None,
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor | None,
token_lora_mapping: torch.Tensor,
max_lora_rank: int,
top_k_num: int,
lora_ids: torch.Tensor,
num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs
adapter_enabled: torch.Tensor,
shrink_block_size_m: int,
shrink_block_size_n: int,
shrink_block_size_k: int,
shrink_group_size_m: int,
shrink_num_warps: int,
shrink_num_stages: int,
shrink_split_k: int,
expand_block_size_m: int,
expand_block_size_n: int,
expand_block_size_k: int,
expand_group_size_m: int,
expand_num_warps: int,
expand_num_stages: int,
expand_split_k: int,
mul_routed_weight: bool = False,
fully_sharded: bool = False,
offset: int = 0,
add_inputs: bool = True,
) -> None:
return
def _fused_moe_lora_shrink_fake(
a_intermediate_cache1: torch.Tensor,
qcurr_hidden_states: torch.Tensor,
lora_a_stacked: list[torch.Tensor],
topk_weights: torch.Tensor,
sorted_token_ids: torch.Tensor | None,
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor | None,
token_lora_mapping: torch.Tensor,
top_k_num: int,
lora_ids: torch.Tensor,
adapter_enabled: torch.Tensor,
device: torch.device,
N: int,
M: int,
EM: int,
K: int,
num_tokens: int,
num_experts: int,
num_slices: int,
block_size_m: int,
block_size_n: int,
block_size_k: int,
group_size_m: int,
num_warps: int,
num_stages: int,
split_k: int,
num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs
mul_routed_weight: bool = False,
use_gdc: bool = False,
use_tma: bool = False,
) -> None:
return
def _fused_moe_lora_expand_fake(
output: torch.Tensor,
a_intermediate_cache1: torch.Tensor,
lora_b_stacked: list[torch.Tensor],
topk_weights: torch.Tensor,
sorted_token_ids: torch.Tensor | None,
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor | None,
token_lora_mapping: torch.Tensor,
top_k_num: int,
lora_ids: torch.Tensor,
adapter_enabled: torch.Tensor,
device: torch.device,
N: int,
M: int,
EM: int,
K: int,
num_tokens: int,
num_experts: int,
num_slices: int,
max_lora_rank: int,
w1_output_dim_size: int,
block_size_m: int,
block_size_n: int,
block_size_k: int,
group_size_m: int,
num_warps: int,
num_stages: int,
split_k: int,
num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs
mul_routed_weight: bool = False,
offset: int = 0,
use_gdc: bool = False,
use_tma: bool = False,
) -> None:
return
try:
direct_register_custom_op(
op_name="fused_moe_lora",
op_func=_fused_moe_lora,
mutates_args=["output"],
fake_impl=_fused_moe_lora_fake,
)
direct_register_custom_op(
op_name="fused_moe_lora_shrink",
op_func=_fused_moe_lora_shrink,
mutates_args=["a_intermediate_cache1"],
fake_impl=_fused_moe_lora_shrink_fake,
)
direct_register_custom_op(
op_name="fused_moe_lora_expand",
op_func=_fused_moe_lora_expand,
mutates_args=["output"],
fake_impl=_fused_moe_lora_expand_fake,
)
fused_moe_lora = torch.ops.vllm.fused_moe_lora
vllm/lora/ops/triton_ops/lora_expand_fp8_op.py
-24
View File
@@ -367,35 +367,11 @@ def _lora_expand_fp8(
return
def _lora_expand_fp8_fake(
inputs: torch.Tensor,
lora_b_weights: list[torch.Tensor],
output_tensor: torch.Tensor,
token_lora_mapping: torch.Tensor,
token_indices_sorted_by_lora_ids: torch.Tensor,
num_tokens_per_lora: torch.Tensor,
lora_token_start_loc: torch.Tensor,
lora_ids: torch.Tensor,
no_lora_flag_cpu: torch.Tensor,
num_active_loras: int,
b_scale: list[torch.Tensor],
a_scale: torch.Tensor | None = None,
offset_start: int = 0,
add_inputs: bool = False,
group_k: int = 0,
group_n: int = 0,
use_fp8_w8a8: bool = False,
per_channel_quant: bool = False,
) -> None:
return
try:
direct_register_custom_op(
op_name="lora_expand_fp8",
op_func=_lora_expand_fp8,
mutates_args=["output_tensor"],
fake_impl=_lora_expand_fp8_fake,
)
lora_expand_fp8 = torch.ops.vllm.lora_expand_fp8
vllm/lora/ops/triton_ops/lora_expand_op.py
-18
View File
@@ -284,29 +284,11 @@ def _lora_expand(
return
def _lora_expand_fake(
inputs: torch.Tensor,
lora_b_weights: list[torch.Tensor],
output_tensor: torch.Tensor,
token_lora_mapping: torch.Tensor,
token_indices_sorted_by_lora_ids: torch.Tensor,
num_tokens_per_lora: torch.Tensor,
lora_token_start_loc: torch.Tensor,
lora_ids: torch.Tensor,
no_lora_flag_cpu: torch.Tensor,
num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs
offset_start: int = 0,
add_inputs: bool = False,
) -> None:
return
try:
direct_register_custom_op(
op_name="lora_expand",
op_func=_lora_expand,
mutates_args=["output_tensor"],
fake_impl=_lora_expand_fake,
)
lora_expand = torch.ops.vllm.lora_expand
vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py
-23
View File
@@ -394,34 +394,11 @@ def _lora_shrink_fp8(
return
def _lora_shrink_fp8_fake(
inputs: torch.Tensor,
lora_a_weights: list[torch.Tensor],
output_tensor: torch.Tensor,
token_lora_mapping: torch.Tensor,
token_indices_sorted_by_lora_ids: torch.Tensor,
num_tokens_per_lora: torch.Tensor,
lora_token_start_loc: torch.Tensor,
lora_ids: torch.Tensor,
no_lora_flag_cpu: torch.Tensor,
num_active_loras: int,
scaling: float,
b_scale: list[torch.Tensor], # LoRA weight scale per slice
a_scale: torch.Tensor | None = None, # Activation scale - per-token or block-wise
group_k: int = 0, # Block size for K in block-wise quantization (0 = tensor-wise)
group_n: int = 0, # Block size for N in block-wise quantization
use_fp8_w8a8: bool = False,
per_channel_quant: bool = False,
) -> None:
return
try:
direct_register_custom_op(
op_name="lora_shrink_fp8",
op_func=_lora_shrink_fp8,
mutates_args=["output_tensor"],
fake_impl=_lora_shrink_fp8_fake,
)
lora_shrink_fp8 = torch.ops.vllm.lora_shrink_fp8
vllm/lora/ops/triton_ops/lora_shrink_op.py
-17
View File
@@ -265,28 +265,11 @@ def _lora_shrink(
return
def _lora_shrink_fake(
inputs: torch.Tensor,
lora_a_weights: list[torch.Tensor],
output_tensor: torch.Tensor,
token_lora_mapping: torch.Tensor,
token_indices_sorted_by_lora_ids: torch.Tensor,
num_tokens_per_lora: torch.Tensor,
lora_token_start_loc: torch.Tensor,
lora_ids: torch.Tensor,
no_lora_flag_cpu: torch.Tensor,
num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs
scaling: float,
) -> None:
return
try:
direct_register_custom_op(
op_name="lora_shrink",
op_func=_lora_shrink,
mutates_args=["output_tensor"],
fake_impl=_lora_shrink_fake,
)
lora_shrink = torch.ops.vllm.lora_shrink
vllm/lora/ops/xpu_ops/lora_ops.py
-35
View File
@@ -85,57 +85,22 @@ def _bgmv_expand_slice_impl(
)
def _bgmv_shrink_fake(
inputs: torch.Tensor,
lora_a_weights: torch.Tensor,
output_tensor: torch.Tensor,
lora_indices_tensor: torch.Tensor,
scaling: float,
) -> None:
return None
def _bgmv_expand_fake(
inputs: torch.Tensor,
lora_b_weights: torch.Tensor,
output_tensor: torch.Tensor,
lora_indices_tensor: torch.Tensor,
add_inputs: bool,
) -> None:
return None
def _bgmv_expand_slice_fake(
inputs: torch.Tensor,
lora_b_weights: torch.Tensor,
output_tensor: torch.Tensor,
lora_indices_tensor: torch.Tensor,
slice_offset: int,
slice_size: int,
add_inputs: bool,
) -> None:
return None
direct_register_custom_op(
op_name="xpu_bgmv_shrink",
op_func=_bgmv_shrink_impl,
mutates_args=["output_tensor"],
fake_impl=_bgmv_shrink_fake,
)
direct_register_custom_op(
op_name="xpu_bgmv_expand",
op_func=_bgmv_expand_impl,
mutates_args=["output_tensor"],
fake_impl=_bgmv_expand_fake,
)
direct_register_custom_op(
op_name="xpu_bgmv_expand_slice",
op_func=_bgmv_expand_slice_impl,
mutates_args=["output_tensor"],
fake_impl=_bgmv_expand_slice_fake,
)
vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py
-12
View File
@@ -139,20 +139,8 @@ def _fp8_gemm_nt_op(
)
def _fp8_gemm_nt_op_fake(
q_input: torch.Tensor,
input_scale: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
output: torch.Tensor,
use_deep_gemm_e8m0: bool,
) -> None:
return None
direct_register_custom_op(
"fp8_gemm_nt_op",
_fp8_gemm_nt_op,
mutates_args=["output"],
fake_impl=_fp8_gemm_nt_op_fake,
)
vllm/model_executor/layers/attention/attention.py
-14
View File
@@ -777,22 +777,8 @@ def unified_attention_with_output(
)
def unified_attention_with_output_fake(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
output: torch.Tensor,
layer_name: LayerNameType,
output_scale: torch.Tensor | None = None,
output_block_scale: torch.Tensor | None = None,
kv_cache_dummy_dep: torch.Tensor | None = None,
) -> None:
return
direct_register_custom_op(
op_name="unified_attention_with_output",
op_func=unified_attention_with_output,
mutates_args=["output", "output_block_scale"],
fake_impl=unified_attention_with_output_fake,
)
vllm/model_executor/layers/attention/mla_attention.py
-19
View File
@@ -1424,29 +1424,10 @@ def unified_mla_attention_with_output(
)
def unified_mla_attention_with_output_fake(
q: torch.Tensor,
kv_c_normed: torch.Tensor,
k_pe: torch.Tensor,
output: torch.Tensor,
layer_name: LayerNameType,
output_scale: torch.Tensor | None = None,
output_block_scale: torch.Tensor | None = None,
kv_cache_dummy_dep: torch.Tensor | None = None,
quant_group_size: int | None = None,
quant_scale_ue8m0: bool | None = None,
quant_col_major: bool | None = None,
quant_tma_aligned: bool | None = None,
q_dcp_replicated: torch.Tensor | None = None,
) -> None:
return
direct_register_custom_op(
op_name="unified_mla_attention_with_output",
op_func=unified_mla_attention_with_output,
mutates_args=["output", "output_block_scale"],
fake_impl=unified_mla_attention_with_output_fake,
dispatch_key=current_platform.dispatch_key,
tags=(torch.Tag.flexible_layout,),
)
vllm/model_executor/layers/attention/static_sink_attention.py
-8
View File
@@ -241,16 +241,8 @@ def maybe_populate_sink(
self.populate_sink_kv(self_kv_cache)
def maybe_populate_sink_fake(
self_kv_cache: torch.Tensor,
layer_name: LayerNameType,
) -> None:
return
direct_register_custom_op(
op_name="maybe_populate_sink",
op_func=maybe_populate_sink,
mutates_args=["self_kv_cache"],
fake_impl=maybe_populate_sink_fake,
)
vllm/model_executor/layers/hpc/rope_norm.py
-10
View File
@@ -78,20 +78,10 @@ def hpc_rope_norm_forward(
rope_norm._forward_impl(qkv, kv_cache, attn_metadata, attn_layer, output)
def hpc_rope_norm_forward_fake(
qkv: torch.Tensor,
output: torch.Tensor,
layer_name: str,
) -> None:
"""Fake impl for torch.compile trace; output is a mutated arg."""
return
direct_register_custom_op(
op_name="hpc_rope_norm_forward",
op_func=hpc_rope_norm_forward,
mutates_args=["output"],
fake_impl=hpc_rope_norm_forward_fake,
)
vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py
-10
View File
@@ -546,20 +546,10 @@ def olmo_hybrid_gdn_full_forward(
)
def olmo_hybrid_gdn_full_forward_fake(
hidden_states: torch.Tensor,
output: torch.Tensor,
layer_name: str,
) -> None:
"""Fake implementation for torch.compile."""
return
direct_register_custom_op(
op_name="olmo_hybrid_gdn_full_forward",
op_func=olmo_hybrid_gdn_full_forward,
mutates_args=["output"],
fake_impl=olmo_hybrid_gdn_full_forward_fake,
)
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
-23
View File
@@ -1949,23 +1949,10 @@ def qwen_gdn_attention_core(
)
def gdn_attention_core_fake(
qkv_or_qkvz: torch.Tensor,
b_or_ba: torch.Tensor,
a_or_z_out: torch.Tensor,
core_attn_out: torch.Tensor,
layer_name: LayerNameType,
use_aiter: bool = False,
) -> None:
"""Fake implementation for torch.compile."""
return
direct_register_custom_op(
op_name="qwen_gdn_attention_core",
op_func=qwen_gdn_attention_core,
mutates_args=["a_or_z_out", "core_attn_out"],
fake_impl=gdn_attention_core_fake,
)
@@ -1985,20 +1972,10 @@ def qwen_gdn_attention_core_fused_norm_packed(
)
def gdn_attention_core_fused_norm_packed_fake(
mixed_qkvz: torch.Tensor,
ba: torch.Tensor,
core_attn_out: torch.Tensor,
layer_name: LayerNameType,
) -> None:
return
direct_register_custom_op(
op_name="qwen_gdn_attention_core_fused_norm_packed",
op_func=qwen_gdn_attention_core_fused_norm_packed,
mutates_args=["core_attn_out"],
fake_impl=gdn_attention_core_fused_norm_packed_fake,
)
vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py
-10
View File
@@ -336,18 +336,8 @@ def linear_attention(
self._forward(hidden_states=hidden_states, output=output, positions=positions)
def linear_attention_fake(
hidden_states: torch.Tensor,
output: torch.Tensor,
positions: torch.Tensor,
layer_name: str,
) -> None:
return
direct_register_custom_op(
op_name="linear_attention",
op_func=linear_attention,
mutates_args=["output"],
fake_impl=linear_attention_fake,
)
vllm/model_executor/layers/mamba/mamba_mixer.py
-9
View File
@@ -532,17 +532,8 @@ def mamba_mixer(
self.forward_impl(hidden_states=hidden_states, output=output)
def mamba_mixer_fake(
hidden_states: torch.Tensor,
output: torch.Tensor,
layer_name: LayerNameType,
) -> None:
return
direct_register_custom_op(
op_name="mamba_mixer",
op_func=mamba_mixer,
mutates_args=["output"],
fake_impl=mamba_mixer_fake,
)
vllm/model_executor/layers/mamba/mamba_mixer2.py
-9
View File
@@ -1271,17 +1271,8 @@ def mamba_mixer2(
self.conv_ssm_forward(projected_states=projected_states, output=output)
def mamba_mixer2_fake(
projected_states: torch.Tensor,
output: torch.Tensor,
layer_name: LayerNameType,
) -> None:
return
direct_register_custom_op(
op_name="mamba_mixer2",
op_func=mamba_mixer2,
mutates_args=["output"],
fake_impl=mamba_mixer2_fake,
)
vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py
-12
View File
@@ -718,17 +718,6 @@ def _spec_aware_nonspec_subset(
return attn_out.squeeze(0)
def cpu_gdn_attention_core_fake(
mixed_qkv: torch.Tensor,
b: torch.Tensor,
a: torch.Tensor,
core_attn_out: torch.Tensor,
layer_name: LayerNameType,
) -> None:
"""Fake implementation for torch.compile."""
return
def register_cpu_gdn_attention_ops() -> None:
global _CPU_GDN_ATTENTION_OPS_REGISTERED
if _CPU_GDN_ATTENTION_OPS_REGISTERED:
@@ -738,6 +727,5 @@ def register_cpu_gdn_attention_ops() -> None:
op_name="cpu_gdn_attention_core",
op_func=cpu_gdn_attention_core,
mutates_args=["core_attn_out"],
fake_impl=cpu_gdn_attention_core_fake,
)
_CPU_GDN_ATTENTION_OPS_REGISTERED = True
vllm/model_executor/layers/mamba/short_conv.py
-9
View File
@@ -373,17 +373,8 @@ def short_conv(
self.forward_native(hidden_states=hidden_states, output=output)
def short_conv_fake(
hidden_states: torch.Tensor,
output: torch.Tensor,
layer_name: str,
) -> None:
return
direct_register_custom_op(
op_name="short_conv",
op_func=short_conv,
mutates_args=["output"],
fake_impl=short_conv_fake,
)
vllm/model_executor/layers/rotary_embedding/common.py
-12
View File
@@ -100,23 +100,11 @@ def _flashinfer_rotary_embedding(
)
def _flashinfer_rotary_embedding_fake(
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor,
head_size: int,
cos_sin_cache: torch.Tensor,
is_neox: bool,
) -> None:
return
# Register flashinfer rotary embedding custom op
direct_register_custom_op(
op_name="flashinfer_rotary_embedding",
op_func=_flashinfer_rotary_embedding,
mutates_args=["query", "key"], # These tensors are modified in-place
fake_impl=_flashinfer_rotary_embedding_fake,
)
vllm/model_executor/models/bailing_moe_v3.py
-13
View File
@@ -115,23 +115,10 @@ def bailing_v3_kda_attention(
)
def bailing_v3_kda_attention_fake(
q_proj_states: torch.Tensor,
k_proj_states: torch.Tensor,
v_proj_states: torch.Tensor,
g1: torch.Tensor,
beta: torch.Tensor,
core_attn_out: torch.Tensor,
layer_name: str,
) -> None:
return
direct_register_custom_op(
op_name="bailing_v3_kda_attention",
op_func=bailing_v3_kda_attention,
mutates_args=["core_attn_out"],
fake_impl=bailing_v3_kda_attention_fake,
)
vllm/model_executor/offloader/prefetch_ops.py
-18
View File
@@ -33,14 +33,6 @@ def _wait_prefetch_impl(
get_offloader()._wait_for_layer(layer_idx)
def _wait_prefetch_fake(
input_tensor: torch.Tensor,
layer_idx: int,
) -> None:
"""Fake implementation for torch.compile tracing."""
return
# --- start_prefetch op ---
@@ -61,14 +53,6 @@ def _start_prefetch_impl(
get_offloader()._start_prefetch(layer_idx)
def _start_prefetch_fake(
output_tensor: torch.Tensor,
layer_idx: int,
) -> None:
"""Fake implementation for torch.compile tracing."""
return
def register_prefetch_offloader_ops() -> None:
"""Register custom ops for prefetch offloader.
@@ -79,14 +63,12 @@ def register_prefetch_offloader_ops() -> None:
op_name="wait_prefetch",
op_func=_wait_prefetch_impl,
mutates_args=["input_tensor"],
fake_impl=_wait_prefetch_fake,
)
direct_register_custom_op(
op_name="start_prefetch",
op_func=_start_prefetch_impl,
mutates_args=["output_tensor"],
fake_impl=_start_prefetch_fake,
)
vllm/models/deepseek_v4/xpu/model.py
-13
View File
@@ -578,23 +578,10 @@ def _deepseek_v4_mega_moe_experts_op(
)
def _deepseek_v4_mega_moe_experts_op_fake(
hidden_states: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
out: torch.Tensor,
layer_name: str,
activation_clamp: float | None,
fast_math: bool,
) -> None:
return None
direct_register_custom_op(
op_name="deepseek_v4_mega_moe_experts",
op_func=_deepseek_v4_mega_moe_experts_op,
mutates_args=["out"],
fake_impl=_deepseek_v4_mega_moe_experts_op_fake,
)
vllm/models/qwen4_exp/amd/ple_layer.py
-18
View File
@@ -1082,14 +1082,6 @@ def qwen4_exp_amd_ple_ngram_embedding(
output.copy_(result)
def qwen4_exp_amd_ple_ngram_embedding_fake(
ngram_ids: torch.Tensor,
output: torch.Tensor,
layer_name: str,
) -> None:
return
def qwen4_exp_ple_short_conv(
inputs: torch.Tensor,
output: torch.Tensor,
@@ -1100,19 +1092,10 @@ def qwen4_exp_ple_short_conv(
output[: result.shape[0]].copy_(result)
def qwen4_exp_ple_short_conv_fake(
inputs: torch.Tensor,
output: torch.Tensor,
layer_name: str,
) -> None:
return
direct_register_custom_op(
op_name="qwen4_exp_amd_ple_ngram_embedding",
op_func=qwen4_exp_amd_ple_ngram_embedding,
mutates_args=["output"],
fake_impl=qwen4_exp_amd_ple_ngram_embedding_fake,
)
@@ -1120,7 +1103,6 @@ direct_register_custom_op(
op_name="qwen4_exp_ple_short_conv",
op_func=qwen4_exp_ple_short_conv,
mutates_args=["output"],
fake_impl=qwen4_exp_ple_short_conv_fake,
)
vllm/models/qwen4_exp/amd/qsa.py
-13
View File
@@ -445,23 +445,10 @@ def qwen4_exp_qsa_with_output(
)
def qwen4_exp_qsa_with_output_fake(
hidden_states: torch.Tensor,
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
output: torch.Tensor,
layer_name: LayerNameType,
) -> None:
del hidden_states, positions, query, key, value, output, layer_name
direct_register_custom_op(
op_name="qwen4_exp_qsa_with_output",
op_func=qwen4_exp_qsa_with_output,
mutates_args=["output"],
fake_impl=qwen4_exp_qsa_with_output_fake,
)
vllm/models/qwen4_exp/nvidia/ple_layer.py
-20
View File
@@ -877,16 +877,6 @@ def qwen4_exp_compute_ple_ngram_ids(
)
def qwen4_exp_compute_ple_ngram_ids_fake(
input_ids: torch.Tensor,
query_start_loc: torch.Tensor,
ngram_context: torch.Tensor,
output: torch.Tensor,
layer_name: str,
) -> None:
return
def qwen4_exp_ple_short_conv(
inputs: torch.Tensor,
residual_output: torch.Tensor,
@@ -896,19 +886,10 @@ def qwen4_exp_ple_short_conv(
layer._short_conv(inputs, residual_output)
def qwen4_exp_ple_short_conv_fake(
inputs: torch.Tensor,
residual_output: torch.Tensor,
layer_name: str,
) -> None:
return
direct_register_custom_op(
op_name="qwen4_exp_compute_ple_ngram_ids",
op_func=qwen4_exp_compute_ple_ngram_ids,
mutates_args=["output"],
fake_impl=qwen4_exp_compute_ple_ngram_ids_fake,
)
@@ -916,7 +897,6 @@ direct_register_custom_op(
op_name="qwen4_exp_ple_short_conv",
op_func=qwen4_exp_ple_short_conv,
mutates_args=["residual_output"],
fake_impl=qwen4_exp_ple_short_conv_fake,
)
vllm/models/qwen4_exp/nvidia/qsa.py
-13
View File
@@ -465,23 +465,10 @@ def qwen4_exp_qsa_with_output(
)
def qwen4_exp_qsa_with_output_fake(
hidden_states: torch.Tensor,
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
output: torch.Tensor,
layer_name: LayerNameType,
) -> None:
del hidden_states, positions, query, key, value, output, layer_name
direct_register_custom_op(
op_name="qwen4_exp_qsa_with_output",
op_func=qwen4_exp_qsa_with_output,
mutates_args=["output"],
fake_impl=qwen4_exp_qsa_with_output_fake,
)
vllm/utils/flashinfer.py
-8
View File
@@ -708,19 +708,11 @@ if has_flashinfer():
concat_mla_k(k, k_nope, k_pe)
def _flashinfer_concat_mla_k_fake(
k: torch.Tensor,
k_nope: torch.Tensor,
k_pe: torch.Tensor,
) -> None:
return
# Register flashinfer concat_mla_k custom op
direct_register_custom_op(
op_name="flashinfer_concat_mla_k",
op_func=_flashinfer_concat_mla_k,
mutates_args=["k"], # k tensor is modified in-place
fake_impl=_flashinfer_concat_mla_k_fake,
)
@torch.library.custom_op(