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

[fkuner]

Summary

• preserve the physical K_INTER SMEM byte image for the SM90 KDA qd, kd, and kr workspaces
• replace layout-aware TensorMap workspace transfers with direct cp.async.bulk S2G/G2S copies
• keep inv and mqk on layout-aware Tensor TMA
• use the same raw G2S workspace ABI in the intracard-CP pre_scan consumer
• refresh the Hopper benchmark data in README.md and BENCHMARK_H200.md

Acknowledgment

The raw workspace transport idea comes from Shikhar Mishra's Flash-Flash-KDA, which identified the same fragmented TensorMap workspace boundary and preserved the physical shared-memory byte image with non-tensor cp.async.bulk copies. This PR adapts that idea to cuLA's SM90 CuTeDSL implementation and its CHUNK=16 workspace ABI. The source is also credited next to the K1 raw workspace stores in code.

Motivation

K1 and K2 use the same K_INTER shared-memory layout for qd, kd, and kr. The previous path stored these tensors through TensorMap/TMA into a logical linear workspace and reconstructed the layout in K2. Flash-Flash-KDA identified this layout traversal as a source of fragmented TMA work for the same workspace boundary; cuLA's NCU results show the corresponding reduction in observable store/load-class executions, XBAR traffic, and kernel duration. NCU does not expose the internal TensorMap segment count directly.

Because both kernels agree on the physical layout:

1. K1 copies each physical SMEM image contiguously to global workspace with raw bulk S2G.
2. K2 and intracard pre_scan copy the same image back with raw bulk G2S.
3. No unswizzle, repacking, or logical coordinate traversal is required for these tensors.

The implementation deliberately leaves the other workspace tensors on Tensor TMA where layout-aware transport remains appropriate.

_{Figure 1. Raw bulk copies preserve three 4 KiB physical K_INTER byte images for qd/kd/kr; inv/mqk remain on layout-aware Tensor TMA.}

Workspace transport

The optimization applies to qd, kd, and kr, whose physical K_INTER shared-memory layout is identical across their producers and consumers. Each cuLA tensor is a contiguous 4 KiB byte image (16 × 128 × sizeof(bf16)).

The baseline uses layout-aware Tensor TMA to convert between the K_INTER SMEM layout and a logical linear GMEM workspace. The optimized path preserves each tensor as an opaque byte image: K1 writes three contiguous payloads with raw bulk S2G copies, while K2 restores all three and intracard pre_scan restores kd/kr with raw bulk G2S copies.

inv and mqk are intentionally unchanged and continue to use layout-aware Tensor TMA. The diagram is conceptual and does not claim an exact internal TensorMap segment count; the observable NCU evidence is reported below.

CuTeDSL compatibility

CuTeDSL 4.6.0 added the missing automatic elect_one inside cute.copy for async bulk-copy atoms; nesting a manual election around those calls is invalid in 4.6+. The raw workspace calls use a runtime version check in cula/ops/kda/sm90/_common.py: CuTeDSL <4.6 gets an explicit elect_one, while 4.6+ relies on the generated election. The retained Tensor TMA stores keep their existing manual election. This preserves the declared nvidia-cutlass-dsl>=4.4.2 compatibility range; validation used CuTeDSL 4.6.1.

Performance

Benchmark provenance

• Hardware: one NVIDIA H200 141GB, 132 SMs
• CUDA 12.9, PyTorch 2.9.1+cu129, Triton 3.5.1, FLA v0.5.0
• Baseline: origin/main 5161546581cd6fecf2ad276dfa7a23a812208cf3; the KDA source baseline is d78f0075c69135fe191358fe8daf10e7383009d4 (the later main commit only refreshed benchmark/docs files)
• Optimized implementation: 4e67e4f55d241af600235898b55cde09bd90c165; current PR head: 988693a31ecaf0e331d92ea5e56f355e11f8e470
• BF16 inputs, H=64, D=128, safe_gate=True, no initial state, 25 warmup iterations, 100 measured iterations, IQR mean
• FLA was imported from the installed v0.5.0 package with PYTHONPATH unset to avoid accidentally using an older source checkout

Absolute optimized result vs FLA

The refreshed 28-shape fixed/variable-length sweep reports 2.72x arithmetic-mean speedup (2.43x geometric mean), with a range of 1.56–7.56x. These are the values now recorded in BENCHMARK_H200.md and summarized in README.md.

Workspace A/B against the latest main KDA baseline

The same 28 configurations were run with the raw workspace change disabled/enabled on the same H200 setup. The benchmark's default use_intracard_cp=None selects the serial K1+K2 path, so this isolates the workspace transport from intracard-CP work.

Metric	Baseline	Optimized
cuLA total latency over 28 rows	33.4064 ms	28.1467 ms
Total speedup	-	1.1869x
Total latency reduction	-	15.74%
Row-wise geometric-mean speedup	-	1.1824x

27/28 rows were faster; the single B=1/T=512 row was within small-kernel timing noise (0.990x). The fastest row improved by 1.2505x (20.04%).

Representative rows:

Workload	Baseline (ms)	Optimized (ms)	Speedup
B=1, T=8192	1.2785	1.1150	1.147x
B=2, T=16384	3.5229	2.8817	1.223x
Uniform 20-seq, T=4096	0.4643	0.3768	1.232x
Skewed 20-seq, T=4096	0.6720	0.5940	1.135x

The baseline and optimized runs produced the same FLA-relative accuracy at printed precision (rel_rmse=0.004568–0.004924, rel_max=0.007772–0.015000, mean_diff=0.000011–0.000013). The dedicated deterministic A/B checks also reported bitwise-identical output and final state.

Intracard CP

Intracard CP is reported separately and is not included in the workspace A/B headline. On the optimized implementation, use_intracard_cp="auto" versus False engages on 28/32 configurations and gives a 4.29x geometric-mean speedup, with a range of 1.93–7.83x. This measures CP planning, pre_scan, merge, and segmented K2 behavior in addition to the workspace transport.

Comparison with Flash-Flash-KDA

Flash-Flash-KDA's deep dive reports isolated workspace-transport latency reductions of 23%, 34%, and 37% on H100 for fixed, uneven packed, and uniform packed inputs with CHUNK=16. cuLA uses the same CHUNK=16 ABI. The cuLA H200 full non-intracard A/B above reduces total latency by 15.74% overall, with a best row reduction of 20.04%. This is comparable in direction and end-to-end scale, but is not claimed as an exact reproduction because the hardware and the measured scope differ.

Nsight Compute evidence

NCU 2025.2.1, --set full, B=1 / T=4096 / H=64:

Metric	Baseline	Optimized	Change
K1 duration	247.840 us	172.736 us	-30.3%
K2 duration	538.528 us	528.224 us	-1.9%
K1 + K2	786.368 us	700.960 us	-10.9%
K1 TMA-store-class executions	819,200	81,920	-90.0%
K1 TMA-store XBAR bytes	419.43 MB	218.10 MB	-48.0%
K2 TMA-load-class executions	1,114,112	376,832	-66.2%
K2 TMA-load XBAR bytes	562.56 MB	361.23 MB	-35.8%

K1 accounts for about 88% of the combined NCU duration reduction, matching the hypothesis that workspace stores were the dominant opportunity.

NCU groups both UTMASTG and UBLKCP.G under its TMA-store-class metric. The optimized K1 SASS adds three UBLKCP.G.S workspace operations per CTA, but the aggregate dynamic count falls from 50 to 5 operations per CTA: three raw workspace copies plus the two retained Tensor TMA stores. Internal TensorMap layout-walk segments are not exposed as a standalone NCU counter, so the evidence is the observable instruction/request count, XBAR traffic, and kernel duration rather than an inferred segment count.

Validation

• pre-commit run --all-files
• python -m py_compile cula/ops/kda/sm90/k1.py cula/ops/kda/sm90/k2.py cula/ops/kda/sm90/cp/pre_scan.py benchmarks/bench_kda_sm90_prefill.py
• python -m pytest tests/test_kda_sm90_prefill_vs_fla.py tests/test_kda_sm90_intracard_cp.py -q (28 passed, 15 skipped)
• full fixed + varlen benchmarks/bench_kda_sm90_prefill.py --mode both sweep on the H200/CUDA 12.9/FLA 0.5.0 environment
• benchmarks/bench_kda_sm90_cp.py on the same optimized environment
• baseline/optimized NCU full reports for B=1, T=4096, H=64 on H200 (K1/K2 kernel-name filtered)

[icavan]

Thanks for working on this optimization. Could you explicitly credit flash-flash-kda as the source of the raw workspace transport idea in the PR description and/or code comments?
Also, before merging, please provide an apples-to-apples performance comparison for the non-intracard path and confirm that the improvement is comparable to what flash-flash-kda reports. Please include the tested shapes, hardware, baseline commit, latency results, and numerical accuracy. This would help separate the benefit of this optimization from any intracard-CP effects.

[icavan]

The SM90 intracard path's pre_scan appears to still load ws_kd and ws_kr through TensorMap/TMA (make_ws_qkd_atom). Since this change stores the physical K_INTER byte images and updates K2 to restore them via raw bulk G2S copies, how is the pre_scan path handled?

Could you please make sure this workspace-format change is also applied correctly to the intracard path—for example, by migrating the corresponding pre_scan loads to raw bulk G2S or otherwise preserving format compatibility—and validate it with tests/test_kda_sm90_intracard_cp.py? As written, it looks possible for pre_scan to reinterpret the raw byte image as the previous logical workspace layout.

[fkuner]

The SM90 intracard path's pre_scan appears to still load ws_kd and ws_kr through TensorMap/TMA (make_ws_qkd_atom). Since this change stores the physical K_INTER byte images and updates K2 to restore them via raw bulk G2S copies, how is the pre_scan path handled?

Could you please make sure this workspace-format change is also applied correctly to the intracard path—for example, by migrating the corresponding pre_scan loads to raw bulk G2S or otherwise preserving format compatibility—and validate it with tests/test_kda_sm90_intracard_cp.py? As written, it looks possible for pre_scan to reinterpret the raw byte image as the previous logical workspace layout.

Thanks for pointing this out. I’ll take a closer look at the intracard pre_scan path and make sure the workspace format is handled consistently.

[fkuner]

Hi, I have one dependency question before updating the package metadata.

This PR's raw workspace path calls cute.copy with CopyBulkS2GOp / CopyBulkG2SOp without an explicit outer elect_one. The CUTLASS 4.6.0 changelog says that automatic elect_one handling for async bulk copies was added in 4.6.0. The package currently declares nvidia-cutlass-dsl>=4.4.2, while the H200 measurements for this PR were run with CuTeDSL 4.6.1.

For context, the previous Tensor TMA store path did manually elect one issuing lane:

if warp_idx == 0:
    with cute.arch.elect_one():
        cute.copy(tma_atom_ws_qd, ...)
        ...

The optimized code keeps that manual elect_one around the retained inv/mqk Tensor TMA stores, but intentionally leaves the raw cute.copy(CopyBulkS2GOp/G2SOp, ...) calls outside it. On CuTeDSL >=4.6.0, cute.copy performs the async-bulk lane election itself. In CuTeDSL 4.4/4.5, the same raw calls would need an explicit elect_one; the current if warp_idx == 0 only selects the warp and does not select one lane.

We therefore cannot simply wrap the raw copies in elect_one unconditionally: on >=4.6.0 that would nest the generated and manual elections, which the changelog warns can cause functionality issues. Supporting both ranges would require a version-specific branch or bypassing cute.copy with a direct async-bulk instruction path.

Would you prefer that we:

1. raise the minimum dependency to nvidia-cutlass-dsl>=4.6.0 (matching the current implementation), or
2. keep >=4.4.2 and add a version-specific compatibility path for older CuTeDSL releases?

I have not changed pyproject.toml yet and will follow the project’s preference.

[icavan]

Hi, I have one dependency question before updating the package metadata.

This PR's raw workspace path calls cute.copy with CopyBulkS2GOp / CopyBulkG2SOp without an explicit outer elect_one. The CUTLASS 4.6.0 changelog says that automatic elect_one handling for async bulk copies was added in 4.6.0. The package currently declares nvidia-cutlass-dsl>=4.4.2, while the H200 measurements for this PR were run with CuTeDSL 4.6.1.

For context, the previous Tensor TMA store path did manually elect one issuing lane:

if warp_idx == 0:
    with cute.arch.elect_one():
        cute.copy(tma_atom_ws_qd, ...)
        ...

The optimized code keeps that manual elect_one around the retained inv/mqk Tensor TMA stores, but intentionally leaves the raw cute.copy(CopyBulkS2GOp/G2SOp, ...) calls outside it. On CuTeDSL >=4.6.0, cute.copy performs the async-bulk lane election itself. In CuTeDSL 4.4/4.5, the same raw calls would need an explicit elect_one; the current if warp_idx == 0 only selects the warp and does not select one lane.

We therefore cannot simply wrap the raw copies in elect_one unconditionally: on >=4.6.0 that would nest the generated and manual elections, which the changelog warns can cause functionality issues. Supporting both ranges would require a version-specific branch or bypassing cute.copy with a direct async-bulk instruction path.

Would you prefer that we:

1. raise the minimum dependency to nvidia-cutlass-dsl>=4.6.0 (matching the current implementation), or
2. keep >=4.4.2 and add a version-specific compatibility path for older CuTeDSL releases?

I have not changed pyproject.toml yet and will follow the project’s preference.

@fkuner #119
We might need to apply dynamic api version check at the runtime to make sure your PR works in both versions.

Please also remember to update the perf in README.md & BENCH_xxx.md after your optimizations.

[icavan]

LGTM

[zhouaihui]

LGTM