feat(kda/sm90): FlashKDA prefill port + intracard-CP prefill#103
feat(kda/sm90): FlashKDA prefill port + intracard-CP prefill#103cherhh wants to merge 45 commits into
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 reorganizes the repository layout by moving KDA public APIs to cula/kda/ and grouping backend kernels by architecture under cula/ops/kda/sm90/ and cula/ops/kda/sm100/. It also introduces a new SM90 (Hopper) two-kernel FlashKDA prefill backend with intracard context-parallel (CP) support, and consolidates the CUDA extensions into a single cula.cudac module.
Feedback on these changes highlights three key areas for improvement: a critical race condition in the SM90 CP driver's scratch buffer cache (_get_scratch) due to a lack of stream-safety, an installation failure on Hopper-only systems because Blackwell build environment validation is called unconditionally in setup.py, and an excessively slow determinism test suite that runs 10,000 iterations instead of a more reasonable number like 100.
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.
| def _get_scratch(key_name: str, shape: tuple, dtype, device) -> torch.Tensor: | ||
| key = (key_name, shape, dtype, str(device)) | ||
| cached = _SCRATCH_CACHE.get(key) | ||
| if cached is None: | ||
| if len(_SCRATCH_CACHE) >= _SCRATCH_CACHE_MAXSIZE: | ||
| _SCRATCH_CACHE.pop(next(iter(_SCRATCH_CACHE))) | ||
| cached = torch.empty(shape, dtype=dtype, device=device) | ||
| _SCRATCH_CACHE[key] = cached | ||
| return cached |
There was a problem hiding this comment.
The _get_scratch function uses a cache key that only includes the device ID, ignoring the current CUDA stream. In concurrent multi-stream environments (such as multi-stream inference servers or pipeline parallel setups), different streams executing concurrently on the same device with the same shape will share the same physical scratch tensors (b_seg, m_seg, carries, seg_final). This leads to a critical race condition and silent data corruption as the asynchronous kernels overwrite each other's intermediate states.\n\nIncluding the current stream pointer in the cache key ensures stream safety.
def _get_scratch(key_name: str, shape: tuple, dtype, device) -> torch.Tensor:\n stream_ptr = int(torch.cuda.current_stream(device).cuda_stream)\n key = (key_name, shape, dtype, str(device), stream_ptr)\n cached = _SCRATCH_CACHE.get(key)\n if cached is None:\n if len(_SCRATCH_CACHE) >= _SCRATCH_CACHE_MAXSIZE:\n _SCRATCH_CACHE.pop(next(iter(_SCRATCH_CACHE)))\n cached = torch.empty(shape, dtype=dtype, device=device)\n _SCRATCH_CACHE[key] = cached\n return cached
There was a problem hiding this comment.
@cherhh Actually it is fine since the caller always try to call k1 & k2 in the same cuda stream.
Adding stream_ptr as a part of the key is a good suggestion. We shall take this suggestion.
| # Validate Blackwell build environment | ||
| assert_blackwell_build_env() |
There was a problem hiding this comment.
The Blackwell build environment validation (assert_blackwell_build_env()) is called unconditionally inside get_arch_flags(). If a user is installing cuLA on a Hopper-only system with an older CUDA toolchain (e.g., CUDA 12.2/12.3) and explicitly disables Blackwell builds via CULA_DISABLE_SM100=1 and CULA_DISABLE_SM103=1, the installation will still fail during setup.\n\nWrapping the validation in a conditional check ensures that the Blackwell toolchain requirements are only enforced when actually building for Blackwell architectures.
| # Validate Blackwell build environment | |
| assert_blackwell_build_env() | |
| # Validate Blackwell build environment if building for Blackwell\n if not DISABLE_SM100 or not DISABLE_SM103:\n assert_blackwell_build_env() |
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.
2 participants
📌 Description
This PR adds the SM90 (Hopper) KDA prefill backend to cuLA, ported from MoonshotAI's [FlashKDA](https://github.com/MoonshotAI/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.
What's included:
cula_kda_prefill, a FLA-compatible API. Each call launches two kernels; workspaces are reused across calls.torch.equal).use_intracard_cp="auto" | True | False, default off. In auto mode the planner turns CP on only when it predicts a real win, from tile counts and the device SM count. Its two tuned constants are chain-cost ratios (fitted on H100); both are env-overridable (cula/ops/kda/sm90/cp/plan.py).🔍 Related Issues
N/A
🚀 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
Full pytest output
test_kda_sm90_prefill_vs_fla.py— serial prefill accuracy against FLA's Tritonchunk_kda(rel err asserted at 5e-3, bf16 tolerance), dense + varlen + initial/final state, plus a workspace-arena reuse guard across shape changes.test_kda_sm90_intracard_cp.py— CP vs serial bit-exactness (torch.equal, incl. unaligned varlen and ragged auto-planned batches), CP determinism, and forced-CP accuracy against FLA through the public API.⚡ Performance
Intracard CP vs the serial path, same kernel binary, auto dispatch (
benchmarks/bench_kda_sm90_cp.py, H100 80GB SXM, median of 100 iters):The serial path is unaffected: CP is off by default and, in auto mode, the planner declines shapes where splitting cannot pay for its re-run cost (short sequences, machines already saturated by batch x heads).
Reviewer Notes
cula/ops/kda/sm90/cp/plan.py(plan_prefill→ trivial plan = serial path); the executor (cp/driver.py) runs a given plan and contains no decisions.B=1;num_kv_heads == num_qk_heads(GVA is a follow-up).