feat(kda/sm90): FlashKDA prefill port + intracard-CP prefill#103
Open
cherhh wants to merge 69 commits into
Open
feat(kda/sm90): FlashKDA prefill port + intracard-CP prefill#103cherhh wants to merge 69 commits into
cherhh wants to merge 69 commits into
Conversation
added 30 commits
June 26, 2026 14:28
…a/ layout
- Move SM100 (Blackwell) modular-chunk backends, decode, and the unwired
fully-fused WIP from flat cula/ops/*.py into cula/ops/kda/{sm100,decode,experimental}/.
- Move the non-KDA lightning/linear prototypes under cula/ops/.
- Add a central CP dispatch policy at cula/ops/kda/policy.py.
- Make cula / cula.ops / cula.kda imports lazy (PEP 562) so `import cula` no
longer eagerly pulls the CuTeDSL/CUDA-heavy modules.
- Repoint all in-repo imports, benchmarks, tests, and docs.
Pure reorganization, no kernel behavior change. The SM90 (Hopper) prefill stays
the existing C++ kernel under csrc/kda/sm90.
- Drop kda_prefill_blackwell from the cula.kda public exports; the fully-fused Blackwell prefill (cula/ops/kda/experimental/sm100_fused/) is unwired WIP. - get_kda_fused_fwd now raises NotImplementedError on SM100/SM103 instead of returning that experimental kernel. - Production Blackwell prefill stays the modular chunk_kda path.
- intracard_fwd_h now raises NotSplittableError when the shape cannot be meaningfully split, instead of silently falling back. - Drop the allow_fallback / skip_precheck flags, the two duplicated _no_cp fallback blocks, and the redundant pre-split heuristic recheck that the dispatch policy already performed. - chunk_gated_delta_rule_fwd_h now owns the fallback: re-raise for forced CP, fall through to the serial body for auto. - NotSplittableError subclasses ValueError for backward compatibility. Behavior-preserving: force -> raise and auto -> serial fallback are unchanged.
…d stale REPO_LAYOUT sections
…st_intracard_cp.py
- cula/__init__.py: drop trailing blank line - cula/ops/kda/__init__.py: dedent module docstring, add final newline - cula/ops/kda/sm100/delta_h.py: drop extra blank lines
- Replace the C++ SM90 KDA kernels (csrc/kda/sm90) with a CuTeDSL K1+K2 prefill
under cula/ops/kda/sm90/{k1,k2,fwd}.py.
- Handle dense and packed-varlen input, including non-CHUNK-aligned lengths.
- Consolidate the CUDA build into a single cula.cudac extension (drop the per-arch
split and the cudac.py shim).
- Select the CuTeDSL arch from compute capability; Hopper verified, Blackwell wired
but unverified.
- Forward-only: compare o and ht vs FLA chunk_kda (assert_close 0.005). - Covers tail chunks, B>1, init_state, and packed varlen. - Handles the SM90 VK-transposed state layout (initial_state/ht transposed).
…port - Stop faking GVA by repeat_interleave-ing q/k up to num_kv_heads. - Reject num_kv_heads != num_qk_heads with NotImplementedError. - Native GVA (kernel-side Q/K vs V indexing, like SM100) is a follow-up.
- Its cases no longer apply: most use use_gate_in_kernel=False, the rest are GVA -- both unsupported on SM90. - The supported path is covered by the SM90 prefill tests.
- Drop benchmarks/bench_kda_fused_fwd.py: misnamed (now a K1+K2 pipeline), routed through the get_kda_fused_fwd dispatcher, and mixed in the SM100 fully-fused path that raises NotImplementedError here. - Add benchmarks/bench_kda_sm90_prefill.py calling cula_kda_prefill directly: Hopper-only, MHA, accuracy + perf vs FLA Triton.
- No remaining callers after the SM90 tests and benchmark stopped using it. - Drop it from cula/utils.py and the two README references.
- flashkda -> kda_sm90; la -> lightning; SM100 KDA gets explicit arch plus recurrence/output/chunk and vs_fla/vs_naive suffixes. - Update the README test commands and stale cross-references.
- The SM90 prefill is a K1+K2 pipeline, not a fused kernel; "fused" was a C++-era leftover. - Repoint the lazy import and drop the now-stale "fused"/layering notes.
- generate_benchmark_hopper_md.py imported the deleted bench_kda_fused_fwd; point it at bench_kda_sm90_prefill and drop the --hv/GVA plumbing (SM90 is MHA-only). - Fix the stale reproduce command in BENCHMARK_H200.md.
- Move the shared kernel helpers (movm_t_b16, add_f16x2_u32, _wrap_input) out of fwd.py into a leaf module cula/ops/kda/sm90/_common.py. - k1/k2 import them from _common, so they no longer import the fwd orchestrator; the dependency is now one-way (fwd -> k1/k2).
- cu_seqlens: int32-only (was int32/int64) - initial_state: fp32-only (was bf16/fp32) - update the with-state smoke test to pass an fp32 initial_state
…mbing - final_state: fp32-only validation (was bf16/fp32), matching initial_state - drop the host-side bf16<->fp32 state adaptation in fwd dispatch + k2 launcher - K2's gmem state ABI was already fp32-only; bf16 was a dead boundary path
- _PrefillProblem.state_fp32 was computed and propagated but never consumed - with fp32-only state the initial/final dtype-match check is also dead
- final_state is written in-place via reshape(-1); a non-contiguous buffer would silently get a copy and drop the kernel's result - validate both initial_state and final_state are contiguous fp32
…la with SM100 pattern
- Add a single-card CP path for the SM90 prefill (cula/ops/kda/sm90/cp/): split long sequences, run K1 once, pre-scan + merge chunk states, re-run K2 per segment. - Recovers sequence-dim parallelism (~3-7x on long sequences that under-fill the SMs). - Route via kda_prefill_hopper(use_intracard_cp="auto"|True|False); the auto-router (policy.py) enables CP only when profitable, else falls back to the serial path. - Handle any sequence length (pad-before-CP), dense or varlen. Inference-only.
… outliers) - benchmark_cuda_mode_fn gains an aggregate param (default "mean", others unchanged) - SM90 prefill bench reduces per-iter times with iqr_mean so a stray slow iteration can't poison the mean; serial is reproducibly ~1.8x vs FLA
- bench_kda_sm90_cp.py targets CP's regime: few heads, small batch, long seq - matched paths: cuLA serial vs FLA serial, and cuLA CP vs FLA intracard CP (FLA_INTRACARD_CP + inference_mode); reports cuLA auto decision + rrmse
- auto-split each long sequence into ~SM_count/H segments (one full SM wave) via a 32-tile floor + one-wave target, instead of staying coarse (T=16384 was capped at 8 segments / 32 CTAs on a 132-SM part) - decouple the CP engage threshold (ENGAGE_MIN_TILES=640) from the split floor so short sequences (where serial is already fast) keep falling back to serial - H=4/8 long-seq: cuLA-CP now beats FLA's intracard CP at T=16384 (~0.9x -> 1.15x) and beats FLA serial too; T>=32768 unchanged (1.8-2.2x); high-H falls back
Delete ~32 assert statements that duplicate upstream validation: - fwd.py: 6 asserts (including 1 bug where assert would incorrectly fire in the needs_padding path where varlen_meta is None) - k2.py: 10 asserts (including 1 bug where assert v_tile_actual_lens is None would incorrectly fire because caller passes dummy tensor, not None) - k1.py: 6 asserts - cp/driver.py: 3 asserts - cp/pre_scan.py: 4 asserts - cp/merge.py: 1 assert - bwd_wy_dqkg.py: 1 assert - wrapper.py: 1 redundant assert (duplicate of the raise below it)
K1 now emits raw beta into a 7th compact workspace (ws_beta, wt_l tile layout, tail rows -80); K2 reads it like the other K1 outputs instead of a host-padded buffer. This removes the per-call host-side beta gather (beta_pad alloc + index_select + index_fill) and the second transpose for non-aligned varlen; q/k/g/v/output were already handled in-kernel via tile_starts. The CP driver adopts the same path (ws_beta is byte-identical to the transposed beta_flat for the CHUNK-aligned data its impl sees). Non-aligned varlen ~10-12% faster at H=4, T<=8192 (where the fixed host overhead is a real fraction of the sub-ms K1+K2); aligned/dense neutral. Prefill (varlen 3/3, dense 7/7) and SM90 intracard CP (17 pass/1 skip) tests pass with unchanged tolerances. k2.py unchanged. 🤖 Generated with [Claude Code](https://claude.com/claude-code)
added 25 commits
July 4, 2026 09:53
Under nvidia-cutlass-dsl 4.4.2 (the pyproject floor), PEP 563 turns the
cutlass.Constexpr[...] annotations into strings the DSL decorators cannot
classify, so Constexpr params (is_varlen, H, scale, ...) degrade to dynamic
values and compilation fails with:
DSLRuntimeError: const_expr(?) received a dynamic expression
Newer DSL releases resolve string annotations; 4.4.2 does not. Drop the
future import from every file defining @cute.kernel/@cute.jit/@dsl_user_op
(k1, k2, _common, cp/merge, cp/pre_scan) so the declared minimum version
actually works. Runtime annotation evaluation is fine on Python >= 3.10.
_compile_k2 faked initial_state and final_state with the SAME cute.sym_int(). Their runtime sizes genuinely differ whenever exactly one of them is present (the absent one is a 1-element dummy, the real one is N*H*D*D), so one sym was being bound to two different values. The JitExecutor launch path never checked the shared-sym constraint, so this stayed latent. The tvm-ffi launch path (next commit) binds each sym to a single value per call, and the inconsistency produced garbage output for every has_initial_state != has_final_state variant. Give each state tensor its own sym; sizes are then independently bound.
Compile run_k1/run_k2 with --enable-tvm-ffi and pass torch tensors (and plain int scalars) straight to the compiled function, the same pattern the decode and SM100 kernels already use. This removes the per-call JitExecutor cost: 17 from_dlpack tensor wrappings plus two rounds of execution-arg marshaling per prefill call (~130 us of the ~320 us host floor; a 2-tensor microbenchmark measures the launch path at 27 us -> 2.9 us per call). Dynamic shapes keep working: sym_int extents and runtime Int32 dims are bound per call (verified bit-exact against the JitExecutor path across changing shapes), so one compiled kernel still serves all batch shapes and there is no per-shape recompile storm. Together with the workspace arena (next commit), the raw-op host cost drops 321 -> 71 us/call on H100, matching the FlashKDA C++ extension (69 us in the same process), and the e2e gap at small prefill sizes disappears (bench T=1024: 0.56x -> 1.01x vs FlashKDA; H=16 T=1024: 0.40x -> 1.02x). T>=4096 configs are unchanged (0.98-1.08x). CUDA graph capture of K1+K2 keeps working.
Each prefill call allocated 9 tensors (7 K1/K2 workspaces, the beta_flat staging buffer) and built a fresh cu_seqlens_tiles arange on the fixed-length path (one extra tiny GPU kernel per call). Replace them with: - a grow-only uint8 arena per (device, stream), carved into cached 256B-aligned views keyed by the workspace sizes; reuse across calls is stream-ordered (the next K1 cannot overwrite before this K2 finished), and different streams get independent arenas. clear_workspace_cache() drops the pinned memory. - a cached cu_seqlens_tiles per (B, tiles_per_seq, device) in launch_k2. The CP driver shares the same arena (beta_flat now comes from it too).
The intracard-CP pipeline was CPU-bound at moderate sizes (H=16 T=16k: e2e 985 us vs 846 us of GPU work). Host profile per call: - launch_merge built firsts/nsegs with torch.tensor(list, device=cuda) every call — two synchronous pageable H2D copies, ~375 us. They are plan-derived constants; cache them per (per_seq, device), same pattern as the driver's _get_plan_tensor. Also cache the has_init=0 dummy state (was a fresh new_zeros + memset kernel per call) and reuse k2's cached CUstream wrapper. - pre_scan/merge still launched through JitExecutor with 17 from_dlpack wrappings per call; compile both with --enable-tvm-ffi and pass torch tensors directly, like K1/K2. pre_scan's dense-path dummy tile tensor (fresh torch.zeros per call) now reuses k2's cached dummy. Driver host cost drops 968 -> 288 us/call and is now fully hidden behind GPU work (e2e - GPU gap 139 -> ~11 us). CP e2e on H100: H=16 T=16384 985 -> 853 us (speedup vs serial 1.86x -> 2.14x), H=16 T=65536 3271 -> 3143 us (2.21x -> 2.30x), H=32 T=32768 3309 -> 3127 us (1.25x -> 1.32x). Kernel times are unchanged; CP output stays bit-exact vs the serial path (dense, aligned/unaligned varlen, with/without initial state) and the unaligned-varlen determinism loop is clean (0/20).
…stages The M-chain reused the S-chain's register fragments, so its leading kd@M GEMMs carried a false dependency on the whole S-chain tail (sigmoid -> movmatrix -> INV -> state update) and could not be scheduled early. Give kd@M its own accumulator fragment (+16 regs/thread) and issue it right after the S-chain's kd@state, letting the two independent chains' MMAs overlap the dependent elementwise stages. pre_scan on H100: 1291 -> 1252 us (H=16 T=65536, 512-tile chains) and 1292 -> 1255 us (H=32 T=32768), ~-3%; neutral at short chains (T=16384). Output stays bit-exact vs the serial path across dense/varlen/init configs; unaligned-varlen determinism loop clean (0/20). Notes from the same investigation, for the record: an output-correction restructuring (compute o with S0=0, correct with the merged carry) is a net loss for KDA — delta-rule tile transitions are Diag - rank16, so the carry needs its own per-tile homogeneous chain and linearity-splitting duplicates every S-consuming GEMM (7 -> 8 big MMAs/tile). Skipping the M-chain for first(no-init)/last segments saves work but not wall time: pre_scan is a single wave, so wall time is the slowest CTA's chain.
The auto router gated CP behind two fixed thresholds: longest sequence >= ENGAGE_MIN_TILES (640) and >= MIN_BENEFICIAL_SEG (5) segments/seq. Both misfire: the tile floor rejects every T < 10240 even where CP wins comfortably, and the segment floor rejects H=32 (s_split=4) outright — auto mode never used CP at 32 heads despite a measured 1.32x at T=32k. Replace them with a calibrated cost model (plan.estimate_cp_speedup): each chain kernel is modeled as per_tile*tiles + fixed us (fitted on H100: serial K2 1.48/tile+84, interleaved pre_scan 2.39/tile+27, segment K2 1.69/tile+19, merge 8+3.3/seg) with a makespan lower bound max(critical chain, load/SM) for waves and imbalance. K1 and driver host time are identical on both sides and cancel. Auto engages when the predicted serial/CP ratio clears CP_ENGAGE_MARGIN (1.10). Constants are env-overridable for other SM90 parts; the decision depends mostly on ratios, not absolute calibration. Forced mode (True) still runs any splittable shape; allow_fallback=True driver callers get the same check. Newly enabled shapes, all measured faster and bit-exact vs serial: H=16 T=8192/4096/2048 -> 1.98x/1.76x/1.15x e2e; H=32 T=32768/8192/2048 -> 1.32x/1.26x/1.11x. Structural rejections (H=64, tiny T, batches that already fill the SM array) are unchanged. Wrapper-level auto routing verified end-to-end (dense + varlen), outputs bit-identical to the serial path.
Same cutlass-dsl 4.4.2 incompatibility as the SM90 kernel files fixed in a65a64b (string annotations defeat Constexpr classification); this was the last @cute.kernel/@cute.jit file still carrying the import. Import-level check only on this box (SM90); the kernel itself needs Blackwell to exercise.
Wrapper cost over the raw op at tiny sizes was ~61 us/call. Cut it to ~22 us (and to ~0 from T>=512, where it fully hides behind GPU work): - use_beta_sigmoid_in_kernel=True (FLA-aligned name): beta is passed as pre-sigmoid logits straight to the kernel, skipping the per-call torch.logit(beta.float()).to(bf16) round-trip (3 elementwise kernels + dispatches). The model produces logits natively; the sigmoid->logit detour existed only for the FLA post-sigmoid convention, which remains the default. Bit-identical to the raw op on the same logits. - optional out= / final_state= preallocated buffers (validated, written in place and returned) instead of per-call allocations. - replace fla.utils input_guard/autocast decorators with a local lightweight guard (contiguous fixup + device ctx only when needed) and torch.amp custom_fwd/custom_bwd: hopper_prefill no longer imports fla (previously an undeclared dependency). - forward-only fast path: skip torch.autograd.Function.apply when no input requires grad (backward raises NotImplementedError anyway); grad-enabled inputs still route through apply. - _cute_arch_for_device: check-and-set CUTE_DSL_ARCH (cached device cc) instead of set+pop per dispatch; compiles are lazy first-call only. Measured (H100, H=64): T=64 wrapper 125 -> 86.5 us vs raw 64.4; T=512 103.5 vs raw 103.4 (hidden). Default FLA-style path also drops 125 -> 109 us and its numerics are unchanged (rel_rmse 6e-3 vs FlashKDA); full test sweep 23 passed.
Every prefill call transposed beta [T_total, H] -> [H, T_total] on the host (an aten elementwise kernel, 3-17 us scaling with T*H, plus a dispatch and an arena slot) purely so K1 could read 16 contiguous values per (head, tile). That layout requirement is inherited from FlashKDA, whose K1 fetches beta through a 1D TMA and genuinely needs contiguous windows; this port already replaced that TMA with plain per-thread loads during porting, so the transpose had become vestigial. Read the original packed layout directly (beta[(token + tidx) * H + h], H is Constexpr; 16 tiny strided loads per CTA) and drop _copy_beta_flat plus the beta_flat arena slot in both the serial dispatch and the CP driver. ws_beta emission and its consumers (K2, pre_scan) are unchanged. cuLA now launches exactly two kernels per prefill call vs FlashKDA's three (prepare + recurrence + its still-present beta transpose): T=1024 H=64 GPU time 174.0 vs 180.4 us. K1 itself is unchanged (64.8 vs 65.1 us). Outputs verified BIT-IDENTICAL to the pre-change code on t-pad, init, aligned/unaligned varlen and forced-CP cases; full test sweep 23 passed.
The auto planner budgeted SM SLOTS, not SM TIME: short sequences were counted as occupying their H CTAs for the whole kernel, and every splittable sequence received the same segment COUNT. Both assumptions break on ragged batches — the documented skew case (1x14336 + 16x128 tokens, H=16) computed s_split=1 and bypassed CP entirely, forfeiting the win; mixed lengths left the critical path on the longest sequence's oversized segments. Add a duration-balanced candidate plan: pick a global target segment length (total tiles over one SM wave's segment slots, floored at AUTO_MIN_SEG_TILES) and give each sequence ceil(tiles/target) segments, so segment LENGTHS equalize across the batch and short sequences drain early for the scheduler to backfill. auto_plan_segments builds both the legacy and the balanced plan and keeps whichever the calibrated cost model predicts faster; ties keep the legacy plan, so uniform batches plan exactly as before (regression-tested over a 96-combo sweep). plan_cp(s_split=None) now routes through auto_plan_segments so the policy decision and the executed plan can never diverge; explicit s_split is untouched. Measured (H100, H=16, bit-exact vs serial): skew 1x14336+16x128 1.42x (previously bypassed, 1.0x); mixed [8192, 2048] 1.46x. Known limitation, documented for follow-up: on heavily oversubscribed ragged grids the cost model's makespan lower bound is optimistic for the segment-K2 pass (measured +55% over prediction from block-scheduling order — long segments launched late; pre_scan matched prediction). A head-major grid rasterization for K2/pre_scan would recover most of that and is a kernel-side follow-up.
Thread blocks dispatch in linear index order, so on ragged CP plans the long and short segment chains interleave in the launch queue; with more CTAs than resident slots, some long chains only get an SM after several generations of short ones drain, and the makespan grows by the late start (skew case: last long K2 segment started ~175 us in; segment-K2 measured 410 us vs the 263 us scheduling lower bound). Add a seq_order/seg_order indirection to k2_kernel and pre_scan_kernel (seg = order[blockIdx.x]); the CP driver launches segments sorted by length descending (stable sort -> identity for uniform splits). The serial path passes a cached identity order, so its behavior is unchanged. Pure launch reordering — outputs stay bit-identical (golden-output check + full 27-test sweep). Measured on the skew case (H=16, 1x14336+16x128, healthy clocks): segment-K2 410 -> 360 us; e2e speedup vs serial 1.42x -> 1.49x; mixed [8192, 2048] 1.46x -> 1.51x. Smaller than the naive estimate because a new bottleneck appears: with long chains holding most slots, the short- chain tail drains through the few remaining slots. The scheduling lower bound in the cost model therefore stays optimistic (~+35%) for heavily oversubscribed ragged grids; the engage decisions remain safe (predicted 2.12x vs actual 1.49x, well above the 1.10 margin).
- intracard CP: ragged varlen batches (one-long-many-short, mixed lengths) planned by the duration-balanced auto planner must stay bit-identical to the serial path; a splittability precondition keeps the cases meaningful. - prefill vs FLA: back-to-back shape changes share the grow-only workspace arena; earlier shapes must keep matching FLA after the arena is re-carved (covers the tail-chunk path too).
Drop repeated tvm-ffi/compile-cache boilerplate, compress cost-model and longest-first launch order notes, and slim docstrings that restated code intent. No runtime changes.
These writeups were never meant to ship with the codebase.
…PMode Introduce cula/ops/kda/cp_mode.py -- CPMode (OFF/AUTO/FORCE, one parse() covering the public "auto"/True/False argument and the use_cp alias) and NotSplittableError -- as the CP dispatch vocabulary shared by both backends. The SM100 decision moves from the mixed cula/ops/kda/policy.py into cula/ops/kda/sm100/policy.py, behavior unchanged; the _no_cp recursion guard becomes an explicit use_intracard_cp=False at the intracard_fwd_h reentry, and forced-mode rejections now raise NotSplittableError (a ValueError subclass, so existing callers keep working).
Serial is CP with nothing split, so the engage decision IS the planner's output: plan.py returns the plan the executor runs verbatim, and a trivial plan means "take the serial path". The wrapper routes on plan_prefill(); driver.py becomes a pure executor (run_cp) with intracard_prefill kept as a plan-and-run entry for direct callers. cula/ops/kda/policy.py is deleted -- each backend owns its dispatch, and planning can no longer diverge from execution. The calibrated cost model is replaced by a dimensionless engage rule: split when max(critical chain, per-slot load) shrinks by more than RERUN_RATIO (3, from the ~2.76x per-tile cost of pre_scan + segment-K2 vs serial K2). Both decision constants are ratios of same-species MMA chains, so they need no per-SKU calibration; derivations live next to the constants. The legacy count-based planner and the two-candidate selection are deleted -- the duration-balanced planner is the only one. Decision sweep vs the old cost model over 287 shapes: 98.6% agreement, strictly conservative, identical segmentation everywhere both engage. E2e boundary probes: declined shapes measure 0.75-0.77x under forced CP, the first engaged size measures 1.03x, deep engage 2.15x. Env knobs: CULA_KDA_CP_COST_*, _ENGAGE_MARGIN, _MIN_SEG, _ENGAGE_MIN_TILES, _MIN_SEG_TILES are removed (setting them now warns); CULA_KDA_CP_RERUN_RATIO is added. Semantic changes: use_intracard_cp=True now runs any splittable shape and raises NotSplittableError only when nothing can split; manual s_split=2 actually runs 2-segment CP instead of silently falling back. Both paths remain bit-exact vs the serial kernel.
Replace the two-line "Copyright + SPDX-License-Identifier" form with the full Apache-2.0 notice (32 files) and unify the "ANTGROUP. All rights reserved." copyright variant to "Ant Group Co., Ltd." (13 files). Third-party attributions (FLA authors, MoonshotAI) are untouched.
Comments argued for the design instead of informing the next reader: derivation arithmetic on the decision constants, portability hedging, and the same trivial-plan semantics restated at three levels. Constants now state their role and that they were tuned on H100; docstrings say each thing once.
Undo the full Apache-2.0 notice expansion from a805f54; the copyright-line unification from that commit stays.
ruff-format line wrapping and whitespace across the SM90/SM100 KDA files, plus two dead k1_T_total assignments in fwd.py left behind when the beta_flat staging copy was removed (F841).
The old cost-model env vars only ever existed on this branch; nothing to migrate.
slots -> parallel_segs, target -> target_seg_tiles, load -> load_bound, r -> tiles; docstring example updated to match.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
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.
Comment on lines
+49
to
+57
| 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 |
Contributor
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
Collaborator
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.
Comment on lines
+118
to
+119
| # Validate Blackwell build environment | ||
| assert_blackwell_build_env() |
Contributor
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.
Suggested change
| # 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() |
…istic The auto decision materialized cu_seqlens on host for the engage heuristic and dropped it, so an enabled decision paid a second D2H sync inside intracard_fwd_h. Return it on CPDecision and reuse it downstream -- one sync when the heuristic runs, zero when CP is off (review feedback on inclusionAI#100).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 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).