[Feature Request] Flexible MLA Kernel Supporting More `d`, `dv`, and `n_heads` Values

## Background

Hi TileLang team!
Thank you for the fantastic repo — the performance of the [MLA](https://github.com/tile-ai/tilelang/tree/main/examples/deepseek_mla) kernel is truly impressive.

We recently proposed a framework called [MHA2MLA](https://github.com/JT-Ushio/MHA2MLA), which enables converting MHA or GQA-based models (such as Llama, Mistral, Qwen, etc.) to an MLA architecture at a very low cost (e.g., continued fine-tuning with only 0.3% of the original pretraining tokens). Due to constraints from the original MHA/GQA dimensions, the converted MLA cannot be perfectly aligned with the deepseek MLA (for instance, [Llama2-7B-d_kv_16](https://huggingface.co/fnlp/Llama-2-7B-MLA-d_kv_16) has `n_heads=32`, `dv=512`, and `d=512+16`), meaning it cannot directly use high-efficiency kernels like FlashMLA that are strictly dimension-bound.

We noticed that TileLang_MLA offers comparable inference efficiency to FlashMLA while allowing flexible dimension adjustments. In our initial experiments, it runs stably with configurations such as `n_heads=64`, `dv=512`, and `d=512+32`. To match these supported shapes, we had to zero-pad models like Llama2-7B-d_kv_16 (`n_heads=32 → 64`, `d=512+16 → 512+32`), and as a result, our TileLang+MLA version achieved faster inference than FlashAttention2+MHA on H100-80G (please see table below).

| LLaMA2-7B(d_kv_16@MLA) | Bsz | Seqlen | Attention's Latency [ms] |
|-------------------------|-----|--------|--------------------------|
|MHA w. FlashAttention2 | 8 | 2K | 5.95 |
| | 8 | 4K | 10.66 |
| | 64 | 2K | 43.74 |
|MLA w. TileLang Kernel | 8 | 2K | 3.52 (-41%) |
| | 8 | 4K | 6.08 (-43%) |
| | 64 | 2K | 4.10 (-91%) |

However, zero-padding may waste memory and compute. Since our kernel programming experience is limited, we would greatly appreciate your help in making the MLA kernel more flexible.

## Feature Requests

**Feature 1: Support for More Flexible Values of `n_heads`, `dv`, and `d`**
- `n_heads`: e.g., 16, 24, 32, etc.
- `dv`: e.g., 64, 128, 256, 512, etc.
- `d`: e.g., `dv+8`, `dv+16`, `dv+32`, `dv+64`, etc.

**Feature 2: Support for Non-shared `k_r` Across Attention Heads**
- In the original MLA, the `k_r` (the portion of the key vector with positional encoding) has the shape `(bsz*seqlen, 1, d_pe)`.
- To align with pretrained MHA parameters in MHA2MLA, `k_r` should be modified to `(bsz*seqlen, n_heads, d_pe)`.

Would the TileLang team be willing to help us develop a more flexible MLA kernel? If not, could you offer some guidance on how to approach this? We would be happy to acknowledge your contributions in our [Paper](https://arxiv.org/abs/2502.14837), GitHub [repo](https://github.com/JT-Ushio/MHA2MLA), and HuggingFace [Models](https://huggingface.co/collections/fnlp/mha2mla-67c51287dfc6cd46127e1b92).

Thanks again for your excellent work, and we look forward to your thoughts!

Best regards