Skip to content

[Fix] Fix FLOPs accounting for non-MLA attention #1980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhuzilin merged 1 commit into from
May 29, 2026
Merged

[Fix] Fix FLOPs accounting for non-MLA attention #1980

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 15 additions & 7 deletions slime/utils/flops_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,19 +6,24 @@ def calculate_lm_head_flops(seqlen, hidden_size, vocab_size):
return 2 * seqlen * hidden_size * vocab_size


def _is_multi_latent_attention(args):
return bool(getattr(args, "multi_latent_attention", False))


def calculate_qkv_projection_flops(args, seqlen, hidden_size, num_attention_heads, num_query_groups):
if args.q_lora_rank is None:
q_flops = 2 * seqlen * hidden_size * num_attention_heads * args.kv_channels
else:
is_mla = _is_multi_latent_attention(args)
if is_mla and args.q_lora_rank is not None:
q_flops = (
2
* seqlen
* args.q_lora_rank
* (args.hidden_size + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim))
)
if args.kv_lora_rank is None:
kv_flops = 2 * 2 * seqlen * hidden_size * num_query_groups * args.kv_channels
else:
q_head_dim = args.qk_head_dim + args.qk_pos_emb_head_dim if is_mla else args.kv_channels
q_flops = 2 * seqlen * hidden_size * num_attention_heads * q_head_dim

if is_mla and args.kv_lora_rank is not None:
kv_flops = (
2
* seqlen
Expand All @@ -28,18 +33,21 @@ def calculate_qkv_projection_flops(args, seqlen, hidden_size, num_attention_head
+ args.hidden_size * args.qk_pos_emb_head_dim
)
)
else:
kv_flops = 2 * 2 * seqlen * hidden_size * num_query_groups * args.kv_channels

return q_flops + kv_flops


def calculate_attention_flops(args, seqlen, num_attention_heads):
is_mla = _is_multi_latent_attention(args)
# QK^T with causal
if args.qk_pos_emb_head_dim:
if is_mla:
flops = 2 * num_attention_heads * seqlen * seqlen * (args.qk_head_dim + args.qk_pos_emb_head_dim) / 2
else:
flops = 2 * num_attention_heads * seqlen * seqlen * args.kv_channels / 2
# A*V
if args.v_head_dim:
if is_mla:
flops += num_attention_heads * seqlen * seqlen * args.v_head_dim
else:
flops += num_attention_heads * seqlen * seqlen * args.kv_channels
Expand Down
Loading