[KDA] Add fused BT=16 inference kernels for KDA prefill#915
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 introduces a fused inference fast path for KDA prefill by implementing two specialized Triton kernels that combine gate activation, L2 normalization, and state propagation to reduce memory overhead. Review feedback highlights the need to enforce a chunk size of 16 throughout the fast path to match kernel expectations. Other suggestions include improving numerical stability by using higher precision for matrix inversions and adjusting normalization epsilons, as well as fixing a redundant memory allocation issue for intermediate states in variable-length sequences.
| if cu_seqlens is not None: | ||
| chunk_indices = prepare_chunk_indices( | ||
| cu_seqlens, | ||
| chunk_size, |
There was a problem hiding this comment.
The inference fast path is designed for fused kernels that are hardcoded for a chunk size of 16 (BT=16). However, chunk_indices are currently prepared using the chunk_size argument passed to the function, which defaults to 64. This mismatch will cause the kernels to process the sequence incorrectly. You should force the chunk size to 16 when preparing indices for this path.
| chunk_size, | |
| 16, |
| use_gate_in_kernel=use_gate_in_kernel, | ||
| A_log=A_log, | ||
| dt_bias=dt_bias, | ||
| chunk_size=chunk_size, |
There was a problem hiding this comment.
As the fused inference kernels require a chunk size of 16, the chunk_size argument passed to chunk_kda_fwd should be explicitly set to 16 here to ensure consistency with the prepared indices and kernel expectations.
| chunk_size=chunk_size, | |
| chunk_size=16, |
| b_L = b_Akk.to(tl.float16) | ||
| b_Ai = m_I.to(tl.float16) - b_L | ||
| b_L2 = tl.dot(b_L, b_L, out_dtype=tl.float16) | ||
| b_Ai = b_Ai + tl.dot(b_Ai, b_L2, out_dtype=tl.float16) | ||
| b_L4 = tl.dot(b_L2, b_L2, out_dtype=tl.float16) | ||
| b_Ai = b_Ai + tl.dot(b_Ai, b_L4, out_dtype=tl.float16) | ||
| b_L8 = tl.dot(b_L4, b_L4, out_dtype=tl.float16) | ||
| b_Ai = b_Ai + tl.dot(b_Ai, b_L8, out_dtype=tl.float16) |
There was a problem hiding this comment.
The parallel prefix expansion for solving tl.float16. Since this involves multiple matrix multiplications (float32 to maintain accuracy, especially since 16x16 matrix operations are very efficient on modern GPUs.
| b_L = b_Akk.to(tl.float16) | |
| b_Ai = m_I.to(tl.float16) - b_L | |
| b_L2 = tl.dot(b_L, b_L, out_dtype=tl.float16) | |
| b_Ai = b_Ai + tl.dot(b_Ai, b_L2, out_dtype=tl.float16) | |
| b_L4 = tl.dot(b_L2, b_L2, out_dtype=tl.float16) | |
| b_Ai = b_Ai + tl.dot(b_Ai, b_L4, out_dtype=tl.float16) | |
| b_L8 = tl.dot(b_L4, b_L4, out_dtype=tl.float16) | |
| b_Ai = b_Ai + tl.dot(b_Ai, b_L8, out_dtype=tl.float16) | |
| b_L = b_Akk.to(tl.float32) | |
| b_Ai = m_I.to(tl.float32) - b_L | |
| b_L2 = tl.dot(b_L, b_L) | |
| b_Ai = b_Ai + tl.dot(b_Ai, b_L2) | |
| b_L4 = tl.dot(b_L2, b_L2) | |
| b_Ai = b_Ai + tl.dot(b_Ai, b_L4) | |
| b_L8 = tl.dot(b_L4, b_L4) | |
| b_Ai = b_Ai + tl.dot(b_Ai, b_L8) |
There was a problem hiding this comment.
Inversion can be done using fp16, provided that qk must l2norm
There was a problem hiding this comment.
The use_qk_l2norm flag controls whether the kernel does norm internally. When False, it assumes q/k are already L2-normalized externally — the fp16 inversion is still safe since the spectral radius is bounded either way.
| h = kg.new_zeros(NT * N, HV, V, K, dtype=kg.dtype) | ||
| else: | ||
| h = kg.new_zeros(NT * N, HV, K, V, dtype=kg.dtype) |
There was a problem hiding this comment.
The memory allocation for intermediate states h is incorrect for variable-length sequences. When cu_seqlens is provided, NT (calculated as len(chunk_indices)) already represents the total number of chunks across all sequences in the batch. Multiplying by N results in redundant memory allocation.
| h = kg.new_zeros(NT * N, HV, V, K, dtype=kg.dtype) | |
| else: | |
| h = kg.new_zeros(NT * N, HV, K, V, dtype=kg.dtype) | |
| if state_v_first: | |
| h = kg.new_zeros(NT if cu_seqlens is not None else NT * N, HV, V, K, dtype=kg.dtype) | |
| else: | |
| h = kg.new_zeros(NT if cu_seqlens is not None else NT * N, HV, K, V, dtype=kg.dtype) |
zhiyuan1i
left a comment
zhiyuan1i
left a comment
There was a problem hiding this comment.
Thank you for your contribution. Looks similar to FLASHKDA. I am very curious about the complete performance, I am curious about the performance of your contributed kernel for D=128, different varlen lengths, and shorter sequence lengths. Also: I hope you can use the dispatch backend and set the default to False, or set the priority to a lower level. You can check the class signatures and comments related to dispatch.
| b_L = b_Akk.to(tl.float16) | ||
| b_Ai = m_I.to(tl.float16) - b_L | ||
| b_L2 = tl.dot(b_L, b_L, out_dtype=tl.float16) | ||
| b_Ai = b_Ai + tl.dot(b_Ai, b_L2, out_dtype=tl.float16) | ||
| b_L4 = tl.dot(b_L2, b_L2, out_dtype=tl.float16) | ||
| b_Ai = b_Ai + tl.dot(b_Ai, b_L4, out_dtype=tl.float16) | ||
| b_L8 = tl.dot(b_L4, b_L4, out_dtype=tl.float16) | ||
| b_Ai = b_Ai + tl.dot(b_Ai, b_L8, out_dtype=tl.float16) |
There was a problem hiding this comment.
Inversion can be done using fp16, provided that qk must l2norm
b5a810f to
bbd6c36
Compare
2 participants
[Ops] Add fused BT=16 inference kernels for KDA prefill
Fuse L2 norm, beta sigmoid, gate activation, intra-chunk attention,
state propagation, and output into two Triton kernels to reduce
intermediate tensor global memory round-trips during inference.
The fused path activates automatically when grad is disabled and
use_gate_in_kernel is set, with no change to the training path.
Testing
Environment: NVIDIA H800 | CUDA 12.9 | PyTorch 2.10.0+cu129
Performance
Fused BT=16 inference kernels vs baseline (multi-kernel path):
Average speedup: ~1.6x, up to 4.25x on small shapes where kernel launch overhead dominates.
Correctness
All tests pass with max output diff < 0.001 and ratio < 0.004:
data：
tle speed
(tle-0.5.1) ╭─root at p-lenovo-h800-01 in ~/workspace/fla-4 on fla-4✘✘✘ 26-06-11 - 10:19:54
╰─(tle-0.5.1) ⠠⠵ FLA_FLASH_KDA=0 FLA_FUSED_INFER=0 /root/workspace/bench_fwd.sh fla-4
shape=[8192,96,128] warmup=30 iters=200 repeats=5
flash_kda (bf16 state) : mean=1.4954 ms, min=1.4884 ms, max=1.5457 ms
flash_kda (no state) : mean=1.5099 ms, min=1.5032 ms, max=1.6497 ms
flash_kda (fp32 state) : mean=1.5123 ms, min=1.5070 ms, max=1.6644 ms
chunk_kda : mean=3.4993 ms, min=3.4796 ms, max=4.2836 ms
chunk_kda (tle) : mean=1.3300 ms, min=1.3122 ms, max=1.4548 ms
chunk_gated_delta_rule : mean=1.8883 ms, min=1.8765 ms, max=2.1909 ms
varlen shape=[8192,96,128] seq_lens=[1300, 547, 2048, 963, 271, 3063] warmup=30 iters=200 repeats=5
flash_kda (bf16 state) : mean=1.5303 ms, min=1.4744 ms, max=1.7063 ms
flash_kda (no state) : mean=1.5097 ms, min=1.4402 ms, max=1.5537 ms
flash_kda (fp32 state) : mean=1.5615 ms, min=1.5022 ms, max=1.7084 ms
chunk_kda : mean=3.5166 ms, min=3.4974 ms, max=3.8342 ms
chunk_kda (tle) : mean=1.2946 ms, min=1.2893 ms, max=1.4155 ms
chunk_gated_delta_rule : mean=1.8844 ms, min=1.8757 ms, max=2.0089 ms
varlen shape=[8192,96,128] seq_lens=[1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024] warmup=30 iters=200 repeats=5
flash_kda (bf16 state) : mean=1.3645 ms, min=1.3537 ms, max=1.3924 ms
flash_kda (no state) : mean=1.3343 ms, min=1.3247 ms, max=1.4495 ms
flash_kda (fp32 state) : mean=1.3804 ms, min=1.3710 ms, max=1.4020 ms
chunk_kda : mean=3.4866 ms, min=3.4748 ms, max=3.7514 ms
chunk_kda (tle) : mean=1.2626 ms, min=1.2533 ms, max=1.3676 ms
chunk_gated_delta_rule : mean=1.8723 ms, min=1.8648 ms, max=2.0289 ms
triton speed
(tle-0.5.1) ╭─root at p-lenovo-h800-01 in ~/workspace/fla-4 on fla-4✘✘✘ 26-06-11 - 10:23:17
╰─(tle-0.5.1) ⠠⠵ FLA_FLASH_KDA=0 FLA_TLE=0 /root/workspace/bench_fwd.sh fla-4
shape=[8192,96,128] warmup=30 iters=200 repeats=5
flash_kda (bf16 state) : mean=1.4959 ms, min=1.4882 ms, max=1.5461 ms
flash_kda (no state) : mean=1.5103 ms, min=1.5043 ms, max=1.5353 ms
flash_kda (fp32 state) : mean=1.5133 ms, min=1.5073 ms, max=1.5482 ms
chunk_kda : mean=3.4994 ms, min=3.4802 ms, max=3.6512 ms
chunk_kda (tle) : mean=1.9965 ms, min=1.9694 ms, max=2.1630 ms
chunk_gated_delta_rule : mean=1.8866 ms, min=1.8741 ms, max=2.0010 ms
varlen shape=[8192,96,128] seq_lens=[1300, 547, 2048, 963, 271, 3063] warmup=30 iters=200 repeats=5
flash_kda (bf16 state) : mean=1.5310 ms, min=1.4749 ms, max=1.5844 ms
flash_kda (no state) : mean=1.5095 ms, min=1.4454 ms, max=1.5972 ms
flash_kda (fp32 state) : mean=1.5613 ms, min=1.4979 ms, max=1.6116 ms
chunk_kda : mean=3.5100 ms, min=3.4975 ms, max=3.6572 ms
chunk_kda (tle) : mean=1.9560 ms, min=1.9217 ms, max=2.2188 ms
chunk_gated_delta_rule : mean=1.8811 ms, min=1.8726 ms, max=2.0651 ms
varlen shape=[8192,96,128] seq_lens=[1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024] warmup=30 iters=200 repeats=5
flash_kda (bf16 state) : mean=1.3645 ms, min=1.3551 ms, max=1.4614 ms
flash_kda (no state) : mean=1.3341 ms, min=1.3247 ms, max=1.4173 ms
flash_kda (fp32 state) : mean=1.3820 ms, min=1.3726 ms, max=1.4370 ms
chunk_kda : mean=3.4890 ms, min=3.4809 ms, max=3.7547 ms
chunk_kda (tle) : mean=1.8754 ms, min=1.8514 ms, max=2.1503 ms
chunk_gated_delta_rule : mean=1.8708 ms, min=1.8628 ms, max=2.0563 ms