[Feature Request] Support q_len > 1 for FP8 decode (Multi-Token Prediction)

[ZZBoom]

Great work on this project! 🎉

I benchmarked the FP8 decode attention on H200 and the results are impressive:

============================================================
GPU: NVIDIA H200
Dtype: FP8
Batch: 128, Q_len: 1, KV_len: 2048
Heads: 64 (Q) / 8 (KV), Dim: 128
Warmup: 10, Iters: 100
============================================================

[hpc-ops] Allocating buffers...
[hpc-ops] Block size: 64

[FA3] Allocating buffers...
[FA3] Block size: 448

============================================================
Results:
============================================================

hpc-ops:
Avg time: 0.147 ms
Throughput: 6800.0 iter/s
Bandwidth: 3672.1 GB/s
TFLOPS: 58.41

FA3:
Avg time: 0.263 ms
Throughput: 3804.7 iter/s
Bandwidth: 2054.6 GB/s
TFLOPS: 32.68

------------------------------------------------------------
hpc-ops is 1.79x faster than FA3
------------------------------------------------------------

Feature Request:

Would it be possible to support q_len > 1 (e.g., q_len = 4) in attention_decode_fp8?

This would enable support for Multi-Token Prediction (MTP) models like DeepSeek-V3, which predicts multiple tokens per forward pass.

Currently returns: RuntimeError: num_seq_q must be 1

Ref: src/attention/entry.cc:310

Thanks!

[reed-lau]

Thank you for the excellent benchmark and the feature request.

Support for q_len > 1 (e.g., for MTP models) is planned for a future update. The current attention_decode_fp8 kernel only supports q_len = 1, and extending it requires additional development.

We appreciate your patience and welcome further discussion or contributions on this feature.

[ZZBoom]

Thanks for the reply!

In src/attention/decode/smallm_splitk_fp8_kernels.cuh, the QK^T computation uses K as the A matrix and Q as the B matrix:

// Line 426-427
auto tKs4r = thr_mma_qk.partition_A(sK);
auto tQs4r = thr_mma_qk.partition_B(sQ);

This computes Score = K × Q^T with output shape [64, 8] (kv_seq × query_heads), using SM90_64x8x32. In contrast, Flash Attention 3 uses the standard Score = Q × K^T.

I have a follow-up question: what's the design rationale behind this choice?

My guess is that for decode (q_len=1), using Q × K^T would result in a very small M dimension, underutilizing the WGMMA instruction. Swapping to K × Q^T puts kv_seq on M=64, which is well utilized, and N=8 fits the typical GQA config where heads_per_group = 8.

Any clarification would be immensely appreciated!