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

BENCHMARK_SM90_RAW_WORKSPACE.md

@@ -0,0 +1,98 @@
+# SM90 Raw Workspace Transport A/B
+
+This benchmark isolates the FlashKDA workspace-transport change from
+intracard context parallelism.  Both runs use the same process configuration,
+inputs, and GPU; only the cuLA checkout changes.
+
+| Item | Value |
+|---|---|
+| Baseline | `d78f0075c69135fe191358fe8daf10e7383009d4` |
+| Optimized | `4e67e4f55d241af600235898b55cde09bd90c165` |
+| GPU | NVIDIA H200, SM90, 143,771 MiB device memory |
+| Driver / CUDA | 570.148.08 / CUDA 13.0.48 |
+| PyTorch | `2.8.0a0+34c6371d24.nv25.8` |
+| CuTeDSL | `nvidia-cutlass-dsl==4.6.1` |
+| FLA | `flash-linear-attention==0.5.0` |
+| Inputs | BF16 Q/K/V/gate, FP32 state, H=64, D=128, `safe_gate=True` |
+| CP mode | Serial K1+K2 (`use_intracard_cp=None`, planner resolves to a trivial plan) |
+| Timing | 25 warmup + 100 measured iterations, IQR-mean CUDA events |
+| Command | `python benchmarks/bench_kda_sm90_prefill.py --mode both` |
+
+The test machine is an **H200**, and these results belong to the H200 SM90
+benchmark set.
+
+## Fixed length
+
+| B | T | Baseline cuLA (ms) | Optimized cuLA (ms) | A/B |
+|---:|---:|---:|---:|---:|
+| 1 | 512 | 0.1241 | 0.1276 | 0.973x |
+| 1 | 1024 | 0.1831 | 0.1583 | **1.157x** |
+| 1 | 4096 | 0.6480 | 0.5704 | **1.136x** |
+| 1 | 8192 | 1.2716 | 1.1108 | **1.145x** |
+| 1 | 16384 | 2.5193 | 2.2082 | **1.141x** |
+| 2 | 512 | 0.1365 | 0.1294 | **1.055x** |
+| 2 | 1024 | 0.2431 | 0.1972 | **1.233x** |
+| 2 | 4096 | 0.8974 | 0.7174 | **1.251x** |
+| 2 | 8192 | 1.7648 | 1.4348 | **1.230x** |
+| 2 | 16384 | 3.5085 | 2.8760 | **1.220x** |
+
+## Variable length
+
+| Distribution | Sequences / total T | Baseline cuLA (ms) | Optimized cuLA (ms) | A/B |
+|---|---:|---:|---:|---:|
+| uniform | 10 / 4096 | 0.4746 | 0.3842 | **1.235x** |
+| random | 10 / 4096 | 0.5537 | 0.4550 | **1.217x** |
+| skewed | 10 / 4096 | 0.6640 | 0.5754 | **1.154x** |
+| uniform | 20 / 4096 | 0.4621 | 0.3752 | **1.232x** |
+| random | 20 / 4096 | 0.5487 | 0.4406 | **1.245x** |
+| skewed | 20 / 4096 | 0.6717 | 0.5904 | **1.138x** |
+| uniform | 10 / 8192 | 0.9028 | 0.7349 | **1.229x** |
+| random | 10 / 8192 | 1.0406 | 0.8582 | **1.213x** |
+| skewed | 10 / 8192 | 1.2627 | 1.0924 | **1.156x** |
+| uniform | 20 / 8192 | 0.8594 | 0.6886 | **1.248x** |
+| random | 20 / 8192 | 1.0135 | 0.8260 | **1.227x** |
+| skewed | 20 / 8192 | 1.2605 | 1.1314 | **1.114x** |
+| uniform | 10 / 16384 | 1.7411 | 1.4488 | **1.202x** |
+| random | 10 / 16384 | 2.0374 | 1.6995 | **1.199x** |
+| skewed | 10 / 16384 | 2.4653 | 2.1679 | **1.137x** |
+| uniform | 20 / 16384 | 1.6536 | 1.3662 | **1.210x** |
+| random | 20 / 16384 | 1.9532 | 1.5986 | **1.222x** |
+| skewed | 20 / 16384 | 2.4597 | 2.2102 | **1.113x** |
+
+## Summary and numerical accuracy
+
+- 27/28 configurations are faster; the only regression is B=1, T=512,
+  where the difference is within small-kernel launch noise.
+- Geometric-mean A/B speedup is **1.1778x**; summing all measured cuLA
+  latencies gives a **15.45%** reduction.
+- The optimized run reports `relative_rms_error` 0.004573–0.004924,
+  `rel_max` 0.007772–0.015000, and `mean_diff` 1.1e-5–1.3e-5 against
+  FLA.  The baseline run reports the same accuracy values to the printed
+  precision, so the raw byte transport does not change numerical behavior.
+- Flash-Flash-KDA reports isolated H100 workspace-transport reductions of
+  23%, 34%, and 37% for fixed, uneven packed, and uniform packed CHUNK=16
+  inputs.  This cuLA result is a complete CP-off prefill A/B on a different
+  GPU and workload mix; it is therefore a comparable direction-of-gain
+  check, not an exact reproduction of those percentages.
+
+## Nsight Compute spot check
+
+NCU 2025.3.0 (`--set full`, `--launch-count 2`, filtering the generated K1/K2
+names) profiled the first fixed shape, B=1, T=512, H=64.  The six `ctc__*`
+metrics were unavailable on this H200 setup, but the kernel duration and
+L1-TEX/XBAR counters below were collected in both reports.
+
+| Metric | Baseline | Optimized | Change |
+|---|---:|---:|---:|
+| K1 GPU duration | 41.152 us | 27.680 us | **-32.7%** |
+| K2 GPU duration | 76.416 us | 75.072 us | **-1.8%** |
+| K1 + K2 duration | 117.568 us | 102.752 us | **-12.6%** |
+| K1 global-op-TMA-store XBAR bytes | 52.43 MB | 27.26 MB | **-48.0%** |
+| K2 global-op-TMA-load XBAR bytes | 70.32 MB | 45.15 MB | **-35.8%** |
+
+These counters show the same signature as the latency A/B: the main saving is
+in K1's workspace stores, while K2's total duration changes little because
+the recurrence and retained TMA loads remain unchanged.  NCU does not expose
+the internal TensorMap segment count directly, so the evidence is the
+observable duration and transport counters rather than an inferred segment
+number.

README.md

@@ -129,11 +129,19 @@ See [BENCHMARK_GB200_CUDA_130.md](BENCHMARK_GB200_CUDA_130.md) tested with CUDA
 
 See [BENCHMARK_H200.md](BENCHMARK_H200.md) tested with CUDA 12.9 for detailed results.
 
+The raw workspace transport A/B is documented separately in
+[BENCHMARK_SM90_RAW_WORKSPACE.md](BENCHMARK_SM90_RAW_WORKSPACE.md).  On the
+H200 validation container it makes 27/28 serial (non-intracard-CP) configurations
+faster, with a 1.18x geometric-mean speedup and 15.45% latency reduction.
+
 **Highlights:**
 - **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).
 - **KDA Fused Forward (Hopper):** **avg 1.58x** speedup across fixed-length and variable-length sequences.
+- **SM90 raw workspace transport:** **1.18x** geometric-mean A/B speedup
+  (15.45% weighted latency reduction, 27/28 configurations) with intracard CP
+  disabled; see [BENCHMARK_SM90_RAW_WORKSPACE.md](BENCHMARK_SM90_RAW_WORKSPACE.md).
 
 To regenerate benchmarks:
 

cula/ops/kda/sm90/_common.py

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

cula/ops/kda/sm90/cp/pre_scan.py

@@ -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,
@@ -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,
@@ -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,
@@ -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:
@@ -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)
@@ -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,
@@ -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))
@@ -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,