[KDA] Refactor csrc kernel selection, add tests and sanity check

[KevinZeng08]

📌 Description

• Refactor sm100 csrc kernel selection code, removing hardcoded kernel objects
• Add tests for comparing with FLA with bf16 beta
• Add sanity check of 128 head dim and bf16 dtype

🔍 Related Issues

#41

🚀 Pull Request Checklist

Thank you for contributing to cuLA! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing.

⚡ Performance

Reviewer Notes

[gemini-code-assist]

Code Review

This pull request refactors the SM100 kernel dispatch logic by introducing BOOL_SWITCH and BETA_TYPE_SWITCH macros, which significantly reduces boilerplate code and removes redundant kernel aliases. It also enforces a head dimension of 128 in the Python wrappers and adds bfloat16 support for beta tensors in tests. Review feedback identifies critical bugs in the benchmark utilities where gating tensors were incorrectly cast to bfloat16, which would cause kernel failures. Additionally, there are suggestions to resolve redundant dimension checks and to consider runtime dispatch for the RoundingTF32 parameter.

[icavan]

LGTM

[cherhh]

LGTM