From 709d7a87a6a0eeceeccfec75d2d81f14e8b3df09 Mon Sep 17 00:00:00 2001 From: Glen Cao Date: Mon, 29 Jun 2026 09:25:45 -0700 Subject: [PATCH] Deleted TDM related code that is inaccurate, incomplete, or incorrect --- test/inductor/test_flex_attention.py | 14 -- test/inductor/test_max_autotune.py | 107 -------------- torch/_inductor/config.py | 24 ---- torch/_inductor/kernel/flex/common.py | 7 +- torch/_inductor/kernel/flex/flex_attention.py | 18 +-- torch/_inductor/kernel/flex/flex_decoding.py | 12 +- torch/_inductor/kernel/mm.py | 34 ----- torch/_inductor/template_heuristics/triton.py | 131 +----------------- torch/_inductor/utils.py | 112 --------------- 9 files changed, 5 insertions(+), 454 deletions(-) diff --git a/test/inductor/test_flex_attention.py b/test/inductor/test_flex_attention.py index c1fe3047fdd3f..a57d4da63c942 100644 --- a/test/inductor/test_flex_attention.py +++ b/test/inductor/test_flex_attention.py @@ -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): diff --git a/test/inductor/test_max_autotune.py b/test/inductor/test_max_autotune.py index 2eccaf1d592d5..65071a2628ec6 100644 --- a/test/inductor/test_max_autotune.py +++ b/test/inductor/test_max_autotune.py @@ -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) diff --git a/torch/_inductor/config.py b/torch/_inductor/config.py index d379050857403..cd2b5aa585174 100644 --- a/torch/_inductor/config.py +++ b/torch/_inductor/config.py @@ -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" @@ -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" diff --git a/torch/_inductor/kernel/flex/common.py b/torch/_inductor/kernel/flex/common.py index 764d7cd386691..bf4006de8399a 100644 --- a/torch/_inductor/kernel/flex/common.py +++ b/torch/_inductor/kernel/flex/common.py @@ -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 diff --git a/torch/_inductor/kernel/flex/flex_attention.py b/torch/_inductor/kernel/flex/flex_attention.py index 6277d26f43c4a..c6017fadb0387 100644 --- a/torch/_inductor/kernel/flex/flex_attention.py +++ b/torch/_inductor/kernel/flex/flex_attention.py @@ -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, @@ -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 @@ -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) diff --git a/torch/_inductor/kernel/flex/flex_decoding.py b/torch/_inductor/kernel/flex/flex_decoding.py index 8a24b71e398f4..4111915c26082 100644 --- a/torch/_inductor/kernel/flex/flex_decoding.py +++ b/torch/_inductor/kernel/flex/flex_decoding.py @@ -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, @@ -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): diff --git a/torch/_inductor/kernel/mm.py b/torch/_inductor/kernel/mm.py index 51d922deec82f..79831ed9feffe 100644 --- a/torch/_inductor/kernel/mm.py +++ b/torch/_inductor/kernel/mm.py @@ -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, @@ -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) diff --git a/torch/_inductor/template_heuristics/triton.py b/torch/_inductor/template_heuristics/triton.py index fb6740e48b37f..2117483aec594 100644 --- a/torch/_inductor/template_heuristics/triton.py +++ b/torch/_inductor/template_heuristics/triton.py @@ -33,7 +33,6 @@ get_backend_num_stages, get_num_sms, get_tma_workspace_arg, - is_gfx1250_arch, TMA_DESCRIPTOR_SIZE, using_b200, ) @@ -51,42 +50,6 @@ from torch._inductor.runtime.triton_compat import Config as TritonConfig -def _is_gfx1250_device() -> bool: - """Detect whether the current device is AMD gfx1250. - - gfx1250 has 320 KB LDS per CU, 512 B/clk/CU LDS bandwidth (R/W combined), - and TDM for async descriptor-based global->LDS copies. When num_stages > 1, - Triton's AMD backend automatically converts tl.load into TDM instructions. - - TDM hardware constraints: - - 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 - - Requests target 128B or 256B aligned contiguous regions - in both global memory and LDS - """ - if not torch.version.hip: - return False - if not inductor_config.enable_tdm_configs: - return False - try: - props = torch.cuda.get_device_properties(torch.cuda.current_device()) - arch = getattr(props, "gcnArchName", "") - return is_gfx1250_arch(arch) - except Exception: - return False - - -def _filter_tdm_block_k_configs( - configs: list[BaseConfig], - dtype_size: int, -) -> list[BaseConfig]: - if dtype_size <= 0: - return configs - block_k_multiple = max(inductor_config.tdm.alignment_bytes // dtype_size, 1) - return [c for c in configs if c.block_k % block_k_multiple == 0] - - # Gemm Configs @dataclasses.dataclass class BaseConfig: @@ -1419,7 +1382,6 @@ def __init__(self) -> None: super().__init__() self.default_num_stages = get_backend_num_stages() - self.uses_tdm_configs = False self.mm_configs: list[BaseConfig] = [ ROCmGemmConfig( @@ -1481,42 +1443,6 @@ def __init__(self) -> None: ROCmGemmConfig(256, 256, 64, self.default_num_stages, 8, group_m=4), ] - # TDM-optimized persistent MM configs for gfx1250. - # - # Hardware constraints on gfx1250: - # - 320 KB LDS per CU, 512 B/clk/CU bandwidth - # - Max 4 outstanding TDM translations per wave (bounds num_stages) - # - Max 6 outstanding TDM translations per SIMD - # - TDM requests target 128B/256B aligned contiguous regions - # - 1 TDM unit per SIMD pair -> recommended 1 wave/SIMD pair issuing - # TDM instructions (num_warps=4 gives 1 wave/SIMD on a 4-SIMD CU) - # - # BLOCK_K alignment requirements (128B / dtype_bytes): - # FP16/BF16 (2B): BLOCK_K multiple of 64 - # FP8/INT8 (1B): BLOCK_K multiple of 128 - # FP32 (4B): BLOCK_K multiple of 32 - # - # LDS usage estimate: (BM*BK + BK*BN)*dtype_size*num_stages - # Must stay under 320 KB. - self.tdm_persistent_mm_configs: list[BaseConfig] = [ - # Small tiles: TDM most beneficial here (LDS-BW-limited to 1 wave/SIMD). - # These are where TDM eliminates cluster load issue overhead. - # `waves_per_eu=0` means Triton compiler should figure out the best value. - ROCmGemmConfig(128, 64, 64, 4, 4, group_m=8, waves_per_eu=0), - ROCmGemmConfig( 64, 128, 64, 4, 4, group_m=8, waves_per_eu=0), - ROCmGemmConfig(128, 64, 128, 4, 4, group_m=8), - ROCmGemmConfig( 64, 128, 128, 4, 4, group_m=8), - # Medium tiles: benefit from TDM prologue overhead reduction. - ROCmGemmConfig(128, 128, 64, 4, 4, group_m=8), - ROCmGemmConfig(128, 128, 64, 4, 8, group_m=16), - ROCmGemmConfig(128, 128, 128, 4, 4, group_m=8), - ROCmGemmConfig(128, 128, 128, 3, 8, group_m=16), - # Larger tiles: exploit the 320 KB LDS with deep pipelines. - ROCmGemmConfig(256, 128, 64, 4, 8, group_m=16), - ROCmGemmConfig(128, 256, 64, 4, 8, group_m=16), - ROCmGemmConfig(256, 128, 128, 3, 8, group_m=16), - ] - # Exhaustive search for mm configs self.exhaustive_configs: list[BaseConfig] = [ ROCmGemmConfig( @@ -1638,57 +1564,12 @@ def _prune_exhaustive_configs( ] return pruned_configs - def preprocess_mm_configs( - self, - m: int, - n: int, - k: int, - configs: list[BaseConfig], - has_int8_tensor: bool = False, - scale: float = 1.0, - exclude: Callable[ - [sympy.Integer, sympy.Integer, sympy.Integer], bool - ] = lambda m, n, k: False, - dtype_size: int = 0, - op_name: str = "mm", - **kwargs, - ) -> Generator[TritonConfig, None, None]: - if self.uses_tdm_configs: - configs = _filter_tdm_block_k_configs(configs, dtype_size) - return super().preprocess_mm_configs( - m, - n, - k, - configs, - has_int8_tensor, - scale, - exclude, - dtype_size, - op_name, - **kwargs, - ) - def _filter_configs(self, configs: list[BaseConfig]) -> list[BaseConfig]: """ ROCm specific filtering - - Normally we force num_stages to `default_num_stages` because AMD's - Triton backend historically didn't benefit from multi-stage pipelining. - However, on gfx1250 with TDM enabled, num_stages > 1 is exactly - what triggers TDM async copies in the StreamPipeliner pass, so we - must preserve the requested num_stages for TDM configs. """ - tdm_enabled = _is_gfx1250_device() for c in configs: - # On gfx1250, preserve num_stages as provided (but cap at 4 to - # respect TDM's per-wave outstanding limit). Otherwise, fall - # back to the legacy behavior of forcing default_num_stages. - if tdm_enabled: - c.num_stages = ( - min(c.num_stages, 4) if c.num_stages > 1 else c.num_stages - ) - else: - c.num_stages = self.default_num_stages + c.num_stages = self.default_num_stages return super()._filter_configs(configs) def _finalize_mm_configs( @@ -3050,15 +2931,7 @@ class PersistentMMTemplateConfigHeuristic( def __init__(self) -> None: super().__init__() - if _is_gfx1250_device(): - # On gfx1250, use TDM-optimized configs that exploit: - # - 320 KB LDS (larger tiles possible) - # - TDM async global->LDS copies (num_stages > 1) - # - Wave specialization across SIMD pairs - self.mm_configs = self.tdm_persistent_mm_configs - self.uses_tdm_configs = True - else: - self.mm_configs = self.persistent_mm_configs + self.mm_configs = self.persistent_mm_configs def _get_template_configs_impl( self, diff --git a/torch/_inductor/utils.py b/torch/_inductor/utils.py index 5184b16001a35..c9c60aa71eb0e 100644 --- a/torch/_inductor/utils.py +++ b/torch/_inductor/utils.py @@ -1977,118 +1977,6 @@ def use_triton_tma_template( ) -def is_gfx1250_arch(arch: str) -> bool: - """Return True only for the gfx1250 target, including feature-suffixed names.""" - return arch.split(":", 1)[0] == "gfx1250" - - -def use_triton_tdm_template(*matrices: IRNode, output_layout=None, add_guards=False): - """Coarsely check whether AMD TDM-optimized persistent template should be used. - Despite the fact that TDM has some restrictions on direct LDS loads ( - when unrolling requests, the requests are to 128B or 256B aligned and - continuous regions in global memory, and writing back to 128B or 256B - aligned and continuous regions in LDS), this gate doesn't check per-element - alignment, instead delegating to Triton compiler's AMD backend. - That said, for GeMM specifically, it checks the leading dimensions produce - 128B-aligned rows, which is the common case where TDM is beneficial. - - Args: - - output_layout: Output layout (used for device detection) - - add_guards: If True, add shape guards to the graph - Returns True when: - - Running on ROCm (torch.version.hip is not None) - - Device is gfx1250 - - config.enable_tdm_configs (defined elsewhere) is True - - Triton version supports gfx1250 TDM (>= 3.6.0) - """ - if not torch.version.hip: - return False - if not config.enable_tdm_configs: - return False - device = matrices[0].get_device() - if device.type != "cuda": - return False - try: - # props = torch.cuda.get_device_properties(torch.cuda.current_device()) - props = torch.cuda.get_device_properties(device) - arch = getattr(props, "gcnArchName", "") - if not is_gfx1250_arch(arch): - return False - except Exception: - return False - - # Check Triton version supports TDM (3.6.0+) - try: - from torch.torch_version import TorchVersion - import triton - if TorchVersion(triton.__version__) < "3.6.0": - return False - except (ImportError, Exception): - return False - - # TDM can lower the same matrix element types as standard Triton GeMM. - # The persistent-template config heuristic still filters BLOCK_K by dtype - # size so each TDM request covers an aligned 128B region. - # FIXME: torch.float4 - triton_tdm_dtypes = { - torch.float16, - torch.bfloat16, - torch.float32, - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float8_e4m3fnuz, - torch.float8_e5m2fnuz, - torch.int8, - } - # TDM unrolls requests to 128B or 256B aligned contiguous regions - # in both global memory and LDS. For GeMM, this means the leading - # dimension (stride) of each matrix should produce rows that start - # at 128B-aligned addresses. If strides are not 128B-aligned, TDM - # may still work (Triton's backend can fall back per-load), but - # performance benefit is reduced. We don't hard-gate on this because: - # 1. The autotuner will benchmark and reject TDM configs if they're slower. - # 2. Triton's pipeliner handles misaligned cases gracefully. - # 3. Most real GeMM workloads have contiguous row-major or col-major layouts. - # For now, only log for debugging purposes. May be upgraded to a hard gate later. - for mat in matrices: - mat_dtype = mat.get_dtype() - if mat_dtype not in triton_tdm_dtypes: - return False - strides = mat.get_stride() - dtype_bytes = mat_dtype.itemsize - # Check if the innermost stride is 1 (contiguous) and - # the outer stride produces 128B-aligned rows. - if len(strides) >= 2: - inner_stride = strides[-1] - outer_stride = strides[-2] - if hasattr(inner_stride, '__int__'): - inner_val = int(inner_stride) - if inner_val != 1: - # TODO: Should be a hard rejection, not just a warning. - log.debug( - "TDM: matrix has non-unit inner stride %d, " - "TDM requests may not be contiguous", - inner_val, - ) - if hasattr(outer_stride, '__int__'): - row_bytes = int(outer_stride) * dtype_bytes - if row_bytes % config.tdm.alignment_bytes != 0: - log.debug( - "TDM: matrix row stride %d bytes is not %dB-aligned, " - "TDM performance may be suboptimal", - row_bytes, - config.tdm.alignment_bytes, - ) - - if add_guards and output_layout is not None: - # `SizeVarAllocator` does not guard device properties. Keep this as a - # static consistency check for the layout/device pair used at compile time. - if (output_layout.device.index or 0) != (device.index or 0): - return False - - return True - - def use_triton_blackwell_tma_template( *matrices: IRNode, output_layout: Layout, add_guards: bool = False ) -> bool: