@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 435c96f9db

[V1][Metrics] Support Sliding Window Attention (SWA) and hybrid layers in MFU/MBU estimation (#55624)

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

Thillai Chithambaram committed Sep 14, 2026 at 12:04 UTC 435c96f9dbdd29258cb8e0f433c5b54a00cf6b16
2 files changed +640 -48
tests/v1/metrics/test_perf_metrics.py
+372 -2
@@ -8,6 +8,7 @@ import types
8 from types import SimpleNamespace
9
10 import pytest
11 +from transformers import Gemma2Config, MistralConfig
12 from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config
13 from transformers.models.llama4.configuration_llama4 import (
14 Llama4Config,
@@ -31,6 +32,7 @@ from vllm.v1.metrics.perf import (
32 MLAAttentionMetrics,
33 ModelMetrics,
34 ParsedArgs,
35 + SlidingWindowAttentionParser,
36 UnembedMetrics,
37 )
38
@@ -38,7 +40,8 @@ from vllm.v1.metrics.perf import (
40 class MockModelConfig:
41 """Mock ModelConfig that implements the getter methods used by parsers."""
42
41 - def __init__(self, hf_config, dtype):
43 + def __init__(self, hf_config, dtype, disable_sliding_window=False):
44 + """Initialize MockModelConfig with HF config and options."""
45 self.hf_config = hf_config
46 self.hf_text_config = get_hf_text_config(hf_config)
47 convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS.get(
@@ -49,8 +52,10 @@ class MockModelConfig:
52 ).convert()
53 self.dtype = dtype
54 self.is_attention_free = False
55 + self.disable_sliding_window = disable_sliding_window
56
57 def __getattr__(self, name):
58 + """Delegate attribute lookups to ModelConfig."""
59 # 1. Check if ModelConfig actually has this attribute
60 if not hasattr(ModelConfig, name):
61 raise AttributeError(
@@ -86,9 +91,13 @@ def create_mock_vllm_config(
91 tensor_parallel_size=1,
92 pipeline_parallel_size=1,
93 enable_expert_parallel=False,
94 + disable_sliding_window=False,
95 ) -> SimpleNamespace:
96 + """Create a mock VllmConfig object for unit tests."""
97 vllm_config = SimpleNamespace()
91 - vllm_config.model_config = MockModelConfig(hf_config, model_dtype)
98 + vllm_config.model_config = MockModelConfig(
99 + hf_config, model_dtype, disable_sliding_window=disable_sliding_window
100 + )
101
102 vllm_config.cache_config = SimpleNamespace()
103 vllm_config.cache_config.cache_dtype = cache_dtype
@@ -1336,3 +1345,364 @@ def test_mla_attention_scaling_with_layers():
1345 assert double_metrics.get_num_flops(ctx) == 2 * base_metrics.get_num_flops(ctx)
1346 assert double_metrics.get_read_bytes(ctx) == 2 * base_metrics.get_read_bytes(ctx)
1347 assert double_metrics.get_write_bytes(ctx) == 2 * base_metrics.get_write_bytes(ctx)
1348 +
1349 +
1350 +#### Sliding Window Attention Tests ####
1351 +
1352 +
1353 +def test_sliding_window_parser_no_swa():
1354 + """Test parser defaults to no SWA when sliding_window is unset."""
1355 + hf_config = Qwen3Config(
1356 + hidden_size=2048,
1357 + num_attention_heads=16,
1358 + num_hidden_layers=12,
1359 + )
1360 + vllm_config = create_mock_vllm_config(hf_config)
1361 + parser = SlidingWindowAttentionParser()
1362 + args = ParsedArgs()
1363 + args.num_hidden_layers = 12
1364 + result = parser.parse(args, vllm_config)
1365 +
1366 + assert result.sliding_window is None
1367 + assert result.num_swa_layers == 0
1368 +
1369 +
1370 +def test_sliding_window_parser_uniform_swa():
1371 + """Test SlidingWindowAttentionParser for full SWA model (e.g. Mistral)."""
1372 + hf_config = MistralConfig(
1373 + hidden_size=4096,
1374 + num_attention_heads=32,
1375 + num_key_value_heads=8,
1376 + num_hidden_layers=32,
1377 + sliding_window=4096,
1378 + )
1379 + vllm_config = create_mock_vllm_config(hf_config)
1380 + parser = SlidingWindowAttentionParser()
1381 + args = ParsedArgs()
1382 + args.num_hidden_layers = 32
1383 + result = parser.parse(args, vllm_config)
1384 +
1385 + assert result.sliding_window == 4096
1386 + assert result.num_swa_layers == 32
1387 +
1388 +
1389 +def test_sliding_window_parser_explicit_layer_types():
1390 + """Test SlidingWindowAttentionParser with explicit layer_types list."""
1391 + hf_config = Qwen3Config(
1392 + hidden_size=2048,
1393 + num_attention_heads=16,
1394 + num_hidden_layers=6,
1395 + )
1396 + hf_config.sliding_window = 2048
1397 + hf_config.layer_types = [
1398 + "sliding_attention",
1399 + "full_attention",
1400 + "sliding_attention",
1401 + "full_attention",
1402 + "sliding_attention",
1403 + "full_attention",
1404 + ]
1405 + vllm_config = create_mock_vllm_config(hf_config)
1406 + parser = SlidingWindowAttentionParser()
1407 + args = ParsedArgs()
1408 + args.num_hidden_layers = 6
1409 + result = parser.parse(args, vllm_config)
1410 +
1411 + assert result.sliding_window == 2048
1412 + assert result.num_swa_layers == 3
1413 +
1414 +
1415 +def test_sliding_window_parser_gemma2():
1416 + """Test SlidingWindowAttentionParser detects alternating SWA for Gemma 2."""
1417 + hf_config = Gemma2Config(
1418 + hidden_size=2304,
1419 + num_attention_heads=8,
1420 + num_key_value_heads=4,
1421 + num_hidden_layers=26,
1422 + sliding_window=4096,
1423 + )
1424 + vllm_config = create_mock_vllm_config(hf_config)
1425 + parser = SlidingWindowAttentionParser()
1426 + args = ParsedArgs()
1427 + args.num_hidden_layers = 26
1428 + result = parser.parse(args, vllm_config)
1429 +
1430 + assert result.sliding_window == 4096
1431 + # Gemma 2 alternates: even layers are sliding window (26 layers -> 13 SWA)
1432 + assert result.num_swa_layers == 13
1433 +
1434 +
1435 +def test_sliding_window_parser_disabled():
1436 + """Test SlidingWindowAttentionParser respects disable_sliding_window."""
1437 + hf_config = MistralConfig(
1438 + hidden_size=4096,
1439 + num_attention_heads=32,
1440 + num_key_value_heads=8,
1441 + num_hidden_layers=32,
1442 + sliding_window=4096,
1443 + )
1444 + vllm_config = create_mock_vllm_config(hf_config, disable_sliding_window=True)
1445 + parser = SlidingWindowAttentionParser()
1446 + args = ParsedArgs()
1447 + args.num_hidden_layers = 32
1448 + result = parser.parse(args, vllm_config)
1449 +
1450 + assert result.sliding_window is None
1451 + assert result.num_swa_layers == 0
1452 +
1453 +
1454 +def test_attention_metrics_sliding_window_flops_bounds():
1455 + """Test that SWA clamps attention core FLOPs when context_len > sliding_window."""
1456 + sw = 1024
1457 + dense_config = MistralConfig(
1458 + hidden_size=2048,
1459 + num_attention_heads=16,
1460 + num_key_value_heads=16,
1461 + num_hidden_layers=8,
1462 + sliding_window=None,
1463 + )
1464 + swa_config = MistralConfig(
1465 + hidden_size=2048,
1466 + num_attention_heads=16,
1467 + num_key_value_heads=16,
1468 + num_hidden_layers=8,
1469 + sliding_window=sw,
1470 + )
1471 +
1472 + dense_vllm = create_mock_vllm_config(dense_config)
1473 + swa_vllm = create_mock_vllm_config(swa_config)
1474 +
1475 + dense_metrics = AttentionMetrics.from_vllm_config(dense_vllm)
1476 + swa_metrics = AttentionMetrics.from_vllm_config(swa_vllm)
1477 +
1478 + # 1. When context <= sliding_window: FLOPs should be identical
1479 + short_ctx = ExecutionContext.from_single_request(
1480 + num_tokens=1, context_len=512, is_prefill=False
1481 + )
1482 + assert swa_metrics.get_num_flops(short_ctx) == dense_metrics.get_num_flops(
1483 + short_ctx
1484 + )
1485 +
1486 + # 2. When context > sliding_window: FLOPs for SWA must be strictly smaller
1487 + long_ctx = ExecutionContext.from_single_request(
1488 + num_tokens=1, context_len=8192, is_prefill=False
1489 + )
1490 + dense_flops = dense_metrics.get_num_flops_breakdown(long_ctx)
1491 + swa_flops = swa_metrics.get_num_flops_breakdown(long_ctx)
1492 +
1493 + # Projections (qkv_proj, out_proj) are unaffected by sliding window
1494 + assert swa_flops["qkv_proj"] == dense_flops["qkv_proj"]
1495 + assert swa_flops["out_proj"] == dense_flops["out_proj"]
1496 +
1497 + # Attention core FLOPs are clamped: 8192 vs 1024 (8x difference)
1498 + assert swa_flops["attn_qk"] < dense_flops["attn_qk"]
1499 + assert swa_flops["attn_av"] < dense_flops["attn_av"]
1500 + assert swa_flops["attn_qk"] * 8 == dense_flops["attn_qk"]
1501 + assert swa_flops["attn_av"] * 8 == dense_flops["attn_av"]
1502 +
1503 +
1504 +def test_attention_metrics_sliding_window_read_bytes():
1505 + """Test that SWA bounds decode KV-cache read bytes to sliding_window."""
1506 + sw = 2048
1507 + dense_config = MistralConfig(
1508 + hidden_size=2048,
1509 + num_attention_heads=16,
1510 + num_key_value_heads=4,
1511 + num_hidden_layers=4,
1512 + sliding_window=None,
1513 + )
1514 + swa_config = MistralConfig(
1515 + hidden_size=2048,
1516 + num_attention_heads=16,
1517 + num_key_value_heads=4,
1518 + num_hidden_layers=4,
1519 + sliding_window=sw,
1520 + )
1521 +
1522 + dense_vllm = create_mock_vllm_config(dense_config)
1523 + swa_vllm = create_mock_vllm_config(swa_config)
1524 +
1525 + dense_metrics = AttentionMetrics.from_vllm_config(dense_vllm)
1526 + swa_metrics = AttentionMetrics.from_vllm_config(swa_vllm)
1527 +
1528 + # Long decode context: 8192 tokens
1529 + ctx = ExecutionContext.from_single_request(
1530 + num_tokens=1, context_len=8192, is_prefill=False
1531 + )
1532 + dense_reads = dense_metrics.get_read_bytes_breakdown(ctx)
1533 + swa_reads = swa_metrics.get_read_bytes_breakdown(ctx)
1534 +
1535 + # Attention input read traffic is lower for SWA
1536 + assert swa_reads["attn_input"] < dense_reads["attn_input"]
1537 +
1538 + # Weight reads and projection inputs are identical
1539 + assert swa_reads["qkv_weight"] == dense_reads["qkv_weight"]
1540 + assert swa_reads["out_weight"] == dense_reads["out_weight"]
1541 + assert swa_reads["qkv_input"] == dense_reads["qkv_input"]
1542 +
1543 +
1544 +def test_attention_metrics_hybrid_alternating_layers():
1545 + """Test hybrid model where half layers are full attention and half are SWA."""
1546 + sw = 1024
1547 + hybrid_config = Qwen3Config(
1548 + hidden_size=2048,
1549 + num_attention_heads=16,
1550 + num_key_value_heads=16,
1551 + num_hidden_layers=4,
1552 + )
1553 + hybrid_config.sliding_window = sw
1554 + hybrid_config.layer_types = [
1555 + "sliding_attention",
1556 + "full_attention",
1557 + "sliding_attention",
1558 + "full_attention",
1559 + ]
1560 +
1561 + vllm_config = create_mock_vllm_config(hybrid_config)
1562 + metrics = AttentionMetrics.from_vllm_config(vllm_config)
1563 +
1564 + assert metrics.num_swa_layers == 2
1565 + assert metrics.sliding_window == sw
1566 +
1567 + ctx = ExecutionContext.from_single_request(
1568 + num_tokens=1, context_len=4096, is_prefill=False
1569 + )
1570 + flops = metrics.get_num_flops_breakdown(ctx)
1571 +
1572 + # 2 full layers attend to 4096; 2 SWA layers attend to 1024
1573 + # Expected TC total = 2 * 4096 + 2 * 1024 = 8192 + 2048 = 10240
1574 + # Formula: 2 * q * d * (L_full * 4096 + L_swa * 1024)
1575 + q = metrics.num_attention_heads
1576 + d = metrics.head_dim
1577 + expected_core_flops = 2 * q * d * (2 * 4096 + 2 * 1024)
1578 + assert flops["attn_qk"] == expected_core_flops
1579 + assert flops["attn_av"] == expected_core_flops
1580 +
1581 +
1582 +def test_attention_metrics_sliding_window_per_gpu():
1583 + """Test that SWA metrics scale properly with tensor and pipeline parallelism."""
1584 + sw = 1024
1585 + hf_config = MistralConfig(
1586 + hidden_size=4096,
1587 + num_attention_heads=32,
1588 + num_key_value_heads=8,
1589 + num_hidden_layers=16,
1590 + sliding_window=sw,
1591 + )
1592 + single_vllm = create_mock_vllm_config(hf_config)
1593 + tp_vllm = create_mock_vllm_config(hf_config, tensor_parallel_size=2)
1594 + pp_vllm = create_mock_vllm_config(hf_config, pipeline_parallel_size=2)
1595 +
1596 + single_metrics = AttentionMetrics.from_vllm_config(single_vllm)
1597 + tp_metrics = AttentionMetrics.from_vllm_config(tp_vllm)
1598 + pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm)
1599 +
1600 + ctx = ExecutionContext.from_single_request(
1601 + num_tokens=1, context_len=4096, is_prefill=False
1602 + )
1603 +
1604 + # Both TP=2 and PP=2 should halve the per-GPU attention core FLOPs
1605 + single_flops = single_metrics.get_num_flops_breakdown(ctx, per_gpu=True)
1606 + tp_flops = tp_metrics.get_num_flops_breakdown(ctx, per_gpu=True)
1607 + pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True)
1608 +
1609 + assert tp_flops["attn_qk"] == single_flops["attn_qk"] // 2
1610 + assert pp_flops["attn_qk"] == single_flops["attn_qk"] // 2
1611 +
1612 +
1613 +def test_attention_metrics_hybrid_sliding_window_per_gpu():
1614 + """Test that hybrid SWA models preserve L_full + L_swa == L under PP."""
1615 + sw = 1024
1616 + hybrid_config = Gemma2Config(
1617 + hidden_size=2048,
1618 + num_attention_heads=16,
1619 + num_key_value_heads=16,
1620 + num_hidden_layers=30,
1621 + sliding_window=sw,
1622 + )
1623 + pp_vllm = create_mock_vllm_config(hybrid_config, pipeline_parallel_size=2)
1624 + pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm)
1625 +
1626 + ctx = ExecutionContext.from_single_request(
1627 + num_tokens=1, context_len=4096, is_prefill=False
1628 + )
1629 + pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True)
1630 + pp_read_bytes = pp_metrics.get_read_bytes_breakdown(ctx, per_gpu=True)
1631 +
1632 + # Under PP=2, each global count is halved: 30 layers -> 15 per GPU, and the
1633 + # 15 SWA / 15 full split becomes 7.5 each. Gemma 2 alternates, so an even
1634 + # split is the correct average stage; deriving it after reducing L instead
1635 + # produced a lopsided 8/7.
1636 + q = pp_metrics.num_attention_heads
1637 + d = pp_metrics.head_dim
1638 + expected_pp_attn = int(2 * q * d * (7.5 * 4096 + 7.5 * sw))
1639 + assert pp_flops["attn_qk"] == expected_pp_attn
1640 + assert pp_flops["attn_av"] == expected_pp_attn
1641 + assert "attn_input" in pp_read_bytes
1642 +
1643 +
1644 +def test_attention_metrics_minority_swa_survives_pipeline_parallelism():
1645 + """A minority SWA sublayer must not round away under PP.
1646 +
1647 + Regression test for a 6-layer model with a single sliding-window layer at
1648 + pp_size=2: deriving the split after reducing L gave
1649 + round(3 * 1 / 6) == 0, billing all three per-GPU layers as full attention
1650 + and dropping the sliding-window contribution entirely.
1651 + """
1652 + sw = 512
1653 + context_len = 4096
1654 + hf_config = Qwen3Config(
1655 + hidden_size=2048,
1656 + num_attention_heads=16,
1657 + num_key_value_heads=16,
1658 + num_hidden_layers=6,
1659 + )
1660 + hf_config.sliding_window = sw
1661 + hf_config.layer_types = ["sliding_attention"] + ["full_attention"] * 5
1662 +
1663 + pp_vllm = create_mock_vllm_config(hf_config, pipeline_parallel_size=2)
1664 + pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm)
1665 + assert pp_metrics.num_swa_layers == 1
1666 +
1667 + ctx = ExecutionContext.from_single_request(
1668 + num_tokens=1, context_len=context_len, is_prefill=False
1669 + )
1670 + pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True)
1671 +
1672 + # 5 full and 1 SWA layer, each halved by PP=2: 2.5 full and 0.5 SWA.
1673 + q = pp_metrics.num_attention_heads
1674 + d = pp_metrics.head_dim
1675 + expected = int(2 * q * d * (2.5 * context_len + 0.5 * sw))
1676 + assert pp_flops["attn_qk"] == expected
1677 + assert pp_flops["attn_av"] == expected
1678 +
1679 + # The sliding-window layer must change the result, i.e. it is not being
1680 + # billed at the full context length.
1681 + all_full = int(2 * q * d * 3 * context_len)
1682 + assert pp_flops["attn_qk"] < all_full
1683 +
1684 +
1685 +def test_attention_metrics_per_gpu_layers_sum_to_whole_model():
1686 + """Per-GPU counts times pp_size must account for every layer.
1687 +
1688 + Flooring the layer count per stage silently drops the remainder, so a model
1689 + whose layers do not divide evenly across stages under-reports its FLOPs.
1690 + """
1691 + hf_config = Qwen3Config(
1692 + hidden_size=2048,
1693 + num_attention_heads=16,
1694 + num_key_value_heads=16,
1695 + num_hidden_layers=7,
1696 + )
1697 + pp_size = 2
1698 + vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=pp_size)
1699 + metrics = AttentionMetrics.from_vllm_config(vllm_config)
1700 +
1701 + ctx = ExecutionContext.from_single_request(
1702 + num_tokens=8, context_len=1024, is_prefill=True
1703 + )
1704 + per_gpu = metrics.get_num_flops_breakdown(ctx, per_gpu=True)
1705 + whole = metrics.get_num_flops_breakdown(ctx, per_gpu=False)
1706 +
1707 + for key, value in whole.items():
1708 + assert per_gpu[key] * pp_size == pytest.approx(value, rel=1e-9)
vllm/v1/metrics/perf.py
+268 -46
@@ -10,7 +10,7 @@ import json
10 import time
11 from abc import ABC, abstractmethod
12 from collections.abc import Iterable
13 -from dataclasses import asdict, dataclass
13 +from dataclasses import asdict, dataclass, field
14 from typing import Any, Protocol
15
16 import prometheus_client
@@ -124,8 +124,12 @@ class ExecutionContext:
124 decode_context_len: int = 0 # sum of context_len for decode requests
125 decode_token_context_product: int = 0 # sum of (num_tokens * context_len)
126
127 + # Per-request records: (num_tokens, context_len, is_prefill)
128 + requests: list[tuple[int, int, bool]] = field(default_factory=list)
129 +
130 def add(self, num_tokens: int, context_len: int, is_prefill: bool) -> None:
131 """Add a single request's statistics to this batch context."""
132 + self.requests.append((num_tokens, context_len, is_prefill))
133 if is_prefill:
134 self.num_prefill_requests += 1
135 self.prefill_num_tokens += num_tokens
@@ -141,9 +145,40 @@ class ExecutionContext:
145 """Total number of tokens across all requests in the batch."""
146 return self.prefill_num_tokens + self.decode_num_tokens
147
144 - def total_token_context_product(self) -> int:
145 - """Total sum of (num_tokens * context_len) across all requests."""
146 - return self.prefill_token_context_product + self.decode_token_context_product
148 + def total_token_context_product(self, sliding_window: int | None = None) -> int:
149 + """Total sum of (num_tokens * context_len) across all requests.
150 +
151 + If sliding_window is provided, context length per request is clamped to
152 + sliding_window.
153 + """
154 + if sliding_window is None:
155 + return (
156 + self.prefill_token_context_product + self.decode_token_context_product
157 + )
158 + return sum(
159 + num_tokens * min(context_len, sliding_window)
160 + for num_tokens, context_len, _ in self.requests
161 + )
162 +
163 + def prefill_context_len_for_window(self, sliding_window: int | None = None) -> int:
164 + """Sum of context_len for prefill requests, clamped to sliding_window if set."""
165 + if sliding_window is None:
166 + return self.prefill_context_len
167 + return sum(
168 + min(context_len, sliding_window)
169 + for _, context_len, is_prefill in self.requests
170 + if is_prefill
171 + )
172 +
173 + def decode_context_len_for_window(self, sliding_window: int | None = None) -> int:
174 + """Sum of context_len for decode requests, clamped to sliding_window if set."""
175 + if sliding_window is None:
176 + return self.decode_context_len
177 + return sum(
178 + min(context_len, sliding_window)
179 + for _, context_len, is_prefill in self.requests
180 + if not is_prefill
181 + )
182
183 def num_logits_tokens(self) -> int:
184 """Number of tokens that require logits computation (unembedding).
@@ -214,6 +249,7 @@ class ParserChain:
249 self.parsers.append(parser)
250
251 def parse(self, vllm_config: VllmConfig) -> ParsedArgs:
252 + """Apply all parsers in sequence and return the accumulated ParsedArgs."""
253 args = ParsedArgs()
254 for parser in self.parsers:
255 args = parser.parse(args, vllm_config)
@@ -234,7 +270,9 @@ class ComponentMetrics(BaseModel, ABC):
270
271 @classmethod
272 @abstractmethod
237 - def component_type(cls) -> str: ...
273 + def component_type(cls) -> str:
274 + """Return the unique string key identifying this component in the registry."""
275 + ...
276
277 @classmethod
278 @abstractmethod
@@ -273,17 +311,23 @@ class ComponentMetrics(BaseModel, ABC):
311 @abstractmethod
312 def get_num_flops_breakdown(
313 self, ctx: ExecutionContext, per_gpu: bool = True
276 - ) -> dict[str, int]: ...
314 + ) -> dict[str, int]:
315 + """Return per-operation FLOPs breakdown for this component."""
316 + ...
317
318 @abstractmethod
319 def get_read_bytes_breakdown(
320 self, ctx: ExecutionContext, per_gpu: bool = True
281 - ) -> dict[str, int]: ...
321 + ) -> dict[str, int]:
322 + """Return per-operation read-memory-traffic breakdown for this component."""
323 + ...
324
325 @abstractmethod
326 def get_write_bytes_breakdown(
327 self, ctx: ExecutionContext, per_gpu: bool = True
286 - ) -> dict[str, int]: ...
328 + ) -> dict[str, int]:
329 + """Return per-operation write-memory-traffic breakdown for this component."""
330 + ...
331
332 def get_num_flops(self, ctx: ExecutionContext, per_gpu: bool = True) -> int:
333 return sum(self.get_num_flops_breakdown(ctx, per_gpu).values())
@@ -306,6 +350,7 @@ class BaseConfigParser(Parser):
350 """
351
352 def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
353 + """Parse base model dimensions, dtype byte sizes, and parallelism config."""
354 model_config = vllm_config.model_config
355
356 args.vocab_size = model_config.get_vocab_size()
@@ -357,6 +402,7 @@ class BaseAttentionConfigParser(Parser):
402 """
403
404 def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
405 + """Parse KV head count, head dimension, and KV-cache dtype byte size."""
406 model_config = vllm_config.model_config
407
408 args.num_key_value_heads = model_config.get_total_num_kv_heads()
@@ -371,6 +417,54 @@ class BaseAttentionConfigParser(Parser):
417 return args
418
419
420 +class SlidingWindowAttentionParser(Parser):
421 + """
422 + Parses sliding window attention configuration and layer breakdown.
423 + Provides: sliding_window, num_swa_layers
424 + """
425 +
426 + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
427 + """Determine sliding_window size and count of SWA layers.
428 +
429 + Detects per-layer type lists (e.g. Llama-4), Gemma 2 alternating
430 + pattern, and uniform SWA models (Mistral, Qwen2.5).
431 + """
432 + model_config = vllm_config.model_config
433 +
434 + sliding_window: int | None = None
435 + if not getattr(model_config, "disable_sliding_window", False):
436 + sliding_window = model_config.get_sliding_window()
437 + if sliding_window is None and hasattr(vllm_config, "cache_config"):
438 + sliding_window = getattr(
439 + vllm_config.cache_config, "sliding_window", None
440 + )
441 +
442 + args.sliding_window = sliding_window
443 +
444 + num_layers = args.num_hidden_layers
445 + if sliding_window is not None and sliding_window > 0:
446 + hf_text_config = getattr(model_config, "hf_text_config", None)
447 + layer_types = getattr(hf_text_config, "layer_types", None)
448 + if layer_types is not None and isinstance(layer_types, list):
449 + num_swa = sum(
450 + 1 for lt in layer_types[:num_layers] if "sliding" in lt.lower()
451 + )
452 + args.num_swa_layers = num_swa
453 + elif (
454 + getattr(getattr(model_config, "hf_config", None), "model_type", None)
455 + == "gemma2"
456 + ):
457 + # Gemma 2 alternates sliding window: even layers (0, 2, ...) are sliding
458 + args.num_swa_layers = (num_layers + 1) // 2
459 + else:
460 + # Full SWA model (e.g. Mistral, Qwen2.5 with sliding_window)
461 + args.num_swa_layers = num_layers
462 + else:
463 + args.num_swa_layers = 0
464 +
465 + return args
466 +
467 +
468 class AttentionQuantizationConfigParser(Parser):
469 """
470 Parses quantization configuration for attention layers.
@@ -378,6 +472,7 @@ class AttentionQuantizationConfigParser(Parser):
472 """
473
474 def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
475 + """Override weight_byte_size based on the active quantization method."""
476 cfg = vllm_config.quant_config
477
478 if cfg is None:
@@ -401,6 +496,7 @@ class AttentionDetectionParser(Parser):
496 """
497
498 def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
499 + """Raise InvalidComponent for MLA models so AttentionMetrics is skipped."""
500 if vllm_config.model_config.is_deepseek_mla:
501 raise InvalidComponent(
502 "Model uses MLA attention; use MLAAttentionMetrics instead"
@@ -425,53 +521,129 @@ class AttentionMetrics(ComponentMetrics):
521 # From BaseConfig Parser, overridden by AttentionQuantizationConfigParser
522 weight_byte_size: int | float = Field(..., gt=0)
523
428 - # TODO: discern cases where we have mixture of different attention layer types
429 - # such as SWA, MLA, etc.
524 + # From SlidingWindowAttentionParser
525 + sliding_window: int | None = Field(None)
526 + num_swa_layers: int = Field(0, ge=0)
527
528 @classmethod
529 def component_type(cls) -> str:
530 + """Return the component registry key for standard (non-MLA) attention."""
531 return "attn"
532
533 @classmethod
534 def get_parser(cls) -> ParserChain:
535 + """Return the parser chain for AttentionMetrics.
536 +
537 + Parsers run in order: MLA detection guard, base model config,
538 + attention-specific config, sliding-window breakdown, quantization
539 + weight-size override.
540 + """
541 return ParserChain(
542 AttentionDetectionParser(),
543 BaseConfigParser(),
544 BaseAttentionConfigParser(),
545 + SlidingWindowAttentionParser(),
546 AttentionQuantizationConfigParser(),
547 )
548
549 + def _layer_counts(self, per_gpu: bool) -> tuple[float, float, float]:
550 + """Return ``(L, L_full, L_swa)``, the layer counts to bill for.
551 +
552 + Per-GPU counts divide each global count by ``pp_size`` independently.
553 + Deriving the sublayer split after reducing ``L`` instead rounds a
554 + minority sublayer away: a 6-layer model with 1 sliding-window layer at
555 + ``pp_size=2`` yields ``round(3 * 1 / 6) == 0``, billing all three layers
556 + as full attention. Dividing each count separately keeps
557 + ``L_full + L_swa == L`` and keeps the per-GPU values summing back to the
558 + whole model across stages.
559 +
560 + The counts are fractional because these metrics are built from config
561 + alone (see ``from_vllm_config``); no pipeline rank is available, so a
562 + per-GPU figure is the average stage rather than a specific one. A model
563 + whose layers do not divide evenly across stages has no single integer
564 + answer here.
565 + """
566 + L = float(self.num_hidden_layers)
567 + L_swa = float(self.num_swa_layers)
568 + L_full = L - L_swa
569 +
570 + if per_gpu:
571 + L /= self.pp_size
572 + L_swa /= self.pp_size
573 + L_full /= self.pp_size
574 +
575 + return L, L_full, L_swa
576 +
577 def get_num_flops_breakdown(
578 self, ctx: ExecutionContext, per_gpu: bool = True
579 ) -> dict[str, int]:
447 - L, D, q, kv, d = (
448 - self.num_hidden_layers,
580 + """Compute FLOPs breakdown for attention layers.
581 +
582 + Accounts for hybrid models by splitting total layers into full-attention
583 + (L_full) and sliding-window (L_swa) sublayers. See ``_layer_counts`` for
584 + how those are scaled under pipeline parallelism.
585 +
586 + Args:
587 + ctx: Execution context describing the current batch.
588 + per_gpu: If True, scale counts down by tensor/pipeline parallelism.
589 +
590 + Returns:
591 + Dict with keys: qkv_proj, out_proj, attn_qk, attn_av.
592 + """
593 + D, q, kv, d = (
594 self.hidden_size,
595 self.num_attention_heads,
596 self.num_key_value_heads,
597 self.head_dim,
598 )
599 T = ctx.total_num_tokens()
455 - TC = ctx.total_token_context_product()
600 + TC_full = ctx.total_token_context_product()
601 +
602 + L, L_full, L_swa = self._layer_counts(per_gpu)
603
604 if per_gpu:
458 - L //= self.pp_size
605 # tensor parallel along heads
606 q = max(1, q // self.tp_size)
607 kv = max(1, kv // self.tp_size)
608
463 - return {
464 - "qkv_proj": 2 * T * D * (q + 2 * kv) * d * L,
465 - "attn_qk": 2 * q * TC * d * L,
466 - "attn_av": 2 * q * TC * d * L,
467 - "out_proj": 2 * T * D * q * d * L,
609 + flops = {
610 + "qkv_proj": int(2 * T * D * (q + 2 * kv) * d * L),
611 + "out_proj": int(2 * T * D * q * d * L),
612 }
613
614 + attn_qk = 0.0
615 + attn_av = 0.0
616 + if L_full > 0:
617 + attn_qk += 2 * q * TC_full * d * L_full
618 + attn_av += 2 * q * TC_full * d * L_full
619 +
620 + if L_swa > 0 and self.sliding_window is not None:
621 + TC_swa = ctx.total_token_context_product(self.sliding_window)
622 + attn_qk += 2 * q * TC_swa * d * L_swa
623 + attn_av += 2 * q * TC_swa * d * L_swa
624 +
625 + flops["attn_qk"] = int(attn_qk)
626 + flops["attn_av"] = int(attn_av)
627 +
628 + return flops
629 +
630 def get_read_bytes_breakdown(
631 self, ctx: ExecutionContext, per_gpu: bool = True
632 ) -> dict[str, int]:
473 - L, D, q, kv, d = (
474 - self.num_hidden_layers,
633 + """Compute read-memory-traffic breakdown for attention layers.
634 +
635 + Splits layers into full-attention and sliding-window sublayers using the
636 + same counts as get_num_flops_breakdown; see ``_layer_counts``.
637 +
638 + Args:
639 + ctx: Execution context describing the current batch.
640 + per_gpu: If True, scale counts down by tensor/pipeline parallelism.
641 +
642 + Returns:
643 + Dict with keys: qkv_input, qkv_weight, attn_input (conditional),
644 + out_input, out_weight.
645 + """
646 + D, q, kv, d = (
647 self.hidden_size,
648 self.num_attention_heads,
649 self.num_key_value_heads,
@@ -479,35 +651,63 @@ class AttentionMetrics(ComponentMetrics):
651 )
652 T = ctx.total_num_tokens()
653
654 + L, L_full, L_swa = self._layer_counts(per_gpu)
655 +
656 if per_gpu:
483 - L //= self.pp_size
657 # tensor parallel along heads
658 q = max(1, q // self.tp_size)
659 kv = max(1, kv // self.tp_size)
660
661 read_bytes = {}
662
490 - read_bytes["qkv_input"] = T * D * self.activation_byte_size * L
663 + read_bytes["qkv_input"] = int(T * D * self.activation_byte_size * L)
664 read_bytes["qkv_weight"] = int(D * (q + 2 * kv) * d * self.weight_byte_size * L)
665
493 - # Attention input reads differ between prefill and decode
494 - # Prefill: read Q, K, V activations (all in activation_byte_size)
495 - if ctx.prefill_num_tokens > 0:
496 - read_bytes["attn_input"] = (
497 - (ctx.prefill_num_tokens * q + 2 * ctx.prefill_context_len * kv)
498 - * d
499 - * self.activation_byte_size
500 - * L
501 - )
666 + attn_input = 0.0
667
503 - # Decode: read Q activations + read K, V from cache (in cache_byte_size)
504 - if ctx.decode_num_tokens > 0:
505 - read_bytes["attn_input"] = read_bytes.get("attn_input", 0) + (
506 - ctx.decode_num_tokens * q * d * self.activation_byte_size * L
507 - + 2 * ctx.decode_context_len * kv * d * self.cache_byte_size * L
508 - )
668 + # Full attention layers
669 + if L_full > 0:
670 + if ctx.prefill_num_tokens > 0:
671 + attn_input += (
672 + (ctx.prefill_num_tokens * q + 2 * ctx.prefill_context_len * kv)
673 + * d
674 + * self.activation_byte_size
675 + * L_full
676 + )
677 + if ctx.decode_num_tokens > 0:
678 + attn_input += (
679 + ctx.decode_num_tokens * q * d * self.activation_byte_size * L_full
680 + + 2
681 + * ctx.decode_context_len
682 + * kv
683 + * d
684 + * self.cache_byte_size
685 + * L_full
686 + )
687 +
688 + # Sliding window attention layers
689 + if L_swa > 0:
690 + sw = self.sliding_window
691 + prefill_sw_len = ctx.prefill_context_len_for_window(sw)
692 + decode_sw_len = ctx.decode_context_len_for_window(sw)
693 +
694 + if ctx.prefill_num_tokens > 0:
695 + attn_input += (
696 + (ctx.prefill_num_tokens * q + 2 * prefill_sw_len * kv)
697 + * d
698 + * self.activation_byte_size
699 + * L_swa
700 + )
701 + if ctx.decode_num_tokens > 0:
702 + attn_input += (
703 + ctx.decode_num_tokens * q * d * self.activation_byte_size * L_swa
704 + + 2 * decode_sw_len * kv * d * self.cache_byte_size * L_swa
705 + )
706 +
707 + if ctx.prefill_num_tokens > 0 or ctx.decode_num_tokens > 0:
708 + read_bytes["attn_input"] = int(attn_input)
709
510 - read_bytes["out_input"] = T * q * d * self.activation_byte_size * L
710 + read_bytes["out_input"] = int(T * q * d * self.activation_byte_size * L)
711 read_bytes["out_weight"] = int(q * d * D * self.weight_byte_size * L)
712
713 return read_bytes
@@ -515,9 +715,13 @@ class AttentionMetrics(ComponentMetrics):
715 def get_write_bytes_breakdown(
716 self, ctx: ExecutionContext, per_gpu: bool = True
717 ) -> dict[str, int]:
518 - """Calculate write memory traffic for attention layers."""
519 - L, D, q, kv, d = (
520 - self.num_hidden_layers,
718 + """Calculate write memory traffic for attention layers.
719 +
720 + Writes do not depend on the attention window, so only the total layer
721 + count is needed; it comes from ``_layer_counts`` so all three breakdowns
722 + bill the same number of layers.
723 + """
724 + D, q, kv, d = (
725 self.hidden_size,
726 self.num_attention_heads,
727 self.num_key_value_heads,
@@ -525,16 +729,17 @@ class AttentionMetrics(ComponentMetrics):
729 )
730 T = ctx.total_num_tokens()
731
732 + L, _, _ = self._layer_counts(per_gpu)
733 +
734 if per_gpu:
529 - L //= self.pp_size
735 # tensor parallel along heads
736 q = max(1, q // self.tp_size)
737 kv = max(1, kv // self.tp_size)
738
739 return {
535 - "qkv_output": T * (q + 2 * kv) * d * self.activation_byte_size * L,
536 - "kv_cache": 2 * T * kv * d * self.cache_byte_size * L,
537 - "out_output": T * D * self.activation_byte_size * L,
740 + "qkv_output": int(T * (q + 2 * kv) * d * self.activation_byte_size * L),
741 + "kv_cache": int(2 * T * kv * d * self.cache_byte_size * L),
742 + "out_output": int(T * D * self.activation_byte_size * L),
743 }
744
745
@@ -549,6 +754,7 @@ class MLADetectionParser(Parser):
754 """
755
756 def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
757 + """Raise InvalidComponent if the model does not use MLA attention."""
758 if not vllm_config.model_config.is_deepseek_mla:
759 raise InvalidComponent("Model does not use MLA attention")
760 return args
@@ -562,6 +768,7 @@ class MLAConfigParser(Parser):
768 """
769
770 def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
771 + """Parse MLA-specific compression dimensions and KV-cache dtype byte size."""
772 model_config = vllm_config.model_config
773 cfg = model_config.hf_text_config
774
@@ -616,10 +823,16 @@ class MLAAttentionMetrics(ComponentMetrics):
823
824 @classmethod
825 def component_type(cls) -> str:
826 + """Return the component registry key for MLA attention."""
827 return "mla_attn"
828
829 @classmethod
830 def get_parser(cls) -> ParserChain:
831 + """Return the parser chain for MLAAttentionMetrics.
832 +
833 + Parsers run in order: MLA detection guard, base model config,
834 + MLA-specific config, quantization weight-size override.
835 + """
836 return ParserChain(
837 MLADetectionParser(),
838 BaseConfigParser(),
@@ -923,6 +1136,7 @@ class FfnQuantizationConfigParser(Parser):
1136 """
1137
1138 def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
1139 + """Override weight_byte_size based on the active FFN quantization method."""
1140 cfg = vllm_config.quant_config
1141
1142 if cfg is None:
@@ -980,10 +1194,16 @@ class FfnMetrics(ComponentMetrics):
1194
1195 @classmethod
1196 def component_type(cls) -> str:
1197 + """Return the component registry key for FFN layers."""
1198 return "ffn"
1199
1200 @classmethod
1201 def get_parser(cls) -> ParserChain:
1202 + """Return the parser chain for FfnMetrics.
1203 +
1204 + Parsers run in order: base model config, FFN parallelism, base FFN
1205 + config, MoE interleave step, MoE layer frequency, quantization override.
1206 + """
1207 return ParserChain(
1208 BaseConfigParser(),
1209 FfnParallelParser(),
@@ -1207,10 +1427,12 @@ class UnembedMetrics(ComponentMetrics):
1427
1428 @classmethod
1429 def component_type(cls) -> str:
1430 + """Return the component registry key for the unembedding (LM head) layer."""
1431 return "unembed"
1432
1433 @classmethod
1434 def get_parser(cls) -> ParserChain:
1435 + """Return the parser chain for UnembedMetrics (base model config only)."""
1436 return ParserChain(
1437 BaseConfigParser(),
1438 )