From 31dc22ebf83925ce934da29ededcb3396cd59572 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 26 Aug 2026 23:01:21 +0000 Subject: [PATCH 1/2] cuda: use Triton TMA for long-context causal prefill Route SM90+ global causal attention with a device-resident KV bound and L_kv >= 16K through shape-specific Triton tensor-descriptor kernels. Cover common head dimensions 64 and 128 with separately profiled resource configurations; unsupported or non-beneficial shapes retain the portable fallback. --- backends/cuda/tests/test_triton_sdpa.py | 128 ++++++++++++++++++ backends/cuda/triton/kernels/sdpa.py | 168 +++++++++++++++++++++++- 2 files changed, 295 insertions(+), 1 deletion(-) diff --git a/backends/cuda/tests/test_triton_sdpa.py b/backends/cuda/tests/test_triton_sdpa.py index 8842128b7ff..576ebe5a934 100644 --- a/backends/cuda/tests/test_triton_sdpa.py +++ b/backends/cuda/tests/test_triton_sdpa.py @@ -13,8 +13,10 @@ Test parametrization adapted from FlashAttention (tests/cute/test_flash_attn.py). """ +import importlib import itertools import unittest +from unittest import mock import torch import torch.nn.functional as F @@ -646,6 +648,132 @@ def test_explicit_mask_composes_with_causal_kv_len(self): self.assertFalse(torch.isnan(out).any()) self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL) + @unittest.skipIf( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9, + "TMA requires SM90+", + ) + def test_tma_causal_prefill_common_head_dims(self): + """TMA causal prefill supports common power-of-two head dimensions.""" + import triton + from triton.runtime._allocation import _allocator + + sdpa_module = importlib.import_module( + "executorch.backends.cuda.triton.kernels.sdpa" + ) + + # Tensor descriptors require a small runtime descriptor workspace in + # eager mode. Inductor supplies this allocator in the production path. + self.addCleanup(triton.set_allocator, _allocator.get()) + triton.set_allocator( + lambda size, alignment, stream: torch.empty( + size, dtype=torch.int8, device="cuda" + ) + ) + B, H_q, H_kv = 1, 4, 2 + Lq, kv_len, Lk = 512, 4096, 16384 + + for D in (64, 128): + with self.subTest(D=D): + self.assertIsNotNone(sdpa_module._tma_prefill_config(D, Lq)) + torch.manual_seed(D) + q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda") + k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda") + v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda") + kv_len_t = torch.tensor([kv_len], dtype=torch.int32, device="cuda") + dense = self._dense_bottom_right_causal_mask(B, Lq, kv_len, Lk, "cuda") + + with mock.patch.object( + sdpa_module, "cuda_targets_are_sm90_or_newer", return_value=True + ): + out_tma = self.sdpa( + q, + k, + v, + attn_mask=None, + enable_gqa=True, + kv_len=kv_len_t, + is_causal=True, + ) + with mock.patch.object( + sdpa_module, "cuda_targets_are_sm90_or_newer", return_value=False + ): + out_existing = self.sdpa( + q, + k, + v, + attn_mask=None, + enable_gqa=True, + kv_len=kv_len_t, + is_causal=True, + ) + out_dense = self.sdpa( + q, + k, + v, + attn_mask=dense, + enable_gqa=True, + kv_len=kv_len_t, + ) + + self.assertFalse(torch.isnan(out_tma).any()) + self.assertLess(_max_abs_error(out_tma, out_dense), MAX_ABS_TOL) + self.assertLess(_max_abs_error(out_tma, out_existing), MAX_ABS_TOL) + + @unittest.skipIf( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9, + "TMA requires SM90+", + ) + def test_tma_falls_back_for_noncontiguous_head_dim(self): + """TMA descriptors are used only when Q/K/V have unit inner stride.""" + B, H_q, H_kv, Lq, kv_len, Lk, D = 1, 4, 2, 512, 4096, 16384, 128 + torch.manual_seed(6) + q = torch.randn(B, H_q, Lq, D * 2, dtype=torch.bfloat16, device="cuda")[ + ..., ::2 + ] + k = torch.randn(B, H_kv, Lk, D * 2, dtype=torch.bfloat16, device="cuda")[ + ..., ::2 + ] + v = torch.randn(B, H_kv, Lk, D * 2, dtype=torch.bfloat16, device="cuda")[ + ..., ::2 + ] + kv_len_t = torch.tensor([kv_len], dtype=torch.int32, device="cuda") + dense = self._dense_bottom_right_causal_mask(B, Lq, kv_len, Lk, "cuda") + + out = self.sdpa(q, k, v, enable_gqa=True, kv_len=kv_len_t, is_causal=True) + ref = _reference_sdpa(q, k, v, attn_mask=dense) + + self.assertFalse(torch.isnan(out).any()) + self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL) + + @unittest.skipIf( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9, + "TMA requires SM90+", + ) + def test_tma_fully_masked_rows_are_finite(self): + """Rows before the beginning of a short KV prefix return zero, not NaN.""" + import triton + from triton.runtime._allocation import _allocator + + self.addCleanup(triton.set_allocator, _allocator.get()) + triton.set_allocator( + lambda size, alignment, stream: torch.empty( + size, dtype=torch.int8, device="cuda" + ) + ) + B, H_q, H_kv, Lq, kv_len, Lk, D = 1, 4, 2, 512, 128, 16384, 128 + torch.manual_seed(7) + q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda") + k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda") + v = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda") + kv_len_t = torch.tensor([kv_len], dtype=torch.int32, device="cuda") + dense = self._dense_bottom_right_causal_mask(B, Lq, kv_len, Lk, "cuda") + + out = self.sdpa(q, k, v, enable_gqa=True, kv_len=kv_len_t, is_causal=True) + ref = _reference_sdpa(q, k, v, attn_mask=dense) + + self.assertFalse(torch.isnan(out).any()) + self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL) + def test_mask_is_causal_matches_dense_decode(self): """is_causal + kv_len is a no-op vs dense for L_q==1 decode over a KV cache.""" D, B, H_q, H_kv = 128, 1, 16, 2 diff --git a/backends/cuda/triton/kernels/sdpa.py b/backends/cuda/triton/kernels/sdpa.py index afc79c86379..f43a1564a9e 100644 --- a/backends/cuda/triton/kernels/sdpa.py +++ b/backends/cuda/triton/kernels/sdpa.py @@ -38,6 +38,7 @@ import torch import triton import triton.language as tl +from executorch.backends.cuda.target_arch import cuda_targets_are_sm90_or_newer from torch.library import triton_op, wrap_triton @@ -52,6 +53,23 @@ def _is_power_of_2(n: int) -> bool: _SPLITK_LKV_THRESHOLD = 256 +_TMA_PREFILL_LKV_THRESHOLD = 16384 + + +def _tma_prefill_config( + head_dim: int, query_len: int +) -> Optional[tuple[int, int, int, int]]: + """Return a validated TMA config for common transformer head dimensions.""" + if head_dim == 64: + return 64, 64, 3, 4 + if head_dim == 128: + if query_len >= 1024: + return 128, 64, 3, 8 + if query_len >= 512: + return 64, 64, 3, 4 + return None + + # Decode split-K occupancy target. A sweep across both production attention # families showed that targeting 16/9 waves gives a good balance between split # kernel occupancy and reduction/empty-CTA overhead. Keep the ratio integral so @@ -389,7 +407,8 @@ def _sdpa_fwd_kernel_non_pow2( l_i = (l_i * alpha + l_ij).to(tl.float32) m_i = m_ij - out = acc / l_i[:, None] + inv_l = tl.where(l_i > 0, 1.0 / l_i, 0.0) + out = acc * inv_l[:, None] if PACK_GQA: o_ptrs = ( @@ -780,6 +799,102 @@ def _sdpa_fwd_kernel( ) +@triton.jit +def _sdpa_prefill_tma_kernel( + Q_ptr, + K_ptr, + V_ptr, + O_ptr, + KV_LEN_ptr, + H_grid, + Lq, + Lk, + stride_qb, + stride_qh, + stride_qm, + stride_kb, + stride_kh, + stride_kn, + stride_vb, + stride_vh, + stride_vn, + stride_ob, + stride_oh, + stride_om, + stride_od, + sm_scale, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, + NUM_GROUPS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + """TMA global-causal prefill for SM90+ and a device-resident KV bound.""" + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + b = pid_bh // H_grid + h_q = pid_bh % H_grid + h_kv = h_q // NUM_GROUPS + + q_desc = tl.make_tensor_descriptor( + Q_ptr + b * stride_qb + h_q * stride_qh, + shape=[Lq, HEAD_DIM], + strides=[stride_qm, 1], + block_shape=[BLOCK_M, HEAD_DIM], + ) + k_desc = tl.make_tensor_descriptor( + K_ptr + b * stride_kb + h_kv * stride_kh, + shape=[Lk, HEAD_DIM], + strides=[stride_kn, 1], + block_shape=[BLOCK_N, HEAD_DIM], + ) + v_desc = tl.make_tensor_descriptor( + V_ptr + b * stride_vb + h_kv * stride_vh, + shape=[Lk, HEAD_DIM], + strides=[stride_vn, 1], + block_shape=[BLOCK_N, HEAD_DIM], + ) + + q_start = pid_m * BLOCK_M + offs_m = q_start + tl.arange(0, BLOCK_M) + offs_n_base = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, HEAD_DIM) + q = tl.load_tensor_descriptor(q_desc, [q_start, 0]) + kv_len = tl.minimum(tl.load(KV_LEN_ptr), Lk) + absolute_q = (kv_len - Lq) + offs_m + m_i = tl.full([BLOCK_M], -float("inf"), tl.float32) + l_i = tl.zeros([BLOCK_M], tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32) + scale = sm_scale.to(tl.float32) + + for start_n in tl.range(0, kv_len, BLOCK_N, num_stages=NUM_STAGES): + offs_n = start_n + offs_n_base + k = tl.load_tensor_descriptor(k_desc, [start_n, 0]) + qk = tl.dot(q, tl.trans(k)).to(tl.float32) * scale + valid = (offs_n[None, :] < kv_len) & (offs_n[None, :] <= absolute_q[:, None]) + qk = tl.where(valid, qk, -float("inf")) + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + safe_m = tl.where(m_ij == -float("inf"), 0.0, m_ij) + alpha = tl.exp(m_i - safe_m) + p = tl.exp(qk - safe_m[:, None]) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + v = tl.load_tensor_descriptor(v_desc, [start_n, 0]) + acc = tl.dot(p.to(tl.bfloat16), v, acc).to(tl.float32) + m_i = m_ij + + inv_l = tl.where(l_i > 0, 1.0 / l_i, 0.0) + out = acc * inv_l[:, None] + o_ptrs = ( + O_ptr + + b * stride_ob + + h_q * stride_oh + + offs_m[:, None] * stride_om + + offs_d[None, :] * stride_od + ) + tl.store(o_ptrs, out.to(tl.bfloat16), mask=offs_m[:, None] < Lq) + + def _validate_sdpa_inputs( query: torch.Tensor, key: torch.Tensor, @@ -875,6 +990,57 @@ def _launch_pow2_kernel( stride_vb, stride_vh, stride_vn, stride_vd = value.stride() stride_ob, stride_oh, stride_om, stride_od = out.stride() + # Hopper/Blackwell: TMA materially improves the long-context global + # causal path. Sliding-window attention keeps the existing kernel, as do + # pre-SM90 GPUs. The same TMA kernel is safe for short chunks and avoids a + # host sync on the device-resident kv_len scalar. + tma_config = _tma_prefill_config(D, L_q) + if ( + cuda_targets_are_sm90_or_newer() + and HAS_KV_LEN + and mask_is_causal + and not HAS_MASK + and query.stride(-1) == 1 + and key.stride(-1) == 1 + and value.stride(-1) == 1 + and tma_config is not None + and L_kv >= _TMA_PREFILL_LKV_THRESHOLD + and L_q > 4 + ): + block_m, block_n, tma_num_stages, tma_num_warps = tma_config + wrap_triton(_sdpa_prefill_tma_kernel)[(triton.cdiv(L_q, block_m), B * H_q)]( + query, + key, + value, + out, + kv_len_ptr, + H_q, + L_q, + L_kv, + stride_qb, + stride_qh, + stride_qm, + stride_kb, + stride_kh, + stride_kn, + stride_vb, + stride_vh, + stride_vn, + stride_ob, + stride_oh, + stride_om, + stride_od, + sm_scale, + BLOCK_M=block_m, + BLOCK_N=block_n, + HEAD_DIM=D, + NUM_GROUPS=num_groups, + NUM_STAGES=tma_num_stages, + num_warps=tma_num_warps, + num_stages=tma_num_stages, + ) + return + if pack_gqa: H_grid = H_kv Lq_packed = L_q * num_groups From 8da3f29323b3900eedd99ec03b8e9d1417cec1cb Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Tue, 1 Sep 2026 03:28:26 -0700 Subject: [PATCH 2/2] cuda: gate TMA causal prefill behind compile spec --- backends/cuda/BUCK | 29 ++++++++++ backends/cuda/cuda_backend.py | 14 +++++ backends/cuda/optimization_config.py | 29 ++++++++++ backends/cuda/target_arch.py | 42 ++++++++++++++ backends/cuda/tests/test_cuda_export.py | 43 ++++++++++++++ backends/cuda/tests/test_triton_sdpa.py | 11 +++- backends/cuda/triton/kernels/sdpa.py | 4 +- .../muse-glimmer/export/export_dflash.py | 57 ++++++++++++++----- .../models/muse-glimmer/export/export_solo.py | 38 +++++++++---- 9 files changed, 240 insertions(+), 27 deletions(-) create mode 100644 backends/cuda/optimization_config.py create mode 100644 backends/cuda/target_arch.py diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 3fe33473472..654b9225fdb 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -3,6 +3,31 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") +fbcode_target( + _kind = runtime.python_library, + name = "target_arch", + srcs = [ + "target_arch.py", + ], + visibility = [ + "//executorch/backends/cuda/...", + ], + deps = [ + "//caffe2:torch", + ], +) + +fbcode_target( + _kind = runtime.python_library, + name = "optimization_config", + srcs = [ + "optimization_config.py", + ], + visibility = [ + "//executorch/backends/cuda/...", + ], +) + fbcode_target( _kind = runtime.python_library, name = "coalesced_int4_tensor", @@ -99,6 +124,8 @@ fbcode_target( visibility = ["PUBLIC"], deps = [ ":cuda_passes", + ":optimization_config", + ":target_arch", ":triton_replacement_pass", "//caffe2:torch", "//executorch/backends/aoti/passes:passes", @@ -137,6 +164,8 @@ fbcode_target( "//executorch/backends/cuda/...", ], deps = [ + ":optimization_config", + ":target_arch", "//caffe2:torch", ], ) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 854ebf8f952..cbc32fbad72 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -24,12 +24,14 @@ CudaWeightCollector, trim_host_memory, ) +from executorch.backends.cuda.optimization_config import cuda_optimization_context from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import ( MoveCondPredicateToCpuPass, ) from executorch.backends.cuda.passes.replace_int64_floordiv import ( ReplaceInt64FloorDivWithFloatPass, ) +from executorch.backends.cuda.target_arch import cuda_targets_are_sm90_or_newer from executorch.backends.cuda.target_smem import target_smem_context from executorch.backends.cuda.triton.replacement_pass import ( ReplaceEdgeOpWithTritonOpPass, @@ -866,6 +868,7 @@ def get_extra_aoti_compile_context_manager( # Parse compile_specs for low_memory_mode (default OFF). compile_specs # may be None when called without specs (parity with base default). low_memory_mode = "OFF" + tma_causal_prefill = False for spec in compile_specs or []: if spec.key == "low_memory_mode": mode = spec.value.decode("utf-8").upper() @@ -874,6 +877,14 @@ def get_extra_aoti_compile_context_manager( f"Invalid low_memory_mode: {mode}. Expected 'ON' or 'OFF'." ) low_memory_mode = mode + elif spec.key == "enable_tma_causal_prefill": + tma_causal_prefill = _on_off_compile_spec_value(spec) + + if tma_causal_prefill and not cuda_targets_are_sm90_or_newer(): + logging.warning( + "enable_tma_causal_prefill requires an SM90+ CUDA target; disabling it" + ) + tma_causal_prefill = False @contextlib.contextmanager def _combined(): @@ -885,6 +896,9 @@ def _combined(): # only the fallback for the `triton_kernel_mode="OFF"` path. stack.enter_context(torch.nn.attention.sdpa_kernel([SDPBackend.MATH])) stack.enter_context(target_smem_context()) + stack.enter_context( + cuda_optimization_context(tma_causal_prefill=tma_causal_prefill) + ) if low_memory_mode == "ON": # Force AOTI's mutated-buffer clones onto CPU during # compile so we stay under tight GPU memory caps (e.g. diff --git a/backends/cuda/optimization_config.py b/backends/cuda/optimization_config.py new file mode 100644 index 00000000000..8b9ee3a062e --- /dev/null +++ b/backends/cuda/optimization_config.py @@ -0,0 +1,29 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Scoped CUDA optimization choices used while AOTInductor traces kernels.""" + +import contextlib +import contextvars +from typing import Iterator + + +_TMA_CAUSAL_PREFILL_ENABLED = contextvars.ContextVar( + "tma_causal_prefill_enabled", default=False +) + + +def tma_causal_prefill_enabled() -> bool: + return _TMA_CAUSAL_PREFILL_ENABLED.get() + + +@contextlib.contextmanager +def cuda_optimization_context(*, tma_causal_prefill: bool) -> Iterator[None]: + tma_token = _TMA_CAUSAL_PREFILL_ENABLED.set(tma_causal_prefill) + try: + yield + finally: + _TMA_CAUSAL_PREFILL_ENABLED.reset(tma_token) diff --git a/backends/cuda/target_arch.py b/backends/cuda/target_arch.py new file mode 100644 index 00000000000..d0aed194d8f --- /dev/null +++ b/backends/cuda/target_arch.py @@ -0,0 +1,42 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Helpers for selecting CUDA export paths from the requested target.""" + +import os +import re + +import torch + + +def cuda_targets_are_sm90_or_newer() -> bool: + """Return whether every requested NVIDIA CUDA target is SM90 or newer. + + Explicit ``TORCH_CUDA_ARCH_LIST`` targets take precedence over the local + export device. This keeps per-architecture AOT exports deterministic while + retaining local-device detection for the usual native-export workflow. + """ + if torch.version.hip is not None: + return False + + arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST") + if arch_list: + target_majors = [] + for target in re.split(r"[;,\s]+", arch_list): + target = target.strip().lower().removeprefix("sm_").removeprefix("compute_") + target = target.removesuffix("+ptx").removesuffix("a") + if not target: + continue + match = re.fullmatch(r"(\d+)(?:\.(\d+))?", target) + if match is None: + return False + major = int(match.group(1)) + if match.group(2) is None and major >= 10: + major //= 10 + target_majors.append(major) + return bool(target_majors) and min(target_majors) >= 9 + + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9 diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index eda5e46de41..04c56c3243d 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -104,6 +104,49 @@ def test_invalid_autotune_at_compile_time_compile_spec(self): [CompileSpec(key="autotune_at_compile_time", value=b"MAYBE")] ) + def test_tma_causal_prefill_defaults_off(self): + from executorch.backends.cuda.optimization_config import ( + tma_causal_prefill_enabled, + ) + + with CudaBackend.get_extra_aoti_compile_context_manager([]): + self.assertFalse(tma_causal_prefill_enabled()) + + def test_tma_causal_prefill_compile_spec(self): + from executorch.backends.cuda.optimization_config import ( + tma_causal_prefill_enabled, + ) + + with patch( + "executorch.backends.cuda.cuda_backend.cuda_targets_are_sm90_or_newer", + return_value=True, + ), CudaBackend.get_extra_aoti_compile_context_manager( + [CompileSpec(key="enable_tma_causal_prefill", value=b"ON")] + ): + self.assertTrue(tma_causal_prefill_enabled()) + + def test_invalid_tma_causal_prefill_compile_spec(self): + with self.assertRaisesRegex(ValueError, "Invalid enable_tma_causal_prefill"): + CudaBackend.get_extra_aoti_compile_context_manager( + [CompileSpec(key="enable_tma_causal_prefill", value=b"MAYBE")] + ) + + def test_tma_causal_prefill_unsupported_target_is_disabled(self): + from executorch.backends.cuda.optimization_config import ( + tma_causal_prefill_enabled, + ) + + with patch( + "executorch.backends.cuda.cuda_backend.cuda_targets_are_sm90_or_newer", + return_value=False, + ), self.assertLogs( + level="WARNING" + ) as logs, CudaBackend.get_extra_aoti_compile_context_manager( + [CompileSpec(key="enable_tma_causal_prefill", value=b"ON")] + ): + self.assertFalse(tma_causal_prefill_enabled()) + self.assertIn("requires an SM90+ CUDA target", "\n".join(logs.output)) + def test_target_smem_context_is_applied(self): with patch( "executorch.backends.cuda.cuda_backend.target_smem_context", diff --git a/backends/cuda/tests/test_triton_sdpa.py b/backends/cuda/tests/test_triton_sdpa.py index 576ebe5a934..4e648ca61da 100644 --- a/backends/cuda/tests/test_triton_sdpa.py +++ b/backends/cuda/tests/test_triton_sdpa.py @@ -20,6 +20,7 @@ import torch import torch.nn.functional as F +from executorch.backends.cuda.optimization_config import cuda_optimization_context def _skip_if_no_cuda(): @@ -682,7 +683,9 @@ def test_tma_causal_prefill_common_head_dims(self): kv_len_t = torch.tensor([kv_len], dtype=torch.int32, device="cuda") dense = self._dense_bottom_right_causal_mask(B, Lq, kv_len, Lk, "cuda") - with mock.patch.object( + with cuda_optimization_context( + tma_causal_prefill=True + ), mock.patch.object( sdpa_module, "cuda_targets_are_sm90_or_newer", return_value=True ): out_tma = self.sdpa( @@ -739,7 +742,8 @@ def test_tma_falls_back_for_noncontiguous_head_dim(self): kv_len_t = torch.tensor([kv_len], dtype=torch.int32, device="cuda") dense = self._dense_bottom_right_causal_mask(B, Lq, kv_len, Lk, "cuda") - out = self.sdpa(q, k, v, enable_gqa=True, kv_len=kv_len_t, is_causal=True) + with cuda_optimization_context(tma_causal_prefill=True): + out = self.sdpa(q, k, v, enable_gqa=True, kv_len=kv_len_t, is_causal=True) ref = _reference_sdpa(q, k, v, attn_mask=dense) self.assertFalse(torch.isnan(out).any()) @@ -768,7 +772,8 @@ def test_tma_fully_masked_rows_are_finite(self): kv_len_t = torch.tensor([kv_len], dtype=torch.int32, device="cuda") dense = self._dense_bottom_right_causal_mask(B, Lq, kv_len, Lk, "cuda") - out = self.sdpa(q, k, v, enable_gqa=True, kv_len=kv_len_t, is_causal=True) + with cuda_optimization_context(tma_causal_prefill=True): + out = self.sdpa(q, k, v, enable_gqa=True, kv_len=kv_len_t, is_causal=True) ref = _reference_sdpa(q, k, v, attn_mask=dense) self.assertFalse(torch.isnan(out).any()) diff --git a/backends/cuda/triton/kernels/sdpa.py b/backends/cuda/triton/kernels/sdpa.py index f43a1564a9e..e6f7825b517 100644 --- a/backends/cuda/triton/kernels/sdpa.py +++ b/backends/cuda/triton/kernels/sdpa.py @@ -38,6 +38,7 @@ import torch import triton import triton.language as tl +from executorch.backends.cuda.optimization_config import tma_causal_prefill_enabled from executorch.backends.cuda.target_arch import cuda_targets_are_sm90_or_newer from torch.library import triton_op, wrap_triton @@ -996,7 +997,8 @@ def _launch_pow2_kernel( # host sync on the device-resident kv_len scalar. tma_config = _tma_prefill_config(D, L_q) if ( - cuda_targets_are_sm90_or_newer() + tma_causal_prefill_enabled() + and cuda_targets_are_sm90_or_newer() and HAS_KV_LEN and mask_is_causal and not HAS_MASK diff --git a/examples/models/muse-glimmer/export/export_dflash.py b/examples/models/muse-glimmer/export/export_dflash.py index 8920e7e9b94..1e4ee8b4fec 100644 --- a/examples/models/muse-glimmer/export/export_dflash.py +++ b/examples/models/muse-glimmer/export/export_dflash.py @@ -58,6 +58,7 @@ def export_dflash( backend: str = "mlx", mmproj: str | None = None, max_vision_patches: int = 16384, + enable_tma_causal_prefill: bool = False, ) -> None: """Export DFlash target + draft to one CUDA or MLX .pte. @@ -138,19 +139,38 @@ def export_dflash( activation_dtype=activation_dtype, ) - backend_export = _export_dflash_mlx if backend == "mlx" else _export_dflash_cuda - backend_export( - target_model, - target_config, - draft_model, - draft_config, - vision_model, - pos_embed_table, - output_dir, - max_seq_len, - activation_dtype, - max_vision_patches, - ) + if backend == "mlx": + _export_dflash_mlx( + target_model, + target_config, + draft_model, + draft_config, + vision_model, + pos_embed_table, + output_dir, + max_seq_len, + activation_dtype, + max_vision_patches, + ) + else: + from executorch.backends.cuda.optimization_config import ( + cuda_optimization_context, + ) + + with cuda_optimization_context(tma_causal_prefill=enable_tma_causal_prefill): + _export_dflash_cuda( + target_model, + target_config, + draft_model, + draft_config, + vision_model, + pos_embed_table, + output_dir, + max_seq_len, + activation_dtype, + max_vision_patches, + enable_tma_causal_prefill=enable_tma_causal_prefill, + ) def _dflash_constant_methods( @@ -422,6 +442,7 @@ def _export_dflash_cuda( max_seq_len: int, activation_dtype: torch.dtype, max_vision_patches: int, + enable_tma_causal_prefill: bool, ) -> None: """Export the CUDA DFlash contract. @@ -627,6 +648,10 @@ def cuda_partitioner(method: str) -> CudaPartitioner: [ CudaBackend.generate_method_name_compile_spec(method), CompileSpec("low_memory_mode", b"ON"), + CompileSpec( + "enable_tma_causal_prefill", + b"ON" if enable_tma_causal_prefill else b"OFF", + ), ] ) @@ -771,6 +796,11 @@ def main() -> None: help="Activation / KV-cache / unquantized-weight dtype. Defaults to " "float16 for MLX and bfloat16 for CUDA.", ) + parser.add_argument( + "--enable-tma-causal-prefill", + action="store_true", + help="Enable the experimental SM90+ TMA causal prefill attention path.", + ) args = parser.parse_args() try: @@ -799,6 +829,7 @@ def main() -> None: backend=args.backend, mmproj=args.mmproj, max_vision_patches=args.max_vision_patches, + enable_tma_causal_prefill=args.enable_tma_causal_prefill, ) diff --git a/examples/models/muse-glimmer/export/export_solo.py b/examples/models/muse-glimmer/export/export_solo.py index a7f06ff6275..9576264e70d 100644 --- a/examples/models/muse-glimmer/export/export_solo.py +++ b/examples/models/muse-glimmer/export/export_solo.py @@ -84,19 +84,26 @@ def export_and_lower( pos_embed_table: torch.Tensor | None = None, max_vision_patches: int = 16384, vision_fp32_mm: str = "none", + enable_tma_causal_prefill: bool = False, ) -> None: if backend == "cuda": - _export_cuda( - model, - config, - output_dir, - sample=sample, - use_turboquant=use_turboquant, - vision_model=vision_model, - pos_embed_table=pos_embed_table, - max_vision_patches=max_vision_patches, - vision_fp32_mm=vision_fp32_mm, + from executorch.backends.cuda.optimization_config import ( + cuda_optimization_context, ) + + with cuda_optimization_context(tma_causal_prefill=enable_tma_causal_prefill): + _export_cuda( + model, + config, + output_dir, + sample=sample, + use_turboquant=use_turboquant, + vision_model=vision_model, + pos_embed_table=pos_embed_table, + max_vision_patches=max_vision_patches, + vision_fp32_mm=vision_fp32_mm, + enable_tma_causal_prefill=enable_tma_causal_prefill, + ) elif backend == "mlx": _export_mlx( model, @@ -158,6 +165,7 @@ def _export_cuda( pos_embed_table: torch.Tensor | None = None, max_vision_patches: int = 16384, vision_fp32_mm: str = "none", + enable_tma_causal_prefill: bool = False, ) -> None: import torch._inductor.config as inductor_config from executorch.backends.cuda.cuda_backend import CudaBackend @@ -279,6 +287,10 @@ def _partitioner_for(name: str) -> "CudaPartitioner": [ CudaBackend.generate_method_name_compile_spec(name), CompileSpec("low_memory_mode", b"ON"), + CompileSpec( + "enable_tma_causal_prefill", + b"ON" if enable_tma_causal_prefill else b"OFF", + ), ] ) @@ -611,6 +623,11 @@ def main() -> None: help="Optional FP32-output linear implementation for vision blocks " "0-34. The default preserves the original all-BF16 encoder.", ) + parser.add_argument( + "--enable-tma-causal-prefill", + action="store_true", + help="Enable the experimental SM90+ TMA causal prefill attention path.", + ) args = parser.parse_args() if args.backend == "cuda" and not torch.cuda.is_available(): @@ -685,6 +702,7 @@ def main() -> None: pos_embed_table=pos_embed_table, max_vision_patches=args.max_vision_patches, vision_fp32_mm=args.vision_fp32_mm, + enable_tma_causal_prefill=args.enable_tma_causal_prefill, )