feat: intracard cp for sm90 - #86
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an optimized Hopper (SM90) KDA prefill path featuring fused gate and L2-norm preprocessing, along with intra-card CP (chunk-parallel) scheduling. Key feedback includes optimizing cp_context.py to avoid a synchronous D2H copy by computing sequence mappings on the CPU, adding device validation checks in the C++ API to prevent illegal memory accesses, and using an if/else block in the fused L2-norm Triton kernel to eliminate redundant load instructions.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
65ae759 to
fbc6487
Compare
|
@Hyaloid Thanks a lot for this contribution! Could you also add some performance numbers comparing this SM90 intra-card CP path against FLA with intra-card CP enabled? |
|
Hi @cherhh, I've added the comparison with the FLA intracard-CP implementation and updated the performance table above. The results show that cuLA isn't always faster than FLA's intracard-CP implementation. One likely reason is the relatively high register pressure in cuLA. The KDA forward kernel keeps q/k/v/g/beta/A_log/dt_bias in registers, and safe_gate=True increases the register usage even further. On Hopper, this doesn't usually show up as traditional register spilling because WGMMA accumulators use dedicated registers, but the high register usage can still reduce occupancy (sometimes down to one CTA per SM), making memory latency harder to hide. If nvcc -Xptxas -v reports a non-zero stack frame, that confirms actual register spill. To consistently outperform FLA across all workloads, I think we'd probably need to rethink parts of cuLA's algorithm rather than just keep tuning the current implementation. That's a much bigger piece of work. And honestly, I don't think it's realistic to expect a single library to be the best choice for every workload and every shape. |
|
@Hyaloid Could you help resolve the confliction? |
|
Hi @icavan, I've resolved the conflicts. |
The CuTeDSL port no longer replaces the C++ backend: csrc/kda/sm90, its API/pybind/build wiring, cula/kda/hopper_fused_fwd.py and the fused bench/test all stay exactly as on main (including the inclusionAI#86 intracard-CP work built on them). The two backends split the public names: kda_prefill_hopper (+_opt/_auto) -> CUDA C++ fused (unchanged) kda_prefill_hopper_cutedsl -> CuTeDSL K1+K2 with intracard CP Docs re-merged to describe both; SM90 CuTeDSL tests/benches import the new name.
csrc/kda/sm90 and hopper_fused_fwd stay exactly as on main (incl. the inclusionAI#86 intra-card CP). Names: kda_prefill_hopper (+_opt/_auto) = C++ fused, kda_prefill_hopper_cutedsl = this CuTeDSL backend.
csrc/kda/sm90 and hopper_fused_fwd stay exactly as on main (incl. the inclusionAI#86 intra-card CP). Names: kda_prefill_hopper (+_opt/_auto) = C++ fused, kda_prefill_hopper_cutedsl = this CuTeDSL backend.
…106) * [KDA] port the SM90 fused prefill to a CuTeDSL two-kernel pipeline K1 prepares decayed q/k, beta and the WY inverse; K2 runs the inter-chunk recurrence. CHUNK=16, D=128, forward only. * test sm90 prefill accuracy against FLA triton * sm90: reject GVA inputs until native support lands * reorganize sm90 tests and benches * rename hopper_fused_fwd.py -> hopper_prefill.py * clean up the sm90 wrapper: import cycle, validation, fp32-only state * consolidate sm90 test files * [KDA] add intra-card CP for the sm90 prefill Split long sequences into segments and run the recurrence per segment: K1 once -> pre_scan -> merge -> segment-K2. Bit-identical to serial. * bench: iqr-mean timing + sm90 cp vs fla comparison * fill the whole SM array when splitting long sequences * remove unreachable asserts in internal launchers * handle varlen beta in-kernel via K1-emitted ws_beta * tidy cp comments and pad sentinels * rewrite cula/kda README for users * test cp against fla ground truth + determinism * simplify arch-context and varlen metadata helpers * fix ws_beta sizing in intracard cp * split the cp driver into plan + pipeline steps * simplify cp tests, add cp-on/off speedup bench * assert varlen args in sm100 bwd_wy_dqkg_fused * fix varlen jit compile storm (config-keyed cache) * add sm90 benchmark numbers and cp debug notes * drop from-future-annotations in sm90 kernel files (cutlass 4.4.2 Constexpr) * fix K2 init/final state fakes sharing one dynamic sym The shape binds wrong whenever only one of the two states is present. * launch K1/K2 through tvm-ffi instead of JitExecutor 27us -> 3us per launch; torch tensors pass straight through. * pool sm90 workspaces in a grow-only arena * cut cp driver host overhead (968 -> 288 us) * pre_scan: issue kd@M ahead of the S-chain * cost-model engage decision for cp auto mode * same future-annotations fix for the sm100 cp merge kernel * slim the prefill wrapper hot path * K1 reads beta straight from its packed [T,H] layout * duration-balanced segment planning for ragged batches * launch longest segments first * test ragged auto-plan cp + workspace arena reuse * trim redundant comments and stale docs * rework cp dispatch: fold the policy layer into the planner plan_prefill returns the plan the executor runs; a trivial plan means serial. SM100 keeps its own decision under sm100/policy.py. * unify copyright headers * slim cp dispatch comments * style: apply ruff formatting * planner polish: clearer names, worked example, drop the env-warning shim * reuse the cpu cu_seqlens copy in sm100 cp dispatch * [KDA] keep the CUDA C++ fused prefill as a separate backend csrc/kda/sm90 and hopper_fused_fwd stay exactly as on main (incl. the #86 intra-card CP). Names: kda_prefill_hopper (+_opt/_auto) = C++ fused, kda_prefill_hopper_cutedsl = this CuTeDSL backend. * rename the cutedsl backend export to flashkda_prefill * drop dead workspace helpers and decorative comments allocate_workspace/clear_workspace_cache lost their last caller when the arena landed; banners and restating comments go per the new comment rules. * Dispatch KDA prefill across SM90 backends FlashKDA remains first choice; unsupported calls fall through to the fully-fused CUDA path with per-backend rejection reasons. * Harden SM90 prefill launches across devices and streams Validate CP inputs before planning, key reusable buffers by CUDA stream, and bound dynamic caches at 32/64 entries. * Benchmark fixed-length SM90 batches as dense inputs * Update repository layout for SM90 dispatch and CUDA extensions * clean comments * remove some test from the test suite * remove test_kda_backend_dispatch * remove redundant doc * enhance doc and bench Co-authored-by: Cursor <cursoragent@cursor.com> * remove some test * flashkda: error on non-contiguous or fp32 inputs instead of converting Review feedback on #106. Also drops the use_cp plumbing here ahead of removing the alias everywhere. * drop the deprecated use_cp alias never shipped in a released tag; use_intracard_cp is the only spelling * bench_kda_sm90_cp: irregular varlen configs, reuse benchmark_cuda_mode_fn * use torch.cuda.device_of for the device guards * fix k2_kernel init bug * refactor pre_scan segment order handling * fix and clean up * remove obsolete tests --------- Co-authored-by: cheheng.ch <cheheng.ch@antgroup.com> Co-authored-by: Cursor <cursoragent@cursor.com>
3 participants
📌 Description
The serial bottleneck
kda_prefill_hopper (cuLA's SM90 KDA prefill) launches one CTA per (seq, head) and runs a strictly
sequential chunk recurrence inside each sequence: h_t = decay(g_t) · h_{t-1} + k_t^T @ (u_t − w_t·h_{t-1}).
Within one sequence, work cannot parallelize across chunks — only across the (raw_batch × H) grid.
This becomes a bottleneck when both:
H=8 occupies only 8 CTAs on a 132‑SM H100 (~6% occupancy). The per‑SM work is so small that most of the card is idle waiting on 8 serial chains.
dominates wall time while short seqs finish in microseconds and leave SMs idle.
Approach
Mirroring FLA's intra‑card CP design (and the SM100 cuLA path in cula/ops/cp/chunk_delta_h.py),
this PR splits long sequences into CP‑chunks on the same card and produces per‑CP‑chunk initial
states so the main C++ kernel can run all CP‑chunks in parallel.
🔍 Related Issues
Similar to this issue #20 , but for SM90.
🚀 Pull Request Checklist
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
python -m pytest tests/test_intracard_cp_sm90.py -v⚡ Performance
python benchmarks/bench_intracard_cp_sm90.py