[KDA] Fuse causal conv1d into MTP verify decode#107
Conversation
…-spec ws)
Fuse causal conv1d (width 4, SiLU) with the KDA MTP verify recurrence into
a single CuTe DSL op, kda_conv_decode_mtp_verify, as a cuLA counterpart to
the Triton fused_kda_conv_gating_verify path.
Kernels:
- vk: single warp/CTA, lane holds a K slice (v-conv Option B + auto-bv) —
small batch.
- warp-spec ws: 4 warps/CTA, warp 0 produces shared q/k conv + gate into
SMEM, the other warps consume disjoint V columns of the recurrence —
large batch.
- auto dispatch: N*HV >= 64 routes to ws (graph-timed vk/ws crossover).
Supports safe/softplus gates, HV >= H (GVA; producer runs H == HV), and
any T >= 1 (Triton verify requires T >= W-1 = 3). Shared q/k conv_state
write-back is owned by the largest-bidx CTA in the group (writes last,
after non-owner reads), keeping outputs aligned with the Triton reference.
Tests (tests/test_kda_conv_mtp.py): conv_window/conv_state bit-exact vs a
torch conv+recurrence+snapshot reference; o/inter_ssm within tolerance;
determinism stable across repeats; {safe,softplus} x {H==HV, GVA} covered.
Benchmarks: benchmarks/bench_fused_kda_conv_mtp.py (single-config torch +
Triton + cuLA) and benchmarks/sweep_conv_mtp.py (T x N x H sweep).
…+ large_batch (warp-spec)
Two CuTe DSL variants of the fused causal-conv1d + KDA MTP verify decode
kernel, selected by an auto dispatch on work_units = N*HV.
small_batch (kda_conv_verify_kernel): pipelined across 2 warps -- warp 0
produces token t's q/k conv + silu + l2norm + v-conv into double-buffered
SMEM; warp 1 consumes token t-1 and runs the delta-rule recurrence. The
two warps are decoupled via 2-slot full/empty named barriers (producer
may run up to 2 tokens ahead), with vectorized (float4) conv-state/weight
preamble and epilogue, gate/beta computed one token ahead on the
consumer, and dead SMEM buffers removed. Per-tier knobs:
- <=32: BV=8 (latency-optimal small batch)
- 64-128: BV=16 (halves CTA count, removes wave spill)
- 256-512: BV=16 + conv weights staged in SMEM (multi-wave occupancy)
large_batch (kda_conv_verify_large_batch_kernel): warp-spec, NWARP=8
warps/CTA -- warps first produce the shared q/k conv + gate into SMEM
(one pass, no per-v-tile redundancy), then each warp consumes BVW v-cols
of recurrence. Low state registers -> high occupancy. Adds producer
state-prefetch, float4 vectorized loads, and mixed_qkv-tap / raw-v double
buffering to hide global-load latency at low occupancy. The auto-bvw rule
{<=8:1, <=16:2, <=32:4, else:8} keeps tile_v (= NWARP*bvw) and grid fixed
per tier; 8 warps/CTA fill the former work=64 valley (0.99 -> 1.00-1.05).
The q/k/g SMEM handoff is vectorized (autovec_copy -> STS.128/LDS.128) and
the K+8 bank-conflict padding dropped (stride K): removes the large_batch
producer store-conflicts (ncu 18432 -> 0) and saves SMEM at bit-identical
output. (Small_batch keeps a residual 164 load-conflict: CuTe vectorizes
the SMEM store but not the load here; negligible, latency-bound.)
Dispatch: large_batch if work_units>=512 else small_batch (graph-timed
crossover on L20X, large_batch reaches parity at work=512). The original
single-warp kernel is removed. Speedup vs the triton fused reference:
1.15-1.43x small batch, 1.24-1.32x large batch, no measured tier below 1.0x.
Correctness: small_batch and large_batch outputs are bit-identical;
conv_window/conv_state bit-exact vs reference, o within the same
tolerance; determinism stress-tested to 1e5 iterations (via KDA_DET_ITERS)
with no race. Unit tests: 42 passed, tolerances and coverage not relaxed.
…batch store dedup + host cleanups + T-aware dispatch
![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 fused causal-conv1d + MTP verify decode kernel implemented in CuTe DSL, supporting both small-batch and large-batch variants with automatic dispatch. It also adds comprehensive benchmarks and unit tests to verify correctness and performance. The review feedback highlights two critical issues: a potential NaN propagation bug due to float32 overflow when unconditionally evaluating softplus exponentials, and a robustness issue where cached arrival counters are not zeroed out before kernel launches, which could lead to corruption if a previous launch fails.
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.
| beta_x = softplus_beta * gx | ||
| sp = (cutlass.Float32(1.0) / softplus_beta) * cute.log( | ||
| cutlass.Float32(1.0) + cute.exp(softplus_beta * gx, fastmath=fast_math), fastmath=fast_math | ||
| ) | ||
| use_sp = cutlass.Float32(1.0) if beta_x <= softplus_threshold else cutlass.Float32(0.0) | ||
| spx = use_sp * sp + (cutlass.Float32(1.0) - use_sp) * gx |
There was a problem hiding this comment.
Evaluating cute.exp(softplus_beta * gx) unconditionally when beta_x > softplus_threshold can cause a float32 overflow to inf. In the subsequent weighted sum spx = use_sp * sp + (1.0 - use_sp) * gx, if use_sp is 0.0 and sp is inf, the multiplication 0.0 * inf results in NaN. This propagates NaN values to the gating term and the final outputs.
Using a proper if/else block (similar to the one on line 625) avoids evaluating the exponential branch when beta_x is large, preventing overflow and NaN propagation.
beta_x = softplus_beta * gx
if beta_x <= softplus_threshold:
sp = (cutlass.Float32(1.0) / softplus_beta) * cute.log(
cutlass.Float32(1.0) + cute.exp(beta_x, fastmath=fast_math), fastmath=fast_math
)
spx = sp
else:
spx = gx| beta_x = softplus_beta * gx | ||
| sp = (cutlass.Float32(1.0) / softplus_beta) * cute.log( | ||
| cutlass.Float32(1.0) + cute.exp(beta_x, fastmath=fast_math), fastmath=fast_math | ||
| ) | ||
| use_sp = cutlass.Float32(1.0) if beta_x <= softplus_threshold else cutlass.Float32(0.0) | ||
| spx = use_sp * sp + (cutlass.Float32(1.0) - use_sp) * gx |
There was a problem hiding this comment.
Evaluating cute.exp(beta_x) unconditionally when beta_x > softplus_threshold can cause a float32 overflow to inf. In the subsequent weighted sum spx = use_sp * sp + (1.0 - use_sp) * gx, if use_sp is 0.0 and sp is inf, the multiplication 0.0 * inf results in NaN. This propagates NaN values to the gating term and the final outputs.
Using a proper if/else block avoids evaluating the exponential branch when beta_x is large, preventing overflow and NaN propagation.
| beta_x = softplus_beta * gx | |
| sp = (cutlass.Float32(1.0) / softplus_beta) * cute.log( | |
| cutlass.Float32(1.0) + cute.exp(beta_x, fastmath=fast_math), fastmath=fast_math | |
| ) | |
| use_sp = cutlass.Float32(1.0) if beta_x <= softplus_threshold else cutlass.Float32(0.0) | |
| spx = use_sp * sp + (cutlass.Float32(1.0) - use_sp) * gx | |
| beta_x = softplus_beta * gx | |
| if beta_x <= softplus_threshold: | |
| sp = (cutlass.Float32(1.0) / softplus_beta) * cute.log( | |
| cutlass.Float32(1.0) + cute.exp(beta_x, fastmath=fast_math), fastmath=fast_math | |
| ) | |
| spx = sp | |
| else: | |
| spx = gx |
| pool_size=intermediate_conv_window.shape[0], | ||
| device=mixed_qkv.device, | ||
| ) | ||
| qk_arrival_counters = _get_qk_arrival_counters(mixed_qkv.device, N, H) |
There was a problem hiding this comment.
The qk_arrival_counters tensor is cached and reused across launches on the same stream. If a previous kernel launch fails, is interrupted, or does not complete all block arrivals (e.g., during autotuning, profiling, or dry runs), the counters will be left in a non-zero state. This will corrupt all subsequent launches on that stream.
Calling qk_arrival_counters.zero_() before launching the kernel ensures the counters always start at zero. This is fully asynchronous and does not introduce any host-device synchronization overhead.
| qk_arrival_counters = _get_qk_arrival_counters(mixed_qkv.device, N, H) | |
| qk_arrival_counters = _get_qk_arrival_counters(mixed_qkv.device, N, H) | |
| qk_arrival_counters.zero_() |
1 participant
[KDA] Fuse causal conv1d into MTP verify decode
📌 Description
Adds a fused CuTe DSL operator for the KDA MTP verify path. It consumes packed pre-convolution
mixed_qkvand executes width-4 depthwise causal conv1d + SiLU, q/k L2 normalization, gate/beta, and the gated delta-rule recurrence in one kernel. This removes the separate conv launch, two transpose copies, and intermediate HBM round trips in the unfused path.The operator follows verify semantics: the persistent SSM state is read-only; per-token SSM and convolution-window snapshots are written for rollback;
conv_stateis rolled to the latest three raw input tokens for the next verify chain or commit.Two variants share the same API and math:
small_batch: a 2-warp producer/consumer pipeline. Warp 0 computes conv + SiLU + q/k normalization into double-buffered shared memory, while warp 1 runs the recurrence.large_batch: an 8-warp warp-specialized kernel. q/k preprocessing is shared across v columns, while each warp owns a v-column tile for higher large-batch occupancy.auto: selectslarge_batchatN × H_V >= 512, or atT >= 8andN × H_V >= 64; otherwise selectssmall_batch.conv_state_indices,cache_indices, andintermediate_state_indicesare normalized independently. Shared q/k convolution-state writeback uses stream-local atomic last-arrival counters, avoiding assumptions about CUDA block completion order while remaining compatible with concurrent streams and CUDA Graph replay.Both safe-gate (
lower_bound) and softplus gate forms are supported, including GVA shapes whereH_V > H.What changed
cula/ops/kda/decode/mtp_conv.py— fused kernels, unified API, and shape-aware dispatch.cula/kda/__init__.py,cula/ops/__init__.py— lazy public API export.tests/test_kda_conv_mtp.py— correctness, state-index, stream/graph, repeated-chain, edge-shape, and determinism tests.benchmarks/bench_fused_kda_conv_mtp.py— CUDA-graph single-shape benchmark and reproducibleH × T × Nsweep, with an optional compatible Triton baseline.🧪 Tests & Accuracy
pytest tests/test_kda_conv_mtp.py— 50 passed in 307.75s.H=H_Vand GVA,T ∈ {2,4,8}, auto dispatch, independent non-contiguous state indices, multistream launches, CUDA Graph replay, and repeated verify chains.H=H_V ∈ {8,16,32,64}×N ∈ {1..128}×T ∈ {3,4,6}× two gates × two seeds × both variants.conv_windowandconv_stateare bit-exact;small_batchandlarge_batchare bit-identical.rel_rmse <= 1.668e-4, intermediate-staterel_rmse <= 5.227e-5. Output vs fp32 stays at the shared bf16 rounding floor (1.529e-3–1.721e-3).compute-sanitizer: memcheck, racecheck, initcheck, and synccheck are clean for both variants (8/8 clean runs).⚡ Performance
NVIDIA H200,
K=V=128, bf16, safe gate (lower_bound=-5.0), CUDA-graph timing. Speedup is compatible Triton baseline / cuLA.The extended sweep contains 192 cuLA measurements and 96 Triton baselines over
H=H_V ∈ {8,16,32,64},T ∈ {4,6,8}, andN ∈ {1..128}. The tables report productionautodispatch speedup.H = H_V = 8
H = H_V = 16
H = H_V = 32
H = H_V = 64
All 96 auto-dispatched shapes are faster than Triton (1.08–1.53×). At
T=4, the variant crossover is consistently nearN × H_V = 512. AtT=6,small_batchis usually best, leaving about 3–5% tuning headroom in part of the mid-batch range. AtT=8, the T-aware dispatch avoids thesmall_batch,BV=16register-spill region. A widerT=4sweep throughN=512measures 1.17–1.41×, with no loss at the large-batch tail.