[Fix] Fix FLOPs accounting for non-MLA attention

[kaysonyu]

Summary

Fix FLOPs accounting so MLA projection/attention formulas are gated by args.multi_latent_attention, not by Megatron's MLA argument defaults.

Megatron registers MLA args globally, so non-MLA models can still have values such as kv_lora_rank=32, qk_pos_emb_head_dim=64, and v_head_dim=128. The previous FLOPs utility treated those non-None/default values as proof that MLA was active, which caused regular MHA/GQA models to use MLA-style accounting.

Evidence

A local scripts/run-qwen3-4B.sh run with temporary branch logging showed a non-MLA GQA model entering MLA-style FLOPs branches:

[FLOPS_DEBUG:qkv] rank=0 multi_latent_attention=False q_branch=standard_q kv_branch=mla_kv_lora kv_lora_rank=32 q_lora_rank=None num_query_groups=8 kv_channels=128 qk_head_dim=128 qk_pos_emb_head_dim=64 v_head_dim=128 seqlen=8391
[FLOPS_DEBUG:attention] rank=0 multi_latent_attention=False qk_branch=mla_qk_pos_emb_dim v_branch=v_head_dim kv_channels=128 qk_head_dim=128 qk_pos_emb_head_dim=64 v_head_dim=128 seqlen=8391

qwen3-4B.sh configures GQA (--group-query-attention, --num-query-groups 8, --kv-channels 128) and does not enable --multi-latent-attention, so this is a metrics-only bug, not a model forward bug.

Changes

• Add an explicit multi_latent_attention gate in slime/utils/flops_utils.py.
• Non-MLA models now always use standard MHA/GQA QKV and attention formulas, ignoring Megatron's MLA defaults.
• MLA models continue to use MLA-specific Q LoRA, KV LoRA, positional QK dim, and V head dim formulas.

Validation

Local-only validation, no test files added:

local flops formula checks passed
python3 -m py_compile slime/utils/flops_utils.py
git diff --check