feat(kda): SM90 FlashKDA prefill, intracard CP, and backend dispatch#106
feat(kda): SM90 FlashKDA prefill, intracard CP, and backend dispatch#106cherhh wants to merge 59 commits into
Conversation
cherhh
changed the title
![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 a new SM90 (Hopper) FlashKDA prefill backend written in CuTeDSL, featuring a two-kernel pipeline and support for intra-card context parallel (CP) execution. It also restructures the repository layout to support multiple backends with a generic registry and dispatch mechanism. The review feedback highlights a few critical issues: a bug in the fully-fused backend where explicitly disabling context parallel incorrectly rejects the call, a potential error in flashkda.py when q.device.index is None, and potential CPU device errors in utils.py when calling torch.cuda.current_stream directly.
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.
K1 prepares decayed q/k, beta and the WY inverse; K2 runs the inter-chunk recurrence. CHUNK=16, D=128, forward only.
Split long sequences into segments and run the recurrence per segment: K1 once -> pre_scan -> merge -> segment-K2. Bit-identical to serial.
The shape binds wrong whenever only one of the two states is present.
27us -> 3us per launch; torch tensors pass straight through.
plan_prefill returns the plan the executor runs; a trivial plan means serial. SM100 keeps its own decision under sm100/policy.py.
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.
allocate_workspace/clear_workspace_cache lost their last caller when the arena landed; banners and restating comments go per the new comment rules.
FlashKDA remains first choice; unsupported calls fall through to the fully-fused CUDA path with per-backend rejection reasons.
Validate CP inputs before planning, key reusable buffers by CUDA stream, and bound dynamic caches at 32/64 entries.
Co-authored-by: Cursor <cursoragent@cursor.com>
| ("32K+1K", [32768, 1024]), | ||
| ("64K+1K", [65536, 1024]), | ||
| ("64K+2x1K", [65536, 1024, 1024]), | ||
| ("64K+5x1K", [65536] + [1024] * 5), |
There was a problem hiding this comment.
we shall add more irregular varlen settings like 65537, 1023, 1025
There was a problem hiding this comment.
Added those irregular lengths (1023, 1025, 64K+1, and a couple ragged mixes). Numbers are in the #106 (comment).
| ] | ||
|
|
||
|
|
||
| def time_kernel(fn, warmup=None, n_iters=None): |
There was a problem hiding this comment.
try reuse benchmark_cuda_mode_fn
There was a problem hiding this comment.
Done — time_kernel now delegates to benchmark_cuda_mode_fn.
| ``CULA_BACKEND_*`` variable makes dispatch try the remaining implementation. | ||
| """ | ||
| device_ctx = torch.cuda.device(q.device) if q.device.type == "cuda" else contextlib.nullcontext() | ||
| with device_ctx: |
There was a problem hiding this comment.
with torch.cuda.device_of(q): is more concise.
Why do we need to specify the device context here? Is it a common pattern in kernel library? When do we have to add ctx guard explicitly?
There was a problem hiding this comment.
Switched to device_of. Needed so launches/allocs run on the same GPU as the input tensors (current device can differ in multi-GPU). FLA's input_guard does the same.
|
|
||
| if cp_context is not None: | ||
| if use_intracard_cp is True: | ||
| if use_intracard_cp is CPMode.FORCE: |
There was a problem hiding this comment.
Why is the intracard cp not compatible with Kimi CP now?
There was a problem hiding this comment.
Kimi CP gives each rank a sequence shard; intracard CP currently assumes a full sequence. Nesting should work in theory, but isn't implemented/tested yet, so we reject for now.
| out: torch.Tensor | None = None, | ||
| final_state: torch.Tensor | None = None, | ||
| ): | ||
| q, k, v, g, beta = _contig(q), _contig(k), _contig(v), _contig(g), _contig(beta) |
There was a problem hiding this comment.
I have a concern about hiding implicitly conversions here. I'd rather return an error to let the caller know they should pass a contiguous tensor or bf16 tensor here.
| @classmethod | ||
| def parse(cls, use_intracard_cp, use_cp=None) -> CPMode | None: | ||
| """Public use_intracard_cp value ("auto"/True/False; use_cp is a | ||
| deprecated alias) -> mode. None stays None (backend default).""" |
There was a problem hiding this comment.
why not drop use_cp if it is a deprecated option
There was a problem hiding this comment.
Removed now. Only use_intracard_cp remains.
| cu_seqlens, | ||
| chunk_indices, | ||
| ) | ||
| with torch.cuda.device(q.device): |
There was a problem hiding this comment.
same question: when do we have to add the device hint?
Review feedback on inclusionAI#106. Also drops the use_cp plumbing here ahead of removing the alias everywhere.
never shipped in a released tag; use_intracard_cp is the only spelling
2 participants
📌 Description
This PR adds the SM90 (Hopper) KDA prefill backend to cuLA, ported from MoonshotAI's FlashKDA CUTLASS C++ kernels to CuTeDSL, and extends it with an intracard context-parallel (CP) mode that is not present upstream.
Huge thanks to the MoonshotAI FlashKDA authors — the kernel design and its numerics are theirs; this port follows the original kernels closely.
Also on top of #105:
kda_prefill(FlashKDA + fused CUTLASS C++, default FlashKDA). Falls back to fused C++ when FlashKDA can’t run (HV > H, fp32g, no CuTeDSL,CULA_BACKEND_*=0, …). Direct entry points (flashkda_prefill,kda_prefill_hopper*, …) stay.use_intracard_cp="auto" | True | False.Related ideas also show up in FlashQLA; we keep an exact
Mand skip their gate-decay warmup.🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to cuLA! Before we review your pull request, please make sure the following items are complete.
✅ 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
⚡ Performance
H100 80GB HBM3. run:
bench_kda_sm90_prefill.py --mode both,bench_kda_sm90_cp.py.FlashKDA (serial; default
use_intracard_cp=None) vs FLA (H=64,D=128, bf16,safe_gate=True, no init state, warmup 25 / iters 100, iqr mean):Same FlashKDA binary, CP
"auto"vs off (warmup 10 / iters 100;cpmarks whether auto actually engaged CP — the irregular-length configs 1023/1025/65537 are included per review feedback):Non-CHUNK-aligned lengths keep the full win (65537: 7.30x/4.33x vs 7.39x/4.39x for 64K). The four non-engaged rows (T=1023/1025) run serial either way; the ~0.9x there is the auto planner's fixed cost on a ~0.14 ms call.