Skip to content

perf(kda): preserve SM90 workspace byte layout with bulk copies #115

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

Merged
fkuner merged 10 commits into from
Aug 12, 2026
Merged

perf(kda): preserve SM90 workspace byte layout with bulk copies #115

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
74af6d6
perf(kda): use raw bulk copies for SM90 workspace
Aug 2, 2026
33f8da1
fix(kda): align intracard pre-scan workspace loads
Aug 3, 2026
6987386
chore(bench): preserve prefill dispatch behavior
Aug 3, 2026
4e67e4f
fix(kda): support bulk-copy election across CuTeDSL versions
Aug 9, 2026
3ad91f2
docs(bench): record SM90 raw workspace A/B
Aug 9, 2026
a3f2e39
docs(bench): add SM90 NCU spot check
Aug 9, 2026
38093d2
style(kda): satisfy ruff import formatting
Aug 9, 2026
988693a
docs(bench): update Hopper benchmark data
Aug 9, 2026
5a9a200
Merge origin/main to resolve PR conflicts
Aug 11, 2026
8c2ca79
docs(bench): refresh H200 FlashKDA results
Aug 12, 2026
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
131 changes: 66 additions & 65 deletions BENCHMARK_H200.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# Benchmark Results — Hopper (SM90)

> Measured on 2026-08-09.
> Measured on 2026-08-12.

> **GPU:** NVIDIA H200 141GB, 132 SMs | **CUDA:** 12.9 | **PyTorch:** 2.9.1+cu129 | **Triton:** 3.5.1

Expand All @@ -19,42 +19,43 @@ iterations, and 100 measured iterations aggregated with the IQR mean.

| B | T | FLA Triton (ms) | cuLA FlashKDA (ms) | Speedup |
|---|---:|---:|---:|---:|
| 1 | 512 | 1.0013 | 0.1284 | **7.80x** |
| 1 | 1024 | 0.9747 | 0.1832 | **5.32x** |
| 1 | 4096 | 0.9994 | 0.6537 | **1.53x** |
| 1 | 8192 | 1.7774 | 1.2831 | **1.39x** |
| 1 | 16384 | 3.5038 | 2.5421 | **1.38x** |
| 2 | 512 | 0.9855 | 0.1386 | **7.11x** |
| 2 | 1024 | 1.0283 | 0.2447 | **4.20x** |
| 2 | 4096 | 1.7690 | 0.9010 | **1.96x** |
| 2 | 8192 | 3.4866 | 1.7810 | **1.96x** |
| 2 | 16384 | 6.9413 | 3.5363 | **1.96x** |
| 1 | 512 | 0.9405 | 0.1244 | **7.56x** |
| 1 | 1024 | 0.9372 | 0.1579 | **5.94x** |
| 1 | 4096 | 0.9344 | 0.5701 | **1.64x** |
| 1 | 8192 | 1.7701 | 1.1150 | **1.59x** |
| 1 | 16384 | 3.4846 | 2.2266 | **1.56x** |
| 2 | 512 | 0.9274 | 0.1262 | **7.35x** |
| 2 | 1024 | 0.9300 | 0.1971 | **4.72x** |
| 2 | 4096 | 1.7628 | 0.7223 | **2.44x** |
| 2 | 8192 | 3.4622 | 1.4306 | **2.42x** |
| 2 | 16384 | 6.8872 | 2.8817 | **2.39x** |

### Variable-Length

| Config | FLA Triton (ms) | cuLA FlashKDA (ms) | Speedup |
|---|---:|---:|---:|
| uniform 10 seqs, T=4096, 409–415 tokens | 1.0295 | 0.4748 | **2.17x** |
| random 10 seqs, T=4096, 24–1201 tokens | 1.0378 | 0.5544 | **1.87x** |
| skewed 10 seqs, T=4096, 227–2053 tokens | 1.0126 | 0.6700 | **1.51x** |
| uniform 20 seqs, T=4096, 204–220 tokens | 1.0510 | 0.4639 | **2.27x** |
| random 20 seqs, T=4096, 5–787 tokens | 1.0250 | 0.5513 | **1.86x** |
| skewed 20 seqs, T=4096, 107–2063 tokens | 1.0048 | 0.6767 | **1.48x** |
| uniform 10 seqs, T=8192, 819–821 tokens | 1.7923 | 0.9051 | **1.98x** |
| random 10 seqs, T=8192, 48–2401 tokens | 1.8277 | 1.0525 | **1.74x** |
| skewed 10 seqs, T=8192, 455–4097 tokens | 1.8457 | 1.2728 | **1.45x** |
| uniform 20 seqs, T=8192, 409–421 tokens | 1.8855 | 0.8628 | **2.19x** |
| random 20 seqs, T=8192, 9–1574 tokens | 1.8785 | 1.0189 | **1.84x** |
| skewed 20 seqs, T=8192, 215–4107 tokens | 1.8875 | 1.2913 | **1.46x** |
| uniform 10 seqs, T=16384, 1638–1642 tokens | 3.5091 | 1.7505 | **2.00x** |
| random 10 seqs, T=16384, 95–4802 tokens | 3.5175 | 2.0473 | **1.72x** |
| skewed 10 seqs, T=16384, 910–8194 tokens | 3.5350 | 2.4892 | **1.42x** |
| uniform 20 seqs, T=16384, 819–823 tokens | 3.5335 | 1.6558 | **2.13x** |
| random 20 seqs, T=16384, 19–3147 tokens | 3.5694 | 1.9728 | **1.81x** |
| skewed 20 seqs, T=16384, 431–8195 tokens | 3.5403 | 2.4949 | **1.42x** |

Across all 28 fixed-length and variable-length configs, FlashKDA averages
**2.39x** over FLA (minimum **1.38x**, maximum **7.80x**).
| uniform 10 seqs, T=4096, 409–415 tokens | 0.9819 | 0.3802 | **2.58x** |
| random 10 seqs, T=4096, 24–1201 tokens | 0.9749 | 0.4574 | **2.13x** |
| skewed 10 seqs, T=4096, 227–2053 tokens | 0.9730 | 0.5753 | **1.69x** |
| uniform 20 seqs, T=4096, 204–220 tokens | 1.0478 | 0.3768 | **2.78x** |
| random 20 seqs, T=4096, 5–787 tokens | 1.0242 | 0.4426 | **2.31x** |
| skewed 20 seqs, T=4096, 107–2063 tokens | 0.9984 | 0.5940 | **1.68x** |
| uniform 10 seqs, T=8192, 819–821 tokens | 1.7868 | 0.7329 | **2.44x** |
| random 10 seqs, T=8192, 48–2401 tokens | 1.8201 | 0.8595 | **2.12x** |
| skewed 10 seqs, T=8192, 455–4097 tokens | 1.8388 | 1.0881 | **1.69x** |
| uniform 20 seqs, T=8192, 409–421 tokens | 1.8781 | 0.6879 | **2.73x** |
| random 20 seqs, T=8192, 9–1574 tokens | 1.8714 | 0.8146 | **2.30x** |
| skewed 20 seqs, T=8192, 215–4107 tokens | 1.8815 | 1.1248 | **1.67x** |
| uniform 10 seqs, T=16384, 1638–1642 tokens | 3.5032 | 1.4460 | **2.42x** |
| random 10 seqs, T=16384, 95–4802 tokens | 3.5105 | 1.6967 | **2.07x** |
| skewed 10 seqs, T=16384, 910–8194 tokens | 3.5259 | 2.1618 | **1.63x** |
| uniform 20 seqs, T=16384, 819–823 tokens | 3.5272 | 1.3552 | **2.60x** |
| random 20 seqs, T=16384, 19–3147 tokens | 3.5600 | 1.5926 | **2.24x** |
| skewed 20 seqs, T=16384, 431–8195 tokens | 3.5296 | 2.2084 | **1.60x** |

Across all 28 fixed-length and variable-length configs, the arithmetic mean of
the per-row speedups is **2.72x** over FLA (geometric mean **2.43x**;
minimum **1.56x**, maximum **7.56x**).

To reproduce:

Expand All @@ -72,41 +73,41 @@ H200's SM array more effectively. Each row uses BF16, `D=128`,

| Config | H | CP engaged | CP off (ms) | CP auto (ms) | Speedup |
|---|---:|:---:|---:|---:|---:|
| T=1023 | 4 | no | 0.1322 | 0.1447 | 0.91x |
| T=1025 | 4 | no | 0.1364 | 0.1388 | 0.98x |
| T=4K | 4 | yes | 0.4374 | 0.2163 | **2.02x** |
| T=8K | 4 | yes | 0.8636 | 0.2820 | **3.06x** |
| T=16K | 4 | yes | 1.7099 | 0.4078 | **4.19x** |
| T=32K | 4 | yes | 3.3888 | 0.5973 | **5.67x** |
| T=64K | 4 | yes | 6.7518 | 0.9658 | **6.99x** |
| T=64K+1 | 4 | yes | 7.3033 | 1.0126 | **7.21x** |
| 2x16K | 4 | yes | 1.7631 | 0.5436 | **3.24x** |
| 32K+4K | 4 | yes | 3.4000 | 0.6682 | **5.09x** |
| 32K+1K | 4 | yes | 3.3923 | 0.6046 | **5.61x** |
| 32K+1023+1025 | 4 | yes | 3.6605 | 0.6821 | **5.37x** |
| 64K+1K | 4 | yes | 6.7713 | 1.0644 | **6.36x** |
| 64K+2x1K | 4 | yes | 6.7412 | 1.0803 | **6.24x** |
| 64K+5x1K | 4 | yes | 6.7648 | 1.1064 | **6.11x** |
| 64K+1+1023+1025 | 4 | yes | 7.3010 | 1.1200 | **6.52x** |
| T=1023 | 8 | no | 0.1365 | 0.1395 | 0.98x |
| T=1025 | 8 | no | 0.1376 | 0.1391 | 0.99x |
| T=4K | 8 | yes | 0.4557 | 0.2356 | **1.93x** |
| T=8K | 8 | yes | 0.8941 | 0.3177 | **2.81x** |
| T=16K | 8 | yes | 1.7600 | 0.5132 | **3.43x** |
| T=32K | 8 | yes | 3.4958 | 0.8774 | **3.98x** |
| T=64K | 8 | yes | 6.9656 | 1.6239 | **4.29x** |
| T=64K+1 | 8 | yes | 7.5212 | 1.7221 | **4.37x** |
| 2x16K | 8 | yes | 1.8764 | 0.8328 | **2.25x** |
| 32K+4K | 8 | yes | 3.5256 | 1.0458 | **3.37x** |
| 32K+1K | 8 | yes | 3.5080 | 0.9871 | **3.55x** |
| 32K+1023+1025 | 8 | yes | 3.7879 | 1.0731 | **3.53x** |
| 64K+1K | 8 | yes | 6.9812 | 1.8345 | **3.81x** |
| 64K+2x1K | 8 | yes | 6.9297 | 1.8368 | **3.77x** |
| 64K+5x1K | 8 | yes | 6.9738 | 1.9318 | **3.61x** |
| 64K+1+1023+1025 | 8 | yes | 7.4816 | 1.9172 | **3.90x** |
| T=1023 | 4 | no | 0.1308 | 0.1456 | 0.90x |
| T=1025 | 4 | no | 0.1312 | 0.1362 | 0.96x |
| T=4K | 4 | yes | 0.4070 | 0.2105 | **1.93x** |
| T=8K | 4 | yes | 0.8141 | 0.2659 | **3.06x** |
| T=16K | 4 | yes | 1.6149 | 0.3792 | **4.26x** |
| T=32K | 4 | yes | 3.2087 | 0.5455 | **5.88x** |
| T=64K | 4 | yes | 6.3907 | 0.8634 | **7.40x** |
| T=64K+1 | 4 | yes | 7.3199 | 0.9349 | **7.83x** |
| 2x16K | 4 | yes | 1.6487 | 0.4960 | **3.32x** |
| 32K+4K | 4 | yes | 3.2272 | 0.6156 | **5.24x** |
| 32K+1K | 4 | yes | 3.2135 | 0.5434 | **5.91x** |
| 32K+1023+1025 | 4 | yes | 3.6700 | 0.6341 | **5.79x** |
| 64K+1K | 4 | yes | 6.4081 | 0.9721 | **6.59x** |
| 64K+2x1K | 4 | yes | 6.4054 | 0.9865 | **6.49x** |
| 64K+5x1K | 4 | yes | 6.4184 | 1.0059 | **6.38x** |
| 64K+1+1023+1025 | 4 | yes | 7.3167 | 1.0531 | **6.95x** |
| T=1023 | 8 | no | 0.1326 | 0.1361 | 0.97x |
| T=1025 | 8 | no | 0.1335 | 0.1362 | 0.98x |
| T=4K | 8 | yes | 0.4246 | 0.2174 | **1.95x** |
| T=8K | 8 | yes | 0.8376 | 0.2854 | **2.94x** |
| T=16K | 8 | yes | 1.6496 | 0.4520 | **3.65x** |
| T=32K | 8 | yes | 3.2839 | 0.7789 | **4.22x** |
| T=64K | 8 | yes | 6.5424 | 1.4323 | **4.57x** |
| T=64K+1 | 8 | yes | 7.4745 | 1.5786 | **4.73x** |
| 2x16K | 8 | yes | 1.7236 | 0.7318 | **2.36x** |
| 32K+4K | 8 | yes | 3.3146 | 0.9499 | **3.49x** |
| 32K+1K | 8 | yes | 3.2958 | 0.8917 | **3.70x** |
| 32K+1023+1025 | 8 | yes | 3.7542 | 0.9980 | **3.76x** |
| 64K+1K | 8 | yes | 6.5675 | 1.6512 | **3.98x** |
| 64K+2x1K | 8 | yes | 6.5733 | 1.6640 | **3.95x** |
| 64K+5x1K | 8 | yes | 6.5782 | 1.7464 | **3.77x** |
| 64K+1+1023+1025 | 8 | yes | 7.4786 | 1.7998 | **4.16x** |

Intracard CP engages for 28 of the 32 configs. On the engaged subset it
delivers a **4.11x geometric-mean** speedup, with a range of **1.93–7.21x**.
delivers a **4.29x geometric-mean** speedup, with a range of **1.93–7.83x**.
The non-CHUNK-aligned `T=64K+1` cases retain essentially the same benefit as
`T=64K`; the four 1023/1025-token rows stay on the serial path and expose only
the auto-planner overhead.
Expand Down
4 changes: 2 additions & 2 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,8 +162,8 @@ See [BENCHMARK_H200.md](BENCHMARK_H200.md) for CuTe DSL FlashKDA results on an H
- **KDA Modular Forward (Blackwell):** **avg 1.33x** speedup on fixed-length, **avg 1.35x** on variable-length (18 configs, uniform/skewed/random).
- **Lightning Attention Prefill (Blackwell):** up to **2.08x** speedup (B=2).
- **Lightning Attention Varlen (Blackwell):** **avg 1.47x** speedup across 126 configs (uniform/skewed/random).
- **FlashKDA Prefill (Hopper):** **avg 2.39x** speedup over FLA across 28 fixed-length and variable-length configs, up to **7.80x**.
- **FlashKDA Intracard CP (Hopper):** **4.11x geo-mean** speedup over serial FlashKDA on 28 CP-engaged long-sequence configs, up to **7.21x**.
- **FlashKDA Prefill (Hopper):** **avg 2.72x** speedup over FLA across 28 fixed-length and variable-length configs, up to **7.56x**.
- **FlashKDA Intracard CP (Hopper):** **4.29x geo-mean** speedup over serial FlashKDA on 28 CP-engaged long-sequence configs, up to **7.83x**.

To reproduce the benchmark suites directly:

Expand Down
42 changes: 42 additions & 0 deletions cula/ops/kda/sm90/_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,55 @@

"""Shared low-level helpers for the SM90 FlashKDA kernels."""

import re
from importlib.metadata import PackageNotFoundError
from importlib.metadata import version as package_version

import cutlass
import cutlass.cute as cute
import torch
from cutlass import Int32
from cutlass._mlir.dialects import llvm as _llvm
from cutlass.cutlass_dsl import T as _T


def _parse_cutedsl_version(raw_version: str) -> tuple[int, int, int]:
"""Return the numeric CuTeDSL version from a release or dev version string."""
match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", raw_version)
if match is None:
raise RuntimeError(f"Unable to parse the installed CuTeDSL version: {raw_version!r}")
return tuple(int(component or 0) for component in match.groups())


def _installed_cutedsl_version() -> tuple[int, int, int]:
"""Read the CuTeDSL version at runtime without relying on a private symbol."""
raw_version = getattr(cutlass, "__version__", None)
if raw_version is None:
try:
raw_version = package_version("nvidia-cutlass-dsl")
except PackageNotFoundError as exc:
raise RuntimeError("nvidia-cutlass-dsl is required by the SM90 FlashKDA backend") from exc
return _parse_cutedsl_version(raw_version)


# CuTeDSL 4.6.0 added the missing elect_one inside cute.copy for async bulk
# atoms. Older releases need an explicit elect_one, while nesting one around
# cute.copy is incorrect in 4.6+. Keep this decision as a compile-time
# constant after detecting the installed runtime version, so the same source
# supports both API behaviours.
_CUTEDSL_VERSION = _installed_cutedsl_version()
_CUTE_COPY_AUTO_ELECTS_BULK = _CUTEDSL_VERSION >= (4, 6, 0)


def copy_async_bulk(atom, src, dst, **kwargs) -> None:
"""Issue a CuTeDSL async bulk copy across supported elect_one APIs."""
if cutlass.const_expr(_CUTE_COPY_AUTO_ELECTS_BULK):
cute.copy(atom, src, dst, **kwargs)
else:
with cute.arch.elect_one():
cute.copy(atom, src, dst, **kwargs)


def _stream_key(device: torch.device) -> tuple[str, int]:
return str(device), int(torch.cuda.current_stream(device).cuda_stream)

Expand Down
56 changes: 20 additions & 36 deletions cula/ops/kda/sm90/cp/pre_scan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@
from cutlass.cute.nvgpu.warpgroup import SmemLayoutAtomKind, make_smem_layout_atom
from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream

from cula.ops.kda.sm90._common import movm_t_b16
from cula.ops.kda.sm90._common import copy_async_bulk, movm_t_b16
from cula.ops.kda.sm90.k2 import (
CHUNK,
D,
Expand All @@ -48,16 +48,14 @@
def pre_scan_kernel(
tma_atom_v: cute.CopyAtom,
tma_tensor_v: cute.Tensor,
tma_atom_kd: cute.CopyAtom,
tma_tensor_kd: cute.Tensor,
tma_atom_kr: cute.CopyAtom,
tma_tensor_kr: cute.Tensor,
tma_atom_inv: cute.CopyAtom,
tma_tensor_inv: cute.Tensor,
tma_atom_gt: cute.CopyAtom,
tma_tensor_gt: cute.Tensor,
tma_atom_beta: cute.CopyAtom,
tma_tensor_beta: cute.Tensor,
ws_kd: cute.Tensor,
ws_kr: cute.Tensor,
H: cutlass.Constexpr[int],
total_tiles: cutlass.Int32,
T_total: cutlass.Int32,
Expand Down Expand Up @@ -118,22 +116,6 @@ def pre_scan_kernel(
cute.group_modes(sV, 0, 2),
cute.group_modes(gSrc_v, 0, 2),
)
gSrc_kd = cute.local_tile(tma_tensor_kd, (CHUNK, D), (None, None, None))
tKDs, tKDg = cpasync.tma_partition(
tma_atom_kd,
0,
cute.make_layout(1),
cute.group_modes(sKd, 0, 2),
cute.group_modes(gSrc_kd, 0, 2),
)
gSrc_kr = cute.local_tile(tma_tensor_kr, (CHUNK, D), (None, None, None))
tKRs, tKRg = cpasync.tma_partition(
tma_atom_kr,
0,
cute.make_layout(1),
cute.group_modes(sKr, 0, 2),
cute.group_modes(gSrc_kr, 0, 2),
)
gSrc_inv = cute.local_tile(tma_tensor_inv, (CHUNK, CHUNK), (None, None, None))
tIs, tIg = cpasync.tma_partition(
tma_atom_inv,
Expand All @@ -158,6 +140,19 @@ def pre_scan_kernel(
cute.group_modes(sBeta, 0, 2),
cute.group_modes(gSrc_beta, 0, 2),
)
raw_copy_atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), cutlass.BFloat16)
raw_stage_layout = cute.make_layout(
(CHUNK * D, STAGES),
stride=(1, CHUNK * D),
)
raw_gmem_layout = cute.make_layout(
(CHUNK * D, total_tiles * H),
stride=(1, CHUNK * D),
)
sKD_raw = cute.make_tensor(sKd.iterator, raw_stage_layout)
sKR_raw = cute.make_tensor(sKr.iterator, raw_stage_layout)
gKD_raw = cute.make_tensor(ws_kd.iterator, raw_gmem_layout)
gKR_raw = cute.make_tensor(ws_kr.iterator, raw_gmem_layout)

# sState=0, sM=I
if tidx < D:
Expand Down Expand Up @@ -273,8 +268,8 @@ def pre_scan_kernel(
cute.copy(tma_atom_v, tVg_seq[(None, t, 0, head_idx)], tVs_seq[(None, s_dyn_l)], tma_bar_ptr=bar_l)
else:
cute.copy(tma_atom_v, tVg[(None, tg_l, 0, head_idx)], tVs[(None, s_dyn_l)], tma_bar_ptr=bar_l)
cute.copy(tma_atom_kd, tKDg[(None, 0, 0, wt_l)], tKDs[(None, s_dyn_l)], tma_bar_ptr=bar_l)
cute.copy(tma_atom_kr, tKRg[(None, 0, 0, wt_l)], tKRs[(None, s_dyn_l)], tma_bar_ptr=bar_l)
copy_async_bulk(raw_copy_atom, gKD_raw[(None, wt_l)], sKD_raw[(None, s_dyn_l)], mbar_ptr=bar_l)
copy_async_bulk(raw_copy_atom, gKR_raw[(None, wt_l)], sKR_raw[(None, s_dyn_l)], mbar_ptr=bar_l)
cute.copy(tma_atom_inv, tIg[(None, 0, 0, wt_l)], tIs[(None, s_dyn_l)], tma_bar_ptr=bar_l)
cute.copy(tma_atom_gt, tGTg[(None, 0, 0, wt_l)], tGTs[(None, s_dyn_l)], tma_bar_ptr=bar_l)
cute.copy(tma_atom_beta, tBg[(None, 0, 0, wt_l)], tBs[(None, s_dyn_l)], tma_bar_ptr=bar_l)
Expand Down Expand Up @@ -487,13 +482,6 @@ def make_thd_atom(t, op):
)
return cpasync.make_tiled_tma_atom(op, view, kinter_smem, (CHUNK, D))

def make_ws_qkd_atom(t):
view = cute.make_tensor(
t.iterator,
cute.make_layout((CHUNK, D, total_tiles * H), stride=(D, 1, CHUNK * D)),
)
return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, kinter_smem, (CHUNK, D))

def make_ws_cc_atom(t):
view = cute.make_tensor(
t.iterator,
Expand All @@ -502,8 +490,6 @@ def make_ws_cc_atom(t):
return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, cc_smem, (CHUNK, CHUNK))

tma_atom_v, tma_tensor_v = make_thd_atom(v, cpasync.CopyBulkTensorTileG2SOp())
tma_atom_kd, tma_tensor_kd = make_ws_qkd_atom(ws_kd)
tma_atom_kr, tma_tensor_kr = make_ws_qkd_atom(ws_kr)
tma_atom_inv, tma_tensor_inv = make_ws_cc_atom(ws_inv)

gt_smem = cute.make_layout((D, 1), stride=(1, D))
Expand Down Expand Up @@ -551,16 +537,14 @@ def make_beta_atom(t):
pre_scan_kernel(
tma_atom_v,
tma_tensor_v,
tma_atom_kd,
tma_tensor_kd,
tma_atom_kr,
tma_tensor_kr,
tma_atom_inv,
tma_tensor_inv,
tma_atom_gt,
tma_tensor_gt,
tma_atom_beta,
tma_tensor_beta,
ws_kd,
ws_kr,
H,
total_tiles,
T_total,
Expand Down
Loading