[CUDA] Tune online-softmax split-K scheduling

[Gasoonjia]

Stack dependency

Depends on #22190.

Summary

Retune split-K scheduling for the numerically stable online-softmax implementation using a hardware-derived policy that adapts to the target GPU and runtime attention dimensions.

• Replace the two-wave target with a hardware-derived 16/9-wave target.
• Cap splits to one per 128 KV elements to avoid excessive reduction and empty-CTA overhead.
• Use contiguous KV chunks, which profile better than persistent striding for the online-softmax kernel.
• Apply the hardware-aware policy to both decode and small-query split-K paths.
• Keep scale arithmetic explicitly in FP32.

The policy depends on the device SM count, L_kv, and B * H_kv; it does not contain model-specific or GPU-specific shape hard-coding.

Test plan

• Triton split-K and FP8 dispatch suite: 47 passed.
• Explicit non-power-of-two split-count coverage.
• 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/22191

• 📄 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.

[Gasoonjia]

Folded into #22190 so the D113 scheduling mechanism, final online-softmax retuning, and FP32 sm_scale handling are reviewed as one coherent split-K optimization.