199cb9b964 · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/_aiter_ops.py

-97

@@ -387,18 +387,6 @@ def _rocm_aiter_topk_softmax_impl(

387 )

388

389

390 -def _rocm_aiter_topk_softmax_fake(

391 - topk_weights: torch.Tensor,

392 - topk_indices: torch.Tensor,

393 - token_expert_indices: torch.Tensor,

394 - gating_output: torch.Tensor,

395 - renormalize: bool,

396 - num_shared_experts: int = 0,

397 - shared_expert_scoring_func: str = "",

398 -) -> None:

399 - pass

400 -

401 -

390 def _rocm_aiter_topk_sigmoid_impl(

391 topk_weights: torch.Tensor,

392 topk_indices: torch.Tensor,

@@ -409,14 +397,6 @@ def _rocm_aiter_topk_sigmoid_impl(

397 topk_sigmoid(topk_weights, topk_indices, gating_output)

398

399

412 -def _rocm_aiter_topk_sigmoid_fake(

413 - topk_weights: torch.Tensor,

414 - topk_indices: torch.Tensor,

415 - gating_output: torch.Tensor,

416 -) -> None:

417 - pass

418 -

419 -

400 def _rocm_aiter_biased_grouped_topk_impl(

401 gating_output: torch.Tensor,

402 correction_bias: torch.Tensor,

@@ -441,19 +421,6 @@ def _rocm_aiter_biased_grouped_topk_impl(

421 )

422

423

444 -def _rocm_aiter_biased_grouped_topk_fake(

445 - gating_output: torch.Tensor,

446 - correction_bias: torch.Tensor,

447 - topk_weights: torch.Tensor,

448 - topk_ids: torch.Tensor,

449 - num_expert_group: int,

450 - topk_group: int,

451 - need_renorm: bool,

452 - routed_scaling_factor: float = 1.0, # mul to topk_weights

453 -) -> None:

454 - pass

455 -

456 -

424 def _rocm_aiter_grouped_topk_impl(

425 gating_output: torch.Tensor,

426 topk_weights: torch.Tensor,

@@ -479,19 +446,6 @@ def _rocm_aiter_grouped_topk_impl(

446 )

447

448

482 -def _rocm_aiter_grouped_topk_fake(

483 - gating_output: torch.Tensor,

484 - topk_weights: torch.Tensor,

485 - topk_ids: torch.Tensor,

486 - num_expert_group: int,

487 - topk_group: int,

488 - need_renorm: bool,

489 - scoring_func: str = "softmax",

490 - routed_scaling_factor: float = 1.0, # mul to topk_weights

491 -) -> None:

492 - pass

493 -

494 -

449 def _rocm_aiter_fused_topk_impl(

450 x: torch.Tensor,

451 router_logits: torch.Tensor,

@@ -618,29 +572,6 @@ def _rocm_aiter_mla_decode_fwd_impl(

572 )

573

574

621 -def _rocm_aiter_mla_decode_fwd_fake(

622 - q: torch.Tensor,

623 - kv_buffer: torch.Tensor,

624 - o: torch.Tensor,

625 - qo_indptr: torch.Tensor,

626 - max_seqlen_qo: int,

627 - kv_indptr: torch.Tensor | None = None,

628 - kv_indices: torch.Tensor | None = None,

629 - kv_last_page_lens: torch.Tensor | None = None,

630 - sm_scale: float = 1.0,

631 - logit_cap: float = 0.0,

632 - q_scale: torch.Tensor | None = None,

633 - kv_scale: torch.Tensor | None = None,

634 - work_meta_data: torch.Tensor | None = None,

635 - work_indptr: torch.Tensor | None = None,

636 - work_info_set: torch.Tensor | None = None,

637 - reduce_indptr: torch.Tensor | None = None,

638 - reduce_final_map: torch.Tensor | None = None,

639 - reduce_partial_map: torch.Tensor | None = None,

640 -) -> None:

641 - pass

642 -

643 -

575 def _rocm_aiter_w8a8_gemm_impl(

576 A: torch.Tensor,

577 B: torch.Tensor,

@@ -1104,15 +1035,6 @@ def _rocm_aiter_per_tensor_quant_impl(

1035 static_per_tensor_quant(out, x, scale)

1036

1037

1107 -def _rocm_aiter_per_tensor_quant_fake(

1108 - out: torch.Tensor,

1109 - x: torch.Tensor,

1110 - scale: torch.Tensor,

1111 - is_dynamic: bool,

1112 -) -> None:

1113 - pass

1114 -

1115 -

1038 def _rocm_aiter_per_token_quant_impl(

1039 x: torch.Tensor, quant_dtype: torch.dtype, scale: torch.Tensor | None = None

1040 ) -> tuple[torch.Tensor, torch.Tensor]:

@@ -1597,18 +1519,6 @@ def _triton_rotary_embedding_impl(

1519 key = key.view(key_shape)

1520

1521

1600 -def _triton_rotary_embedding_fake(

1601 - positions: torch.Tensor,

1602 - query: torch.Tensor,

1603 - key: torch.Tensor,

1604 - head_size: int,

1605 - cos_sin_cache: torch.Tensor,

1606 - is_neox_style: bool,

1607 - offsets: torch.Tensor | None = None,

1608 -) -> None:

1609 - return

1610 -

1611 -

1522 def _rocm_aiter_fp8_attn_impl(

1523 q: torch.Tensor,

1524 k: torch.Tensor,

@@ -2132,7 +2042,6 @@ class rocm_aiter_ops:

2042 op_name="rocm_aiter_topk_softmax",

2043 op_func=_rocm_aiter_topk_softmax_impl,

2044 mutates_args=["topk_weights", "topk_indices", "token_expert_indices"],

2135 - fake_impl=_rocm_aiter_topk_softmax_fake,

2045 dispatch_key=current_platform.dispatch_key,

2046 )

2047

@@ -2140,7 +2049,6 @@ class rocm_aiter_ops:

2049 op_name="rocm_aiter_topk_sigmoid",

2050 op_func=_rocm_aiter_topk_sigmoid_impl,

2051 mutates_args=["topk_weights", "topk_indices"],

2143 - fake_impl=_rocm_aiter_topk_sigmoid_fake,

2052 dispatch_key=current_platform.dispatch_key,

2053 )

2054

@@ -2148,7 +2056,6 @@ class rocm_aiter_ops:

2056 op_name="rocm_aiter_biased_grouped_topk",

2057 op_func=_rocm_aiter_biased_grouped_topk_impl,

2058 mutates_args=["topk_weights", "topk_ids"],

2151 - fake_impl=_rocm_aiter_biased_grouped_topk_fake,

2059 dispatch_key=current_platform.dispatch_key,

2060 )

2061

@@ -2156,7 +2063,6 @@ class rocm_aiter_ops:

2063 op_name="rocm_aiter_grouped_topk",

2064 op_func=_rocm_aiter_grouped_topk_impl,

2065 mutates_args=["topk_weights", "topk_ids"],

2159 - fake_impl=_rocm_aiter_grouped_topk_fake,

2066 dispatch_key=current_platform.dispatch_key,

2067 )

2068

@@ -2172,7 +2078,6 @@ class rocm_aiter_ops:

2078 op_name="rocm_aiter_mla_decode_fwd",

2079 op_func=_rocm_aiter_mla_decode_fwd_impl,

2080 mutates_args=["o"],

2175 - fake_impl=_rocm_aiter_mla_decode_fwd_fake,

2081 )

2082

2083 direct_register_custom_op(

@@ -2266,7 +2171,6 @@ class rocm_aiter_ops:

2171 op_name="rocm_aiter_per_tensor_quant",

2172 op_func=_rocm_aiter_per_tensor_quant_impl,

2173 mutates_args=["out", "scale"],

2269 - fake_impl=_rocm_aiter_per_tensor_quant_fake,

2174 dispatch_key=current_platform.dispatch_key,

2175 )

2176

@@ -2306,7 +2210,6 @@ class rocm_aiter_ops:

2210 op_name="rocm_aiter_triton_rotary_embedding",

2211 op_func=_triton_rotary_embedding_impl,

2212 mutates_args=["query", "key"], # These tensors are modified in-place

2309 - fake_impl=_triton_rotary_embedding_fake,

2213 )

2214

2215 direct_register_custom_op(

vllm/_custom_ops.py

-56

@@ -100,17 +100,6 @@ if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "scaled_fp4_quant"):

100 m = input.numel() // n

101 return create_fp4_output_tensors(m, n, input.device, is_sf_swizzled_layout)

102

103 - @register_fake("_C::scaled_fp4_quant.out")

104 - def _scaled_fp4_quant_out_fake(

105 - input: torch.Tensor,

106 - input_scale: torch.Tensor,

107 - is_sf_swizzled_layout: bool,

108 - *,

109 - output: torch.Tensor,

110 - output_scale: torch.Tensor,

111 - ) -> None:

112 - return None

113 -

103

104 # page attention ops

105 def paged_attention_rocm(

@@ -740,27 +729,6 @@ def moe_gptq_gemm_rdna3(

729 )

730

731

743 -if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3"):

744 -

745 - @register_fake("_rocm_C::moe_gptq_gemm_rdna3")

746 - def _moe_gptq_gemm_rdna3_fake(

747 - a: torch.Tensor,

748 - c: torch.Tensor,

749 - b_q_weight: torch.Tensor,

750 - b_scales: torch.Tensor,

751 - b_qzeros: torch.Tensor,

752 - topk_weights: torch.Tensor,

753 - sorted_token_ids: torch.Tensor,

754 - expert_ids: torch.Tensor,

755 - num_tokens_post_padded: torch.Tensor,

756 - top_k: int,

757 - block_size_m: int,

758 - mul_topk_weight: bool,

759 - output_topk: int = 0,

760 - ) -> None:

761 - return

762 -

763 -

732 if hasattr(torch.ops._C, "allspark_w8a16_gemm"):

733

734 @register_fake("_C::allspark_w8a16_gemm")

@@ -2437,17 +2405,6 @@ def fp32_router_gemm(

2405 return output

2406

2407

2440 -if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"):

2441 -

2442 - @register_fake("_C::fp32_router_gemm")

2443 - def fp32_router_gemm_fake(

2444 - output: torch.Tensor,

2445 - mat_a: torch.Tensor,

2446 - mat_b: torch.Tensor,

2447 - ) -> None:

2448 - return

2449 -

2450 -

2408 def topk_softmax(

2409 topk_weights: torch.Tensor,

2410 topk_ids: torch.Tensor,

@@ -4363,19 +4320,6 @@ def safeFusedQuantizeNv(

4320 return

4321

4322

4366 -if hasattr(torch.ops._qutlass_C, "fusedQuantizeNv"):

4367 -

4368 - @register_fake("vllm::safeFusedQuantizeNv")

4369 - def _fake_fused_quantize_nv(

4370 - a: torch.Tensor,

4371 - b: torch.Tensor,

4372 - xh_e2m1: torch.Tensor,

4373 - xh_e4m3: torch.Tensor,

4374 - global_scale: torch.Tensor,

4375 - ) -> None:

4376 - return

4377 -

4378 -

4323 def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor:

4324 """

4325 Perform Hadamard transforms using [Hadacore](https://arxiv.org/abs/2412.08832)

vllm/_xpu_ops.py

-60

@@ -125,15 +125,6 @@ def _gemma_rms_norm_impl(

125 torch.ops._C.gemma_rms_norm(out, input, weight, epsilon)

126

127

128 -def _gemma_rms_norm_fake(

129 - out: torch.Tensor,

130 - input: torch.Tensor,

131 - weight: torch.Tensor,

132 - epsilon: float,

133 -) -> None:

134 - return None

135 -

136 -

128 def _fused_add_gemma_rms_norm_impl(

129 input: torch.Tensor,

130 residual: torch.Tensor,

@@ -143,15 +134,6 @@ def _fused_add_gemma_rms_norm_impl(

134 torch.ops._C.fused_add_gemma_rms_norm(input, residual, weight, epsilon)

135

136

146 -def _fused_add_gemma_rms_norm_fake(

147 - input: torch.Tensor,

148 - residual: torch.Tensor,

149 - weight: torch.Tensor,

150 - epsilon: float,

151 -) -> None:

152 - return None

153 -

154 -

137 def _gdn_attention_core_xpu_impl(

138 core_attn_out: torch.Tensor,

139 z: torch.Tensor,

@@ -240,16 +222,6 @@ def _gdn_attention_core_xpu_impl(

222 )

223

224

243 -def _gdn_attention_core_xpu_fake(

244 - core_attn_out: torch.Tensor,

245 - z: torch.Tensor,

246 - projected_states_qkvz: torch.Tensor,

247 - projected_states_ba: torch.Tensor,

248 - layer_name: str,

249 -) -> None:

250 - return

251 -

252 -

225 def _xpu_ops_deepseek_scaling_rope_impl(

226 positions: torch.Tensor,

227 query: torch.Tensor,

@@ -471,19 +443,6 @@ def _xpu_deepseek_fused_indexer_q_rope_fp8_impl(

443 )

444

445

474 -def _xpu_deepseek_fused_indexer_q_rope_fp8_fake(

475 - index_q: torch.Tensor,

476 - positions: torch.Tensor,

477 - index_q_cos_sin_cache: torch.Tensor,

478 - index_weights: torch.Tensor,

479 - index_weights_softmax_scale: float,

480 - index_weights_head_scale: float,

481 - index_q_fp8: torch.Tensor,

482 - index_weights_out: torch.Tensor,

483 -) -> None:

484 - return

485 -

486 -

446 def _xpu_deepseek_fused_indexer_q_rope_mxfp4_impl(

447 index_q: torch.Tensor,

448 positions: torch.Tensor,

@@ -528,20 +487,6 @@ def _xpu_deepseek_fused_indexer_q_rope_mxfp4_impl(

487 )

488

489

531 -def _xpu_deepseek_fused_indexer_q_rope_mxfp4_fake(

532 - index_q: torch.Tensor,

533 - positions: torch.Tensor,

534 - index_q_cos_sin_cache: torch.Tensor,

535 - index_weights: torch.Tensor,

536 - index_weights_softmax_scale: float,

537 - index_weights_head_scale: float,

538 - index_q_packed: torch.Tensor,

539 - index_q_scale: torch.Tensor,

540 - index_weights_out: torch.Tensor,

541 -) -> None:

542 - return

543 -

544 -

490 def _xpu_mxfp8_quantize_impl(

491 x: torch.Tensor, dtype: torch.dtype | None = None

492 ) -> tuple[torch.Tensor, torch.Tensor]:

@@ -1262,14 +1207,12 @@ class xpu_ops:

1207 op_name="xpu_gemma_rms_norm",

1208 op_func=_gemma_rms_norm_impl,

1209 mutates_args=["out"],

1265 - fake_impl=_gemma_rms_norm_fake,

1210 )

1211

1212 direct_register_custom_op(

1213 op_name="xpu_fused_add_gemma_rms_norm",

1214 op_func=_fused_add_gemma_rms_norm_impl,

1215 mutates_args=["input", "residual"],

1272 - fake_impl=_fused_add_gemma_rms_norm_fake,

1216 )

1217

1218 direct_register_custom_op(

@@ -1314,7 +1257,6 @@ class xpu_ops:

1257 op_name="gdn_attention_core_xpu",

1258 op_func=eager_break_during_capture(_gdn_attention_core_xpu_impl),

1259 mutates_args=["core_attn_out", "z"],

1317 - fake_impl=_gdn_attention_core_xpu_fake,

1260 )

1261

1262 direct_register_custom_op(

@@ -1327,7 +1269,6 @@ class xpu_ops:

1269 op_name="xpu_deepseek_fused_indexer_q_rope_fp8",

1270 op_func=_xpu_deepseek_fused_indexer_q_rope_fp8_impl,

1271 mutates_args=["index_q_fp8", "index_weights_out"],

1330 - fake_impl=_xpu_deepseek_fused_indexer_q_rope_fp8_fake,

1272 )

1273

1274 direct_register_custom_op(

@@ -1338,7 +1279,6 @@ class xpu_ops:

1279 "index_q_scale",

1280 "index_weights_out",

1281 ],

1341 - fake_impl=_xpu_deepseek_fused_indexer_q_rope_mxfp4_fake,

1282 )

1283

1284 _OPS_REGISTERED = True

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

-19

@@ -293,24 +293,6 @@ if flashinfer_comm is not None:

293 or num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,

294 )

295

296 - def call_trtllm_fused_allreduce_norm_fake(

297 - allreduce_in: torch.Tensor,

298 - residual: torch.Tensor,

299 - rms_gamma: torch.Tensor,

300 - rms_eps: float,

301 - world_size: int,

302 - launch_with_pdl: bool,

303 - fp32_acc: bool,

304 - max_token_num: int,

305 - pattern_code: int,

306 - norm_out: torch.Tensor | None = None,

307 - quant_out: torch.Tensor | None = None,

308 - scale_out: torch.Tensor | None = None,

309 - scale_factor: torch.Tensor | None = None,

310 - weight_bias: float = 0.0,

311 - ) -> None:

312 - pass

313 -

296 direct_register_custom_op(

297 op_name="flashinfer_trtllm_fused_allreduce_norm",

298 op_func=call_trtllm_fused_allreduce_norm,

@@ -321,7 +303,6 @@ if flashinfer_comm is not None:

303 "quant_out",

304 "scale_out",

305 ],

324 - fake_impl=call_trtllm_fused_allreduce_norm_fake,

306 )

307 flashinfer_trtllm_fused_allreduce_norm = (

308 torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default

vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py

-10

@@ -100,15 +100,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |

100 return result

101

102

103 -def fake_impl(

104 - result: torch.Tensor, # [num_tokens, hidden_size]

105 - input: torch.Tensor, # [num_tokens, hidden_size]

106 - scale: torch.Tensor, # [num_tokens, 1]

107 - scale_ub: torch.Tensor | None = None, # scalar tensor

108 -) -> None:

109 - return

110 -

111 -

103 def baseline(

104 result: torch.Tensor, # [num_tokens, hidden_size]

105 input: torch.Tensor, # [num_tokens, hidden_size]

@@ -135,7 +126,6 @@ def baseline(

126 mutates_args=["result", "scale"],

127 config_picker=pick_config,

128 input_generator=generate_inputs,

138 - fake_impl=fake_impl,

129 helion_settings=helion.Settings(

130 autotune_baseline_fn=baseline,

131 ignore_warnings=[helion.exc.TensorOperationInWrapper],

vllm/kernels/helion/ops/fused_qk_norm_rope.py

-18

@@ -159,23 +159,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |

159 return result

160

161

162 -def fake_impl(

163 - qkv: torch.Tensor, # [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]

164 - num_heads_q: int,

165 - num_heads_k: int,

166 - num_heads_v: int,

167 - head_dim: int,

168 - eps: float,

169 - q_weight: torch.Tensor,

170 - k_weight: torch.Tensor,

171 - cos_sin_cache: torch.Tensor, # [max_position, rotary_dim]

172 - is_neox: bool,

173 - position_ids: torch.Tensor, # [num_tokens],

174 - forced_token_heads_per_warp: int = -1, # dummy

175 -) -> None:

176 - return

177 -

178 -

162 def baseline(

163 qkv: torch.Tensor, # [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]

164 num_heads_q: int,

@@ -216,7 +199,6 @@ def baseline(

199 mutates_args=["qkv"],

200 config_picker=pick_config,

201 input_generator=generate_inputs,

219 - fake_impl=fake_impl,

202 helion_settings=helion.Settings(

203 autotune_baseline_fn=baseline,

204 autotune_baseline_atol=5e-2,

vllm/kernels/helion/ops/per_token_group_fp8_quant.py

-18

@@ -135,23 +135,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |

135 return result

136

137

138 -def fake_impl(

139 - input: torch.Tensor, # [num_tokens, hidden_size]

140 - output_q: torch.Tensor, # [num_tokens, hidden_size]

141 - output_s: torch.Tensor, # [num_tokens, groups_per_row]

142 - group_size: int,

143 - eps: float,

144 - fp8_min: float,

145 - fp8_max: float,

146 - scale_ue8m0: bool,

147 - # Unused dummy args

148 - # Kept for consistency with existing kernel interface

149 - dummy_is_scale_transposed: bool = False,

150 - dummy_is_tma_aligned: bool = False,

151 -) -> None:

152 - return

153 -

154 -

138 def baseline(

139 input: torch.Tensor, # [num_tokens, hidden_size]

140 output_q: torch.Tensor, # [num_tokens, hidden_size]

@@ -181,7 +164,6 @@ def baseline(

164 mutates_args=["output_q", "output_s"],

165 config_picker=pick_config,

166 input_generator=generate_inputs,

184 - fake_impl=fake_impl,

167 helion_settings=helion.Settings(

168 autotune_baseline_fn=baseline,

169 ),

vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py

-13

@@ -119,18 +119,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |

119 return result

120

121

122 -def fake_impl(

123 - result: torch.Tensor, # [num_tokens, hidden_size]

124 - input: torch.Tensor, # [num_tokens, hidden_size]

125 - weight: torch.Tensor, # [hidden_size]

126 - scale: torch.Tensor, # [num_tokens, 1]

127 - epsilon: float,

128 - scale_ub: torch.Tensor | None = None, # []

129 - residual: torch.Tensor | None = None, # [num_tokens, hidden_size]

130 -) -> None:

131 - return

132 -

133 -

122 def baseline(

123 result: torch.Tensor, # [num_tokens, hidden_size]

124 input: torch.Tensor, # [num_tokens, hidden_size]

@@ -179,7 +167,6 @@ def baseline(

167 mutates_args=["result", "scale", "residual"],

168 config_picker=pick_config,

169 input_generator=generate_inputs,

182 - fake_impl=fake_impl,

170 helion_settings=helion.Settings(

171 autotune_baseline_fn=baseline,

172 ignore_warnings=[helion.exc.TensorOperationInWrapper],

vllm/kernels/helion/ops/rms_norm_per_block_quant.py

-15

@@ -153,20 +153,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |

153 return result

154

155

156 -def fake_impl(

157 - result: torch.Tensor, # [num_tokens, hidden_size]

158 - input: torch.Tensor, # [num_tokens, hidden_size]

159 - weight: torch.Tensor, # [hidden_size]

160 - scale: torch.Tensor, # [num_tokens, groups_per_row]

161 - epsilon: float,

162 - scale_ub: torch.Tensor | None, # []

163 - residual: torch.Tensor | None, # [num_tokens, hidden_size]

164 - group_size: int,

165 - is_scale_transposed: bool, # dummy

166 -) -> None:

167 - return

168 -

169 -

156 def baseline(

157 result: torch.Tensor, # [num_tokens, hidden_size]

158 input: torch.Tensor, # [num_tokens, hidden_size]

@@ -219,7 +205,6 @@ def baseline(

205 mutates_args=["result", "scale", "residual"],

206 config_picker=pick_config,

207 input_generator=generate_inputs,

222 - fake_impl=fake_impl,

208 helion_settings=helion.Settings(

209 autotune_baseline_fn=baseline,

210 ignore_warnings=[helion.exc.TensorOperationInWrapper],

vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py

-12

@@ -139,17 +139,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |

139 return result

140

141

142 -def fake_impl(

143 - out: torch.Tensor, # [num_tokens, intermediate_size]

144 - input: torch.Tensor, # [num_tokens, 2 * intermediate_size]

145 - scales: torch.Tensor, # [num_tokens, groups_per_row]

146 - group_size: int,

147 - scale_ub: torch.Tensor | None = None, # scalar tensor

148 - is_scale_transposed: bool = False,

149 -) -> None:

150 - return

151 -

152 -

142 def baseline(

143 out: torch.Tensor, # [num_tokens, intermediate_size]

144 input: torch.Tensor, # [num_tokens, 2 * intermediate_size]

@@ -193,7 +182,6 @@ def baseline(

182 mutates_args=["out", "scales"],

183 config_picker=pick_config,

184 input_generator=generate_inputs,

196 - fake_impl=fake_impl,

185 helion_settings=helion.Settings(

186 autotune_baseline_fn=baseline,

187 ignore_warnings=[helion.exc.TensorOperationInWrapper],

vllm/kernels/helion/register.py

+7 -2

@@ -38,6 +38,7 @@ Key Classes

38

39 from __future__ import annotations

40

41 +import inspect

42 from collections.abc import Callable

43 from typing import Any

44

@@ -258,7 +259,7 @@ class HelionKernelWrapper:

259 self,

260 raw_kernel_func: Callable,

261 op_name: str,

261 - fake_impl: Callable,

262 + fake_impl: Callable | None,

263 config_picker: ConfigPicker,

264 mutates_args: list[str] | None = None,

265 helion_settings: helion.Settings | None = None,

@@ -445,8 +446,12 @@ def register_kernel(

446 f"Use a different op_name or check for duplicate registrations."

447 )

448

449 + # A void-returning kernel that mutates its args needs no fake impl:

450 + # PyTorch auto-generates a trivial one for mutable, no-return schemas.

451 + returns_none = inspect.signature(kernel_func).return_annotation is None

452 +

453 final_fake_impl = fake_impl

449 - if final_fake_impl is None:

454 + if final_fake_impl is None and not (mutates_args and returns_none):

455 final_fake_impl = infer_fake_impl(kernel_func, helion_settings)

456 logger.debug(

457 "Auto-generated fake_impl for Helion kernel '%s'",

vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py

-133

@@ -869,156 +869,23 @@ def _fused_moe_lora_fp8(

869 )

870

871

872 -def _fused_moe_lora_fp8_fake(

873 - output: torch.Tensor,

874 - qcurr_hidden_states: torch.Tensor,

875 - lora_a_stacked: list[torch.Tensor],

876 - lora_b_stacked: list[torch.Tensor],

877 - topk_weights: torch.Tensor,

878 - sorted_token_ids: torch.Tensor | None,

879 - expert_ids: torch.Tensor,

880 - num_tokens_post_padded: torch.Tensor | None,

881 - token_lora_mapping: torch.Tensor,

882 - max_lora_rank: int,

883 - top_k_num: int,

884 - lora_ids: torch.Tensor,

885 - num_active_loras: int,

886 - adapter_enabled: torch.Tensor,

887 - shrink_block_size_m: int,

888 - shrink_block_size_n: int,

889 - shrink_block_size_k: int,

890 - shrink_group_size_m: int,

891 - shrink_num_warps: int,

892 - shrink_num_stages: int,

893 - shrink_split_k: int,

894 - expand_block_size_m: int,

895 - expand_block_size_n: int,

896 - expand_block_size_k: int,

897 - expand_group_size_m: int,

898 - expand_num_warps: int,

899 - expand_num_stages: int,

900 - expand_split_k: int,

901 - lora_a_scale_stacked: list[torch.Tensor],

902 - lora_b_scale_stacked: list[torch.Tensor],

903 - mul_routed_weight: bool = False,

904 - fully_sharded: bool = False,

905 - offset: int = 0,

906 - shrink_act_scale: torch.Tensor | None = None,

907 - expand_act_scale: torch.Tensor | None = None,

908 - use_fp8_w8a8: bool = False,

909 - use_int8_w8a8: bool = False,

910 - use_int8_w8a16: bool = False,

911 - per_channel_quant: bool = False,

912 - block_shape: List[int] | None = None, # noqa: UP006, UP007

913 -) -> None:

914 - return

915 -

916 -

917 -def _fused_moe_lora_shrink_fp8_fake(

918 - a_intermediate_cache1: torch.Tensor,

919 - qcurr_hidden_states: torch.Tensor,

920 - lora_a_stacked: list[torch.Tensor],

921 - topk_weights: torch.Tensor,

922 - sorted_token_ids: torch.Tensor | None,

923 - expert_ids: torch.Tensor,

924 - num_tokens_post_padded: torch.Tensor | None,

925 - token_lora_mapping: torch.Tensor,

926 - top_k_num: int,

927 - lora_ids: torch.Tensor,

928 - adapter_enabled: torch.Tensor,

929 - device: torch.device,

930 - N: int,

931 - M: int,

932 - EM: int,

933 - K: int,

934 - num_tokens: int,

935 - num_experts: int,

936 - num_slices: int,

937 - block_size_m: int,

938 - block_size_n: int,

939 - block_size_k: int,

940 - group_size_m: int,

941 - num_warps: int,

942 - num_stages: int,

943 - split_k: int,

944 - num_active_loras: int,

945 - lora_a_scale_stacked: list[torch.Tensor],

946 - mul_routed_weight: bool = False,

947 - use_gdc: bool = False,

948 - act_scale: torch.Tensor | None = None,

949 - use_fp8_w8a8: bool = False,

950 - use_int8_w8a8: bool = False,

951 - use_int8_w8a16: bool = False,

952 - per_channel_quant: bool = False,

953 - block_shape: List[int] | None = None, # noqa: UP006, UP007

954 -) -> None:

955 - return

956 -

957 -

958 -def _fused_moe_lora_expand_fp8_fake(

959 - output: torch.Tensor,

960 - a_intermediate_cache1: torch.Tensor,

961 - lora_b_stacked: list[torch.Tensor],

962 - topk_weights: torch.Tensor,

963 - sorted_token_ids: torch.Tensor | None,

964 - expert_ids: torch.Tensor,

965 - num_tokens_post_padded: torch.Tensor | None,

966 - token_lora_mapping: torch.Tensor,

967 - top_k_num: int,

968 - lora_ids: torch.Tensor,

969 - adapter_enabled: torch.Tensor,

970 - device: torch.device,

971 - N: int,

972 - M: int,

973 - EM: int,

974 - K: int,

975 - num_tokens: int,

976 - num_experts: int,

977 - num_slices: int,

978 - max_lora_rank: int,

979 - w1_output_dim_size: int,

980 - block_size_m: int,

981 - block_size_n: int,

982 - block_size_k: int,

983 - group_size_m: int,

984 - num_warps: int,

985 - num_stages: int,

986 - split_k: int,

987 - num_active_loras: int,

988 - act_scale: torch.Tensor,

989 - lora_b_scale_stacked: list[torch.Tensor],

990 - mul_routed_weight: bool = False,

991 - offset: int = 0,

992 - use_fp8_w8a8: bool = False,

993 - use_int8_w8a8: bool = False,

994 - use_int8_w8a16: bool = False,

995 - per_channel_quant: bool = False,

996 - block_shape: List[int] | None = None, # noqa: UP006, UP007

997 - use_gdc: bool = False,

998 -) -> None:

999 - return

1000 -

1001 -

872 try:

873 direct_register_custom_op(

874 op_name="fused_moe_lora_fp8",

875 op_func=_fused_moe_lora_fp8,

876 mutates_args=["output"],

1007 - fake_impl=_fused_moe_lora_fp8_fake,

877 )

878

879 direct_register_custom_op(

880 op_name="fused_moe_lora_shrink_fp8",

881 op_func=_fused_moe_lora_shrink_fp8,

882 mutates_args=["a_intermediate_cache1"],

1014 - fake_impl=_fused_moe_lora_shrink_fp8_fake,

883 )

884

885 direct_register_custom_op(

886 op_name="fused_moe_lora_expand_fp8",

887 op_func=_fused_moe_lora_expand_fp8,

888 mutates_args=["output"],

1021 - fake_impl=_fused_moe_lora_expand_fp8_fake,

889 )

890

891 fused_moe_lora_fp8 = torch.ops.vllm.fused_moe_lora_fp8

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

-113

@@ -1662,136 +1662,23 @@ def _fused_moe_lora(

1662 )

1663

1664

1665 -def _fused_moe_lora_fake(

1666 - output: torch.Tensor,

1667 - qcurr_hidden_states: torch.Tensor,

1668 - lora_a_stacked: list[torch.Tensor],

1669 - lora_b_stacked: list[torch.Tensor],

1670 - topk_weights: torch.Tensor,

1671 - sorted_token_ids: torch.Tensor | None,

1672 - expert_ids: torch.Tensor,

1673 - num_tokens_post_padded: torch.Tensor | None,

1674 - token_lora_mapping: torch.Tensor,

1675 - max_lora_rank: int,

1676 - top_k_num: int,

1677 - lora_ids: torch.Tensor,

1678 - num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

1679 - adapter_enabled: torch.Tensor,

1680 - shrink_block_size_m: int,

1681 - shrink_block_size_n: int,

1682 - shrink_block_size_k: int,

1683 - shrink_group_size_m: int,

1684 - shrink_num_warps: int,

1685 - shrink_num_stages: int,

1686 - shrink_split_k: int,

1687 - expand_block_size_m: int,

1688 - expand_block_size_n: int,

1689 - expand_block_size_k: int,

1690 - expand_group_size_m: int,

1691 - expand_num_warps: int,

1692 - expand_num_stages: int,

1693 - expand_split_k: int,

1694 - mul_routed_weight: bool = False,

1695 - fully_sharded: bool = False,

1696 - offset: int = 0,

1697 - add_inputs: bool = True,

1698 -) -> None:

1699 - return

1700 -

1701 -

1702 -def _fused_moe_lora_shrink_fake(

1703 - a_intermediate_cache1: torch.Tensor,

1704 - qcurr_hidden_states: torch.Tensor,

1705 - lora_a_stacked: list[torch.Tensor],

1706 - topk_weights: torch.Tensor,

1707 - sorted_token_ids: torch.Tensor | None,

1708 - expert_ids: torch.Tensor,

1709 - num_tokens_post_padded: torch.Tensor | None,

1710 - token_lora_mapping: torch.Tensor,

1711 - top_k_num: int,

1712 - lora_ids: torch.Tensor,

1713 - adapter_enabled: torch.Tensor,

1714 - device: torch.device,

1715 - N: int,

1716 - M: int,

1717 - EM: int,

1718 - K: int,

1719 - num_tokens: int,

1720 - num_experts: int,

1721 - num_slices: int,

1722 - block_size_m: int,

1723 - block_size_n: int,

1724 - block_size_k: int,

1725 - group_size_m: int,

1726 - num_warps: int,

1727 - num_stages: int,

1728 - split_k: int,

1729 - num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

1730 - mul_routed_weight: bool = False,

1731 - use_gdc: bool = False,

1732 - use_tma: bool = False,

1733 -) -> None:

1734 - return

1735 -

1736 -

1737 -def _fused_moe_lora_expand_fake(

1738 - output: torch.Tensor,

1739 - a_intermediate_cache1: torch.Tensor,

1740 - lora_b_stacked: list[torch.Tensor],

1741 - topk_weights: torch.Tensor,

1742 - sorted_token_ids: torch.Tensor | None,

1743 - expert_ids: torch.Tensor,

1744 - num_tokens_post_padded: torch.Tensor | None,

1745 - token_lora_mapping: torch.Tensor,

1746 - top_k_num: int,

1747 - lora_ids: torch.Tensor,

1748 - adapter_enabled: torch.Tensor,

1749 - device: torch.device,

1750 - N: int,

1751 - M: int,

1752 - EM: int,

1753 - K: int,

1754 - num_tokens: int,

1755 - num_experts: int,

1756 - num_slices: int,

1757 - max_lora_rank: int,

1758 - w1_output_dim_size: int,

1759 - block_size_m: int,

1760 - block_size_n: int,

1761 - block_size_k: int,

1762 - group_size_m: int,

1763 - num_warps: int,

1764 - num_stages: int,

1765 - split_k: int,

1766 - num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

1767 - mul_routed_weight: bool = False,

1768 - offset: int = 0,

1769 - use_gdc: bool = False,

1770 - use_tma: bool = False,

1771 -) -> None:

1772 - return

1773 -

1774 -

1665 try:

1666 direct_register_custom_op(

1667 op_name="fused_moe_lora",

1668 op_func=_fused_moe_lora,

1669 mutates_args=["output"],

1780 - fake_impl=_fused_moe_lora_fake,

1670 )

1671

1672 direct_register_custom_op(

1673 op_name="fused_moe_lora_shrink",

1674 op_func=_fused_moe_lora_shrink,

1675 mutates_args=["a_intermediate_cache1"],

1787 - fake_impl=_fused_moe_lora_shrink_fake,

1676 )

1677

1678 direct_register_custom_op(

1679 op_name="fused_moe_lora_expand",

1680 op_func=_fused_moe_lora_expand,

1681 mutates_args=["output"],

1794 - fake_impl=_fused_moe_lora_expand_fake,

1682 )

1683

1684 fused_moe_lora = torch.ops.vllm.fused_moe_lora

vllm/lora/ops/triton_ops/lora_expand_fp8_op.py

-24

@@ -367,35 +367,11 @@ def _lora_expand_fp8(

367 return

368

369

370 -def _lora_expand_fp8_fake(

371 - inputs: torch.Tensor,

372 - lora_b_weights: list[torch.Tensor],

373 - output_tensor: torch.Tensor,

374 - token_lora_mapping: torch.Tensor,

375 - token_indices_sorted_by_lora_ids: torch.Tensor,

376 - num_tokens_per_lora: torch.Tensor,

377 - lora_token_start_loc: torch.Tensor,

378 - lora_ids: torch.Tensor,

379 - no_lora_flag_cpu: torch.Tensor,

380 - num_active_loras: int,

381 - b_scale: list[torch.Tensor],

382 - a_scale: torch.Tensor | None = None,

383 - offset_start: int = 0,

384 - add_inputs: bool = False,

385 - group_k: int = 0,

386 - group_n: int = 0,

387 - use_fp8_w8a8: bool = False,

388 - per_channel_quant: bool = False,

389 -) -> None:

390 - return

391 -

392 -

370 try:

371 direct_register_custom_op(

372 op_name="lora_expand_fp8",

373 op_func=_lora_expand_fp8,

374 mutates_args=["output_tensor"],

398 - fake_impl=_lora_expand_fp8_fake,

375 )

376 lora_expand_fp8 = torch.ops.vllm.lora_expand_fp8

377

vllm/lora/ops/triton_ops/lora_expand_op.py

-18

@@ -284,29 +284,11 @@ def _lora_expand(

284 return

285

286

287 -def _lora_expand_fake(

288 - inputs: torch.Tensor,

289 - lora_b_weights: list[torch.Tensor],

290 - output_tensor: torch.Tensor,

291 - token_lora_mapping: torch.Tensor,

292 - token_indices_sorted_by_lora_ids: torch.Tensor,

293 - num_tokens_per_lora: torch.Tensor,

294 - lora_token_start_loc: torch.Tensor,

295 - lora_ids: torch.Tensor,

296 - no_lora_flag_cpu: torch.Tensor,

297 - num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

298 - offset_start: int = 0,

299 - add_inputs: bool = False,

300 -) -> None:

301 - return

302 -

303 -

287 try:

288 direct_register_custom_op(

289 op_name="lora_expand",

290 op_func=_lora_expand,

291 mutates_args=["output_tensor"],

309 - fake_impl=_lora_expand_fake,

292 )

293 lora_expand = torch.ops.vllm.lora_expand

294

vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py

-23

@@ -394,34 +394,11 @@ def _lora_shrink_fp8(

394 return

395

396

397 -def _lora_shrink_fp8_fake(

398 - inputs: torch.Tensor,

399 - lora_a_weights: list[torch.Tensor],

400 - output_tensor: torch.Tensor,

401 - token_lora_mapping: torch.Tensor,

402 - token_indices_sorted_by_lora_ids: torch.Tensor,

403 - num_tokens_per_lora: torch.Tensor,

404 - lora_token_start_loc: torch.Tensor,

405 - lora_ids: torch.Tensor,

406 - no_lora_flag_cpu: torch.Tensor,

407 - num_active_loras: int,

408 - scaling: float,

409 - b_scale: list[torch.Tensor], # LoRA weight scale per slice

410 - a_scale: torch.Tensor | None = None, # Activation scale - per-token or block-wise

411 - group_k: int = 0, # Block size for K in block-wise quantization (0 = tensor-wise)

412 - group_n: int = 0, # Block size for N in block-wise quantization

413 - use_fp8_w8a8: bool = False,

414 - per_channel_quant: bool = False,

415 -) -> None:

416 - return

417 -

418 -

397 try:

398 direct_register_custom_op(

399 op_name="lora_shrink_fp8",

400 op_func=_lora_shrink_fp8,

401 mutates_args=["output_tensor"],

424 - fake_impl=_lora_shrink_fp8_fake,

402 )

403 lora_shrink_fp8 = torch.ops.vllm.lora_shrink_fp8

404

vllm/lora/ops/triton_ops/lora_shrink_op.py

-17

@@ -265,28 +265,11 @@ def _lora_shrink(

265 return

266

267

268 -def _lora_shrink_fake(

269 - inputs: torch.Tensor,

270 - lora_a_weights: list[torch.Tensor],

271 - output_tensor: torch.Tensor,

272 - token_lora_mapping: torch.Tensor,

273 - token_indices_sorted_by_lora_ids: torch.Tensor,

274 - num_tokens_per_lora: torch.Tensor,

275 - lora_token_start_loc: torch.Tensor,

276 - lora_ids: torch.Tensor,

277 - no_lora_flag_cpu: torch.Tensor,

278 - num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs

279 - scaling: float,

280 -) -> None:

281 - return

282 -

283 -

268 try:

269 direct_register_custom_op(

270 op_name="lora_shrink",

271 op_func=_lora_shrink,

272 mutates_args=["output_tensor"],

289 - fake_impl=_lora_shrink_fake,

273 )

274 lora_shrink = torch.ops.vllm.lora_shrink

275

vllm/lora/ops/xpu_ops/lora_ops.py

-35

@@ -85,57 +85,22 @@ def _bgmv_expand_slice_impl(

85 )

86

87

88 -def _bgmv_shrink_fake(

89 - inputs: torch.Tensor,

90 - lora_a_weights: torch.Tensor,

91 - output_tensor: torch.Tensor,

92 - lora_indices_tensor: torch.Tensor,

93 - scaling: float,

94 -) -> None:

95 - return None

96 -

97 -

98 -def _bgmv_expand_fake(

99 - inputs: torch.Tensor,

100 - lora_b_weights: torch.Tensor,

101 - output_tensor: torch.Tensor,

102 - lora_indices_tensor: torch.Tensor,

103 - add_inputs: bool,

104 -) -> None:

105 - return None

106 -

107 -

108 -def _bgmv_expand_slice_fake(

109 - inputs: torch.Tensor,

110 - lora_b_weights: torch.Tensor,

111 - output_tensor: torch.Tensor,

112 - lora_indices_tensor: torch.Tensor,

113 - slice_offset: int,

114 - slice_size: int,

115 - add_inputs: bool,

116 -) -> None:

117 - return None

118 -

119 -

88 direct_register_custom_op(

89 op_name="xpu_bgmv_shrink",

90 op_func=_bgmv_shrink_impl,

91 mutates_args=["output_tensor"],

124 - fake_impl=_bgmv_shrink_fake,

92 )

93

94 direct_register_custom_op(

95 op_name="xpu_bgmv_expand",

96 op_func=_bgmv_expand_impl,

97 mutates_args=["output_tensor"],

131 - fake_impl=_bgmv_expand_fake,

98 )

99

100 direct_register_custom_op(

101 op_name="xpu_bgmv_expand_slice",

102 op_func=_bgmv_expand_slice_impl,

103 mutates_args=["output_tensor"],

138 - fake_impl=_bgmv_expand_slice_fake,

104 )

105

106

vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py

-12

@@ -139,20 +139,8 @@ def _fp8_gemm_nt_op(

139 )

140

141

142 -def _fp8_gemm_nt_op_fake(

143 - q_input: torch.Tensor,

144 - input_scale: torch.Tensor,

145 - weight: torch.Tensor,

146 - weight_scale: torch.Tensor,

147 - output: torch.Tensor,

148 - use_deep_gemm_e8m0: bool,

149 -) -> None:

150 - return None

151 -

152 -

142 direct_register_custom_op(

143 "fp8_gemm_nt_op",

144 _fp8_gemm_nt_op,

145 mutates_args=["output"],

157 - fake_impl=_fp8_gemm_nt_op_fake,

146 )

vllm/model_executor/layers/attention/attention.py

-14

@@ -777,22 +777,8 @@ def unified_attention_with_output(

777 )

778

779

780 -def unified_attention_with_output_fake(

781 - query: torch.Tensor,

782 - key: torch.Tensor,

783 - value: torch.Tensor,

784 - output: torch.Tensor,

785 - layer_name: LayerNameType,

786 - output_scale: torch.Tensor | None = None,

787 - output_block_scale: torch.Tensor | None = None,

788 - kv_cache_dummy_dep: torch.Tensor | None = None,

789 -) -> None:

790 - return

791 -

792 -

780 direct_register_custom_op(

781 op_name="unified_attention_with_output",

782 op_func=unified_attention_with_output,

783 mutates_args=["output", "output_block_scale"],

797 - fake_impl=unified_attention_with_output_fake,

784 )

vllm/model_executor/layers/attention/mla_attention.py

-19

@@ -1424,29 +1424,10 @@ def unified_mla_attention_with_output(

1424 )

1425

1426

1427 -def unified_mla_attention_with_output_fake(

1428 - q: torch.Tensor,

1429 - kv_c_normed: torch.Tensor,

1430 - k_pe: torch.Tensor,

1431 - output: torch.Tensor,

1432 - layer_name: LayerNameType,

1433 - output_scale: torch.Tensor | None = None,

1434 - output_block_scale: torch.Tensor | None = None,

1435 - kv_cache_dummy_dep: torch.Tensor | None = None,

1436 - quant_group_size: int | None = None,

1437 - quant_scale_ue8m0: bool | None = None,

1438 - quant_col_major: bool | None = None,

1439 - quant_tma_aligned: bool | None = None,

1440 - q_dcp_replicated: torch.Tensor | None = None,

1441 -) -> None:

1442 - return

1443 -

1444 -

1427 direct_register_custom_op(

1428 op_name="unified_mla_attention_with_output",

1429 op_func=unified_mla_attention_with_output,

1430 mutates_args=["output", "output_block_scale"],

1449 - fake_impl=unified_mla_attention_with_output_fake,

1431 dispatch_key=current_platform.dispatch_key,

1432 tags=(torch.Tag.flexible_layout,),

1433 )

vllm/model_executor/layers/attention/static_sink_attention.py

-8

@@ -241,16 +241,8 @@ def maybe_populate_sink(

241 self.populate_sink_kv(self_kv_cache)

242

243

244 -def maybe_populate_sink_fake(

245 - self_kv_cache: torch.Tensor,

246 - layer_name: LayerNameType,

247 -) -> None:

248 - return

249 -

250 -

244 direct_register_custom_op(

245 op_name="maybe_populate_sink",

246 op_func=maybe_populate_sink,

247 mutates_args=["self_kv_cache"],

255 - fake_impl=maybe_populate_sink_fake,

248 )

vllm/model_executor/layers/hpc/rope_norm.py

-10

@@ -78,20 +78,10 @@ def hpc_rope_norm_forward(

78 rope_norm._forward_impl(qkv, kv_cache, attn_metadata, attn_layer, output)

79

80

81 -def hpc_rope_norm_forward_fake(

82 - qkv: torch.Tensor,

83 - output: torch.Tensor,

84 - layer_name: str,

85 -) -> None:

86 - """Fake impl for torch.compile trace; output is a mutated arg."""

87 - return

88 -

89 -

81 direct_register_custom_op(

82 op_name="hpc_rope_norm_forward",

83 op_func=hpc_rope_norm_forward,

84 mutates_args=["output"],

94 - fake_impl=hpc_rope_norm_forward_fake,

85 )

86

87

vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py

-10

@@ -546,20 +546,10 @@ def olmo_hybrid_gdn_full_forward(

546 )

547

548

549 -def olmo_hybrid_gdn_full_forward_fake(

550 - hidden_states: torch.Tensor,

551 - output: torch.Tensor,

552 - layer_name: str,

553 -) -> None:

554 - """Fake implementation for torch.compile."""

555 - return

556 -

557 -

549 direct_register_custom_op(

550 op_name="olmo_hybrid_gdn_full_forward",

551 op_func=olmo_hybrid_gdn_full_forward,

552 mutates_args=["output"],

562 - fake_impl=olmo_hybrid_gdn_full_forward_fake,

553 )

554

555

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

-23

@@ -1949,23 +1949,10 @@ def qwen_gdn_attention_core(

1949 )

1950

1951

1952 -def gdn_attention_core_fake(

1953 - qkv_or_qkvz: torch.Tensor,

1954 - b_or_ba: torch.Tensor,

1955 - a_or_z_out: torch.Tensor,

1956 - core_attn_out: torch.Tensor,

1957 - layer_name: LayerNameType,

1958 - use_aiter: bool = False,

1959 -) -> None:

1960 - """Fake implementation for torch.compile."""

1961 - return

1962 -

1963 -

1952 direct_register_custom_op(

1953 op_name="qwen_gdn_attention_core",

1954 op_func=qwen_gdn_attention_core,

1955 mutates_args=["a_or_z_out", "core_attn_out"],

1968 - fake_impl=gdn_attention_core_fake,

1956 )

1957

1958

@@ -1985,20 +1972,10 @@ def qwen_gdn_attention_core_fused_norm_packed(

1972 )

1973

1974

1988 -def gdn_attention_core_fused_norm_packed_fake(

1989 - mixed_qkvz: torch.Tensor,

1990 - ba: torch.Tensor,

1991 - core_attn_out: torch.Tensor,

1992 - layer_name: LayerNameType,

1993 -) -> None:

1994 - return

1995 -

1996 -

1975 direct_register_custom_op(

1976 op_name="qwen_gdn_attention_core_fused_norm_packed",

1977 op_func=qwen_gdn_attention_core_fused_norm_packed,

1978 mutates_args=["core_attn_out"],

2001 - fake_impl=gdn_attention_core_fused_norm_packed_fake,

1979 )

1980

1981

vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py

-10

@@ -336,18 +336,8 @@ def linear_attention(

336 self._forward(hidden_states=hidden_states, output=output, positions=positions)

337

338

339 -def linear_attention_fake(

340 - hidden_states: torch.Tensor,

341 - output: torch.Tensor,

342 - positions: torch.Tensor,

343 - layer_name: str,

344 -) -> None:

345 - return

346 -

347 -

339 direct_register_custom_op(

340 op_name="linear_attention",

341 op_func=linear_attention,

342 mutates_args=["output"],

352 - fake_impl=linear_attention_fake,

343 )

vllm/model_executor/layers/mamba/mamba_mixer.py

-9

@@ -532,17 +532,8 @@ def mamba_mixer(

532 self.forward_impl(hidden_states=hidden_states, output=output)

533

534

535 -def mamba_mixer_fake(

536 - hidden_states: torch.Tensor,

537 - output: torch.Tensor,

538 - layer_name: LayerNameType,

539 -) -> None:

540 - return

541 -

542 -

535 direct_register_custom_op(

536 op_name="mamba_mixer",

537 op_func=mamba_mixer,

538 mutates_args=["output"],

547 - fake_impl=mamba_mixer_fake,

539 )

vllm/model_executor/layers/mamba/mamba_mixer2.py

-9

@@ -1271,17 +1271,8 @@ def mamba_mixer2(

1271 self.conv_ssm_forward(projected_states=projected_states, output=output)

1272

1273

1274 -def mamba_mixer2_fake(

1275 - projected_states: torch.Tensor,

1276 - output: torch.Tensor,

1277 - layer_name: LayerNameType,

1278 -) -> None:

1279 - return

1280 -

1281 -

1274 direct_register_custom_op(

1275 op_name="mamba_mixer2",

1276 op_func=mamba_mixer2,

1277 mutates_args=["output"],

1286 - fake_impl=mamba_mixer2_fake,

1278 )

vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py

-12

@@ -718,17 +718,6 @@ def _spec_aware_nonspec_subset(

718 return attn_out.squeeze(0)

719

720

721 -def cpu_gdn_attention_core_fake(

722 - mixed_qkv: torch.Tensor,

723 - b: torch.Tensor,

724 - a: torch.Tensor,

725 - core_attn_out: torch.Tensor,

726 - layer_name: LayerNameType,

727 -) -> None:

728 - """Fake implementation for torch.compile."""

729 - return

730 -

731 -

721 def register_cpu_gdn_attention_ops() -> None:

722 global _CPU_GDN_ATTENTION_OPS_REGISTERED

723 if _CPU_GDN_ATTENTION_OPS_REGISTERED:

@@ -738,6 +727,5 @@ def register_cpu_gdn_attention_ops() -> None:

727 op_name="cpu_gdn_attention_core",

728 op_func=cpu_gdn_attention_core,

729 mutates_args=["core_attn_out"],

741 - fake_impl=cpu_gdn_attention_core_fake,

730 )

731 _CPU_GDN_ATTENTION_OPS_REGISTERED = True

vllm/model_executor/layers/mamba/short_conv.py

-9

@@ -373,17 +373,8 @@ def short_conv(

373 self.forward_native(hidden_states=hidden_states, output=output)

374

375

376 -def short_conv_fake(

377 - hidden_states: torch.Tensor,

378 - output: torch.Tensor,

379 - layer_name: str,

380 -) -> None:

381 - return

382 -

383 -

376 direct_register_custom_op(

377 op_name="short_conv",

378 op_func=short_conv,

379 mutates_args=["output"],

388 - fake_impl=short_conv_fake,

380 )

vllm/model_executor/layers/rotary_embedding/common.py

-12

@@ -100,23 +100,11 @@ def _flashinfer_rotary_embedding(

100 )

101

102

103 -def _flashinfer_rotary_embedding_fake(

104 - positions: torch.Tensor,

105 - query: torch.Tensor,

106 - key: torch.Tensor,

107 - head_size: int,

108 - cos_sin_cache: torch.Tensor,

109 - is_neox: bool,

110 -) -> None:

111 - return

112 -

113 -

103 # Register flashinfer rotary embedding custom op

104 direct_register_custom_op(

105 op_name="flashinfer_rotary_embedding",

106 op_func=_flashinfer_rotary_embedding,

107 mutates_args=["query", "key"], # These tensors are modified in-place

119 - fake_impl=_flashinfer_rotary_embedding_fake,

108 )

109

110

vllm/model_executor/models/bailing_moe_v3.py

-13

@@ -115,23 +115,10 @@ def bailing_v3_kda_attention(

115 )

116

117

118 -def bailing_v3_kda_attention_fake(

119 - q_proj_states: torch.Tensor,

120 - k_proj_states: torch.Tensor,

121 - v_proj_states: torch.Tensor,

122 - g1: torch.Tensor,

123 - beta: torch.Tensor,

124 - core_attn_out: torch.Tensor,

125 - layer_name: str,

126 -) -> None:

127 - return

128 -

129 -

118 direct_register_custom_op(

119 op_name="bailing_v3_kda_attention",

120 op_func=bailing_v3_kda_attention,

121 mutates_args=["core_attn_out"],

134 - fake_impl=bailing_v3_kda_attention_fake,

122 )

123

124

vllm/model_executor/offloader/prefetch_ops.py

-18

@@ -33,14 +33,6 @@ def _wait_prefetch_impl(

33 get_offloader()._wait_for_layer(layer_idx)

34

35

36 -def _wait_prefetch_fake(

37 - input_tensor: torch.Tensor,

38 - layer_idx: int,

39 -) -> None:

40 - """Fake implementation for torch.compile tracing."""

41 - return

42 -

43 -

36 # --- start_prefetch op ---

37

38

@@ -61,14 +53,6 @@ def _start_prefetch_impl(

53 get_offloader()._start_prefetch(layer_idx)

54

55

64 -def _start_prefetch_fake(

65 - output_tensor: torch.Tensor,

66 - layer_idx: int,

67 -) -> None:

68 - """Fake implementation for torch.compile tracing."""

69 - return

70 -

71 -

56 def register_prefetch_offloader_ops() -> None:

57 """Register custom ops for prefetch offloader.

58

@@ -79,14 +63,12 @@ def register_prefetch_offloader_ops() -> None:

63 op_name="wait_prefetch",

64 op_func=_wait_prefetch_impl,

65 mutates_args=["input_tensor"],

82 - fake_impl=_wait_prefetch_fake,

66 )

67

68 direct_register_custom_op(

69 op_name="start_prefetch",

70 op_func=_start_prefetch_impl,

71 mutates_args=["output_tensor"],

89 - fake_impl=_start_prefetch_fake,

72 )

73

74

vllm/models/deepseek_v4/xpu/model.py

-13

@@ -578,23 +578,10 @@ def _deepseek_v4_mega_moe_experts_op(

578 )

579

580

581 -def _deepseek_v4_mega_moe_experts_op_fake(

582 - hidden_states: torch.Tensor,

583 - topk_weights: torch.Tensor,

584 - topk_ids: torch.Tensor,

585 - out: torch.Tensor,

586 - layer_name: str,

587 - activation_clamp: float | None,

588 - fast_math: bool,

589 -) -> None:

590 - return None

591 -

592 -

581 direct_register_custom_op(

582 op_name="deepseek_v4_mega_moe_experts",

583 op_func=_deepseek_v4_mega_moe_experts_op,

584 mutates_args=["out"],

597 - fake_impl=_deepseek_v4_mega_moe_experts_op_fake,

585 )

586

587

vllm/models/qwen4_exp/amd/ple_layer.py

-18

@@ -1082,14 +1082,6 @@ def qwen4_exp_amd_ple_ngram_embedding(

1082 output.copy_(result)

1083

1084

1085 -def qwen4_exp_amd_ple_ngram_embedding_fake(

1086 - ngram_ids: torch.Tensor,

1087 - output: torch.Tensor,

1088 - layer_name: str,

1089 -) -> None:

1090 - return

1091 -

1092 -

1085 def qwen4_exp_ple_short_conv(

1086 inputs: torch.Tensor,

1087 output: torch.Tensor,

@@ -1100,19 +1092,10 @@ def qwen4_exp_ple_short_conv(

1092 output[: result.shape[0]].copy_(result)

1093

1094

1103 -def qwen4_exp_ple_short_conv_fake(

1104 - inputs: torch.Tensor,

1105 - output: torch.Tensor,

1106 - layer_name: str,

1107 -) -> None:

1108 - return

1109 -

1110 -

1095 direct_register_custom_op(

1096 op_name="qwen4_exp_amd_ple_ngram_embedding",

1097 op_func=qwen4_exp_amd_ple_ngram_embedding,

1098 mutates_args=["output"],

1115 - fake_impl=qwen4_exp_amd_ple_ngram_embedding_fake,

1099 )

1100

1101

@@ -1120,7 +1103,6 @@ direct_register_custom_op(

1103 op_name="qwen4_exp_ple_short_conv",

1104 op_func=qwen4_exp_ple_short_conv,

1105 mutates_args=["output"],

1123 - fake_impl=qwen4_exp_ple_short_conv_fake,

1106 )

1107

1108

vllm/models/qwen4_exp/amd/qsa.py

-13

@@ -445,23 +445,10 @@ def qwen4_exp_qsa_with_output(

445 )

446

447

448 -def qwen4_exp_qsa_with_output_fake(

449 - hidden_states: torch.Tensor,

450 - positions: torch.Tensor,

451 - query: torch.Tensor,

452 - key: torch.Tensor,

453 - value: torch.Tensor,

454 - output: torch.Tensor,

455 - layer_name: LayerNameType,

456 -) -> None:

457 - del hidden_states, positions, query, key, value, output, layer_name

458 -

459 -

448 direct_register_custom_op(

449 op_name="qwen4_exp_qsa_with_output",

450 op_func=qwen4_exp_qsa_with_output,

451 mutates_args=["output"],

464 - fake_impl=qwen4_exp_qsa_with_output_fake,

452 )

453

454

vllm/models/qwen4_exp/nvidia/ple_layer.py

-20

@@ -877,16 +877,6 @@ def qwen4_exp_compute_ple_ngram_ids(

877 )

878

879

880 -def qwen4_exp_compute_ple_ngram_ids_fake(

881 - input_ids: torch.Tensor,

882 - query_start_loc: torch.Tensor,

883 - ngram_context: torch.Tensor,

884 - output: torch.Tensor,

885 - layer_name: str,

886 -) -> None:

887 - return

888 -

889 -

880 def qwen4_exp_ple_short_conv(

881 inputs: torch.Tensor,

882 residual_output: torch.Tensor,

@@ -896,19 +886,10 @@ def qwen4_exp_ple_short_conv(

886 layer._short_conv(inputs, residual_output)

887

888

899 -def qwen4_exp_ple_short_conv_fake(

900 - inputs: torch.Tensor,

901 - residual_output: torch.Tensor,

902 - layer_name: str,

903 -) -> None:

904 - return

905 -

906 -

889 direct_register_custom_op(

890 op_name="qwen4_exp_compute_ple_ngram_ids",

891 op_func=qwen4_exp_compute_ple_ngram_ids,

892 mutates_args=["output"],

911 - fake_impl=qwen4_exp_compute_ple_ngram_ids_fake,

893 )

894

895

@@ -916,7 +897,6 @@ direct_register_custom_op(

897 op_name="qwen4_exp_ple_short_conv",

898 op_func=qwen4_exp_ple_short_conv,

899 mutates_args=["residual_output"],

919 - fake_impl=qwen4_exp_ple_short_conv_fake,

900 )

901

902

vllm/models/qwen4_exp/nvidia/qsa.py

-13

@@ -465,23 +465,10 @@ def qwen4_exp_qsa_with_output(

465 )

466

467

468 -def qwen4_exp_qsa_with_output_fake(

469 - hidden_states: torch.Tensor,

470 - positions: torch.Tensor,

471 - query: torch.Tensor,

472 - key: torch.Tensor,

473 - value: torch.Tensor,

474 - output: torch.Tensor,

475 - layer_name: LayerNameType,

476 -) -> None:

477 - del hidden_states, positions, query, key, value, output, layer_name

478 -

479 -

468 direct_register_custom_op(

469 op_name="qwen4_exp_qsa_with_output",

470 op_func=qwen4_exp_qsa_with_output,

471 mutates_args=["output"],

484 - fake_impl=qwen4_exp_qsa_with_output_fake,

472 )

473

474

vllm/utils/flashinfer.py

-8

@@ -708,19 +708,11 @@ if has_flashinfer():

708

709 concat_mla_k(k, k_nope, k_pe)

710

711 - def _flashinfer_concat_mla_k_fake(

712 - k: torch.Tensor,

713 - k_nope: torch.Tensor,

714 - k_pe: torch.Tensor,

715 - ) -> None:

716 - return

717 -

711 # Register flashinfer concat_mla_k custom op

712 direct_register_custom_op(

713 op_name="flashinfer_concat_mla_k",

714 op_func=_flashinfer_concat_mla_k,

715 mutates_args=["k"], # k tensor is modified in-place

723 - fake_impl=_flashinfer_concat_mla_k_fake,

716 )

717

718 @torch.library.custom_op(

[Kernel] Remove unused fake implementation (#55535)