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 6, 13:15 UTC
[Torch 2.11] Guard torch._C._cpu attribute checks for forward compatibility (#35673)
Signed-off-by: atalman <atalman@fb.com>
Andrey Talman committed
Mar 17, 2026 at 14:47 UTC
68f783a72749c714971af725ce5632b40c29b8cf
9 files changed
+10
-12
benchmarks/kernels/cpu/benchmark_cpu_attn.py
+1
-1
@@ -27,7 +27,7 @@ def get_attn_isa(
27
else:
28
if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
29
return "neon"
30
- elif torch._C._cpu._is_amx_tile_supported():
30
+ elif torch.cpu._is_amx_tile_supported():
31
return "amx"
32
else:
33
return "vec"
benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py
+1
-1
@@ -24,7 +24,7 @@ except (ImportError, AttributeError) as e:
24
sys.exit(1)
25
26
# ISA selection following test_cpu_fused_moe.py pattern
27
-ISA_CHOICES = ["amx", "vec"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]
27
+ISA_CHOICES = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]
28
29
30
@torch.inference_mode()
tests/kernels/attention/test_cpu_attn.py
+2
-4
@@ -48,7 +48,7 @@ def get_attn_isa(
48
else:
49
if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
50
return "neon"
51
- elif torch._C._cpu._is_amx_tile_supported():
51
+ elif torch.cpu._is_amx_tile_supported():
52
return "amx"
53
else:
54
return "vec"
@@ -400,9 +400,7 @@ def test_varlen_with_paged_kv_normal_vec(
400
@pytest.mark.parametrize("use_alibi", [False])
401
@pytest.mark.parametrize("use_sink", [False])
402
@pytest.mark.parametrize("isa", ["amx"])
403
-@pytest.mark.skipif(
404
- not torch._C._cpu._is_amx_tile_supported(), reason="no AMX support."
405
-)
403
+@pytest.mark.skipif(not torch.cpu._is_amx_tile_supported(), reason="no AMX support.")
404
def test_varlen_with_paged_kv_normal_amx(
405
seq_lens: list[tuple[int, int]],
406
num_heads: tuple[int, int],
tests/kernels/moe/test_cpu_fused_moe.py
+1
-1
@@ -22,7 +22,7 @@ INTERMEDIATE_DIM = [128, 2880]
22
BATCH_SIZE = [1, 64, 256]
23
ACT = [MoEActivation.SILU, MoEActivation.SWIGLUOAI]
24
USE_BIAS = [True, False]
25
-ISA = ["amx", "vec"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]
25
+ISA = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]
26
DTYPE = [torch.bfloat16]
27
28
vllm/model_executor/kernels/linear/mixed_precision/cpu.py
+1
-1
@@ -119,7 +119,7 @@ class CPUWNA16LinearKernel(MPLinearKernel):
119
120
121
def _get_isa_hint(dtype: torch.dtype) -> str:
122
- supports_amx = torch._C._cpu._is_amx_tile_supported()
122
+ supports_amx = torch.cpu._is_amx_tile_supported()
123
if supports_amx and dtype in (torch.bfloat16,):
124
return "amx"
125
else:
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+1
-1
@@ -280,7 +280,7 @@ class CPUFusedMOE:
280
if not (w13_output_size % 32 == 0 and w2_output_size % 32 == 0):
281
return False, "none"
282
283
- supports_amx = torch._C._cpu._is_amx_tile_supported()
283
+ supports_amx = torch.cpu._is_amx_tile_supported()
284
285
if (
286
supports_amx
vllm/model_executor/layers/quantization/cpu_wna16.py
+1
-1
@@ -292,7 +292,7 @@ class CPUAWQLinearMethod(LinearMethodBase):
292
293
294
def _get_isa_hint(dtype: torch.dtype) -> str:
295
- supports_amx = torch._C._cpu._is_amx_tile_supported()
295
+ supports_amx = torch.cpu._is_amx_tile_supported()
296
if supports_amx and dtype in (torch.bfloat16,):
297
return "amx"
298
else:
vllm/model_executor/layers/utils.py
+1
-1
@@ -212,7 +212,7 @@ direct_register_custom_op(
212
213
def check_cpu_sgl_kernel(n: int, k: int, dtype: torch.dtype) -> bool:
214
return (
215
- torch._C._cpu._is_amx_tile_supported()
215
+ torch.cpu._is_amx_tile_supported()
216
and (dtype in (torch.bfloat16, torch.int8))
217
and k % 32 == 0
218
and n % 16 == 0
vllm/v1/attention/backends/cpu_attn.py
+1
-1
@@ -482,7 +482,7 @@ def _get_attn_isa(
482
) -> str:
483
if head_size is not None and head_size % 32 != 0 and head_size % 16 == 0:
484
return "vec16"
485
- supports_amx = torch._C._cpu._is_amx_tile_supported()
485
+ supports_amx = torch.cpu._is_amx_tile_supported()
486
supports_arm = current_platform.get_cpu_architecture() == CpuArchEnum.ARM
487
supports_vxe = current_platform.get_cpu_architecture() == CpuArchEnum.S390X
488
if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0: