feat: support head_dim=256 for SM90 BF16 decode

[Carlos779988]

Summary

Add head_dim=256 support to the existing SM90 BF16 decode paths:

• static split-K
• dynamic task-map scheduling

This PR covers decode only; D256 prefill is not included. Refs #70.

Since each D256 K/V stage is twice as large as D128, D256 uses one pipeline stage while D128 keeps two. This keeps their staged K/V shared-memory footprints comparable. The static reducer and dynamic combine kernel now iterate over output-dimension tiles instead of assuming D128, so the same code covers both D128 and D256.

While extending the tests, I also found that dynamic decode produced NaNs for q_len=5 at both D128 and D256. The cross-warp softmax reduction handled 2/4/6/8 values per lane but not 10; this adds the missing 10-value layout.

Test plan

• make format-check
• pytest tests/test_attention_decode_bf16.py tests/test_attention_decode_qkpertoken_perhead_vperhead_fp8.py tests/test_attention_decode_qpertoken_perhead_kvpertensor_fp8.py (582 passed)
  • 198 BF16 cases, including 53 D256 cases and a D128 q_len=5 regression case
  • 384 FP8 regression cases across the two existing quantization paths
• D256 coverage: static q_len=1..3, dynamic q_len=1..5, and selected GQA 4/8, NHD/HND, and page-size 16/32/64 configurations
• Output checks for all nine benchmark cases: dynamic, FlashAttention-3, and FlashInfer against static (atol=0.016, rtol=1e-5)

Performance

H20, BF16, Hq/Hkv=16/4, head_dim=256, q_len=1, CUDA Graph, 10 warmup iterations and 100 measured iterations. The software versions are CUDA 12.8, PyTorch 2.8.0, FlashAttention 3.0.0, and FlashInfer 0.6.17. Times are in microseconds; lower is better. K denotes 1024 tokens.

KV-length distribution	static split-K	dynamic task-map	FlashInfer (tensor core)	FA3 (best tuned)
64 x 512	51.49	55.62	66.94	178.88
64 x 4K	344.13	321.12	374.34	1218.53
32 x 128 + 32 x 4K	164.77	175.17	207.68	614.91
1 x 16K + 15 x 64	42.85	37.73	46.56	91.26
1 x 64K + 7 x 4K	147.23	125.41	133.60	411.97
1 x 64K + 15 x 4K	192.86	166.50	184.90	613.66
1 x 64K + 31 x 4K	281.31	246.50	365.18	915.42
1 x 128K + 31 x 4K	372.83	316.35	685.25	1069.98
2 x 32K + 30 x 4K	252.13	233.63	360.54	907.74

Dynamic timings include the attention and combine kernels using a prebuilt task map. Task-map construction is outside the timed region, matching the exclusion of FlashInfer planning and FA3 KV-cache conversion. HPC-Ops and FlashInfer use page size 64; FA3 uses page size 256, its minimum paged-KV block size.

The FA3 result is the best per-case configuration from 48 combinations of num_splits and GQA packing (432 runs across the nine cases). On these cases, dynamic is 1.07-2.17x faster than FlashInfer's tensor-core path and 2.42-3.89x faster than the best tuned FA3 result. Static remains faster for the short uniform case and the 32 x 128 + 32 x 4K mix.