Enable support for 16 or 32 heads in the SM90 Sparse Attention

[princepride]

Related to vllm-project/vllm#31473

Purposal:

The SM90 kernel now automatically selects the appropriate block size based on h_q:
h_q = 64, 128, 192, ... → Use B_H=64
h_q = 32, 96, 160, ... → Use B_H=32
h_q = 16, 48, 80, ... → Use B_H=16
This allows for more flexible support of different head configurations while maintaining optimal performance (prioritizing larger block sizes).

Test Plan:

@LucasWilkinson This code was generated by Gemini 3.0 Pro. I don't know much about kernel. Is this code modification reasonable? How should I test it?🥹

[LucasWilkinson]

you can test it by using export FLASH_MLA_SRC_DIR=<path-to-modified-flashmla> then rebuild vLLM

NOTE: id recommend installing ccache and using:

VLLM_DISABLE_SCCACHE=1 CCACHE_NOHASHDIR="true" uv pip install --no-build-isolation -e . -v

for faster rebuilds (first build will still be very slow), you can then even do:

VLLM_DISABLE_SCCACHE=1 python setup.py build_ext --inplace

to just rebuild the kernels

[princepride]

@LucasWilkinson It seems much more complicated than I imagined 😂. After trying this kernel and making several more versions, the model's performance has decreased.