feat: relax heads_per_group constraint to support 1~8 for BF16 and FP8 decode attention

[Religious-J]

Summary

Relax the heads_per_group (= num_head_q / num_head_kv) restriction from {4, 8} to [1, 8] for BF16 and FP8 decode attention kernels, enabling models with arbitrary GQA ratios (including MHA where hpg=1).

Background

The constraint existed because kHeadsPerGroup=8 is hardcoded in the SM90 kernel launch path as the TMA Q/Y tile size. For hpg < 8, the tile over-extends the actual tensor — but this is safe due to SM90 TMA hardware guarantees:

• Load OOB: coordinates outside the tensor box are zero-filled by hardware, not read from adjacent memory.
• Store OOB: writes outside the tensor box are silently discarded by hardware.

The zero-filled Q rows (Q=0) produce garbage attention outputs for those positions, but those outputs are always discarded by the TMA Y store clamp.
Softmax correctness is preserved because gMax/gSum are maintained independently per M-row — the ghost rows' max=0 cannot contaminate valid rows.

Changes

• src/attention/entry.cc: relax two TORCH_CHECK guards from hpg == 4 || hpg == 8 to hpg >= 1 && hpg <= 8 (one for BF16, one for FP8)
• smallm_bf16_dim128_static.cu: add hpg > 8 → return false guard in the launch function to prevent silent mis-computation
• 4 FP8 launch files (static/dynamic × 2 quant types): same guard change from != 8 && != 4 to < 1 || > 8
• 3 test files: extend kv_head_q_head to cover hpg ∈ {1,2,3,4,5,6,7,8}

Test Results

# BF16
python3 -m pytest tests/test_attention_decode_bf16.py -v
# FP8
python3 -m pytest tests/test_attention_decode_qkpertoken_perhead_vperhead_fp8.py \
                  tests/test_attention_decode_qpertoken_perhead_kvpertensor_fp8.py -v

Suite	Result
BF16 decode attention	720 / 720 passed
FP8 qkpertoken_perhead_vperhead	864 / 864 passed
FP8 qpertoken_perhead_kvpertensor	864 / 864 passed

Benchmark

Static splitk, qpertoken_perhead_kvpertensor, KV=4 fixed, 100 iterations:

for hpg in 1 2 3 4 5 6 7 8; do
  python3 benchmark/attention_decode/bench_attention_decode_fp8.py \
    --cases uniform_512 uniform_4096 skewed_mix skewed_extreme one_64k_7x4k one_64k_31x4k \
    --quant-types qpertoken_perhead_kvpertensor \
    --num-head-kv 4 --num-head-q $((4 * hpg)) \
    --warmup 10 --iters 100
done

Environment: NVIDIA H20 (SM90, 96 GB) · CUDA 13.0 · PyTorch 2.11.0

Scenario	hpg=1	hpg=2	hpg=3	hpg=4	hpg=5	hpg=6	hpg=7	hpg=8
uniform_512	20.2	20.2	20.4	20.5	20.6	20.6	20.8	20.7
uniform_4096	112.5	114.3	113.0	113.5	113.9	114.6	114.1	114.6
skewed_mix	58.7	58.8	59.0	59.3	59.2	60.1	59.3	59.1
one_64k_31x4k	148.0	147.0	147.6	148.4	147.8	147.8	149.5	148.3

hpg means heads_per_group = num_head_q / num_head_kv，latency in µs (lower is better)