From 0f9a03550ec60db6fc2c373073ad751d12227a30 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 19 Oct 2026 13:07:54 -0700 Subject: [PATCH] [executorch][cuda] Autotune SDPA on a visible mask and tune pack GQA with the tile Pull Request resolved: https://github.com/pytorch/executorch/pull/23727 D123681196 declared representative KV lengths for SDPA kernels that read KV_LEN_ptr. Masked calls still autotuned unrepresentative work: fabricated boolean masks were all false, so candidate kernels skipped visible KV blocks. Pack-GQA selection also relied on a fixed tile-utilization heuristic rather than the candidate tile. Declare Mask_ptr beside KV_LEN_ptr on the forward, decode-split-K and small-query-split-K kernels, using an all-visible mask for each representative KV-length scenario. Each candidate therefore sweeps the declared visible extent. Calls without masks pass zero and remain untouched. Make PACK_GQA a forward-kernel config dimension. Tune packed and unpacked candidates for GQA, prune packing without GQA, and derive the launch-grid head count from the chosen config. The non-power-of-two kernel keeps its fixed-BLOCK_M rule. Existing masked numerical/reference and autotune-input tests cover the contract. ghstack-source-id: 451457287 @exported-using-ghexport Differential Revision: [D124908769](https://our.internmc.facebook.com/intern/diff/D124908769/) --- backends/cuda/tests/test_autotune_inputs.py | 131 +++++++++++++++++--- backends/cuda/tests/test_triton_sdpa.py | 68 ++++++++++ backends/cuda/triton/kernels/sdpa.py | 101 +++++++++------ 3 files changed, 243 insertions(+), 57 deletions(-) diff --git a/backends/cuda/tests/test_autotune_inputs.py b/backends/cuda/tests/test_autotune_inputs.py index 6a0ae763671..e6c1e5a8f6c 100644 --- a/backends/cuda/tests/test_autotune_inputs.py +++ b/backends/cuda/tests/test_autotune_inputs.py @@ -149,14 +149,14 @@ def _config_key(config) -> tuple: @contextlib.contextmanager -def _record_decode_pick(picks: list): +def _record_pick(kernel_name: str, picks: list): from torch._inductor.runtime.triton_heuristics import CachingAutotuner original = CachingAutotuner.benchmark_all_configs def record(self, *args, **kwargs): timings = original(self, *args, **kwargs) - if getattr(self.fn, "__name__", "") == "_sdpa_decode_splitk_kernel": + if getattr(self.fn, "__name__", "") == kernel_name: picks.append(_config_key(min(timings, key=timings.get).config)) return timings @@ -167,28 +167,16 @@ def record(self, *args, **kwargs): CachingAutotuner.benchmark_all_configs = original -def _decode_inputs(): - torch.manual_seed(0) - q = torch.randn( - 1, _H_KV * _GROUPS, 1, _HEAD_DIM, dtype=torch.bfloat16, device="cuda" - ) - k = torch.randn(1, _H_KV, _CACHE, _HEAD_DIM, dtype=torch.bfloat16, device="cuda") - v = torch.randn_like(k) - return q, k, v - - -def _compiled_decode_pick() -> tuple: +def _compile_picks(module, inputs, kernel_name: str) -> list: from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower - q, k, v = _decode_inputs() - kv_len = torch.tensor([_RUNTIME_KV_LENS[0]], dtype=torch.int64, device="cuda") with torch.no_grad(): - program = torch.export.export(_Decode(), (q, k, v, kv_len), strict=True) + program = torch.export.export(module, inputs, strict=True) picks = [] - with torch.compiler.config.patch(force_disable_caches=True), _record_decode_pick( - picks + with torch.compiler.config.patch(force_disable_caches=True), _record_pick( + kernel_name, picks ): to_edge_transform_and_lower( program, @@ -199,6 +187,23 @@ def _compiled_decode_pick() -> tuple: ], compile_config=EdgeCompileConfig(_check_ir_validity=False), ) + return picks + + +def _decode_inputs(): + torch.manual_seed(0) + q = torch.randn( + 1, _H_KV * _GROUPS, 1, _HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + k = torch.randn(1, _H_KV, _CACHE, _HEAD_DIM, dtype=torch.bfloat16, device="cuda") + v = torch.randn_like(k) + return q, k, v + + +def _compiled_decode_pick() -> tuple: + q, k, v = _decode_inputs() + kv_len = torch.tensor([_RUNTIME_KV_LENS[0]], dtype=torch.int64, device="cuda") + picks = _compile_picks(_Decode(), (q, k, v, kv_len), "_sdpa_decode_splitk_kernel") assert len(picks) == 1, picks return picks[0] @@ -247,6 +252,14 @@ def test_sdpa_kernels_reading_kv_len_are_declared(self) -> None: "_sdpa_small_query_splitk_kernel", ): self.assertIn("KV_LEN_ptr", ai._REGISTRY[name].providers, name) + self.assertIn("Mask_ptr", ai._REGISTRY[name].providers, name) + + def test_visible_mask_pairs_with_every_kv_length(self) -> None: + args = {"Lk": 65536, "Lq": 32} + self.assertEqual( + _SDPA._visible_mask_scenarios(args), + [True] * len(_SDPA._kv_len_scenarios(args)), + ) def test_compiled_pick_is_fastest_at_runtime_kv_lengths(self) -> None: pick = _compiled_decode_pick() @@ -258,5 +271,87 @@ def test_compiled_pick_is_fastest_at_runtime_kv_lengths(self) -> None: self.assertLessEqual(regret, self.MAX_REGRET, (kv_len, pick, times)) +# A cell-layout step: T queries of T sequences over one pool, an explicit +# visibility mask and the live extent as kv_len (lower_offgraph_kv.cell_step). +_CELL_T, _CELL_H_Q, _CELL_POOL = 32, 32, 65536 +_CELL_LIVE = (4650, 9300) + + +class _CellStep(torch.nn.Module): + def forward(self, q, k, v, mask, kv_len): + return torch.ops.triton.sdpa(q, k, v, mask, 0.0, False, 0.0, True, kv_len) + + +def _cell_inputs(live: int): + """Each sequence's 128-cell prompt, then decode cells interleaved across + the sequences, filling ``live`` cells.""" + torch.manual_seed(0) + q = torch.randn( + 1, _CELL_H_Q, _CELL_T, _HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + k = torch.randn( + 1, _CELL_POOL, _H_KV, _HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ).transpose(1, 2) + v = torch.randn_like(k) + mask = torch.zeros(1, 1, _CELL_T, _CELL_POOL, dtype=torch.bool, device="cuda") + prompt = 128 + for i in range(_CELL_T): + mask[0, 0, i, i * prompt : (i + 1) * prompt] = True + tail = torch.arange(_CELL_T * prompt, live, device="cuda") + mask[0, 0, (tail - _CELL_T * prompt) % _CELL_T, tail] = True + kv_len = torch.tensor([live], dtype=torch.int64, device="cuda") + return q, k, v, mask, kv_len + + +def _time_every_cell_config(live: int) -> dict: + import triton + from executorch.backends.cuda.autotune.cuda_graph_timing import ( + time_with_cuda_graph_us, + ) + + q, k, v, mask, kv_len = _cell_inputs(live) + kernel = _SDPA._sdpa_fwd_kernel + configs = _SDPA._sdpa_fwd_prune( + kernel.configs, None, HEAD_DIM=_HEAD_DIM, NUM_GROUPS=_CELL_H_Q // _H_KV + ) + times = {} + with torch.cuda.stream(torch.cuda.Stream()): + for config in configs: + single = triton.autotune(configs=[config], key=kernel.keys)(kernel.fn) + with mock.patch.object(_SDPA, "_sdpa_fwd_kernel", single): + try: + times[_config_key(config)] = time_with_cuda_graph_us( + lambda: _CellStep()(q, k, v, mask, kv_len) + ) + except Exception: # noqa: BLE001 - configs that do not fit are skipped + continue + return times + + +class SdpaCellStepAutotuneTest(unittest.TestCase): + """The standard SDPA config (tile and pack GQA) compiled for a masked + cell step is the fastest at live extents, not at the all-False mask + autotuning would otherwise time.""" + + MAX_REGRET = 1.10 + + @classmethod + def setUpClass(cls) -> None: + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA required") + + def test_compiled_pick_is_fastest_at_live_extents(self) -> None: + picks = _compile_picks( + _CellStep(), _cell_inputs(_CELL_LIVE[0]), "_sdpa_fwd_kernel" + ) + self.assertEqual(len(picks), 1, picks) + for live in _CELL_LIVE: + times = _time_every_cell_config(live) + best = min(times.values()) + regret = times[picks[0]] / best + print(f"live={live}: compiled pick {regret:.2f}x of the best") + self.assertLessEqual(regret, self.MAX_REGRET, (live, picks[0], times)) + + if __name__ == "__main__": unittest.main() diff --git a/backends/cuda/tests/test_triton_sdpa.py b/backends/cuda/tests/test_triton_sdpa.py index 3dd31a42c69..1fdb1dedef9 100644 --- a/backends/cuda/tests/test_triton_sdpa.py +++ b/backends/cuda/tests/test_triton_sdpa.py @@ -299,6 +299,74 @@ def test_gqa_prefill(self): self.assertFalse(torch.isnan(out).any()) self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL) + def test_every_tile_packed_or_not_matches(self): + """Each standard-kernel tile over a masked GQA step bounded by kv_len: + packing changes which rows share a tile, not any row's arithmetic, so + packed and unpacked agree exactly, and both match the reference.""" + import triton + + sdpa_module = importlib.import_module( + "executorch.backends.cuda.triton.kernels.sdpa" + ) + kernel = sdpa_module._sdpa_fwd_kernel + B, H_q, H_kv, D, Lk, kv_len = 1, 32, 2, 128, 512, 300 + torch.manual_seed(0) + 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.int64, device="cuda") + configs = sdpa_module._sdpa_fwd_prune( + kernel.configs, None, HEAD_DIM=D, NUM_GROUPS=H_q // H_kv + ) + self.assertEqual({c.kwargs["PACK_GQA"] for c in configs}, {False, True}) + for Lq in (5, 32, 64): + q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda") + # Query i sees an interleaved subset of the live cells, as a cell + # step's mask does. + mask = torch.zeros(B, 1, Lq, Lk, dtype=torch.bool, device="cuda") + cells = torch.arange(kv_len, device="cuda") + mask[0, 0, cells % Lq, cells] = True + ref = _reference_sdpa( + q, k[:, :, :kv_len], v[:, :, :kv_len], attn_mask=mask[..., :kv_len] + ) + # Few cells per row leave outputs near V's magnitude, where bf16 + # rounding of the output alone approaches MAX_ABS_TOL. + tol = max(MAX_ABS_TOL, 2 * _max_abs_error(ref.to(torch.bfloat16), ref)) + outs = {} + for config in configs: + single = triton.autotune(configs=[config], key=kernel.keys)(kernel.fn) + with self.subTest(Lq=Lq, config=str(config)), mock.patch.object( + sdpa_module, "_sdpa_fwd_kernel", single + ): + out = self.sdpa( + q, k, v, attn_mask=mask, enable_gqa=True, kv_len=kv_len_t + ) + self.assertLess(_max_abs_error(out, ref), tol) + tile = ( + tuple( + sorted( + (key, value) + for key, value in config.kwargs.items() + if key != "PACK_GQA" + ) + ), + config.num_warps, + config.num_stages, + ) + outs.setdefault(tile, []).append(out) + for tile, (unpacked, packed) in outs.items(): + with self.subTest(Lq=Lq, tile=tile): + self.assertTrue(torch.equal(unpacked, packed)) + + def test_mha_never_packs(self): + sdpa_module = importlib.import_module( + "executorch.backends.cuda.triton.kernels.sdpa" + ) + configs = sdpa_module._sdpa_fwd_prune( + sdpa_module._sdpa_fwd_kernel.configs, None, HEAD_DIM=128, NUM_GROUPS=1 + ) + self.assertTrue(configs) + self.assertFalse(any(c.kwargs["PACK_GQA"] for c in configs)) + def test_gqa_non_pow2_head_dim(self): """GQA with non-power-of-2 head dimensions.""" for D in HEAD_DIMS_NON_POW2: diff --git a/backends/cuda/triton/kernels/sdpa.py b/backends/cuda/triton/kernels/sdpa.py index 585c709679c..e01c0922cef 100644 --- a/backends/cuda/triton/kernels/sdpa.py +++ b/backends/cuda/triton/kernels/sdpa.py @@ -13,9 +13,10 @@ # # Pack GQA folds multiple Q heads that share the same KV head into the M # (sequence) dimension of a single tile, so K/V are loaded once per KV head -# instead of once per Q head. The tile-utilization heuristic decides when -# packing is beneficial (short seqlen_q, e.g. decode) vs. when simple head -# remapping suffices (long seqlen_q, e.g. prefill). +# instead of once per Q head. The power-of-2 kernel autotunes whether to pack +# along with its tile; the non-power-of-2 kernel keeps the tile-utilization +# heuristic, which packs for short seqlen_q (e.g. decode) and not for long +# seqlen_q (e.g. prefill). """ Triton SDPA Kernel for ExecuTorch CUDA Backend. @@ -51,6 +52,16 @@ def _kv_len_scenarios(args) -> list[int]: return sorted({min(lk, max(lq, length)) for length in (1024, lk // 8, lk)}) +def _visible_mask_scenarios(args) -> list[bool]: + """An all-visible mask for each of ``_kv_len_scenarios``, paired by position. + + The kernels skip KV blocks no row of a tile can see, so the all-False bool + mask autotuning would otherwise see times no attention at all and picks + whichever tile loads the mask fastest. All-visible times the full sweep + up to each KV length.""" + return [True] * len(_kv_len_scenarios(args)) + + def _is_power_of_2(n: int) -> bool: """Check if n is a power of 2.""" return n > 0 and (n & (n - 1)) == 0 @@ -723,9 +734,33 @@ def _sdpa_prefill_prune(configs, nargs, **kwargs): return kept -@autotune_inputs(KV_LEN_ptr=_kv_len_scenarios) +# Pack GQA is tuned with the tile: which wins depends on the query count, the +# group size, the tile and how many CTAs the device runs at once, which a +# utilization rule cannot weigh. +_SDPA_FWD_CONFIGS = [ + triton.Config( + {**config.kwargs, "PACK_GQA": pack}, + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + for config in _SDPA_PREFILL_CONFIGS + for pack in (False, True) +] + + +def _sdpa_fwd_prune(configs, nargs, **kwargs): + """``_sdpa_prefill_prune``, and without GQA only the unpacked tiles.""" + num_groups = kwargs.get("NUM_GROUPS") + if num_groups is None and nargs is not None: + num_groups = nargs.get("NUM_GROUPS") + if num_groups is not None and num_groups <= 1: + configs = [c for c in configs if not c.kwargs["PACK_GQA"]] + return _sdpa_prefill_prune(configs, nargs, **kwargs) + + +@autotune_inputs(KV_LEN_ptr=_kv_len_scenarios, Mask_ptr=_visible_mask_scenarios) @triton.autotune( - configs=_SDPA_PREFILL_CONFIGS, + configs=_SDPA_FWD_CONFIGS, key=[ "Lq", "Lk", @@ -734,9 +769,8 @@ def _sdpa_prefill_prune(configs, nargs, **kwargs): "IS_CAUSAL", "MASK_IS_CAUSAL", "NUM_GROUPS", - "PACK_GQA", ], - prune_configs_by={"early_config_prune": _sdpa_prefill_prune}, + prune_configs_by={"early_config_prune": _sdpa_fwd_prune}, ) @triton.jit def _sdpa_fwd_kernel( @@ -747,7 +781,7 @@ def _sdpa_fwd_kernel( Mask_ptr, KV_LEN_ptr, B, - H_grid, + H_q, Lq, Lk, stride_qb, @@ -780,6 +814,11 @@ def _sdpa_fwd_kernel( BLOCK_N: tl.constexpr, MASK_IS_CAUSAL: tl.constexpr = False, ): + # Packed, the grid walks the KV heads; unpacked, the query heads. + if PACK_GQA: + H_grid = H_q // NUM_GROUPS + else: + H_grid = H_q _sdpa_fwd_kernel_body( Q_ptr, K_ptr, @@ -1003,7 +1042,6 @@ def _launch_pow2_kernel( stride_mk: int, is_causal: bool, num_groups: int, - pack_gqa: bool, kv_len_ptr: Optional[torch.Tensor] = None, HAS_KV_LEN: bool = False, mask_is_causal: bool = False, @@ -1065,21 +1103,16 @@ def _launch_pow2_kernel( ) return - if pack_gqa: - H_grid = H_kv - Lq_packed = L_q * num_groups - else: - H_grid = H_q - Lq_packed = L_q - def grid(meta): - return (triton.cdiv(Lq_packed, meta["BLOCK_M"]), B * H_grid) - - # Single autotuned kernel: the config set spans BLOCK_M in {16..128} and - # `_sdpa_prefill_prune` keeps only non-spilling tiles for this HEAD_DIM, so - # the autotuner picks a high-occupancy tile (small BLOCK_M for large HEAD_DIM, - # larger BLOCK_M / more CTAs for small problems) — subsuming the old - # CTA-count m32/m64 selector. + if meta["PACK_GQA"]: + return (triton.cdiv(L_q * num_groups, meta["BLOCK_M"]), B * H_kv) + return (triton.cdiv(L_q, meta["BLOCK_M"]), B * H_q) + + # Single autotuned kernel: the config set spans BLOCK_M in {16..128}, each + # with and without pack GQA, and `_sdpa_prefill_prune` keeps only + # non-spilling tiles for this HEAD_DIM, so the autotuner picks a + # high-occupancy tile (small BLOCK_M for large HEAD_DIM, larger BLOCK_M / + # more CTAs for small problems). wrap_triton(_sdpa_fwd_kernel)[grid]( query, key, @@ -1088,7 +1121,7 @@ def grid(meta): Mask_ptr if HAS_MASK else 0, kv_len_ptr if HAS_KV_LEN else 0, B, - H_grid, + H_q, L_q, L_kv, stride_qb, @@ -1116,7 +1149,6 @@ def grid(meta): HAS_KV_LEN=HAS_KV_LEN, HEAD_DIM=D, NUM_GROUPS=num_groups, - PACK_GQA=pack_gqa, MASK_IS_CAUSAL=mask_is_causal, ) @@ -1350,17 +1382,6 @@ def sdpa( ) return out - # Decide whether to pack GQA based on tile utilization heuristic. - # Use the actual BLOCK_M that the launched kernel will use: - # - non-pow2 path always uses BLOCK_M=32 - # - pow2 path selects M32 or M64 based on CTA occupancy - if not _is_power_of_2(D): - block_m = 32 - else: - total_ctas_m64 = ((L_q * num_groups + 63) // 64) * (B * H_kv) - block_m = 32 if total_ctas_m64 < 4 * 84 else 64 - pack_gqa = _should_pack_gqa(L_q, num_groups, block_m) - if _is_power_of_2(D): _launch_pow2_kernel( query, @@ -1381,12 +1402,14 @@ def sdpa( stride_mk, kernel_is_causal, num_groups, - pack_gqa, kv_len_t, HAS_KV_LEN, mask_is_causal, ) else: + # The non-pow2 kernel always uses BLOCK_M=32 and is not tuned over pack + # GQA, so a tile-utilization rule decides it. + pack_gqa = _should_pack_gqa(L_q, num_groups, 32) _launch_non_pow2_kernel( query, key, @@ -1447,7 +1470,7 @@ def _sdpa_abstract( # then reduces partial results in a second kernel. -@autotune_inputs(KV_LEN_ptr=_kv_len_scenarios) +@autotune_inputs(KV_LEN_ptr=_kv_len_scenarios, Mask_ptr=_visible_mask_scenarios) @triton.autotune( configs=[ triton.Config({"BLOCK_N": 32}, num_warps=2, num_stages=1), @@ -1904,7 +1927,7 @@ def _sdpa_decode_splitk_abstract( # occupancy, then reduces per-query partial results in a second kernel. -@autotune_inputs(KV_LEN_ptr=_kv_len_scenarios) +@autotune_inputs(KV_LEN_ptr=_kv_len_scenarios, Mask_ptr=_visible_mask_scenarios) @triton.autotune( configs=[ triton.Config({"BLOCK_N": 32}, num_warps=2, num_stages=1),