From 199cb9b964822e59ab9b58d88e7be31eb419a2ae Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Mon, 7 Sep 2026 08:54:32 +0800 Subject: [PATCH] [Kernel] Remove unused fake implementation (#55535) Signed-off-by: Jee Jee Li --- vllm/_aiter_ops.py | 97 ------------- vllm/_custom_ops.py | 56 -------- vllm/_xpu_ops.py | 60 -------- .../passes/fusion/allreduce_rms_fusion.py | 19 --- .../ops/dynamic_per_token_scaled_fp8_quant.py | 10 -- vllm/kernels/helion/ops/fused_qk_norm_rope.py | 18 --- .../helion/ops/per_token_group_fp8_quant.py | 18 --- .../ops/rms_norm_dynamic_per_token_quant.py | 13 -- .../helion/ops/rms_norm_per_block_quant.py | 15 -- .../ops/silu_and_mul_per_block_quant.py | 12 -- vllm/kernels/helion/register.py | 9 +- .../ops/triton_ops/fused_moe_lora_fp8_op.py | 133 ------------------ vllm/lora/ops/triton_ops/fused_moe_lora_op.py | 113 --------------- .../lora/ops/triton_ops/lora_expand_fp8_op.py | 24 ---- vllm/lora/ops/triton_ops/lora_expand_op.py | 18 --- .../lora/ops/triton_ops/lora_shrink_fp8_op.py | 23 --- vllm/lora/ops/triton_ops/lora_shrink_op.py | 17 --- vllm/lora/ops/xpu_ops/lora_ops.py | 35 ----- .../kernels/linear/scaled_mm/deep_gemm.py | 12 -- .../layers/attention/attention.py | 14 -- .../layers/attention/mla_attention.py | 19 --- .../layers/attention/static_sink_attention.py | 8 -- vllm/model_executor/layers/hpc/rope_norm.py | 10 -- .../layers/mamba/gdn/olmo_gdn_linear_attn.py | 10 -- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 23 --- .../mamba/linear/minimax_linear_attn.py | 10 -- .../layers/mamba/mamba_mixer.py | 9 -- .../layers/mamba/mamba_mixer2.py | 9 -- .../layers/mamba/ops/cpu/gdn_attention.py | 12 -- .../model_executor/layers/mamba/short_conv.py | 9 -- .../layers/rotary_embedding/common.py | 12 -- vllm/model_executor/models/bailing_moe_v3.py | 13 -- vllm/model_executor/offloader/prefetch_ops.py | 18 --- vllm/models/deepseek_v4/xpu/model.py | 13 -- vllm/models/qwen4_exp/amd/ple_layer.py | 18 --- vllm/models/qwen4_exp/amd/qsa.py | 13 -- vllm/models/qwen4_exp/nvidia/ple_layer.py | 20 --- vllm/models/qwen4_exp/nvidia/qsa.py | 13 -- vllm/utils/flashinfer.py | 8 -- 39 files changed, 7 insertions(+), 956 deletions(-) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index a68b187598..cde07eb8eb 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -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( diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index ecc3cf3c9f..450a7a57da 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -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) diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index 112f5478f0..df66e64033 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -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 diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index 559a99ddd9..5e7570ff0d 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -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 diff --git a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py index 75f240c2aa..0ebbf72409 100644 --- a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py +++ b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py @@ -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], diff --git a/vllm/kernels/helion/ops/fused_qk_norm_rope.py b/vllm/kernels/helion/ops/fused_qk_norm_rope.py index c97ad1e514..dbbb9e0886 100644 --- a/vllm/kernels/helion/ops/fused_qk_norm_rope.py +++ b/vllm/kernels/helion/ops/fused_qk_norm_rope.py @@ -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, diff --git a/vllm/kernels/helion/ops/per_token_group_fp8_quant.py b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py index 42fd0215f5..789fa8dd65 100644 --- a/vllm/kernels/helion/ops/per_token_group_fp8_quant.py +++ b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py @@ -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, ), diff --git a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py index 8bf352cc51..276c7f9670 100644 --- a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py +++ b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py @@ -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], diff --git a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py index fe995b2c62..8c357386bd 100644 --- a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py +++ b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py @@ -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], diff --git a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py index 31386943e7..7cc0f99ea5 100644 --- a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py +++ b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py @@ -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], diff --git a/vllm/kernels/helion/register.py b/vllm/kernels/helion/register.py index 764022de77..438011b1aa 100644 --- a/vllm/kernels/helion/register.py +++ b/vllm/kernels/helion/register.py @@ -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'", diff --git a/vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py b/vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py index deb34cfe43..033ed1df6c 100644 --- a/vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py +++ b/vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py @@ -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 diff --git a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py b/vllm/lora/ops/triton_ops/fused_moe_lora_op.py index 3133ea6228..1bde7e0845 100644 --- a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py +++ b/vllm/lora/ops/triton_ops/fused_moe_lora_op.py @@ -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 diff --git a/vllm/lora/ops/triton_ops/lora_expand_fp8_op.py b/vllm/lora/ops/triton_ops/lora_expand_fp8_op.py index d5850f1181..2dc33965e3 100644 --- a/vllm/lora/ops/triton_ops/lora_expand_fp8_op.py +++ b/vllm/lora/ops/triton_ops/lora_expand_fp8_op.py @@ -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 diff --git a/vllm/lora/ops/triton_ops/lora_expand_op.py b/vllm/lora/ops/triton_ops/lora_expand_op.py index 7f8ed577ec..40f640e243 100644 --- a/vllm/lora/ops/triton_ops/lora_expand_op.py +++ b/vllm/lora/ops/triton_ops/lora_expand_op.py @@ -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 diff --git a/vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py b/vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py index d58368753d..9b4fe75f3c 100644 --- a/vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py +++ b/vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py @@ -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 diff --git a/vllm/lora/ops/triton_ops/lora_shrink_op.py b/vllm/lora/ops/triton_ops/lora_shrink_op.py index 88c24c740d..4144c16ead 100644 --- a/vllm/lora/ops/triton_ops/lora_shrink_op.py +++ b/vllm/lora/ops/triton_ops/lora_shrink_op.py @@ -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 diff --git a/vllm/lora/ops/xpu_ops/lora_ops.py b/vllm/lora/ops/xpu_ops/lora_ops.py index ca363b94ab..b55419acfd 100644 --- a/vllm/lora/ops/xpu_ops/lora_ops.py +++ b/vllm/lora/ops/xpu_ops/lora_ops.py @@ -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, ) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py b/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py index 70122f7b4a..79458f5cda 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py @@ -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, ) diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index c6ba8ea712..a72d36a718 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -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, ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 00e10e9142..a45f074837 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -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,), ) diff --git a/vllm/model_executor/layers/attention/static_sink_attention.py b/vllm/model_executor/layers/attention/static_sink_attention.py index 8d199be0a5..d27ccc728a 100644 --- a/vllm/model_executor/layers/attention/static_sink_attention.py +++ b/vllm/model_executor/layers/attention/static_sink_attention.py @@ -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, ) diff --git a/vllm/model_executor/layers/hpc/rope_norm.py b/vllm/model_executor/layers/hpc/rope_norm.py index eed19cf0a4..c1503c754c 100644 --- a/vllm/model_executor/layers/hpc/rope_norm.py +++ b/vllm/model_executor/layers/hpc/rope_norm.py @@ -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, ) diff --git a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py index 2c1deec687..a1a494e121 100644 --- a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py @@ -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, ) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 108a73223e..718f8de05c 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -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, ) diff --git a/vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py b/vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py index 3842e3e2cd..e7618b29b7 100644 --- a/vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py +++ b/vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py @@ -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, ) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py index 42f060e816..5d09da4c3b 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer.py @@ -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, ) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 0de22ead83..5bc890d007 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -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, ) diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index 3f28babb44..b5c1d11a16 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -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 diff --git a/vllm/model_executor/layers/mamba/short_conv.py b/vllm/model_executor/layers/mamba/short_conv.py index 3f152623bd..37aef2b5bd 100644 --- a/vllm/model_executor/layers/mamba/short_conv.py +++ b/vllm/model_executor/layers/mamba/short_conv.py @@ -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, ) diff --git a/vllm/model_executor/layers/rotary_embedding/common.py b/vllm/model_executor/layers/rotary_embedding/common.py index 17cf66b025..0a392812b6 100644 --- a/vllm/model_executor/layers/rotary_embedding/common.py +++ b/vllm/model_executor/layers/rotary_embedding/common.py @@ -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, ) diff --git a/vllm/model_executor/models/bailing_moe_v3.py b/vllm/model_executor/models/bailing_moe_v3.py index 0c15f073a1..e788771f74 100644 --- a/vllm/model_executor/models/bailing_moe_v3.py +++ b/vllm/model_executor/models/bailing_moe_v3.py @@ -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, ) diff --git a/vllm/model_executor/offloader/prefetch_ops.py b/vllm/model_executor/offloader/prefetch_ops.py index d1f59b67b4..b7f599210a 100644 --- a/vllm/model_executor/offloader/prefetch_ops.py +++ b/vllm/model_executor/offloader/prefetch_ops.py @@ -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, ) diff --git a/vllm/models/deepseek_v4/xpu/model.py b/vllm/models/deepseek_v4/xpu/model.py index ddf8ae175b..fe50160346 100644 --- a/vllm/models/deepseek_v4/xpu/model.py +++ b/vllm/models/deepseek_v4/xpu/model.py @@ -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, ) diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py index 69215b1aa4..fc494be055 100644 --- a/vllm/models/qwen4_exp/amd/ple_layer.py +++ b/vllm/models/qwen4_exp/amd/ple_layer.py @@ -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, ) diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py index 0e65956e96..c0c766aad1 100644 --- a/vllm/models/qwen4_exp/amd/qsa.py +++ b/vllm/models/qwen4_exp/amd/qsa.py @@ -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, ) diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 9d7091d7a2..7c2ab8a8d5 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -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, ) diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 2e135149af..53ad18362e 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -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, ) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 3b3457c4f0..294e91391b 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -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(