@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4171ae406c

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

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

Signed-off-by: Thillai Chithambaram <thillaichithambaram.a@gmail.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com>

Thillai Chithambaram committed Jun 12, 2026 at 09:28 UTC 4171ae406cdcec1c9952ed6fc00cd9ac91e3e342

2 files changed +600

tests/v1/metrics/test_perf_metrics.py

+315

@@ -28,6 +28,7 @@ from vllm.v1.metrics.perf import (

28 ExecutionContext,

29 FfnMetrics,

30 InvalidComponent,

31 + MLAAttentionMetrics,

32 ModelMetrics,

33 ParsedArgs,

34 UnembedMetrics,

@@ -1021,3 +1022,317 @@ def test_quantized_model_metrics_aggregation():

1022

1023 assert total_flops > 0

1024 assert total_flops == sum(breakdown.values())

1025 +

1026 +

1027 +#### MLA Attention Tests ####

1028 +

1029 +

1030 +def test_mla_config_parser():

1031 + """Test MLAConfigParser extracts MLA-specific fields from DeepseekV3Config."""

1032 + hf_config = DeepseekV3Config(

1033 + hidden_size=7168,

1034 + num_attention_heads=128,

1035 + num_hidden_layers=61,

1036 + kv_lora_rank=512,

1037 + qk_nope_head_dim=128,

1038 + qk_rope_head_dim=64,

1039 + v_head_dim=128,

1040 + q_lora_rank=1536,

1041 + )

1042 + vllm_config = create_mock_vllm_config(hf_config)

1043 +

1044 + parser_chain = MLAAttentionMetrics.get_parser()

1045 + result = parser_chain.parse(vllm_config)

1046 +

1047 + assert result.kv_lora_rank == 512

1048 + assert result.qk_nope_head_dim == 128

1049 + assert result.qk_rope_head_dim == 64

1050 + assert result.v_head_dim == 128

1051 + assert result.q_lora_rank == 1536

1052 + assert result.num_attention_heads == 128

1053 + assert result.hidden_size == 7168

1054 +

1055 +

1056 +def test_mla_attention_metrics_decode():

1057 + """Test MLA decode metrics use compressed KV cache, not standard head_dim."""

1058 + hf_config = DeepseekV3Config(

1059 + hidden_size=7168,

1060 + num_attention_heads=128,

1061 + num_hidden_layers=1,

1062 + kv_lora_rank=512,

1063 + qk_nope_head_dim=128,

1064 + qk_rope_head_dim=64,

1065 + v_head_dim=128,

1066 + q_lora_rank=1536,

1067 + )

1068 + vllm_config = create_mock_vllm_config(hf_config)

1069 + metrics = MLAAttentionMetrics.from_vllm_config(vllm_config)

1070 +

1071 + # Single decode token with 1024 context

1072 + ctx = ExecutionContext.from_single_request(

1073 + num_tokens=1, context_len=1024, is_prefill=False

1074 + )

1075 +

1076 + write_breakdown = metrics.get_write_bytes_breakdown(ctx, per_gpu=False)

1077 +

1078 + # KV cache write should be 1 * (512 + 64) * cache_byte_size * 1 layer

1079 + # = 576 * 2 = 1152 bytes (for bfloat16 cache)

1080 + kv_compressed_dim = 512 + 64 # kv_lora_rank + qk_rope_head_dim

1081 + expected_kv_cache_write = 1 * kv_compressed_dim * 2 * 1 # T * dim * bytes * L

1082 + assert write_breakdown["kv_cache"] == expected_kv_cache_write

1083 +

1084 + # Verify read bytes include compressed KV cache reads for context

1085 + read_breakdown = metrics.get_read_bytes_breakdown(ctx, per_gpu=False)

1086 + assert "attn_input" in read_breakdown

1087 + assert read_breakdown["attn_input"] > 0

1088 +

1089 +

1090 +def test_mla_attention_metrics_prefill():

1091 + """Test MLA prefill metrics account for low-rank Q and KV projections."""

1092 + hf_config = DeepseekV3Config(

1093 + hidden_size=7168,

1094 + num_attention_heads=128,

1095 + num_hidden_layers=1,

1096 + kv_lora_rank=512,

1097 + qk_nope_head_dim=128,

1098 + qk_rope_head_dim=64,

1099 + v_head_dim=128,

1100 + q_lora_rank=1536,

1101 + )

1102 + vllm_config = create_mock_vllm_config(hf_config)

1103 + metrics = MLAAttentionMetrics.from_vllm_config(vllm_config)

1104 +

1105 + ctx = ExecutionContext.from_single_request(

1106 + num_tokens=2048, context_len=2048, is_prefill=True

1107 + )

1108 +

1109 + flops_breakdown = metrics.get_num_flops_breakdown(ctx, per_gpu=False)

1110 +

1111 + # Should have two-stage Q projection (q_a and q_b)

1112 + assert "q_a_proj" in flops_breakdown

1113 + assert "q_b_proj" in flops_breakdown

1114 + assert "q_proj" not in flops_breakdown # Since q_lora_rank is not None

1115 +

1116 + # Should have KV projections

1117 + assert "kv_a_proj" in flops_breakdown

1118 + assert "kv_b_proj" in flops_breakdown

1119 +

1120 + # Should have attention and output

1121 + assert "attn_qk" in flops_breakdown

1122 + assert "attn_av" in flops_breakdown

1123 + assert "out_proj" in flops_breakdown

1124 +

1125 + # Verify q_a_proj: 2 * T * D * q_lora_rank * L

1126 + expected_q_a = 2 * 2048 * 7168 * 1536 * 1

1127 + assert flops_breakdown["q_a_proj"] == expected_q_a

1128 +

1129 + # Verify kv_a_proj: 2 * T * D * (kv_lora_rank + qk_rope_head_dim) * L

1130 + expected_kv_a = 2 * 2048 * 7168 * (512 + 64) * 1

1131 + assert flops_breakdown["kv_a_proj"] == expected_kv_a

1132 +

1133 +

1134 +def test_mla_kv_cache_vs_standard_attention():

1135 + """Test MLA KV cache writes are dramatically smaller than standard MHA."""

1136 + # MLA config (DeepSeek-V3 style)

1137 + mla_config = DeepseekV3Config(

1138 + hidden_size=7168,

1139 + num_attention_heads=128,

1140 + num_hidden_layers=1,

1141 + kv_lora_rank=512,

1142 + qk_nope_head_dim=128,

1143 + qk_rope_head_dim=64,

1144 + v_head_dim=128,

1145 + )

1146 + mla_vllm_config = create_mock_vllm_config(mla_config)

1147 + mla_metrics = MLAAttentionMetrics.from_vllm_config(mla_vllm_config)

1148 +

1149 + # Standard MHA config with same num_heads and head_dim

1150 + standard_config = Qwen3Config(

1151 + hidden_size=7168,

1152 + num_attention_heads=128,

1153 + num_key_value_heads=128, # MHA: same as num_heads

1154 + num_hidden_layers=1,

1155 + head_dim=128,

1156 + )

1157 + standard_vllm_config = create_mock_vllm_config(standard_config)

1158 + standard_metrics = AttentionMetrics.from_vllm_config(standard_vllm_config)

1159 +

1160 + # Compare KV cache write for 100 tokens

1161 + ctx = ExecutionContext.from_single_request(

1162 + num_tokens=100, context_len=100, is_prefill=True

1163 + )

1164 +

1165 + mla_write = mla_metrics.get_write_bytes_breakdown(ctx, per_gpu=False)

1166 + standard_write = standard_metrics.get_write_bytes_breakdown(ctx, per_gpu=False)

1167 +

1168 + # MLA: T * (kv_lora_rank + qk_rope_head_dim) * cache_bytes * L

1169 + # = 100 * 576 * 2 * 1 = 115,200

1170 + mla_kv_cache = mla_write["kv_cache"]

1171 +

1172 + # Standard: 2 * T * num_kv_heads * head_dim * cache_bytes * L

1173 + # = 2 * 100 * 128 * 128 * 2 * 1 = 6,553,600

1174 + standard_kv_cache = standard_write["kv_cache"]

1175 +

1176 + # MLA KV cache should be dramatically smaller (about 57x)

1177 + assert mla_kv_cache < standard_kv_cache

1178 + ratio = standard_kv_cache / mla_kv_cache

1179 + assert ratio > 50 # Should be ~56.9x

1180 +

1181 +

1182 +def test_mla_per_gpu_with_tensor_parallelism():

1183 + """Test MLA metrics with tensor parallelism."""

1184 + hf_config = DeepseekV3Config(

1185 + hidden_size=7168,

1186 + num_attention_heads=128,

1187 + num_hidden_layers=8,

1188 + kv_lora_rank=512,

1189 + qk_nope_head_dim=128,

1190 + qk_rope_head_dim=64,

1191 + v_head_dim=128,

1192 + q_lora_rank=1536,

1193 + )

1194 +

1195 + # Test with TP=8

1196 + vllm_config = create_mock_vllm_config(hf_config, tensor_parallel_size=8)

1197 + metrics = MLAAttentionMetrics.from_vllm_config(vllm_config)

1198 +

1199 + ctx = ExecutionContext.from_single_request(

1200 + num_tokens=64, context_len=1024, is_prefill=True

1201 + )

1202 +

1203 + global_flops = metrics.get_num_flops(ctx, per_gpu=False)

1204 + per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True)

1205 +

1206 + # Both should be positive

1207 + assert global_flops > 0

1208 + assert per_gpu_flops > 0

1209 + # Global should exceed per-GPU

1210 + assert global_flops > per_gpu_flops

1211 +

1212 +

1213 +def test_mla_per_gpu_with_pipeline_parallelism():

1214 + """Test MLA metrics with pipeline parallelism."""

1215 + hf_config = DeepseekV3Config(

1216 + hidden_size=7168,

1217 + num_attention_heads=128,

1218 + num_hidden_layers=16, # Divisible by PP

1219 + kv_lora_rank=512,

1220 + qk_nope_head_dim=128,

1221 + qk_rope_head_dim=64,

1222 + v_head_dim=128,

1223 + q_lora_rank=1536,

1224 + )

1225 +

1226 + vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=4)

1227 + metrics = MLAAttentionMetrics.from_vllm_config(vllm_config)

1228 +

1229 + ctx = ExecutionContext.from_single_request(

1230 + num_tokens=1, context_len=512, is_prefill=False

1231 + )

1232 +

1233 + global_flops = metrics.get_num_flops(ctx, per_gpu=False)

1234 + per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True)

1235 +

1236 + # With PP=4, layers are divided by 4

1237 + assert global_flops == 4 * per_gpu_flops

1238 +

1239 +

1240 +def test_mla_model_metrics_excludes_standard_attention():

1241 + """Test that ModelMetrics uses MLAAttentionMetrics, not AttentionMetrics,

1242 + for DeepSeek MLA models."""

1243 + hf_config = DeepseekV3Config(

1244 + hidden_size=7168,

1245 + num_attention_heads=128,

1246 + num_hidden_layers=4,

1247 + kv_lora_rank=512,

1248 + qk_nope_head_dim=128,

1249 + qk_rope_head_dim=64,

1250 + v_head_dim=128,

1251 + q_lora_rank=1536,

1252 + )

1253 + vllm_config = create_mock_vllm_config(hf_config)

1254 + model_metrics = ModelMetrics(vllm_config)

1255 +

1256 + # Should have MLAAttentionMetrics but NOT standard AttentionMetrics

1257 + component_types = [m.component_type() for m in model_metrics.metrics]

1258 + assert "mla_attn" in component_types

1259 + assert "attn" not in component_types

1260 +

1261 + # Should still have FFN and unembed

1262 + assert "ffn" in component_types

1263 + assert "unembed" in component_types

1264 +

1265 + # Breakdowns should work end-to-end

1266 + ctx = ExecutionContext.from_single_request(

1267 + num_tokens=100, context_len=512, is_prefill=True

1268 + )

1269 + total_flops = model_metrics.get_num_flops(ctx)

1270 + breakdown = model_metrics.get_num_flops_breakdown(ctx)

1271 + assert total_flops == sum(breakdown.values())

1272 + assert total_flops > 0

1273 +

1274 + # Verify MLA-specific keys in breakdown

1275 + assert any(k.startswith("mla_attn.") for k in breakdown)

1276 + assert not any(k.startswith("attn.") for k in breakdown)

1277 +

1278 +

1279 +def test_standard_attention_still_works_for_non_mla():

1280 + """Regression test: non-MLA models still use standard AttentionMetrics."""

1281 + hf_config = Qwen3Config(

1282 + hidden_size=2048,

1283 + num_attention_heads=16,

1284 + num_hidden_layers=12,

1285 + vocab_size=32000,

1286 + intermediate_size=8192,

1287 + )

1288 + vllm_config = create_mock_vllm_config(hf_config)

1289 + model_metrics = ModelMetrics(vllm_config)

1290 +

1291 + component_types = [m.component_type() for m in model_metrics.metrics]

1292 + assert "attn" in component_types

1293 + assert "mla_attn" not in component_types

1294 +

1295 + ctx = ExecutionContext.from_single_request(

1296 + num_tokens=100, context_len=512, is_prefill=True

1297 + )

1298 + total_flops = model_metrics.get_num_flops(ctx)

1299 + assert total_flops > 0

1300 +

1301 +

1302 +def test_mla_attention_scaling_with_layers():

1303 + """Test that MLA attention metrics scale proportionally with layers."""

1304 + base_config = DeepseekV3Config(

1305 + hidden_size=7168,

1306 + num_attention_heads=128,

1307 + num_hidden_layers=8,

1308 + kv_lora_rank=512,

1309 + qk_nope_head_dim=128,

1310 + qk_rope_head_dim=64,

1311 + v_head_dim=128,

1312 + q_lora_rank=1536,

1313 + )

1314 + double_config = DeepseekV3Config(

1315 + hidden_size=7168,

1316 + num_attention_heads=128,

1317 + num_hidden_layers=16, # Double layers

1318 + kv_lora_rank=512,

1319 + qk_nope_head_dim=128,

1320 + qk_rope_head_dim=64,

1321 + v_head_dim=128,

1322 + q_lora_rank=1536,

1323 + )

1324 +

1325 + base_vllm = create_mock_vllm_config(base_config)

1326 + double_vllm = create_mock_vllm_config(double_config)

1327 +

1328 + base_metrics = MLAAttentionMetrics.from_vllm_config(base_vllm)

1329 + double_metrics = MLAAttentionMetrics.from_vllm_config(double_vllm)

1330 +

1331 + ctx = ExecutionContext.from_single_request(

1332 + num_tokens=100, context_len=512, is_prefill=True

1333 + )

1334 +

1335 + # All metrics should double with double layers

1336 + assert double_metrics.get_num_flops(ctx) == 2 * base_metrics.get_num_flops(ctx)

1337 + assert double_metrics.get_read_bytes(ctx) == 2 * base_metrics.get_read_bytes(ctx)

1338 + assert double_metrics.get_write_bytes(ctx) == 2 * base_metrics.get_write_bytes(ctx)

vllm/v1/metrics/perf.py

+285

@@ -396,6 +396,20 @@ class AttentionQuantizationConfigParser(Parser):

396 return args

397

398

399 +class AttentionDetectionParser(Parser):

400 + """

401 + Prevents standard AttentionMetrics from being instantiated for MLA models.

402 + MLA models should use MLAAttentionMetrics instead.

403 + """

404 +

405 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:

406 + if vllm_config.model_config.is_deepseek_mla:

407 + raise InvalidComponent(

408 + "Model uses MLA attention; use MLAAttentionMetrics instead"

409 + )

410 + return args

411 +

412 +

413 class AttentionMetrics(ComponentMetrics):

414 # From BaseConfigParser

415 num_hidden_layers: int = Field(..., gt=0)

@@ -423,6 +437,7 @@ class AttentionMetrics(ComponentMetrics):

437 @classmethod

438 def get_parser(cls) -> ParserChain:

439 return ParserChain(

440 + AttentionDetectionParser(),

441 BaseConfigParser(),

442 BaseAttentionConfigParser(),

443 AttentionQuantizationConfigParser(),

@@ -525,6 +540,276 @@ class AttentionMetrics(ComponentMetrics):

540 }

541

542

543 +#### MLA Attention ####

544 +

545 +

546 +class MLADetectionParser(Parser):

547 + """

548 + Validates that the model uses MLA attention.

549 + Raises InvalidComponent if the model does not use MLA,

550 + so MLAAttentionMetrics is silently skipped for non-MLA models.

551 + """

552 +

553 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:

554 + if not vllm_config.model_config.is_deepseek_mla:

555 + raise InvalidComponent("Model does not use MLA attention")

556 + return args

557 +

558 +

559 +class MLAConfigParser(Parser):

560 + """

561 + Parses MLA-specific configuration fields.

562 + Provides: kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim,

563 + v_head_dim, q_lora_rank

564 + """

565 +

566 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:

567 + model_config = vllm_config.model_config

568 + cfg = model_config.hf_text_config

569 +

570 + args.kv_lora_rank = get_required(cfg, "kv_lora_rank")

571 + args.qk_nope_head_dim = get_required(cfg, "qk_nope_head_dim")

572 + args.qk_rope_head_dim = get_required(cfg, "qk_rope_head_dim")

573 + args.v_head_dim = get_required(cfg, "v_head_dim")

574 + args.q_lora_rank = getattr(cfg, "q_lora_rank", None)

575 +

576 + model_dtype = vllm_config.model_config.dtype

577 + cache_dtype = vllm_config.cache_config.cache_dtype

578 + kv_cache_torch_dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)

579 + args.cache_byte_size = get_dtype_size(kv_cache_torch_dtype)

580 +

581 + return args

582 +

583 +

584 +class MLAAttentionMetrics(ComponentMetrics):

585 + """

586 + Performance metrics for Multi-Latent Attention (MLA) layers.

587 +

588 + MLA uses a compressed latent representation for KV cache:

589 + - KV cache stores a single compressed vector of size

590 + (kv_lora_rank + qk_rope_head_dim) per token per layer,

591 + instead of 2 * num_kv_heads * head_dim as in standard MHA/GQA.

592 + - Q path uses optional low-rank compression:

593 + h -> q_lora_rank -> num_heads * qk_head_dim

594 + - KV path: h -> (kv_lora_rank + qk_rope_head_dim),

595 + then kv_lora_rank -> num_heads * (qk_nope_head_dim + v_head_dim)

596 +

597 + Used by DeepSeek-V2, DeepSeek-V3, DeepSeek-R1, and similar models.

598 + """

599 +

600 + # From BaseConfigParser

601 + num_hidden_layers: int = Field(..., gt=0)

602 + hidden_size: int = Field(..., gt=0)

603 + num_attention_heads: int = Field(..., gt=0)

604 + activation_byte_size: int = Field(..., gt=0)

605 + tp_size: int = Field(..., gt=0)

606 + pp_size: int = Field(..., gt=0)

607 +

608 + # From BaseConfigParser, can be overridden by AttentionQuantizationConfigParser

609 + weight_byte_size: int | float = Field(..., gt=0)

610 +

611 + # From MLAConfigParser

612 + kv_lora_rank: int = Field(..., gt=0)

613 + qk_nope_head_dim: int = Field(..., gt=0)

614 + qk_rope_head_dim: int = Field(..., gt=0)

615 + v_head_dim: int = Field(..., gt=0)

616 + q_lora_rank: int | None = Field(None)

617 + cache_byte_size: int = Field(..., gt=0)

618 +

619 + @classmethod

620 + def component_type(cls) -> str:

621 + return "mla_attn"

622 +

623 + @classmethod

624 + def get_parser(cls) -> ParserChain:

625 + return ParserChain(

626 + MLADetectionParser(),

627 + BaseConfigParser(),

628 + MLAConfigParser(),

629 + AttentionQuantizationConfigParser(),

630 + )

631 +

632 + def get_num_flops_breakdown(

633 + self, ctx: ExecutionContext, per_gpu: bool = True

634 + ) -> dict[str, int]:

635 + """Calculate flops breakdown for MLA attention layers.

636 +

637 + MLA projection structure:

638 + - Q path: h -> q_lora_rank -> num_heads * qk_head_dim

639 + (or h -> num_heads * qk_head_dim if q_lora_rank is None)

640 + - KV path: h -> (kv_lora_rank + qk_rope_head_dim),

641 + then kv_lora_rank -> num_heads * (qk_nope_head_dim + v_head_dim)

642 + - Attention: Q @ K^T and attn @ V

643 + - Output: num_heads * v_head_dim -> h

644 + """

645 + L = self.num_hidden_layers

646 + D = self.hidden_size

647 + q = self.num_attention_heads

648 + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

649 + v_d = self.v_head_dim

650 + c = self.kv_lora_rank

651 + r = self.qk_rope_head_dim

652 + q_rank = self.q_lora_rank

653 +

654 + T = ctx.total_num_tokens()

655 + TC = ctx.total_token_context_product()

656 +

657 + if per_gpu:

658 + L //= self.pp_size

659 + q = max(1, q // self.tp_size)

660 +

661 + flops: dict[str, int] = {}

662 +

663 + # Q projection

664 + if q_rank is not None:

665 + # Two-stage: h -> q_lora_rank -> num_heads * qk_head_dim

666 + flops["q_a_proj"] = 2 * T * D * q_rank * L

667 + flops["q_b_proj"] = 2 * T * q_rank * q * qk_head_dim * L

668 + else:

669 + # Direct: h -> num_heads * qk_head_dim

670 + flops["q_proj"] = 2 * T * D * q * qk_head_dim * L

671 +

672 + # KV projection (always compressed, shared across heads)

673 + # kv_a: h -> (kv_lora_rank + qk_rope_head_dim) [replicated]

674 + flops["kv_a_proj"] = 2 * T * D * (c + r) * L

675 + # kv_b: kv_lora_rank -> num_heads * (qk_nope + v_head_dim)

676 + flops["kv_b_proj"] = 2 * T * c * q * (self.qk_nope_head_dim + v_d) * L

677 +

678 + # Attention core

679 + flops["attn_qk"] = 2 * q * TC * qk_head_dim * L

680 + flops["attn_av"] = 2 * q * TC * v_d * L

681 +

682 + # Output projection: num_heads * v_head_dim -> h

683 + flops["out_proj"] = 2 * T * q * v_d * D * L

684 +

685 + return flops

686 +

687 + def get_read_bytes_breakdown(

688 + self, ctx: ExecutionContext, per_gpu: bool = True

689 + ) -> dict[str, int]:

690 + """Calculate read memory traffic for MLA attention layers."""

691 + L = self.num_hidden_layers

692 + D = self.hidden_size

693 + q = self.num_attention_heads

694 + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

695 + v_d = self.v_head_dim

696 + c = self.kv_lora_rank

697 + r = self.qk_rope_head_dim

698 + q_rank = self.q_lora_rank

699 +

700 + T = ctx.total_num_tokens()

701 + # Compressed KV cache size per token

702 + kv_compressed_dim = c + r

703 +

704 + if per_gpu:

705 + L //= self.pp_size

706 + q = max(1, q // self.tp_size)

707 +

708 + read_bytes: dict[str, int] = {}

709 +

710 + # Q projection weight + input reads

711 + if q_rank is not None:

712 + read_bytes["q_a_input"] = T * D * self.activation_byte_size * L

713 + read_bytes["q_a_weight"] = int(D * q_rank * self.weight_byte_size * L)

714 + read_bytes["q_b_input"] = T * q_rank * self.activation_byte_size * L

715 + read_bytes["q_b_weight"] = int(

716 + q_rank * q * qk_head_dim * self.weight_byte_size * L

717 + )

718 + else:

719 + read_bytes["q_input"] = T * D * self.activation_byte_size * L

720 + read_bytes["q_weight"] = int(

721 + D * q * qk_head_dim * self.weight_byte_size * L

722 + )

723 +

724 + # KV projection weight + input reads

725 + # kv_a is replicated (not TP-sharded)

726 + read_bytes["kv_a_input"] = T * D * self.activation_byte_size * L

727 + read_bytes["kv_a_weight"] = int(

728 + D * kv_compressed_dim * self.weight_byte_size * L

729 + )

730 + # kv_b is TP-sharded along heads

731 + read_bytes["kv_b_input"] = T * c * self.activation_byte_size * L

732 + read_bytes["kv_b_weight"] = int(

733 + c * q * (self.qk_nope_head_dim + v_d) * self.weight_byte_size * L

734 + )

735 +

736 + # Attention input reads

737 + # Prefill: read Q activations + K,V from kv_b_proj output

738 + if ctx.prefill_num_tokens > 0:

739 + read_bytes["attn_input"] = (

740 + ctx.prefill_num_tokens * q * qk_head_dim * self.activation_byte_size * L

741 + + ctx.prefill_context_len

742 + * q

743 + * (qk_head_dim + v_d)

744 + * self.activation_byte_size

745 + * L

746 + )

747 +

748 + # Decode: read Q activations + read compressed KV from cache

749 + if ctx.decode_num_tokens > 0:

750 + read_bytes["attn_input"] = read_bytes.get("attn_input", 0) + (

751 + ctx.decode_num_tokens * q * qk_head_dim * self.activation_byte_size * L

752 + + ctx.decode_context_len * kv_compressed_dim * self.cache_byte_size * L

753 + )

754 +

755 + # Output projection reads

756 + read_bytes["out_input"] = T * q * v_d * self.activation_byte_size * L

757 + read_bytes["out_weight"] = int(q * v_d * D * self.weight_byte_size * L)

758 +

759 + return read_bytes

760 +

761 + def get_write_bytes_breakdown(

762 + self, ctx: ExecutionContext, per_gpu: bool = True

763 + ) -> dict[str, int]:

764 + """Calculate write memory traffic for MLA attention layers."""

765 + L = self.num_hidden_layers

766 + D = self.hidden_size

767 + q = self.num_attention_heads

768 + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

769 + v_d = self.v_head_dim

770 + c = self.kv_lora_rank

771 + r = self.qk_rope_head_dim

772 + q_rank = self.q_lora_rank

773 +

774 + T = ctx.total_num_tokens()

775 + kv_compressed_dim = c + r

776 +

777 + if per_gpu:

778 + L //= self.pp_size

779 + q = max(1, q // self.tp_size)

780 +

781 + write_bytes: dict[str, int] = {}

782 +

783 + # Q projection outputs

784 + if q_rank is not None:

785 + write_bytes["q_a_output"] = T * q_rank * self.activation_byte_size * L

786 + write_bytes["q_b_output"] = (

787 + T * q * qk_head_dim * self.activation_byte_size * L

788 + )

789 + else:

790 + write_bytes["q_output"] = (

791 + T * q * qk_head_dim * self.activation_byte_size * L

792 + )

793 +

794 + # KV projection outputs

795 + write_bytes["kv_a_output"] = (

796 + T * kv_compressed_dim * self.activation_byte_size * L

797 + )

798 + write_bytes["kv_b_output"] = (

799 + T * q * (self.qk_nope_head_dim + v_d) * self.activation_byte_size * L

800 + )

801 +

802 + # KV cache write: one compressed vector per token

803 + # (kv_lora_rank + qk_rope_head_dim) instead of

804 + # 2 * num_kv_heads * head_dim in standard MHA

805 + write_bytes["kv_cache"] = T * kv_compressed_dim * self.cache_byte_size * L

806 +

807 + # Output projection

808 + write_bytes["out_output"] = T * D * self.activation_byte_size * L

809 +

810 + return write_bytes

811 +

812 +

813 #### Ffn ####

814

815