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
• keep one implementation path, with no environment flag or permanent benchmark file

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 below 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, the workspace can instead be treated as an opaque byte image:

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.

Performance

Commit A/B summary

Execution path	Cases improved	Speedup range	Geometric mean	Mean latency reduction	Output / final-state accuracy
Non-intracard (use_intracard_cp=False)	6 / 6	1.134–1.255x	1.197x	16.4%	Bitwise identical
Intracard CP (s_split=8)	6 / 6	1.043–1.081x	1.058x	5.5%	Bitwise identical

Common setup:

• Hardware: one NVIDIA H200, SM90, 143,771 MiB reported device memory
• Baseline: d78f0075c69135fe191358fe8daf10e7383009d4
• Optimized kernel implementation: 33f8da18594af07b2e7ab52103b2c6f280e516e3 (current PR head 698738643134f80d6079360b441d269d58c4a279 only restores the benchmark's original dispatch behavior)
• BF16 inputs, FP32 initial/final state, D=128, CHUNK=16
• 100 measured iterations, interquartile mean
• Run order: baseline A, optimized, baseline B; baseline columns average the two bracketing runs
• Direct baseline/optimized comparison: output and final state are bitwise identical in every case (relative RMSE, relative max error, and maximum absolute error are all 0)

Non-intracard details

These runs explicitly pass use_intracard_cp=False, so they contain no intracard planning, pre_scan, or merge work. Warmup: 25 iterations. Uneven 8K denotes [1300, 547, 2048, 963, 271, 3063].

H	Workload	Baseline (ms)	Optimized (ms)	Reduction	Speedup
64	Fixed 1 × 8192	1.3532	1.1928	11.9%	1.134x
64	Uneven 8K, 6 sequences	1.0764	0.9287	13.7%	1.159x
64	Uniform 8 × 1024	0.8317	0.6629	20.3%	1.255x
96	Fixed 1 × 8192	1.5871	1.3495	15.0%	1.176x
96	Uneven 8K, 6 sequences	1.5232	1.2497	18.0%	1.219x
96	Uniform 8 × 1024	1.2384	0.9939	19.7%	1.246x

Intracard-CP details

These runs force the same manual s_split=8 plan on both commits, so every row executes K1, intracard pre_scan, merge, and segmented K2 rather than relying on an auto heuristic that might fall back to serial. Warmup: 10 iterations.

H	Sequence lengths	CP segments	Baseline (ms)	Optimized (ms)	Reduction	Speedup
4	[8192]	8	0.3716	0.3561	4.1%	1.043x
4	[8192, 2048]	16	0.3817	0.3638	4.7%	1.049x
4	[16384]	8	0.6628	0.6357	4.1%	1.043x
8	[8192]	8	0.4026	0.3768	6.4%	1.068x
8	[8192, 2048]	16	0.4228	0.3911	7.5%	1.081x
8	[16384]	8	0.7207	0.6759	6.2%	1.066x

The CP-on measurements show smaller end-to-end gains for these tested H=4/8 shapes. They are not a direct comparison with the CP-off H=64/96 cases above. A likely reason is that complete CP latency also includes pre_scan recurrence, merge, and segmented recurrence work that is unchanged by the raw workspace transport. Confirming the phase-level contribution would require a dedicated CP-on NCU breakdown.

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 also uses CHUNK=16; its complete CP-off prefill A/B on H200, where raw workspace transport is the only kernel implementation difference, shows 11.9–20.3% latency reductions. This is not claimed as an exact reproduction of Flash-Flash-KDA's isolated percentages.

It is nevertheless comparable in direction and end-to-end scale: the vLLM FlashKDA integration PR reports 1.145–1.240x for complete pure-prefill calls on GB300, while the six CP-off cuLA cases above span 1.134–1.255x. Hardware, kernel scheduling, and the other optimizations included by Flash-Flash-KDA differ, so the raw transport contribution should not be compared as if these were identical kernels.

Broader cuLA shape sweep

The existing benchmarks/bench_kda_sm90_prefill.py sweep also uses the non-intracard path. It calls flashkda_prefill directly; the default use_intracard_cp=None resolves to the serial K1+K2 path, so the benchmark file does not need to override dispatch explicitly. With its normal 25 warmup / 100 measured iterations, baseline d78f007 versus the original optimized commit 74af6d6 produced:

• all 28 fixed and varlen configurations improved
• geometric-mean speedup: 1.175x
• mean latency reduction: 14.8%
• varlen geometric-mean speedup: 1.192x
• output relative RMSE versus FLA remained approximately 0.00457–0.00492

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
• pytest -q -m "not kda_slow" tests/test_kda_sm90_intracard_cp.py
• reviewer-requested CP-off A/B with deterministic output/final-state comparison
• forced-CP (s_split=8) baseline/optimized A/B with deterministic output/final-state comparison
• full fixed + varlen benchmarks/bench_kda_sm90_prefill.py sweep
• baseline/optimized NCU full reports for B=1, T=4096, H=64

[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.