[KDA] KDA MTP decode: recurrent + KVBuffer chunkwise verify + flush#96
Merged
Conversation
Single-kernel recurrent gated-delta-rule multi-token-prediction decode with register-resident state. vk (lane=K butterfly-reduce) and ws (4-warp warp-spec) CuTe ops behind a unified dispatch; single-token T=1 routes to vk regardless of batch.
Chunkwise parallel-verification KVBuffer ops for KDA MTP speculative decoding: tp (token-parallel SIMT) and cute-gemm (sm90 tensor-core) verify emit a compact u-buffer instead of T per-token states; rank-m flush rebuilds the accepted state. Adds unit + determinism tests and the unified decode-mtp benchmark.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
…added CTAs, empty dummies flush kvbuffer: accept_len is now per-request and read at runtime from an [N] int32 buffer (m_buf[i_n]) instead of a compile-time constant. The kernel statically unrolls T and masks i_i < m_n, so it compiles exactly once per (shape, BV) regardless of accept length; b_m uses the per-request token m_n-1. Host API accepts an int (broadcast to all N) or a per-request [N] tensor. small-batch decode (vk + kv): wrap the compute body in `if cache_idx >= 0:` so padded slots (cache_idx < 0) skip the whole T-loop, matching the ws kernel (~1.3x on a half-padded batch). kv hoists its k_split constexpr decisions to top level so they stay python constants inside the guarded block. kvbuffer verify: torch.empty instead of torch.zeros for the write_ubuf=False dummy buffers (only ever written, never read) — drops a per-call memset.
- cg SMEM stride K+8 -> K+4: fixes MMA fragment bank conflict (-16~19% @ N>=4, bit-identical) - drop redundant P4 doubling-chain barrier - 3xTF32 emulation on P3 & GEMM1 GEMMs: cgkvb max|Δ| 2.44e-4 -> ~vk level/bf16 floor (+5~6% cgkvb_v @ large N) - bench: default H=HV=32, graph-calls=20
Longxmas
force-pushed
the
feat/kda-mtp-verify-4ops
branch
from
4fd3268 to
9bf0493
Compare
Model config kda_safe_gate=true with kda_lower_bound set uses the safe gate g = lower_bound * sigmoid(exp(A_log) * x); the MTP decode gates previously only implemented the softplus gate g = -exp(A_log) * softplus(x). Add the lower_bound branch (compile-time const_expr; lower_bound=None keeps the softplus path, bit-identical) to all five gates: vk/ws/kv in kda_decode_mtp and tp/cg in kda_decode_mtp_kvbuffer, threaded through launcher / compile cache key / host / dispatch. Tests: oracle lower_bound branch + test_lower_bound_safe_gate (vk/ws/kv) + test_lower_bound_kvbuffer (tp/cg).
Longxmas
force-pushed
the
feat/kda-mtp-verify-4ops
branch
from
50a60e6 to
cf3c1e7
Compare
dynamic-N: mark the batch (and state-pool) axes dynamic via mark_compact_shape_dynamic(mode=0, stride_order=...) and mark_layout_dynamic() for the index tensors, and drop N + pool_size from the compile cache key, so one cubin serves all batch sizes (removes the per-N JIT; no startup prewarm needed). Applied to all five MTP decode gate kernels: vk/ws/kv in kda_decode_mtp and tp/cg in kda_decode_mtp_kvbuffer. Validated: unit tests bit-exact/bf16-level, e2e gsm8k unchanged at 0.8696 with cuda-graph capture succeeding without prewarm. ruff: fix lint in the touched files (E402 noqa on intentional mid-file imports, SIM comparison order, isort, F841 unused var, E702 multiple statements).
Run the ruff-format pre-commit hook on the files it flagged for pre-existing format drift: cula/ops/kda_decode_mtp.py, cula/ops/kda_decode_mtp_kvbuffer.py, tests/test_kda_decode_mtp.py, benchmarks/bench_kda_decode_mtp.py. Formatting only, no logic change.
Add kda_flush_kvbuffer_all_layers: one launch over all L KDA layers (2D grid, x = single-layer grid, y = layer) instead of the per-layer Python loop the caller previously used for the spec-decode state commit. dyn-N: N (request count) is not a compile constant -- the index tensors are marked layout-dynamic and N is dropped from the cache key, so one cubin serves all batch sizes (no per-N compile storm). cute.compile traces the real, already-allocated tensors directly. Bit-identical to the per-layer kda_flush_kvbuffer (MAXDIFF=0 vs the loop over all layers); the single-layer entry point is left unchanged.
Longxmas
force-pushed
the
feat/kda-mtp-verify-4ops
branch
from
68e8f3f to
0da49c7
Compare
… support) Replace the kinv = k / b_run factorization (divides by the cumulative gate product, overflows to inf/NaN once unbounded softplus gates underflow b_run) with bounded ordered products: - tp/cg verify: T x T scores via per-pair decay-ratio chains r(t,i) = prod_(i<j<=t) g_j <= 1 in fp32 SIMT (cg keeps tensor-core MMA for the S0 projection and state-update GEMMs; state B-operand becomes the suffix-decayed key ksuf_t = kn_t * prod_(j>t) g_j). - scratch now stores raw (u, k, g) per token, the same triplet as ReplaySSM's (d, k, g) ring. - single-layer and all-layers flush rebuild S_m with descending suffix products; every factor stays <= 1, no division anywhere. Exact identity transformation; also valid for safe-gate models.
The flush kernel is DRAM-latency bound with no per-CTA data reuse, so more, smaller CTAs win: a back-to-back (bv, warps-per-CTA) sweep on H200 shows bv=8 beats bv=32 by 17-20% at large N*HV (e.g. N=128 T=6 HV=16: 91.2 vs 109.9 us) and ties at small, while multi-warp CTAs only lose (MLP loss outweighs the saved duplicate k/g reads). _select_vk_bv is tuned for the compute-heavier vk verify kernel and picked 32 exactly where flush wants 8; use a flush-specific selector instead. Also evaluated and rejected: row-grouped running-product Stage-3 for tp verify (op count T^3/6 -> T^2 but the serial ratio chain kills ILP; measured flat to -2%).
The sglang adapter feeds q/k/v as strided views of the fused qkv projection; forcing them contiguous costs ~90 copy kernels per verify step and, at small batch, launch-bound GPU idle. Compile a dynamic- layout kernel variant (mark_layout_dynamic, innermost K axis static) when the inputs are strided with a contiguous K axis, and skip the copies. Contiguous inputs keep the byte-identical compact variant; dyn_stride joins the compile cache key. Bit-exact vs the contiguous path (output, final state, intermediate states) at N=4/16/32, T=4, H=HV=16, K=V=128.
… / shuffle) + d/k/g_buffer - Restore the warp-spec recurrent op as kda_decode_mtp_recurrent_ws; recurrent dispatch now routes state_layout=kv -> kv, T<=4 or N*HV<2048 -> recurrent (vk), and T>4 with N*HV>=2048 -> recurrent_ws (the regime where single-warp vk hits the DRAM-bandwidth wall). - Rename small_batch -> recurrent (vk/kv single-warp host + kernels). - kvbuffer: cg/tp -> tensor_core/shuffle across function/kernel/launcher/compile-cache names and the _kvbuffer_prefer_* helper; verify dispatch picks tensor_core vs shuffle by S = HV*N and T (bench-derived collapse). Scratch params u_buffer/kinv_buffer/b_buffer -> d_buffer/k_buffer/g_buffer (replayssm d/k/g semantics). - recurrent_ws warp-group gate opt: distribute the per-K-channel decay gate across all 4 warps (channel = warp_idx*threads_per_group + lane_in_group) instead of computing all 128 channels on warp 0; bit-exact, q/k l2norm-reduce + beta stay on warp 0. - Condense the kvbuffer flush/rebuild comments and the module docstring. - tests and bench updated for every rename plus the recurrent_ws variant. Validation: pytest tests/test_kda_decode_mtp.py = 131 passed; bench --check reports all variants (vk / recurrent_ws / shuffle / tcore) within ~1e-5..6e-5 of the Triton oracle. Note: tests/test_kda_decode.py switched from `cula.kda` to `cula.ops` imports to avoid the box's incomplete fla install (missing fla.ops.cp); this is an environment workaround and is independent of the naming/op changes above.
… + canonical naming recurrent_ws: in the dispatch regime (T>4 and N*HV>=2048) the verify kernel is write-BW bound by the T*d^2 snapshot writes; use ilp_rows=2 there (lower register pressure, higher occupancy hides the wall) -- ~6% faster, bit-exact. Other verify shapes keep the work_units heuristic. kvbuffer dispatch: add _select_kvb_variant(N,HV,T) -- two threshold lines on the work size wu=N*HV (T<=2 shuffle; T=3 wu<=64; T=4 wu<=32; T>=5 tensor_core); _kvbuffer_prefer_tensor_core delegates to it. naming: rename bench op-selector flags + output columns and the two kvbuffer test functions to canonical recurrent / recurrent_ws / tensor_core / shuffle (drop vk/tp/cg/tcore from selectors, flags, and labels; kernel variant params kept).
… q/k/v (dyn-stride)
Extend the vk dyn-stride path to recurrent_ws and both kvbuffer
(tensor_core/shuffle) verify ops. When q/k/v arrive as K-contiguous strided
views the host skips the .contiguous() copy and compiles the
mark_layout_dynamic(leading_dim=3) kernel variant (dyn_stride enters the
compile cache key); contiguous inputs keep the compact byte-identical kernel.
Kernel bodies are layout-agnostic and unchanged (a small _dlp_qkv helper picks
the descriptor).
Bit-exact vs the contiguous-copy path: 18/18 (max abs diff 0) across
recurrent_ws / tensor_core / shuffle x softplus/safe-gate x N{4,8,2} T=4
HV{16,32} K=V=128. Removes the adapter's per-layer input copies on the
ws/kvbuffer dispatch paths.
Longxmas
force-pushed
the
feat/kda-mtp-verify-4ops
branch
3 times, most recently
from
5097eec to
cbd3800
Compare
…ctors; determinism 100000; drop sm90 wording; ruff Import decode ops from the public cula.kda surface; name the single-warp recurrent test selectors recurrent_vk / recurrent_kv; bump the determinism iteration default to 100000 and make the kvbuffer verify/flush determinism tests honor the same KDA_MTP_DET_ITERS knob; drop sm_90 wording from the tensor_core kvbuffer comments/docstrings and benchmark; ruff import-sort + format.
Longxmas
force-pushed
the
feat/kda-mtp-verify-4ops
branch
from
cbd3800 to
5fec6eb
Compare
…acle-tolerance line
kvbuffer: write r(t,i) as the ordered product prod_{i<j<=t} g_j (drop the b_t/b_i
division form) so the formula matches the no-division numerical note. recurrent:
drop the validation-tolerance sentence from the module docstring.
Relax the recurrent dispatch guard from `T <= 4` to `T < 4` so T==4 uses the same N*HV >= _WS_WORK_UNIT_THRESHOLD work-unit threshold as T > 4. Under the production dyn-stride (K-contiguous strided q/k/v) path the single-warp vk kernel loses its vectorized-load fast path at large batch (~+25% vs contiguous at N=128/HV=16, dropping below both triton and warp-spec); recurrent_ws stages through SMEM and is ~stride-insensitive (~+4%). Contiguous vk and ws tie there, so the change is non-regressive.
Longxmas
force-pushed
the
feat/kda-mtp-verify-4ops
branch
from
1c81fa8 to
9e69e4e
Compare
Resolve conflicts against the arch-first KDA backend reorganization: - move MTP decode ops into the new layout: cula/ops/kda_decode_mtp.py -> cula/ops/kda/decode/mtp.py cula/ops/kda_decode_mtp_kvbuffer.py -> cula/ops/kda/decode/mtp_kvbuffer.py and update imports in tests/benchmarks accordingly - export kda_decode_mtp / kda_decode_mtp_recurrent / kda_decode_mtp_recurrent_ws via the lazy-import tables in cula/kda and cula/ops - drop the opt_level plumbing previously added to the single-token decode op (cula/ops/kda/decode/cute.py stays identical to upstream); the MTP ops keep their own internal opt_level compile knobs - list the MTP files in REPO_LAYOUT.md
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
3 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
Adds KDA (Kimi Delta Attention) multi-token-prediction (MTP) decode — the target-side gated-delta-rule recurrence for speculative decoding — in two complementary forms:
kda_decode_mtp): a single register-resident CuTe kernel threading the recurrence over theTdraft tokens per (batch, head).recurrent(single-warp, lane=K / lane=V) covers small batch;recurrent_ws(4-warp warp-spec) covers the large-batch / long-Tcorner where the single warp turns DRAM-bandwidth bound. A unified dispatch picks per shape.kda_decode_mtp_kvbuffer): the parallel-verification path — verify emits a compact scratch (~2·T·d) instead of theTper-token states (T·d²), and a rank-mflush rebuilds the accepted state. Two variants —shuffle(token-parallel SIMT) /tensor_core(CuTe tensor-core MMA, flat-in-T). Scores are built from bounded ordered decay-ratio products (no division), so they stay finite under unbounded softplus gates.Both gate forms are supported: softplus and safe-gate (
lower_bound,g = exp(lower_bound · σ(exp(A_log) · x))).What changed
cula/ops/kda_decode_mtp.py—recurrent/recurrent_wsops + unified dispatch.cula/ops/kda_decode_mtp_kvbuffer.py—shuffle/tensor_corechunkwise verify + rank-mflush.tests/test_kda_decode_mtp.py— unit (vs fp32 oracle) + bit-exact determinism.benchmarks/bench_kda_decode_mtp.py— unified verify-chain benchmark.🔍 Related Issues
Closes #17
🧪 Tests & Accuracy
pytest tests/test_kda_decode_mtp.py— 131 passed. recurrent (lane=K / lane=V) / recurrent_ws / shuffle / tensor_core verify output + rank-mflush vs the fp32 single-token recurrence oracle, plus bit-exact determinism (torch.equal), for both softplus and safe-gate gates.Kernel-level accuracy sweep (
HV=H ∈ {8,16,32,64}×N ∈ {1..128}×T ∈ {2,3,4,6},K=V=128, bf16, safe-gatelower_bound=-5.0) vs Tritonfused_sigmoid_gating_delta_rule_update:rel_rmse ≤ 3.6e-5,rel_max ≤ 3.3e-3; shuffle / tensor_corerel_rmse ≤ 3.4e-5,rel_max ≤ 3.3e-3— within bf16 rounding noise.rel_rmse ~1e-7,|mean_diff| ~1e-11(bit-clean, unbiased).torch.equal) — output, state, u-buffer.compute-sanitizer(memcheck / racecheck / initcheck / synccheck) reports 0 errors / 0 hazards across small + large shapes and both contiguous and dyn-stride (strided q/k/v) input paths.⚡ Performance
H200,
K=V=128, bf16, safe-gatelower_bound=-5.0(both baseline and cuLA), commit = official sglang scatter, acceptm=full, CUDA-graph kernel-chain. Speedup =triton_specverify chain / op chain (>1= faster than Triton). recurrent-class =max(recurrent, recurrent_ws), kvbuffer-class =max(tensor_core, shuffle); for each shape the faster class is bold.HV = H = 16
recurrent
kvbuffer
HV = H = 32
recurrent
kvbuffer
Takeaways:
T(up to ~1.9× atN ≤ 4) and tapers toward ~1.0× asNgrows (the single warp turns bandwidth-bound;recurrent_wsholds the floor at ~1.0× in the large-batch / long-Tcorner).T(up to ~2.3× atN=32, T=6) and grows withT(flat-in-Ttensor_core).