Skip to content

For adding TDM in separate PRs, deleted TDM related code that is inaccurate, incomplete, or incorrect #3377

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
glen-amd wants to merge 1 commit into
base: release/2.11
Choose a base branch
from
Open

For adding TDM in separate PRs, deleted TDM related code that is inaccurate, incomplete, or incorrect #3377

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 0 additions & 14 deletions test/inductor/test_flex_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -446,20 +446,6 @@ def batch_reserve(paged_attention: PagedAttention, target_seq_len: Tensor):
)


class TestFlexAttentionTDMOptions(InductorTestCase):
def test_apply_tdm_num_stages_uses_triton_launch_option(self):
from torch._inductor.kernel.flex.common import apply_tdm_num_stages

kernel_options = {"num_stages": 1, "NUM_STAGES": 4}

apply_tdm_num_stages(kernel_options)

self.assertEqual(
kernel_options["num_stages"], config.tdm.max_outstanding_per_wave
)
self.assertNotIn("NUM_STAGES", kernel_options)


@large_tensor_test_class("2GB", device=test_device[0])
class TestFlexAttention(InductorTestCase):
def setUp(self):
Expand Down
107 changes: 0 additions & 107 deletions test/inductor/test_max_autotune.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -328,113 +328,6 @@ def mock_get_tma_workspace_arg(*args, **kwargs):
mm_tma_heuristic.mm_configs = original_tma_configs
mm_heuristic.mm_configs = original_mm_configs

def test_tdm_arch_gate_accepts_only_gfx1250(self):
from torch._inductor.utils import is_gfx1250_arch

self.assertTrue(is_gfx1250_arch("gfx1250"))
self.assertTrue(is_gfx1250_arch("gfx1250:sramecc+:xnack-"))
self.assertFalse(is_gfx1250_arch("gfx1251"))
self.assertFalse(is_gfx1250_arch("amd-gfx1250"))

def test_tdm_persistent_template_precedes_rocm_tma_fallback(self):
from torch._inductor.kernel import mm as mm_kernel

for persistent_tma_enabled in (False, True):
templates = []
with (
config.patch(
{
"triton.enable_persistent_tma_matmul": persistent_tma_enabled
}
),
mock.patch.object(
mm_kernel,
"use_triton_blackwell_tma_template",
return_value=False,
),
mock.patch.object(
mm_kernel,
"use_triton_tdm_template",
return_value=True,
),
mock.patch.object(
mm_kernel,
"use_triton_tma_template",
side_effect=AssertionError("TMA fallback should not run"),
),
):
selected = mm_kernel._append_persistent_mm_template(
templates, mock.Mock(), mock.Mock(), mock.Mock()
)

self.assertEqual(selected, "tdm")
self.assertEqual(templates, [mm_kernel.persistent_mm_template])

def test_tdm_template_add_guards_checks_compile_time_device(self):
from torch._inductor.utils import use_triton_tdm_template

class FakeMatrix:
def __init__(self, device):
self.device = device

def get_device(self):
return self.device

def get_dtype(self):
return torch.float16

def get_stride(self):
return (128, 1)

mat = FakeMatrix(torch.device("cuda", 0))
with (
config.patch({"enable_tdm_configs": True}),
mock.patch.object(torch.version, "hip", "7.2.0"),
mock.patch.object(
torch.cuda,
"get_device_properties",
return_value=mock.Mock(gcnArchName="gfx1250"),
),
):
self.assertTrue(
use_triton_tdm_template(
mat,
output_layout=mock.Mock(device=torch.device("cuda", 0)),
add_guards=True,
)
)
self.assertFalse(
use_triton_tdm_template(
mat,
output_layout=mock.Mock(device=torch.device("cuda", 1)),
add_guards=True,
)
)

def test_tdm_block_k_filter_is_dtype_size_aware(self):
from torch._inductor.template_heuristics.triton import (
_filter_tdm_block_k_configs,
ROCmGemmConfig,
)

configs = [
ROCmGemmConfig(128, 64, 64, 4, 4, group_m=8),
ROCmGemmConfig(128, 64, 128, 4, 4, group_m=8),
]

self.assertEqual(
[c.block_k for c in _filter_tdm_block_k_configs(configs, 2)],
[64, 128],
)
self.assertEqual(
[c.block_k for c in _filter_tdm_block_k_configs(configs, 1)],
[128],
)
self.assertEqual(
[c.block_k for c in _filter_tdm_block_k_configs(configs, 4)],
[64, 128],
)

def test_shared_memory_estimation_counts_num_stages_once(self):
heuristic = ROCmMMTemplateConfigHeuristic()
gemm_config = GemmConfig(128, 64, 64, 4, 4, group_m=8)
Expand Down
24 changes: 0 additions & 24 deletions torch/_inductor/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -734,19 +734,6 @@ def use_autoheuristic(name: str) -> bool:

force_layout_optimization = os.environ.get("TORCHINDUCTOR_FORCE_LAYOUT_OPT", "0") == "1"

# AMD TDM config flag for TDM support for gfx1250, is a configuration gate
# rather than something that causes TDM instructions to be emitted directly.
# It's set to control if gfx1250-specific autotuning configs
# (with larger tile sizes, more pipeline stages) are included
# in the candidate set during `max_autotune`.
# When enabled and running on a gfx1250 device, Inductor may emit Triton kernels
# that leverage TDM for asynchronous global->LDS tensor tile copies.
# Requires Triton >= 3.6.0 with AMD TDM backend support.
# Disabled by default on ROCm; TDM is further gated by device arch (gfx1250)
# in `use_triton_tdm_template()` at codegen time.
# Set to "0" to disable TDM even on gfx1250 hardware.
enable_tdm_configs = os.environ.get("TORCHINDUCTOR_ENABLE_TDM_CONFIGS", "0") == "1"

# Whether to keep the output strides the same as eager after layout optimization.
keep_output_stride = os.environ.get("TORCHINDUCTOR_KEEP_OUTPUT_STRIDE", "1") == "1"

Expand Down Expand Up @@ -2281,17 +2268,6 @@ class rocm:
contiguous_threshold: int = 16


class tdm:
# Maximum outstanding TDM address translations per wave is 4.
# For small tiles (e.g., 128x64 FP16), this is the binding constraint.
# For larger tiles with 2 waves/SIMD, the SIMD limit of 6 applies.
max_outstanding_per_wave: int = 4
max_outstanding_per_simd: int = 6
# TDM requires 128B/256B aligned contiguous regions
# in both global memory and LDS.
alignment_bytes: int = 128


# Backend to use for CPU codegen either "cpp" or "triton" (experimental) or "halide" (experimental) or "pallas" (experimental)
cpu_backend: Literal["cpp", "triton", "halide", "pallas"] = "cpp"

Expand Down
7 changes: 1 addition & 6 deletions torch/_inductor/kernel/flex/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,17 +43,12 @@
to_dtype,
)
from ...select_algorithm import realize_inputs
from ...utils import config, load_template
from ...utils import load_template


SubgraphResults = Union[list[Optional[ComputedBuffer]], Optional[ComputedBuffer]]


def apply_tdm_num_stages(kernel_options: dict[str, Any]) -> None:
kernel_options["num_stages"] = config.tdm.max_outstanding_per_wave
kernel_options.pop("NUM_STAGES", None)


def zeros_and_scatter_lowering(shape: list[int], indices, values):
"""To support backwards on captured buffers we register a specific lowering for our specific custom up"""
# Always accumulate into fp32 then cast
Expand Down
18 changes: 1 addition & 17 deletions torch/_inductor/kernel/flex/flex_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,8 @@
SymbolicGridFn,
TritonTemplate,
)
from ...utils import can_use_tma, use_triton_tdm_template
from ...utils import can_use_tma
from .common import (
apply_tdm_num_stages,
build_subgraph_buffer,
create_indices_fake,
create_num_blocks_fake_generator,
Expand Down Expand Up @@ -411,15 +410,6 @@ def flex_attention(
if cur_kernel_options["USE_TMA"] and not can_use_tma(query, key, value):
cur_kernel_options["USE_TMA"] = False

# For gfx1250 TDM, ensure the non-TMA path uses enough pipeline stages
# to trigger TDM async copies in Triton's AMD backend (TTGIR pipelining pass).
# Standard attention block sizes (64, 128, 256) with FP16/BF16
# produce 128B aligned tiles and are compatible.
if not cur_kernel_options["USE_TMA"] and use_triton_tdm_template(
query, key, value
):
apply_tdm_num_stages(cur_kernel_options)

cur_kernel_options.setdefault("BLOCK_M", conf.block_m)
cur_kernel_options.setdefault("BLOCK_N", conf.block_n)
# Blocksparse options
Expand Down Expand Up @@ -937,12 +927,6 @@ def flex_attention_backward(*args, **kwargs):
if cur_kernel_options["USE_TMA"] and not can_use_tma(query, key, value):
cur_kernel_options["USE_TMA"] = False

# See the comments at the corresponding place in function `flex_attention` above.
if not cur_kernel_options["USE_TMA"] and use_triton_tdm_template(
query, key, value
):
apply_tdm_num_stages(cur_kernel_options)

cur_kernel_options.setdefault("BLOCK_M1", conf.block_m1)
cur_kernel_options.setdefault("BLOCK_N1", conf.block_n1)
cur_kernel_options.setdefault("BLOCK_M2", conf.block_m2)
Expand Down
12 changes: 1 addition & 11 deletions torch/_inductor/kernel/flex/flex_decoding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,8 @@
SymbolicGridFn,
TritonTemplate,
)
from ...utils import can_use_tma, use_triton_tdm_template
from ...utils import can_use_tma
from .common import (
apply_tdm_num_stages,
create_indices_fake,
create_num_blocks_fake_generator,
freeze_irnodes,
Expand Down Expand Up @@ -355,15 +354,6 @@ def create_flex_decoding_kernel(*args, **kwargs):
if cur_kernel_options["USE_TMA"] and not can_use_tma(query, key, value):
cur_kernel_options["USE_TMA"] = False

# For gfx1250 TDM, ensure the non-TMA path uses enough pipeline stages
# to trigger TDM async copies in Triton's AMD backend (TTGIR pipelining pass).
# Standard attention block sizes (64, 128, 256) with FP16/BF16
# produce 128B aligned tiles and are compatible.
if not cur_kernel_options["USE_TMA"] and use_triton_tdm_template(
query, key, value
):
apply_tdm_num_stages(cur_kernel_options)

# Add ROCm-specific parameters if they exist in the config
for attrib in ["kpack", "matrix_instr_nonkdim", "waves_per_eu"]:
if hasattr(conf, attrib):
Expand Down
34 changes: 0 additions & 34 deletions torch/_inductor/kernel/mm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,6 @@
use_triton_scaling_template,
use_triton_template,
use_triton_tma_template,
use_triton_tdm_template,
)
from .mm_common import (
_is_static_problem,
Expand Down Expand Up @@ -141,39 +140,6 @@ def _append_persistent_mm_template(
):
templates_to_use.append(blackwell_ws_persistent_device_tma_mm_template)
return "blackwell_tma"
if use_triton_tdm_template(mat1, mat2, output_layout=layout, add_guards=True):
# GFX1250 TDM: use the non-TMA persistent template. Triton's AMD backend
# automatically inserts TDM instructions when compiling with num_stages > 1.
#
# TDM hardware constraints (gfx1250):
# Documentation-only for now because Triton's pipeliner handles the
# wave assignment internally. If Triton exposes wave-specialization
# knobs in the future, these comments identify where to hook them.
# - 1 TDM unit per SIMD pair (2 SIMDs share 1 TDM unit)
# - Max 4 outstanding TDM address translations per wave
# - Max 6 outstanding TDM address translations per SIMD
# - Recommended: waves specialized to load A or B
# - 8 waves/WG (num_warps=8): 4 waves for A, 4 for B
# - 4 waves/WG (num_warps=4): 2 waves for A, 2 for B
# - All waves can load both A & B. This is preferable for uniform
# unrolling across CUs and helps multicast load latency.
# - Do not interleave different TDMs within a wave. Complete one TDM
# instruction's unrolling before switching.
#
# The persistent template heuristic supplies gfx1250 TDM configs.
# Key constraints:
# - TDM requests target 128B or 256B aligned contiguous regions in both
# global memory and LDS.
# - For FP16: BLOCK_K must be a multiple of 64 (64 * 2B = 128B).
# - For FP8/FP4: BLOCK_K must be a multiple of 128 (128 * 1B = 128B).
# - For FP32: BLOCK_K must be a multiple of 32 (32 * 4B = 128B).
# - BLOCK_M and BLOCK_N should also produce 128B-aligned LDS rows.
# TDM is most beneficial for small tiles such as 128x64/64x128 FP16
# and 128x128/128x256 FP8 with num_stages=4 (matching the
# 4-outstanding-per-wave TDM address translation limit) and
# num_warps=4 or 8 (for wave-specialized A/B loading).
templates_to_use.append(persistent_mm_template)
return "tdm"
if use_triton_tma_template(mat1, mat2, output_layout=layout, add_guards=True):
if torch.version.hip is None:
templates_to_use.append(persistent_tma_mm_template)
Expand Down
Loading