feat: support 192/128 and 256/256 ffa_fa4 in magi

[NVIDIA-JerryChen]

Summary

• Add head_dim_v propagation and support SM100 configurations:
  • Q/K dim = 192, V dim = 128
  • Q/K/V dim = 256
• Align forward q-stage selection and backward K2Q CSR metadata with the FFA_FA4 kernel tile contract.

KV Communication

• Symmetric (256, 256) path keeps the existing sequence-axis fusion with packed_times=2.
• Asymmetric (192, 128) path uses MLA-style K/V packing:
  • concatenate K and V along the last head-dimension axis;
  • communicate them as one fused token stream with packed_times=1;
  • split the received buffer back into K and V.
• Backward dK/dV communication follows the same fused representation and is split after reduction.
• The asymmetric path does not use the symmetric concat_kv/concat_dkv sequence-axis packing.

CSR Tile Mapping

• (192, 128) backward K2Q CSR metadata uses 128x256 blocks.
• (256, 256) backward K2Q CSR metadata uses 256x128 blocks.
• Forward q-stage and CSR block sizes are selected to match the FFA_FA4 kernel.

Validation

• Ran selected SM100 FFA_FA4 tests for configurations with n_func <= 13.

Related PR

• Supersedes MagiAttention#326.
• After this PR is merged, PR [WIP] head dim 256 & 192/128 MLA support for ffa4 #326 can be closed.

[lijinnn]

Hello! Excited to work together. 🙌

I noticed that the pre-commit checks failed for this PR: #331.

Could you please refer to our contributing guide (https://github.com/SandAI-org/MagiAttention/blob/main/CONTRIBUTING.md) to run the pre-commit checks locally, fix the formatting, and update your commit?

Thank you so much!

[codecov]

Codecov Report

❌ Patch coverage is 45.05495% with 50 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
magi_attention/functional/dist_attn.py	36.17%	30 Missing ⚠️
magi_attention/meta/collection/calc_meta.py	41.66%	14 Missing ⚠️
magi_attention/functional/fa4_utils.py	0.00%	4 Missing ⚠️
magi_attention/dist_attn_runtime_mgr.py	60.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[lijinnn]

LGTM