diff --git a/tests/v1/metrics/test_perf_metrics.py b/tests/v1/metrics/test_perf_metrics.py index bd77fbe91f..ab30f1bb9e 100644 --- a/tests/v1/metrics/test_perf_metrics.py +++ b/tests/v1/metrics/test_perf_metrics.py @@ -28,6 +28,7 @@ from vllm.v1.metrics.perf import ( ExecutionContext, FfnMetrics, InvalidComponent, + MLAAttentionMetrics, ModelMetrics, ParsedArgs, UnembedMetrics, @@ -1021,3 +1022,317 @@ def test_quantized_model_metrics_aggregation(): assert total_flops > 0 assert total_flops == sum(breakdown.values()) + + +#### MLA Attention Tests #### + + +def test_mla_config_parser(): + """Test MLAConfigParser extracts MLA-specific fields from DeepseekV3Config.""" + hf_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=61, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + vllm_config = create_mock_vllm_config(hf_config) + + parser_chain = MLAAttentionMetrics.get_parser() + result = parser_chain.parse(vllm_config) + + assert result.kv_lora_rank == 512 + assert result.qk_nope_head_dim == 128 + assert result.qk_rope_head_dim == 64 + assert result.v_head_dim == 128 + assert result.q_lora_rank == 1536 + assert result.num_attention_heads == 128 + assert result.hidden_size == 7168 + + +def test_mla_attention_metrics_decode(): + """Test MLA decode metrics use compressed KV cache, not standard head_dim.""" + hf_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=1, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + vllm_config = create_mock_vllm_config(hf_config) + metrics = MLAAttentionMetrics.from_vllm_config(vllm_config) + + # Single decode token with 1024 context + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=1024, is_prefill=False + ) + + write_breakdown = metrics.get_write_bytes_breakdown(ctx, per_gpu=False) + + # KV cache write should be 1 * (512 + 64) * cache_byte_size * 1 layer + # = 576 * 2 = 1152 bytes (for bfloat16 cache) + kv_compressed_dim = 512 + 64 # kv_lora_rank + qk_rope_head_dim + expected_kv_cache_write = 1 * kv_compressed_dim * 2 * 1 # T * dim * bytes * L + assert write_breakdown["kv_cache"] == expected_kv_cache_write + + # Verify read bytes include compressed KV cache reads for context + read_breakdown = metrics.get_read_bytes_breakdown(ctx, per_gpu=False) + assert "attn_input" in read_breakdown + assert read_breakdown["attn_input"] > 0 + + +def test_mla_attention_metrics_prefill(): + """Test MLA prefill metrics account for low-rank Q and KV projections.""" + hf_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=1, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + vllm_config = create_mock_vllm_config(hf_config) + metrics = MLAAttentionMetrics.from_vllm_config(vllm_config) + + ctx = ExecutionContext.from_single_request( + num_tokens=2048, context_len=2048, is_prefill=True + ) + + flops_breakdown = metrics.get_num_flops_breakdown(ctx, per_gpu=False) + + # Should have two-stage Q projection (q_a and q_b) + assert "q_a_proj" in flops_breakdown + assert "q_b_proj" in flops_breakdown + assert "q_proj" not in flops_breakdown # Since q_lora_rank is not None + + # Should have KV projections + assert "kv_a_proj" in flops_breakdown + assert "kv_b_proj" in flops_breakdown + + # Should have attention and output + assert "attn_qk" in flops_breakdown + assert "attn_av" in flops_breakdown + assert "out_proj" in flops_breakdown + + # Verify q_a_proj: 2 * T * D * q_lora_rank * L + expected_q_a = 2 * 2048 * 7168 * 1536 * 1 + assert flops_breakdown["q_a_proj"] == expected_q_a + + # Verify kv_a_proj: 2 * T * D * (kv_lora_rank + qk_rope_head_dim) * L + expected_kv_a = 2 * 2048 * 7168 * (512 + 64) * 1 + assert flops_breakdown["kv_a_proj"] == expected_kv_a + + +def test_mla_kv_cache_vs_standard_attention(): + """Test MLA KV cache writes are dramatically smaller than standard MHA.""" + # MLA config (DeepSeek-V3 style) + mla_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=1, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ) + mla_vllm_config = create_mock_vllm_config(mla_config) + mla_metrics = MLAAttentionMetrics.from_vllm_config(mla_vllm_config) + + # Standard MHA config with same num_heads and head_dim + standard_config = Qwen3Config( + hidden_size=7168, + num_attention_heads=128, + num_key_value_heads=128, # MHA: same as num_heads + num_hidden_layers=1, + head_dim=128, + ) + standard_vllm_config = create_mock_vllm_config(standard_config) + standard_metrics = AttentionMetrics.from_vllm_config(standard_vllm_config) + + # Compare KV cache write for 100 tokens + ctx = ExecutionContext.from_single_request( + num_tokens=100, context_len=100, is_prefill=True + ) + + mla_write = mla_metrics.get_write_bytes_breakdown(ctx, per_gpu=False) + standard_write = standard_metrics.get_write_bytes_breakdown(ctx, per_gpu=False) + + # MLA: T * (kv_lora_rank + qk_rope_head_dim) * cache_bytes * L + # = 100 * 576 * 2 * 1 = 115,200 + mla_kv_cache = mla_write["kv_cache"] + + # Standard: 2 * T * num_kv_heads * head_dim * cache_bytes * L + # = 2 * 100 * 128 * 128 * 2 * 1 = 6,553,600 + standard_kv_cache = standard_write["kv_cache"] + + # MLA KV cache should be dramatically smaller (about 57x) + assert mla_kv_cache < standard_kv_cache + ratio = standard_kv_cache / mla_kv_cache + assert ratio > 50 # Should be ~56.9x + + +def test_mla_per_gpu_with_tensor_parallelism(): + """Test MLA metrics with tensor parallelism.""" + hf_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=8, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + + # Test with TP=8 + vllm_config = create_mock_vllm_config(hf_config, tensor_parallel_size=8) + metrics = MLAAttentionMetrics.from_vllm_config(vllm_config) + + ctx = ExecutionContext.from_single_request( + num_tokens=64, context_len=1024, is_prefill=True + ) + + global_flops = metrics.get_num_flops(ctx, per_gpu=False) + per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True) + + # Both should be positive + assert global_flops > 0 + assert per_gpu_flops > 0 + # Global should exceed per-GPU + assert global_flops > per_gpu_flops + + +def test_mla_per_gpu_with_pipeline_parallelism(): + """Test MLA metrics with pipeline parallelism.""" + hf_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=16, # Divisible by PP + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + + vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=4) + metrics = MLAAttentionMetrics.from_vllm_config(vllm_config) + + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=512, is_prefill=False + ) + + global_flops = metrics.get_num_flops(ctx, per_gpu=False) + per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True) + + # With PP=4, layers are divided by 4 + assert global_flops == 4 * per_gpu_flops + + +def test_mla_model_metrics_excludes_standard_attention(): + """Test that ModelMetrics uses MLAAttentionMetrics, not AttentionMetrics, + for DeepSeek MLA models.""" + hf_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=4, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + vllm_config = create_mock_vllm_config(hf_config) + model_metrics = ModelMetrics(vllm_config) + + # Should have MLAAttentionMetrics but NOT standard AttentionMetrics + component_types = [m.component_type() for m in model_metrics.metrics] + assert "mla_attn" in component_types + assert "attn" not in component_types + + # Should still have FFN and unembed + assert "ffn" in component_types + assert "unembed" in component_types + + # Breakdowns should work end-to-end + ctx = ExecutionContext.from_single_request( + num_tokens=100, context_len=512, is_prefill=True + ) + total_flops = model_metrics.get_num_flops(ctx) + breakdown = model_metrics.get_num_flops_breakdown(ctx) + assert total_flops == sum(breakdown.values()) + assert total_flops > 0 + + # Verify MLA-specific keys in breakdown + assert any(k.startswith("mla_attn.") for k in breakdown) + assert not any(k.startswith("attn.") for k in breakdown) + + +def test_standard_attention_still_works_for_non_mla(): + """Regression test: non-MLA models still use standard AttentionMetrics.""" + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_hidden_layers=12, + vocab_size=32000, + intermediate_size=8192, + ) + vllm_config = create_mock_vllm_config(hf_config) + model_metrics = ModelMetrics(vllm_config) + + component_types = [m.component_type() for m in model_metrics.metrics] + assert "attn" in component_types + assert "mla_attn" not in component_types + + ctx = ExecutionContext.from_single_request( + num_tokens=100, context_len=512, is_prefill=True + ) + total_flops = model_metrics.get_num_flops(ctx) + assert total_flops > 0 + + +def test_mla_attention_scaling_with_layers(): + """Test that MLA attention metrics scale proportionally with layers.""" + base_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=8, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + double_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=16, # Double layers + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + + base_vllm = create_mock_vllm_config(base_config) + double_vllm = create_mock_vllm_config(double_config) + + base_metrics = MLAAttentionMetrics.from_vllm_config(base_vllm) + double_metrics = MLAAttentionMetrics.from_vllm_config(double_vllm) + + ctx = ExecutionContext.from_single_request( + num_tokens=100, context_len=512, is_prefill=True + ) + + # All metrics should double with double layers + assert double_metrics.get_num_flops(ctx) == 2 * base_metrics.get_num_flops(ctx) + assert double_metrics.get_read_bytes(ctx) == 2 * base_metrics.get_read_bytes(ctx) + assert double_metrics.get_write_bytes(ctx) == 2 * base_metrics.get_write_bytes(ctx) diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index 38135b9b15..3336fca606 100644 --- a/vllm/v1/metrics/perf.py +++ b/vllm/v1/metrics/perf.py @@ -396,6 +396,20 @@ class AttentionQuantizationConfigParser(Parser): return args +class AttentionDetectionParser(Parser): + """ + Prevents standard AttentionMetrics from being instantiated for MLA models. + MLA models should use MLAAttentionMetrics instead. + """ + + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + if vllm_config.model_config.is_deepseek_mla: + raise InvalidComponent( + "Model uses MLA attention; use MLAAttentionMetrics instead" + ) + return args + + class AttentionMetrics(ComponentMetrics): # From BaseConfigParser num_hidden_layers: int = Field(..., gt=0) @@ -423,6 +437,7 @@ class AttentionMetrics(ComponentMetrics): @classmethod def get_parser(cls) -> ParserChain: return ParserChain( + AttentionDetectionParser(), BaseConfigParser(), BaseAttentionConfigParser(), AttentionQuantizationConfigParser(), @@ -525,6 +540,276 @@ class AttentionMetrics(ComponentMetrics): } +#### MLA Attention #### + + +class MLADetectionParser(Parser): + """ + Validates that the model uses MLA attention. + Raises InvalidComponent if the model does not use MLA, + so MLAAttentionMetrics is silently skipped for non-MLA models. + """ + + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + if not vllm_config.model_config.is_deepseek_mla: + raise InvalidComponent("Model does not use MLA attention") + return args + + +class MLAConfigParser(Parser): + """ + Parses MLA-specific configuration fields. + Provides: kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, + v_head_dim, q_lora_rank + """ + + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + model_config = vllm_config.model_config + cfg = model_config.hf_text_config + + args.kv_lora_rank = get_required(cfg, "kv_lora_rank") + args.qk_nope_head_dim = get_required(cfg, "qk_nope_head_dim") + args.qk_rope_head_dim = get_required(cfg, "qk_rope_head_dim") + args.v_head_dim = get_required(cfg, "v_head_dim") + args.q_lora_rank = getattr(cfg, "q_lora_rank", None) + + model_dtype = vllm_config.model_config.dtype + cache_dtype = vllm_config.cache_config.cache_dtype + kv_cache_torch_dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype) + args.cache_byte_size = get_dtype_size(kv_cache_torch_dtype) + + return args + + +class MLAAttentionMetrics(ComponentMetrics): + """ + Performance metrics for Multi-Latent Attention (MLA) layers. + + MLA uses a compressed latent representation for KV cache: + - KV cache stores a single compressed vector of size + (kv_lora_rank + qk_rope_head_dim) per token per layer, + instead of 2 * num_kv_heads * head_dim as in standard MHA/GQA. + - Q path uses optional low-rank compression: + h -> q_lora_rank -> num_heads * qk_head_dim + - KV path: h -> (kv_lora_rank + qk_rope_head_dim), + then kv_lora_rank -> num_heads * (qk_nope_head_dim + v_head_dim) + + Used by DeepSeek-V2, DeepSeek-V3, DeepSeek-R1, and similar models. + """ + + # From BaseConfigParser + num_hidden_layers: int = Field(..., gt=0) + hidden_size: int = Field(..., gt=0) + num_attention_heads: int = Field(..., gt=0) + activation_byte_size: int = Field(..., gt=0) + tp_size: int = Field(..., gt=0) + pp_size: int = Field(..., gt=0) + + # From BaseConfigParser, can be overridden by AttentionQuantizationConfigParser + weight_byte_size: int | float = Field(..., gt=0) + + # From MLAConfigParser + kv_lora_rank: int = Field(..., gt=0) + qk_nope_head_dim: int = Field(..., gt=0) + qk_rope_head_dim: int = Field(..., gt=0) + v_head_dim: int = Field(..., gt=0) + q_lora_rank: int | None = Field(None) + cache_byte_size: int = Field(..., gt=0) + + @classmethod + def component_type(cls) -> str: + return "mla_attn" + + @classmethod + def get_parser(cls) -> ParserChain: + return ParserChain( + MLADetectionParser(), + BaseConfigParser(), + MLAConfigParser(), + AttentionQuantizationConfigParser(), + ) + + def get_num_flops_breakdown( + self, ctx: ExecutionContext, per_gpu: bool = True + ) -> dict[str, int]: + """Calculate flops breakdown for MLA attention layers. + + MLA projection structure: + - Q path: h -> q_lora_rank -> num_heads * qk_head_dim + (or h -> num_heads * qk_head_dim if q_lora_rank is None) + - KV path: h -> (kv_lora_rank + qk_rope_head_dim), + then kv_lora_rank -> num_heads * (qk_nope_head_dim + v_head_dim) + - Attention: Q @ K^T and attn @ V + - Output: num_heads * v_head_dim -> h + """ + L = self.num_hidden_layers + D = self.hidden_size + q = self.num_attention_heads + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + v_d = self.v_head_dim + c = self.kv_lora_rank + r = self.qk_rope_head_dim + q_rank = self.q_lora_rank + + T = ctx.total_num_tokens() + TC = ctx.total_token_context_product() + + if per_gpu: + L //= self.pp_size + q = max(1, q // self.tp_size) + + flops: dict[str, int] = {} + + # Q projection + if q_rank is not None: + # Two-stage: h -> q_lora_rank -> num_heads * qk_head_dim + flops["q_a_proj"] = 2 * T * D * q_rank * L + flops["q_b_proj"] = 2 * T * q_rank * q * qk_head_dim * L + else: + # Direct: h -> num_heads * qk_head_dim + flops["q_proj"] = 2 * T * D * q * qk_head_dim * L + + # KV projection (always compressed, shared across heads) + # kv_a: h -> (kv_lora_rank + qk_rope_head_dim) [replicated] + flops["kv_a_proj"] = 2 * T * D * (c + r) * L + # kv_b: kv_lora_rank -> num_heads * (qk_nope + v_head_dim) + flops["kv_b_proj"] = 2 * T * c * q * (self.qk_nope_head_dim + v_d) * L + + # Attention core + flops["attn_qk"] = 2 * q * TC * qk_head_dim * L + flops["attn_av"] = 2 * q * TC * v_d * L + + # Output projection: num_heads * v_head_dim -> h + flops["out_proj"] = 2 * T * q * v_d * D * L + + return flops + + def get_read_bytes_breakdown( + self, ctx: ExecutionContext, per_gpu: bool = True + ) -> dict[str, int]: + """Calculate read memory traffic for MLA attention layers.""" + L = self.num_hidden_layers + D = self.hidden_size + q = self.num_attention_heads + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + v_d = self.v_head_dim + c = self.kv_lora_rank + r = self.qk_rope_head_dim + q_rank = self.q_lora_rank + + T = ctx.total_num_tokens() + # Compressed KV cache size per token + kv_compressed_dim = c + r + + if per_gpu: + L //= self.pp_size + q = max(1, q // self.tp_size) + + read_bytes: dict[str, int] = {} + + # Q projection weight + input reads + if q_rank is not None: + read_bytes["q_a_input"] = T * D * self.activation_byte_size * L + read_bytes["q_a_weight"] = int(D * q_rank * self.weight_byte_size * L) + read_bytes["q_b_input"] = T * q_rank * self.activation_byte_size * L + read_bytes["q_b_weight"] = int( + q_rank * q * qk_head_dim * self.weight_byte_size * L + ) + else: + read_bytes["q_input"] = T * D * self.activation_byte_size * L + read_bytes["q_weight"] = int( + D * q * qk_head_dim * self.weight_byte_size * L + ) + + # KV projection weight + input reads + # kv_a is replicated (not TP-sharded) + read_bytes["kv_a_input"] = T * D * self.activation_byte_size * L + read_bytes["kv_a_weight"] = int( + D * kv_compressed_dim * self.weight_byte_size * L + ) + # kv_b is TP-sharded along heads + read_bytes["kv_b_input"] = T * c * self.activation_byte_size * L + read_bytes["kv_b_weight"] = int( + c * q * (self.qk_nope_head_dim + v_d) * self.weight_byte_size * L + ) + + # Attention input reads + # Prefill: read Q activations + K,V from kv_b_proj output + if ctx.prefill_num_tokens > 0: + read_bytes["attn_input"] = ( + ctx.prefill_num_tokens * q * qk_head_dim * self.activation_byte_size * L + + ctx.prefill_context_len + * q + * (qk_head_dim + v_d) + * self.activation_byte_size + * L + ) + + # Decode: read Q activations + read compressed KV from cache + if ctx.decode_num_tokens > 0: + read_bytes["attn_input"] = read_bytes.get("attn_input", 0) + ( + ctx.decode_num_tokens * q * qk_head_dim * self.activation_byte_size * L + + ctx.decode_context_len * kv_compressed_dim * self.cache_byte_size * L + ) + + # Output projection reads + read_bytes["out_input"] = T * q * v_d * self.activation_byte_size * L + read_bytes["out_weight"] = int(q * v_d * D * self.weight_byte_size * L) + + return read_bytes + + def get_write_bytes_breakdown( + self, ctx: ExecutionContext, per_gpu: bool = True + ) -> dict[str, int]: + """Calculate write memory traffic for MLA attention layers.""" + L = self.num_hidden_layers + D = self.hidden_size + q = self.num_attention_heads + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + v_d = self.v_head_dim + c = self.kv_lora_rank + r = self.qk_rope_head_dim + q_rank = self.q_lora_rank + + T = ctx.total_num_tokens() + kv_compressed_dim = c + r + + if per_gpu: + L //= self.pp_size + q = max(1, q // self.tp_size) + + write_bytes: dict[str, int] = {} + + # Q projection outputs + if q_rank is not None: + write_bytes["q_a_output"] = T * q_rank * self.activation_byte_size * L + write_bytes["q_b_output"] = ( + T * q * qk_head_dim * self.activation_byte_size * L + ) + else: + write_bytes["q_output"] = ( + T * q * qk_head_dim * self.activation_byte_size * L + ) + + # KV projection outputs + write_bytes["kv_a_output"] = ( + T * kv_compressed_dim * self.activation_byte_size * L + ) + write_bytes["kv_b_output"] = ( + T * q * (self.qk_nope_head_dim + v_d) * self.activation_byte_size * L + ) + + # KV cache write: one compressed vector per token + # (kv_lora_rank + qk_rope_head_dim) instead of + # 2 * num_kv_heads * head_dim in standard MHA + write_bytes["kv_cache"] = T * kv_compressed_dim * self.cache_byte_size * L + + # Output projection + write_bytes["out_output"] = T * D * self.activation_byte_size * L + + return write_bytes + + #### Ffn ####