[executorch][cuda] Autotune SDPA on a visible mask and tune pack GQA with the tile

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()

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: