diff --git a/tests/utils/test_count_flops.py b/tests/utils/test_count_flops.py index aa521786b0..ceba7a9e7e 100644 --- a/tests/utils/test_count_flops.py +++ b/tests/utils/test_count_flops.py @@ -58,6 +58,12 @@ def qwen3_5_counter(): return VeomniFlopsCounter(config) +@pytest.fixture +def qwen3_counter(): + config = _load_toy_config("tests/toy_config/qwen3_toy") + return VeomniFlopsCounter(config) + + @pytest.fixture def qwen3_5_moe_counter(): config = _load_toy_config("tests/toy_config/qwen3_5_moe_toy") @@ -110,13 +116,13 @@ def test_numerical(self, qwen3_5_counter): batch_seqlens = [1024, 1024, 1024, 1024] flops, _ = qwen3_5_counter.estimate_flops(batch_seqlens, delta_time=1.0) # Embedding lookup is not a matmul; only lm_head contributes vocab_size * hidden_size. - assert flops == pytest.approx(105.419032756224, rel=1e-9) + assert flops == pytest.approx(106.965220982784, rel=1e-9) def test_numerical_with_vit(self, qwen3_5_counter): batch_seqlens = [1024, 1024, 1024, 1024] flops, _ = qwen3_5_counter.estimate_flops(batch_seqlens, delta_time=1.0, images_seqlens=[256, 512]) # Embedding lookup is not a matmul; only lm_head contributes vocab_size * hidden_size. - assert flops == pytest.approx(107.650266169344, rel=1e-9) + assert flops == pytest.approx(109.196454395904, rel=1e-9) class TestQwen35MoeFlops: @@ -137,13 +143,40 @@ def test_numerical(self, qwen3_5_moe_counter): batch_seqlens = [1024, 1024, 1024, 1024] flops, _ = qwen3_5_moe_counter.estimate_flops(batch_seqlens, delta_time=1.0) # Embedding lookup is not a matmul; only lm_head contributes vocab_size * hidden_size. - assert flops == pytest.approx(16.68192141312, rel=1e-9) + assert flops == pytest.approx(16.888079843328, rel=1e-9) def test_numerical_with_vit(self, qwen3_5_moe_counter): batch_seqlens = [1024, 1024, 1024, 1024] flops, _ = qwen3_5_moe_counter.estimate_flops(batch_seqlens, delta_time=1.0, images_seqlens=[256, 512]) # Embedding lookup is not a matmul; only lm_head contributes vocab_size * hidden_size. - assert flops == pytest.approx(18.847925010432, rel=1e-9) + assert flops == pytest.approx(19.05408344064, rel=1e-9) + + +class TestQwen3Flops: + pytestmark = pytest.mark.usefixtures("mock_device_flops") + + def test_uses_explicit_head_dim_for_projection_shapes(self, qwen3_counter): + config = qwen3_counter.config + batch_seqlens = [12, 5] + tokens_sum = sum(batch_seqlens) + q_size = config.num_attention_heads * config.head_dim + kv_size = config.num_key_value_heads * config.head_dim + + mlp_N = config.hidden_size * config.intermediate_size * 3 + attn_linear_N = config.hidden_size * (2 * q_size + 2 * kv_size) + lm_head_N = config.hidden_size * config.vocab_size + dense_N = (mlp_N + attn_linear_N) * config.num_hidden_layers + lm_head_N + expected_flops = 6 * dense_N * tokens_sum + expected_flops += ( + 12 + * sum(seqlen * seqlen for seqlen in batch_seqlens) + * config.head_dim + * config.num_attention_heads + * config.num_hidden_layers + ) + + flops, _ = qwen3_counter.estimate_flops(batch_seqlens, delta_time=1.0) + assert flops == pytest.approx(expected_flops / 1e12, rel=1e-9) class TestGptOssFlops: diff --git a/veomni/ops/kernels/gated_delta_rule/_ascend/triton_core/chunk_scaled_dot_kkt.py b/veomni/ops/kernels/gated_delta_rule/_ascend/triton_core/chunk_scaled_dot_kkt.py index 7abd6f42c8..7040ac9fbf 100644 --- a/veomni/ops/kernels/gated_delta_rule/_ascend/triton_core/chunk_scaled_dot_kkt.py +++ b/veomni/ops/kernels/gated_delta_rule/_ascend/triton_core/chunk_scaled_dot_kkt.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang from typing import Optional diff --git a/veomni/ops/kernels/gated_delta_rule/_ascend/triton_core/l2norm.py b/veomni/ops/kernels/gated_delta_rule/_ascend/triton_core/l2norm.py index 1c60635170..fd9deabb9e 100644 --- a/veomni/ops/kernels/gated_delta_rule/_ascend/triton_core/l2norm.py +++ b/veomni/ops/kernels/gated_delta_rule/_ascend/triton_core/l2norm.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang from typing import Optional diff --git a/veomni/ops/kernels/gated_delta_rule/_ascend/triton_core/utils.py b/veomni/ops/kernels/gated_delta_rule/_ascend/triton_core/utils.py index df362ab702..4ca43de1aa 100644 --- a/veomni/ops/kernels/gated_delta_rule/_ascend/triton_core/utils.py +++ b/veomni/ops/kernels/gated_delta_rule/_ascend/triton_core/utils.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang import itertools @@ -95,8 +94,8 @@ def prepare_chunk_indices_list( num_chunks = (length + chunk_size - 1) // chunk_size for chunk_id in range(num_chunks): - indices.append((i)) - indices.append((chunk_id)) + indices.append(i) + indices.append(chunk_id) return indices diff --git a/veomni/ops/kernels/gated_delta_rule/npu_ascendc_gated_delta_rule.py b/veomni/ops/kernels/gated_delta_rule/npu_ascendc_gated_delta_rule.py index d568003046..8c229f6661 100644 --- a/veomni/ops/kernels/gated_delta_rule/npu_ascendc_gated_delta_rule.py +++ b/veomni/ops/kernels/gated_delta_rule/npu_ascendc_gated_delta_rule.py @@ -35,8 +35,6 @@ from __future__ import annotations -from typing import Optional - import torch from ._ascend.flash_gated_delta_rule import flash_gated_delta_rule @@ -48,11 +46,11 @@ def chunk_gated_delta_rule( v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, - scale: Optional[float] = None, - initial_state: Optional[torch.Tensor] = None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, + cu_seqlens: torch.Tensor | None = None, **kwargs, ): # q/k/v arrive as [B, T, H, D]; the AscendC entry wants [B, H, T, D]. diff --git a/veomni/utils/count_flops.py b/veomni/utils/count_flops.py index 8517ac421a..6768c66569 100644 --- a/veomni/utils/count_flops.py +++ b/veomni/utils/count_flops.py @@ -104,6 +104,45 @@ def _compute_lm_head_params(hidden_size, vocab_size): # nn.Embedding is a table lookup, so only the lm_head matmul contributes FLOPs. return vocab_size * hidden_size + def _compute_qwen_dense_text_flops(self, config, tokens_sum, batch_seqlens): + """Compute raw text FLOPs shared by dense Qwen text and VL models.""" + hidden_size = config.hidden_size + num_attention_heads = config.num_attention_heads + num_hidden_layers = config.num_hidden_layers + head_dim = getattr(config, "head_dim", hidden_size // num_attention_heads) + q_size = num_attention_heads * head_dim + kv_size = config.num_key_value_heads * head_dim + + mlp_N = hidden_size * config.intermediate_size * 3 + attn_linear_N = hidden_size * (2 * q_size + 2 * kv_size) + lm_head_N = self._compute_lm_head_params(hidden_size, config.vocab_size) + dense_N = (mlp_N + attn_linear_N) * num_hidden_layers + lm_head_N + dense_N_flops = 6 * dense_N * tokens_sum + + seqlen_square_sum = sum(seqlen * seqlen for seqlen in batch_seqlens) + attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers + return dense_N_flops + attn_qkv_flops + + def _compute_qwen3_moe_text_flops(self, config, tokens_sum, batch_seqlens): + """Compute raw text FLOPs shared by Qwen3 MoE text and VL models.""" + hidden_size = config.hidden_size + num_attention_heads = config.num_attention_heads + num_hidden_layers = config.num_hidden_layers + head_dim = getattr(config, "head_dim", hidden_size // num_attention_heads) + q_size = num_attention_heads * head_dim + kv_size = config.num_key_value_heads * head_dim + + moe_gata_N = hidden_size * config.num_experts + moe_expertmlp_N = hidden_size * config.moe_intermediate_size * config.num_experts_per_tok * 3 + attn_linear_N = hidden_size * (2 * q_size + 2 * kv_size) + lm_head_N = self._compute_lm_head_params(hidden_size, config.vocab_size) + moe_N = (moe_gata_N + moe_expertmlp_N + attn_linear_N) * num_hidden_layers + lm_head_N + dense_N_flops = 6 * moe_N * tokens_sum + + seqlen_square_sum = sum(seqlen * seqlen for seqlen in batch_seqlens) + attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers + return dense_N_flops + attn_qkv_flops + def _estimate_seed_flops(self, tokens_sum, batch_seqlens, delta_time): hidden_size = self.config.hidden_size vocab_size = self.config.vocab_size @@ -307,40 +346,7 @@ def _estimate_deepseek_v4_flops(self, tokens_sum, batch_seqlens, delta_time): return flops_all_token * (1.0 / delta_time) / 1e12 def _estimate_qwen3_moe_flops(self, tokens_sum, batch_seqlens, delta_time): - hidden_size = self.config.hidden_size - vocab_size = self.config.vocab_size - moe_intermediate_size = self.config.moe_intermediate_size - num_hidden_layers = self.config.num_hidden_layers - num_key_value_heads = self.config.num_key_value_heads - num_attention_heads = self.config.num_attention_heads - moe_intermediate_size = self.config.moe_intermediate_size - moe_num_expert = self.config.num_experts - moe_topk = self.config.num_experts_per_tok - - head_dim = getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads) - q_size = num_attention_heads * head_dim - k_size = num_key_value_heads * head_dim - v_size = num_key_value_heads * head_dim - - # non-attn per layer parm - moe_gata_N = hidden_size * moe_num_expert - # moe has gate_proj, up_proj and down_proj using SwiGLU in ExpertMlp layer & shared experts - moe_expertmlp_N = hidden_size * moe_intermediate_size * (moe_topk) * 3 - attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) - lm_head_N = self._compute_lm_head_params(hidden_size, vocab_size) - # non-attn all_layer parm - moe_N = (moe_gata_N + moe_expertmlp_N + attn_linear_N) * (num_hidden_layers) + lm_head_N - # non-attn all_layer & all_token fwd & bwd flops - dense_N_flops = 6 * moe_N * tokens_sum - - # attn all_layer & all_token fwd & bwd flops - seqlen_square_sum = 0 - for seqlen in batch_seqlens: - seqlen_square_sum += seqlen * seqlen - attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers - - # all_layer & all_token fwd & bwk flops - flops_all_token = dense_N_flops + attn_qkv_flops + flops_all_token = self._compute_qwen3_moe_text_flops(self.config, tokens_sum, batch_seqlens) flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12 return flops_achieved @@ -421,36 +427,7 @@ def _estimate_gpt_oss_flops(self, tokens_sum, batch_seqlens, delta_time): return flops_achieved def _estimate_qwen2_flops(self, tokens_sum, batch_seqlens, delta_time): - hidden_size = self.config.hidden_size - vocab_size = self.config.vocab_size - num_hidden_layers = self.config.num_hidden_layers - num_key_value_heads = self.config.num_key_value_heads - num_attention_heads = self.config.num_attention_heads - intermediate_size = self.config.intermediate_size - - head_dim = hidden_size // num_attention_heads - q_size = num_attention_heads * head_dim - k_size = num_key_value_heads * head_dim - v_size = num_key_value_heads * head_dim - - # non-attn per layer parm - # llama use SwiGelu, gate, having up and down linear layer in mlp - mlp_N = hidden_size * intermediate_size * 3 - attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) - lm_head_N = self._compute_lm_head_params(hidden_size, vocab_size) - # non-attn all_layer parm - dense_N = (mlp_N + attn_linear_N) * num_hidden_layers + lm_head_N - # non-attn all_layer & all_token fwd & bwd flops - dense_N_flops = 6 * dense_N * tokens_sum - - # attn all_layer & all_token fwd & bwd flops - seqlen_square_sum = 0 - for seqlen in batch_seqlens: - seqlen_square_sum += seqlen * seqlen - attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers - - # all_layer & all_token fwd & bwd flops - flops_all_token = dense_N_flops + attn_qkv_flops + flops_all_token = self._compute_qwen_dense_text_flops(self.config, tokens_sum, batch_seqlens) flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12 return flops_achieved @@ -489,32 +466,8 @@ def _estimate_llama_flops(self, tokens_sum, batch_seqlens, delta_time): return flops_achieved def _estimate_qwen2_vl_flops(self, tokens_sum, batch_seqlens, delta_time, **kargs): - hidden_size = self.config.text_config.hidden_size - vocab_size = self.config.text_config.vocab_size - num_hidden_layers = self.config.text_config.num_hidden_layers - num_key_value_heads = self.config.text_config.num_key_value_heads - num_attention_heads = self.config.text_config.num_attention_heads - intermediate_size = self.config.text_config.intermediate_size - - head_dim = hidden_size // num_attention_heads - q_size = num_attention_heads * head_dim - k_size = num_key_value_heads * head_dim - v_size = num_key_value_heads * head_dim - - # non-attn per layer parm - mlp_N = hidden_size * intermediate_size * 3 - attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) - lm_head_N = self._compute_lm_head_params(hidden_size, vocab_size) - # non-attn all_layer parm - dense_N = (mlp_N + attn_linear_N) * num_hidden_layers + lm_head_N - # non-attn all_layer & all_token fwd & bwd flops - dense_N_flops = 6 * dense_N * tokens_sum - - # attn all_layer & all_token fwd & bwd flops - seqlen_square_sum = 0 - for seqlen in batch_seqlens: - seqlen_square_sum += seqlen * seqlen - attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers + text_config = self.config.text_config if hasattr(self.config, "text_config") else self.config + text_flops = self._compute_qwen_dense_text_flops(text_config, tokens_sum, batch_seqlens) # vit flops images_seqlens = kargs.get("images_seqlens", None) @@ -524,40 +477,12 @@ def _estimate_qwen2_vl_flops(self, tokens_sum, batch_seqlens, delta_time, **karg vit_flops = 0 # all_layer & all_token fwd & bwd flops - flops_all_token = dense_N_flops + attn_qkv_flops + vit_flops + flops_all_token = text_flops + vit_flops flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12 return flops_achieved def _estimate_qwen3_vl_flops(self, tokens_sum, batch_seqlens, delta_time, **kargs): - # qwen3_vl uses text_config and vision_config to distinguish configs of different parts. - hidden_size = self.config.text_config.hidden_size - vocab_size = self.config.text_config.vocab_size - num_hidden_layers = self.config.text_config.num_hidden_layers - num_key_value_heads = self.config.text_config.num_key_value_heads - num_attention_heads = self.config.text_config.num_attention_heads - intermediate_size = self.config.text_config.intermediate_size - - head_dim = hidden_size // num_attention_heads - q_size = num_attention_heads * head_dim - k_size = num_key_value_heads * head_dim - v_size = num_key_value_heads * head_dim - - # non-attn per layer parm - mlp_N = hidden_size * intermediate_size * 3 - attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) - lm_head_N = self._compute_lm_head_params(hidden_size, vocab_size) - # non-attn all_layer parm - dense_N = (mlp_N + attn_linear_N) * num_hidden_layers + lm_head_N - # non-attn all_layer & all_token fwd & bwd flops - dense_N_flops = 6 * dense_N * tokens_sum - - # qwen3_vl uses deepstack to merge visual embeds and text embeds, but it has no tensor operation. - - # attn all_layer & all_token fwd & bwd flops - seqlen_square_sum = 0 - for seqlen in batch_seqlens: - seqlen_square_sum += seqlen * seqlen - attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers + text_flops = self._compute_qwen_dense_text_flops(self.config.text_config, tokens_sum, batch_seqlens) # vit flops images_seqlens = kargs.get("images_seqlens", None) @@ -567,47 +492,12 @@ def _estimate_qwen3_vl_flops(self, tokens_sum, batch_seqlens, delta_time, **karg vit_flops = 0 # all_layer & all_token fwd & bwd flops - flops_all_token = dense_N_flops + attn_qkv_flops + vit_flops + flops_all_token = text_flops + vit_flops flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12 return flops_achieved def _estimate_qwen3_vl_moe_flops(self, tokens_sum, batch_seqlens, delta_time, **kargs): - # qwen3_vl uses text_config and vision_config to distinguish configs of different parts. - hidden_size = self.config.text_config.hidden_size - vocab_size = self.config.text_config.vocab_size - moe_intermediate_size = self.config.text_config.moe_intermediate_size - num_hidden_layers = self.config.text_config.num_hidden_layers - num_key_value_heads = self.config.text_config.num_key_value_heads - num_attention_heads = self.config.text_config.num_attention_heads - moe_intermediate_size = self.config.text_config.moe_intermediate_size - moe_num_expert = self.config.text_config.num_experts - moe_topk = self.config.text_config.num_experts_per_tok - head_dim = getattr( - self.config.text_config, - "head_dim", - self.config.text_config.hidden_size // self.config.text_config.num_attention_heads, - ) - q_size = num_attention_heads * head_dim - k_size = num_key_value_heads * head_dim - v_size = num_key_value_heads * head_dim - # non-attn per layer parm - moe_gata_N = hidden_size * moe_num_expert - # moe has gate_proj, up_proj and down_proj using SwiGLU in ExpertMlp layer & shared experts - moe_expertmlp_N = hidden_size * moe_intermediate_size * (moe_topk) * 3 - attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) - lm_head_N = self._compute_lm_head_params(hidden_size, vocab_size) - # non-attn all_layer parm - moe_N = (moe_gata_N + moe_expertmlp_N + attn_linear_N) * (num_hidden_layers) + lm_head_N - # non-attn all_layer & all_token fwd & bwd flops - dense_N_flops = 6 * moe_N * tokens_sum - # attn all_layer & all_token fwd & bwd flops - seqlen_square_sum = 0 - for seqlen in batch_seqlens: - seqlen_square_sum += seqlen * seqlen - attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers - # all_layer & all_token fwd & bwk flops - flops_all_token = dense_N_flops + attn_qkv_flops - flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12 + text_flops = self._compute_qwen3_moe_text_flops(self.config.text_config, tokens_sum, batch_seqlens) # vit flops images_seqlens = kargs.get("images_seqlens", None) if images_seqlens is not None: @@ -615,7 +505,7 @@ def _estimate_qwen3_vl_moe_flops(self, tokens_sum, batch_seqlens, delta_time, ** else: vit_flops = 0 # all_layer & all_token fwd & bwd flops - flops_all_token = dense_N_flops + attn_qkv_flops + vit_flops + flops_all_token = text_flops + vit_flops flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12 return flops_achieved @@ -740,14 +630,11 @@ def _compute_hybrid_attn_params(config): by 1 full attention layer. Full attention (Qwen3_5Attention) projections: - q_proj: hidden_size -> num_attention_heads * head_dim (output gate ignored, see note) + q_proj: hidden_size -> 2 * num_attention_heads * head_dim (query + output gate) k_proj: hidden_size -> num_key_value_heads * head_dim v_proj: hidden_size -> num_key_value_heads * head_dim o_proj: num_attention_heads * head_dim -> hidden_size - Note: q_proj actually outputs 2x (half query, half gate via sigmoid), but the gate - contribution is ignored here for consistency with existing qwen3_next estimation. - GatedDeltaNet (Qwen3_5GatedDeltaNet) projections: in_proj_qkv: hidden_size -> 2 * linear_k_size + linear_v_size in_proj_z: hidden_size -> linear_v_size (output gate) @@ -807,8 +694,8 @@ def _compute_hybrid_attn_params(config): "The config may use an unsupported entry in `config.layer_types`." ) - # Full attention: q_proj + k_proj + v_proj + o_proj - full_attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) + # Full attention: q_proj emits query and output-gate halves, while o_proj consumes only query-sized output. + full_attn_linear_N = hidden_size * (2 * q_size + k_size + v_size + q_size) # GatedDeltaNet linear projections and depthwise conv1d linear_k_size = config.linear_num_key_heads * config.linear_key_head_dim