biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Kernel] Apply torch.Tag.needs_fixed_stride_order only for torch==2.6.0 (#19346)

Browse Source 
Signed-off-by: rzou <zou3519@gmail.com>
This commit is contained in:
Richard Zou
2025-07-18 14:10:21 -04:00
committed by GitHub
parent 21274ab476
commit b2eb2b5ad7
3 changed files with 19 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
vllm/model_executor/layers/fused_moe/fused_moe.py
View File

@@ -33,7 +33,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
dequant_mxfp4)
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils import direct_register_custom_op
from vllm.utils import direct_register_custom_op, is_torch_equal_or_newer
from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used
from .rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled
@@ -1056,7 +1056,8 @@ direct_register_custom_op(
op_func=inplace_fused_experts,
mutates_args=["hidden_states"],
fake_impl=inplace_fused_experts_fake,
tags=(torch.Tag.needs_fixed_stride_order, ),
tags=(() if is_torch_equal_or_newer("2.7.0") else
(torch.Tag.needs_fixed_stride_order, )),
)
@@ -1122,7 +1123,8 @@ direct_register_custom_op(
op_func=outplace_fused_experts,
mutates_args=[],
fake_impl=outplace_fused_experts_fake,
tags=(torch.Tag.needs_fixed_stride_order, ),
tags=(() if is_torch_equal_or_newer("2.7.0") else
(torch.Tag.needs_fixed_stride_order, )),
)