A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 01:00 UTC
[Kernel] Apply torch.Tag.needs_fixed_stride_order only for torch==2.6.0 (#19346)
Signed-off-by: rzou <zou3519@gmail.com>
Richard Zou committed
Jul 18, 2025 at 14:10 UTC
b2eb2b5ad7090ad3b3e002b200104a82eeb2fa7f
3 files changed
+19
-9
csrc/torch_bindings.cpp
+8
-4
@@ -20,13 +20,17 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
20
// vLLM custom ops
21
//
22
23
- // The default behavior in PyTorch 2.6 is "requires_contiguous", so we need
23
+ // The default behavior in PyTorch 2.6 was changed to "requires_contiguous",
24
+ // so we need
25
// to override this for many GEMMs with the following tag. Otherwise,
26
// torch.compile will force all input tensors to be contiguous(), which
27
// will break many custom ops that require column-major weight matrices.
27
- // TODO: remove this for PyTorch 2.8, when the default is planned to switch
28
- // to match exact eager-mode strides.
29
- at::Tag stride_tag = at::Tag::needs_fixed_stride_order;
28
+ // This was a bug and PyTorch 2.7 has since fixed this.
29
+#if TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 6
30
+ #define stride_tag at::Tag::needs_fixed_stride_order
31
+#else
32
+ #define stride_tag
33
+#endif
34
35
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
36
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
vllm/attention/ops/rocm_aiter_mla.py
+6
-2
@@ -6,7 +6,7 @@ from typing import Optional
6
import torch
7
8
from vllm.platforms import current_platform
9
-from vllm.utils import direct_register_custom_op
9
+from vllm.utils import direct_register_custom_op, is_torch_equal_or_newer
10
11
12
def get_aiter_mla_metadata(max_batch_size: int, block_size: int,
@@ -93,8 +93,12 @@ def mla_decode_fwd_fake(
93
94
95
if current_platform.is_rocm():
96
+ if is_torch_equal_or_newer("2.7.0"):
97
+ tags = ()
98
+ else:
99
+ tags = (torch.Tag.needs_fixed_stride_order, ),
100
direct_register_custom_op(op_name="rocm_aiter_mla_decode_fwd",
101
op_func=mla_decode_fwd_impl,
102
mutates_args=["o"],
103
fake_impl=mla_decode_fwd_fake,
100
- tags=[torch.Tag.needs_fixed_stride_order])
104
+ tags=tags)
vllm/model_executor/layers/fused_moe/fused_moe.py
+5
-3
@@ -33,7 +33,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
33
dequant_mxfp4)
34
from vllm.platforms import current_platform
35
from vllm.triton_utils import tl, triton
36
-from vllm.utils import direct_register_custom_op
36
+from vllm.utils import direct_register_custom_op, is_torch_equal_or_newer
37
from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used
38
39
from .rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled
@@ -1056,7 +1056,8 @@ direct_register_custom_op(
1056
op_func=inplace_fused_experts,
1057
mutates_args=["hidden_states"],
1058
fake_impl=inplace_fused_experts_fake,
1059
- tags=(torch.Tag.needs_fixed_stride_order, ),
1059
+ tags=(() if is_torch_equal_or_newer("2.7.0") else
1060
+ (torch.Tag.needs_fixed_stride_order, )),
1061
)
1062
1063
@@ -1122,7 +1123,8 @@ direct_register_custom_op(
1123
op_func=outplace_fused_experts,
1124
mutates_args=[],
1125
fake_impl=outplace_fused_experts_fake,
1125
- tags=(torch.Tag.needs_fixed_stride_order, ),
1126
+ tags=(() if is_torch_equal_or_newer("2.7.0") else
1127
+ (torch.Tag.needs_fixed_stride_order, )),
1128
)
1129
1130