iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[V1][Metrics] Add MLA attention metrics for DeepSeek MFU estimation (#39457)

Browse Source 
Signed-off-by: Thillai Chithambaram <thillaichithambaram.a@gmail.com>
Co-authored-by: Mark McLoughlin <markmc@redhat.com>
This commit is contained in:
Thillai Chithambaram
2026-06-12 14:28:40 +01:00
committed by GitHub
co-authored by Mark McLoughlin
parent b7f9b6ab27
commit 4171ae406c
2 changed files with 600 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/metrics/test_perf_metrics.py
+315
View File
@@ -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)
vllm/v1/metrics/perf.py
+285
View File
@@ -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 ####