[Torch 2.11] Guard torch._C._cpu attribute checks for forward compatibility (#35673)

Signed-off-by: atalman <atalman@fb.com>
This commit is contained in:
Andrey Talman
2026-03-17 14:47:59 -04:00
committed by GitHub
parent c5030c439d
commit 68f783a727
9 changed files with 10 additions and 12 deletions

View File

@@ -48,7 +48,7 @@ def get_attn_isa(
else:
if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
return "neon"
elif torch._C._cpu._is_amx_tile_supported():
elif torch.cpu._is_amx_tile_supported():
return "amx"
else:
return "vec"
@@ -400,9 +400,7 @@ def test_varlen_with_paged_kv_normal_vec(
@pytest.mark.parametrize("use_alibi", [False])
@pytest.mark.parametrize("use_sink", [False])
@pytest.mark.parametrize("isa", ["amx"])
@pytest.mark.skipif(
not torch._C._cpu._is_amx_tile_supported(), reason="no AMX support."
)
@pytest.mark.skipif(not torch.cpu._is_amx_tile_supported(), reason="no AMX support.")
def test_varlen_with_paged_kv_normal_amx(
seq_lens: list[tuple[int, int]],
num_heads: tuple[int, int],