[CUDA] Tune hardware-aware split-K decode scheduling

[Gasoonjia]

Stack dependency

Depends on #22189.

Summary

Combine the D113 SM-sized split-K work with the final scheduling policy for the numerically stable online-softmax implementation.

• Derive split counts from the target GPU SM count, L_kv, and B * H_kv.
• Keep the launch size static while reading runtime kv_len on device, preserving CUDA graph compatibility.
• Retune the occupancy target from D113's two waves to 16/9 waves.
• Cap splits to one per 128 KV elements to limit reduction and empty-CTA overhead.
• Use contiguous KV partitions instead of D113's persistent strided traversal; profiling showed that contiguous partitions are faster for the online-softmax kernel.
• Apply the same hardware-aware policy to both decode and small-query split-K paths.
• Keep scale arithmetic explicitly in FP32.

The policy contains no model-specific or GPU-specific shape hard-coding. It adapts from hardware properties and runtime attention dimensions.

Test plan

• Triton split-K and FP8 dispatch suite: 47 passed.
• Standalone D113 decode correctness smoke: 3 passed.
• Covered runtime KV-length tail exclusion, non-power-of-two split counts, and large-logit stability.
• Long-context numerical smoke remained finite with online softmax.
• git diff --check.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22190

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.