| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the analytic estimators in metrics/flops.py.""" |
| 4 | |
| 5 | import types |
| 6 | from types import SimpleNamespace |
| 7 | |
| 8 | import pytest |
| 9 | from transformers import Gemma2Config, MistralConfig |
| 10 | from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config |
| 11 | from transformers.models.llama4.configuration_llama4 import ( |
| 12 | Llama4Config, |
| 13 | Llama4TextConfig, |
| 14 | ) |
| 15 | from transformers.models.qwen3.configuration_qwen3 import Qwen3Config |
| 16 | from transformers.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig |
| 17 | |
| 18 | from vllm.config.model import ModelConfig, get_hf_text_config |
| 19 | from vllm.transformers_utils.model_arch_config_convertor import ( |
| 20 | MODEL_ARCH_CONFIG_CONVERTORS, |
| 21 | ModelArchConfigConvertorBase, |
| 22 | ) |
| 23 | from vllm.v1.metrics.perf import ( |
| 24 | _QUANT_WEIGHT_BYTE_SIZE, |
| 25 | AttentionMetrics, |
| 26 | BaseConfigParser, |
| 27 | ExecutionContext, |
| 28 | FfnMetrics, |
| 29 | InvalidComponent, |
| 30 | MLAAttentionMetrics, |
| 31 | ModelMetrics, |
| 32 | ParsedArgs, |
| 33 | SlidingWindowAttentionParser, |
| 34 | UnembedMetrics, |
| 35 | ) |
| 36 | |
| 37 | |
| 38 | class MockModelConfig: |
| 39 | """Mock ModelConfig that implements the getter methods used by parsers.""" |
| 40 | |
| 41 | def __init__(self, hf_config, dtype, disable_sliding_window=False): |
| 42 | """Initialize MockModelConfig with HF config and options.""" |
| 43 | self.hf_config = hf_config |
| 44 | self.hf_text_config = get_hf_text_config(hf_config) |
| 45 | convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS.get( |
| 46 | self.hf_config.model_type, ModelArchConfigConvertorBase |
| 47 | ) |
| 48 | self.model_arch_config = convertor_cls( |
| 49 | self.hf_config, self.hf_text_config |
| 50 | ).convert() |
| 51 | self.dtype = dtype |
| 52 | self.is_attention_free = False |
| 53 | self.disable_sliding_window = disable_sliding_window |
| 54 | |
| 55 | def __getattr__(self, name): |
| 56 | """Delegate attribute lookups to ModelConfig.""" |
| 57 | # 1. Check if ModelConfig actually has this attribute |
| 58 | if not hasattr(ModelConfig, name): |
| 59 | raise AttributeError( |
| 60 | f"'{type(self).__name__}' object has no attribute '{name}' " |
| 61 | f"and neither does 'ModelConfig'." |
| 62 | ) |
| 63 | |
| 64 | # 2. Fetch the attribute from the ModelConfig CLASS |
| 65 | attr = getattr(ModelConfig, name) |
| 66 | |
| 67 | # 3. Case A: It is a @property |
| 68 | if isinstance(attr, property): |
| 69 | # Manually invoke the property's getter, passing 'self' (this mock instance) |
| 70 | return attr.__get__(self, self.__class__) |
| 71 | |
| 72 | # 4. Case B: It is a standard method (function) |
| 73 | if isinstance(attr, types.FunctionType): |
| 74 | # Bind the function to 'self' so it acts like a method of |
| 75 | # this instance. This creates a bound method where 'self' is |
| 76 | # automatically passed as the first arg. |
| 77 | return types.MethodType(attr, self) |
| 78 | |
| 79 | # 5. Case C: It is a class attribute / static variable |
| 80 | return attr |
| 81 | |
| 82 | |
| 83 | def create_mock_vllm_config( |
| 84 | hf_config, |
| 85 | model_dtype="bfloat16", |
| 86 | cache_dtype="auto", |
| 87 | quant_config=None, |
| 88 | data_parallel_size=1, |
| 89 | tensor_parallel_size=1, |
| 90 | pipeline_parallel_size=1, |
| 91 | enable_expert_parallel=False, |
| 92 | disable_sliding_window=False, |
| 93 | ) -> SimpleNamespace: |
| 94 | """Create a mock VllmConfig object for unit tests.""" |
| 95 | vllm_config = SimpleNamespace() |
| 96 | vllm_config.model_config = MockModelConfig( |
| 97 | hf_config, model_dtype, disable_sliding_window=disable_sliding_window |
| 98 | ) |
| 99 | |
| 100 | vllm_config.cache_config = SimpleNamespace() |
| 101 | vllm_config.cache_config.cache_dtype = cache_dtype |
| 102 | |
| 103 | vllm_config.quant_config = quant_config |
| 104 | |
| 105 | vllm_config.parallel_config = SimpleNamespace() |
| 106 | vllm_config.parallel_config.data_parallel_size = data_parallel_size |
| 107 | vllm_config.parallel_config.tensor_parallel_size = tensor_parallel_size |
| 108 | vllm_config.parallel_config.pipeline_parallel_size = pipeline_parallel_size |
| 109 | vllm_config.parallel_config.enable_expert_parallel = enable_expert_parallel |
| 110 | |
| 111 | return vllm_config |
| 112 | |
| 113 | |
| 114 | #### Parser Tests #### |
| 115 | |
| 116 | |
| 117 | def test_base_config_parser(): |
| 118 | """Test BaseConfigParser extracts base model attributes correctly.""" |
| 119 | hf_config = Qwen3Config( |
| 120 | vocab_size=50000, |
| 121 | hidden_size=2048, |
| 122 | num_attention_heads=16, |
| 123 | num_hidden_layers=24, |
| 124 | ) |
| 125 | vllm_config = create_mock_vllm_config(hf_config, model_dtype="float16") |
| 126 | |
| 127 | parser = BaseConfigParser() |
| 128 | args = ParsedArgs() |
| 129 | result = parser.parse(args, vllm_config) |
| 130 | |
| 131 | assert result.vocab_size == 50000 |
| 132 | assert result.hidden_size == 2048 |
| 133 | assert result.num_attention_heads == 16 |
| 134 | assert result.num_hidden_layers == 24 |
| 135 | assert result.weight_byte_size == 2 # float16 is 2 bytes |
| 136 | assert result.activation_byte_size == 2 # float16 activations are 2 bytes |
| 137 | |
| 138 | |
| 139 | @pytest.mark.parametrize( |
| 140 | "model_dtype,expected_byte_size", |
| 141 | [("float16", 2), ("bfloat16", 2), ("float32", 4)], |
| 142 | ) |
| 143 | def test_base_config_parser_activation_byte_size_follows_dtype( |
| 144 | model_dtype: str, expected_byte_size: int |
| 145 | ): |
| 146 | """Activation traffic is sized from the model's compute dtype. |
| 147 | |
| 148 | Activations are produced in the model dtype, so pinning them to 2 bytes |
| 149 | under-reported memory traffic by half for a float32 model while that same |
| 150 | model's weight traffic was sized correctly. |
| 151 | """ |
| 152 | hf_config = Qwen3Config( |
| 153 | hidden_size=2048, |
| 154 | num_attention_heads=16, |
| 155 | num_hidden_layers=24, |
| 156 | ) |
| 157 | vllm_config = create_mock_vllm_config(hf_config, model_dtype=model_dtype) |
| 158 | |
| 159 | result = BaseConfigParser().parse(ParsedArgs(), vllm_config) |
| 160 | |
| 161 | assert result.activation_byte_size == expected_byte_size |
| 162 | assert result.weight_byte_size == expected_byte_size |
| 163 | |
| 164 | |
| 165 | def test_attention_read_bytes_scale_with_activation_dtype(): |
| 166 | """A float32 model reads twice the activation bytes of a bfloat16 one.""" |
| 167 | hf_config = Qwen3Config( |
| 168 | hidden_size=2048, |
| 169 | num_attention_heads=16, |
| 170 | num_key_value_heads=16, |
| 171 | num_hidden_layers=4, |
| 172 | ) |
| 173 | ctx = ExecutionContext.from_single_request( |
| 174 | num_tokens=128, context_len=128, is_prefill=True |
| 175 | ) |
| 176 | |
| 177 | bf16 = AttentionMetrics.from_vllm_config( |
| 178 | create_mock_vllm_config(hf_config, model_dtype="bfloat16") |
| 179 | ).get_read_bytes_breakdown(ctx) |
| 180 | fp32 = AttentionMetrics.from_vllm_config( |
| 181 | create_mock_vllm_config(hf_config, model_dtype="float32") |
| 182 | ).get_read_bytes_breakdown(ctx) |
| 183 | |
| 184 | # qkv_input is pure activation traffic, so it doubles exactly. |
| 185 | assert fp32["qkv_input"] == 2 * bf16["qkv_input"] |
| 186 | |
| 187 | |
| 188 | def test_base_attention_config_parser_with_gqa(): |
| 189 | """Test BaseAttentionConfigParser with grouped query attention.""" |
| 190 | hf_config = Qwen3Config( |
| 191 | hidden_size=4096, |
| 192 | num_attention_heads=32, |
| 193 | num_key_value_heads=8, # GQA with 4:1 ratio |
| 194 | head_dim=128, |
| 195 | ) |
| 196 | vllm_config = create_mock_vllm_config(hf_config) |
| 197 | |
| 198 | parser_chain = AttentionMetrics.get_parser() |
| 199 | result = parser_chain.parse(vllm_config) |
| 200 | |
| 201 | assert result.num_key_value_heads == 8 |
| 202 | assert result.head_dim == 128 |
| 203 | |
| 204 | |
| 205 | def test_base_attention_config_parser_without_gqa(): |
| 206 | """Test BaseAttentionConfigParser defaults to MHA when num_key_value_heads not |
| 207 | specified. |
| 208 | """ |
| 209 | hf_config = Qwen3Config( |
| 210 | hidden_size=4096, |
| 211 | num_attention_heads=32, |
| 212 | # No num_key_value_heads specified |
| 213 | ) |
| 214 | vllm_config = create_mock_vllm_config(hf_config) |
| 215 | |
| 216 | parser_chain = AttentionMetrics.get_parser() |
| 217 | result = parser_chain.parse(vllm_config) |
| 218 | |
| 219 | # Should default to MHA (num_key_value_heads = num_attention_heads) |
| 220 | assert result.num_key_value_heads == 32 |
| 221 | |
| 222 | |
| 223 | def test_base_ffn_config_parser_dense(): |
| 224 | """Test BaseFfnConfigParser for dense FFN.""" |
| 225 | hf_config = Qwen3Config( |
| 226 | hidden_size=4096, |
| 227 | intermediate_size=11008, |
| 228 | num_hidden_layers=32, |
| 229 | ) |
| 230 | vllm_config = create_mock_vllm_config(hf_config) |
| 231 | |
| 232 | parser_chain = FfnMetrics.get_parser() |
| 233 | result = parser_chain.parse(vllm_config) |
| 234 | |
| 235 | assert result.intermediate_size == 11008 |
| 236 | assert result.num_experts == 0 |
| 237 | assert result.num_experts_per_tok == 0 |
| 238 | assert result.num_moe_layers == 0 # No MoE |
| 239 | |
| 240 | |
| 241 | def test_base_ffn_config_parser_moe(): |
| 242 | """Test BaseFfnConfigParser for MoE FFN.""" |
| 243 | hf_config = Qwen3MoeConfig( |
| 244 | hidden_size=4096, |
| 245 | intermediate_size=11008, |
| 246 | num_hidden_layers=32, |
| 247 | num_experts=64, |
| 248 | num_experts_per_tok=8, |
| 249 | moe_intermediate_size=14336, |
| 250 | n_shared_experts=2, |
| 251 | ) |
| 252 | vllm_config = create_mock_vllm_config(hf_config) |
| 253 | |
| 254 | parser_chain = FfnMetrics.get_parser() |
| 255 | result = parser_chain.parse(vllm_config) |
| 256 | |
| 257 | assert result.num_experts == 64 |
| 258 | assert result.num_experts_per_tok == 8 |
| 259 | assert result.moe_intermediate_size == 14336 |
| 260 | assert result.num_shared_experts == 2 |
| 261 | assert result.num_moe_layers == 32 # All layers are MoE by default |
| 262 | |
| 263 | |
| 264 | def test_interleave_moe_layer_step_parser(): |
| 265 | """Test InterleaveMoeLayerStepParser correctly computes MoE layer count.""" |
| 266 | hf_config = Llama4Config( |
| 267 | text_config=Llama4TextConfig( |
| 268 | num_hidden_layers=32, |
| 269 | num_local_experts=64, |
| 270 | interleave_moe_layer_step=4, # Every 4th layer is MoE |
| 271 | ), |
| 272 | ) |
| 273 | |
| 274 | vllm_config = create_mock_vllm_config(hf_config) |
| 275 | |
| 276 | parser_chain = FfnMetrics.get_parser() |
| 277 | result = parser_chain.parse(vllm_config) |
| 278 | |
| 279 | assert result.num_moe_layers == 8 |
| 280 | |
| 281 | |
| 282 | def test_moe_layer_freq_parser(): |
| 283 | """Test MoeLayerFreqParser correctly computes MoE layer count.""" |
| 284 | hf_config = DeepseekV3Config( |
| 285 | num_hidden_layers=30, |
| 286 | n_routed_experts=64, |
| 287 | moe_layer_freq=3, # Every 3rd layer after first_k_dense_replace |
| 288 | first_k_dense_replace=6, # First 6 layers are dense |
| 289 | ) |
| 290 | vllm_config = create_mock_vllm_config(hf_config) |
| 291 | |
| 292 | parser_chain = FfnMetrics.get_parser() |
| 293 | result = parser_chain.parse(vllm_config) |
| 294 | |
| 295 | # Layers >= 6 and divisible by 3: 6, 9, 12, 15, 18, 21, 24, 27 |
| 296 | expected_moe_layers = len( |
| 297 | [layer for layer in range(30) if layer >= 6 and layer % 3 == 0] |
| 298 | ) |
| 299 | assert expected_moe_layers == 8 |
| 300 | assert result.num_moe_layers == expected_moe_layers |
| 301 | |
| 302 | |
| 303 | #### ComponentMetrics Tests #### |
| 304 | |
| 305 | |
| 306 | def test_attention_metrics_scaling(): |
| 307 | """Test that attention metrics scale proportionally with model dimensions.""" |
| 308 | base_hf_config = Qwen3Config( |
| 309 | hidden_size=2048, |
| 310 | num_attention_heads=16, |
| 311 | num_key_value_heads=16, |
| 312 | num_hidden_layers=12, |
| 313 | head_dim=128, |
| 314 | ) |
| 315 | |
| 316 | base_vllm_config = create_mock_vllm_config(base_hf_config) |
| 317 | base_metrics = AttentionMetrics.from_vllm_config(base_vllm_config) |
| 318 | |
| 319 | # Test scaling with number of layers |
| 320 | double_layers_hf_config = Qwen3Config( |
| 321 | hidden_size=2048, |
| 322 | num_attention_heads=16, |
| 323 | num_key_value_heads=16, |
| 324 | num_hidden_layers=24, # Double the layers |
| 325 | head_dim=128, |
| 326 | ) |
| 327 | double_layers_vllm_config = create_mock_vllm_config(double_layers_hf_config) |
| 328 | double_layers_metrics = AttentionMetrics.from_vllm_config(double_layers_vllm_config) |
| 329 | |
| 330 | ctx = ExecutionContext.from_single_request( |
| 331 | num_tokens=100, context_len=512, is_prefill=True |
| 332 | ) |
| 333 | |
| 334 | # FLOPS should double when layers double |
| 335 | base_flops = base_metrics.get_num_flops(ctx) |
| 336 | double_flops = double_layers_metrics.get_num_flops(ctx) |
| 337 | assert double_flops == 2 * base_flops |
| 338 | |
| 339 | # Read/write bytes should also scale proportionally |
| 340 | base_read = base_metrics.get_read_bytes(ctx) |
| 341 | double_read = double_layers_metrics.get_read_bytes(ctx) |
| 342 | assert double_read == 2 * base_read |
| 343 | |
| 344 | base_write = base_metrics.get_write_bytes(ctx) |
| 345 | double_write = double_layers_metrics.get_write_bytes(ctx) |
| 346 | assert double_write == 2 * base_write |
| 347 | |
| 348 | |
| 349 | def test_attention_metrics_grouped_query(): |
| 350 | """Test attention metrics handle grouped query attention correctly.""" |
| 351 | mha_hf_config = Qwen3Config( |
| 352 | hidden_size=4096, |
| 353 | num_attention_heads=32, |
| 354 | num_key_value_heads=32, # MHA |
| 355 | num_hidden_layers=1, |
| 356 | ) |
| 357 | mha_config = create_mock_vllm_config(mha_hf_config) |
| 358 | |
| 359 | gqa_hf_config = Qwen3Config( |
| 360 | hidden_size=4096, |
| 361 | num_attention_heads=32, |
| 362 | num_key_value_heads=8, # GQA with 4:1 ratio |
| 363 | num_hidden_layers=1, |
| 364 | ) |
| 365 | gqa_config = create_mock_vllm_config(gqa_hf_config) |
| 366 | |
| 367 | mha_metrics = AttentionMetrics.from_vllm_config(mha_config) |
| 368 | gqa_metrics = AttentionMetrics.from_vllm_config(gqa_config) |
| 369 | |
| 370 | ctx = ExecutionContext.from_single_request( |
| 371 | num_tokens=1, context_len=1024, is_prefill=False |
| 372 | ) |
| 373 | |
| 374 | # GQA should have less KV cache reads since fewer KV heads |
| 375 | mha_read = mha_metrics.get_read_bytes(ctx) |
| 376 | gqa_read = gqa_metrics.get_read_bytes(ctx) |
| 377 | assert gqa_read < mha_read |
| 378 | |
| 379 | |
| 380 | def test_ffn_metrics_scaling(): |
| 381 | """Test FFN metrics scale proportionally with model dimensions.""" |
| 382 | base_hf_config = Qwen3Config( |
| 383 | hidden_size=2048, |
| 384 | intermediate_size=8192, |
| 385 | num_hidden_layers=12, |
| 386 | ) |
| 387 | base_vllm_config = create_mock_vllm_config(base_hf_config) |
| 388 | base_metrics = FfnMetrics.from_vllm_config(base_vllm_config) |
| 389 | |
| 390 | # Test scaling with intermediate size |
| 391 | larger_ffn_hf_config = Qwen3Config( |
| 392 | hidden_size=2048, |
| 393 | intermediate_size=16384, # Double intermediate size |
| 394 | num_hidden_layers=12, |
| 395 | ) |
| 396 | larger_ffn_vllm_config = create_mock_vllm_config(larger_ffn_hf_config) |
| 397 | larger_ffn_metrics = FfnMetrics.from_vllm_config(larger_ffn_vllm_config) |
| 398 | |
| 399 | ctx = ExecutionContext.from_single_request( |
| 400 | num_tokens=100, context_len=512, is_prefill=True |
| 401 | ) |
| 402 | |
| 403 | # FLOPS should double when intermediate size doubles |
| 404 | base_flops = base_metrics.get_num_flops(ctx) |
| 405 | larger_flops = larger_ffn_metrics.get_num_flops(ctx) |
| 406 | assert larger_flops == base_flops * 2 |
| 407 | |
| 408 | |
| 409 | def test_moe_metrics_vs_dense(): |
| 410 | """Test MoE metrics versus dense metrics.""" |
| 411 | dense_hf_config = Qwen3Config( |
| 412 | hidden_size=2048, |
| 413 | intermediate_size=8192, |
| 414 | num_hidden_layers=12, |
| 415 | ) |
| 416 | dense_config = create_mock_vllm_config(dense_hf_config) |
| 417 | |
| 418 | moe_hf_config = Qwen3MoeConfig( |
| 419 | hidden_size=2048, |
| 420 | intermediate_size=8192, |
| 421 | num_hidden_layers=12, |
| 422 | num_experts=64, |
| 423 | num_experts_per_tok=2, # 2 routed expert |
| 424 | moe_intermediate_size=8192, |
| 425 | n_shared_experts=0, |
| 426 | ) |
| 427 | moe_config = create_mock_vllm_config(moe_hf_config) |
| 428 | |
| 429 | dense_metrics = FfnMetrics.from_vllm_config(dense_config) |
| 430 | moe_metrics = FfnMetrics.from_vllm_config(moe_config) |
| 431 | |
| 432 | ctx = ExecutionContext.from_single_request( |
| 433 | num_tokens=100, context_len=512, is_prefill=True |
| 434 | ) |
| 435 | |
| 436 | # MoE should have different compute/memory characteristics |
| 437 | dense_flops = dense_metrics.get_num_flops(ctx) |
| 438 | moe_flops = moe_metrics.get_num_flops(ctx) |
| 439 | |
| 440 | # 2 routed experts vs 1 dense. |
| 441 | assert moe_flops == dense_flops * 2 |
| 442 | |
| 443 | |
| 444 | def test_unembed_metrics_scaling(): |
| 445 | """Test unembedding metrics scale with vocab size.""" |
| 446 | small_vocab_hf_config = Qwen3Config( |
| 447 | hidden_size=2048, |
| 448 | vocab_size=32000, |
| 449 | ) |
| 450 | small_vocab_config = create_mock_vllm_config(small_vocab_hf_config) |
| 451 | |
| 452 | large_vocab_hf_config = Qwen3Config( |
| 453 | hidden_size=2048, |
| 454 | vocab_size=64000, # Double vocab size |
| 455 | ) |
| 456 | large_vocab_config = create_mock_vllm_config(large_vocab_hf_config) |
| 457 | |
| 458 | small_vocab_metrics = UnembedMetrics.from_vllm_config(small_vocab_config) |
| 459 | large_vocab_metrics = UnembedMetrics.from_vllm_config(large_vocab_config) |
| 460 | |
| 461 | ctx = ExecutionContext.from_single_request( |
| 462 | num_tokens=100, context_len=512, is_prefill=True |
| 463 | ) |
| 464 | |
| 465 | # FLOPS should double when vocab size doubles |
| 466 | small_flops = small_vocab_metrics.get_num_flops(ctx) |
| 467 | large_flops = large_vocab_metrics.get_num_flops(ctx) |
| 468 | assert large_flops == 2 * small_flops |
| 469 | |
| 470 | |
| 471 | def test_prefill_vs_decode_differences(): |
| 472 | """Test that prefill and decode have different memory access patterns.""" |
| 473 | hf_config = Qwen3Config( |
| 474 | hidden_size=2048, |
| 475 | num_attention_heads=16, |
| 476 | num_key_value_heads=16, |
| 477 | num_hidden_layers=1, |
| 478 | ) |
| 479 | config = create_mock_vllm_config(hf_config) |
| 480 | |
| 481 | metrics = AttentionMetrics.from_vllm_config(config) |
| 482 | |
| 483 | prefill_ctx = ExecutionContext.from_single_request( |
| 484 | num_tokens=512, context_len=512, is_prefill=True |
| 485 | ) |
| 486 | decode_ctx = ExecutionContext.from_single_request( |
| 487 | num_tokens=1, context_len=512, is_prefill=False |
| 488 | ) |
| 489 | |
| 490 | prefill_read = metrics.get_read_bytes(prefill_ctx) |
| 491 | decode_read = metrics.get_read_bytes(decode_ctx) |
| 492 | |
| 493 | assert prefill_read != decode_read |
| 494 | |
| 495 | |
| 496 | def test_model_metrics_aggregation(): |
| 497 | """Test ModelMetrics correctly aggregates across components.""" |
| 498 | hf_config = Qwen3Config( |
| 499 | hidden_size=2048, |
| 500 | num_attention_heads=16, |
| 501 | num_hidden_layers=12, |
| 502 | vocab_size=32000, |
| 503 | intermediate_size=8192, |
| 504 | ) |
| 505 | config = create_mock_vllm_config(hf_config) |
| 506 | |
| 507 | model_metrics = ModelMetrics(config) |
| 508 | ctx = ExecutionContext.from_single_request( |
| 509 | num_tokens=100, context_len=512, is_prefill=True |
| 510 | ) |
| 511 | |
| 512 | # Should have metrics for attention, ffn, and unembed |
| 513 | total_flops = model_metrics.get_num_flops(ctx) |
| 514 | breakdown = model_metrics.get_num_flops_breakdown(ctx) |
| 515 | |
| 516 | # Breakdown should sum to total |
| 517 | assert total_flops == sum(breakdown.values()) |
| 518 | |
| 519 | |
| 520 | def test_moe_expert_activation_proportional_scaling(): |
| 521 | """Test that routed expert metrics scale proportionally with num_experts_per_tok.""" |
| 522 | base_moe_config = Qwen3MoeConfig( |
| 523 | hidden_size=2048, |
| 524 | intermediate_size=8192, |
| 525 | num_hidden_layers=12, |
| 526 | num_experts=64, |
| 527 | num_experts_per_tok=1, # 1 expert per token |
| 528 | moe_intermediate_size=8192, |
| 529 | n_shared_experts=2, |
| 530 | ) |
| 531 | |
| 532 | double_experts_config = Qwen3MoeConfig( |
| 533 | hidden_size=2048, |
| 534 | intermediate_size=8192, |
| 535 | num_hidden_layers=12, |
| 536 | num_experts=64, |
| 537 | num_experts_per_tok=2, # 2 experts per token (double) |
| 538 | moe_intermediate_size=8192, |
| 539 | n_shared_experts=2, # Same shared experts |
| 540 | ) |
| 541 | |
| 542 | triple_experts_config = Qwen3MoeConfig( |
| 543 | hidden_size=2048, |
| 544 | intermediate_size=8192, |
| 545 | num_hidden_layers=12, |
| 546 | num_experts=64, |
| 547 | num_experts_per_tok=3, # 3 experts per token (triple) |
| 548 | moe_intermediate_size=8192, |
| 549 | n_shared_experts=2, # Same shared experts |
| 550 | ) |
| 551 | |
| 552 | base_vllm_config = create_mock_vllm_config(base_moe_config) |
| 553 | double_vllm_config = create_mock_vllm_config(double_experts_config) |
| 554 | triple_vllm_config = create_mock_vllm_config(triple_experts_config) |
| 555 | |
| 556 | base_metrics = FfnMetrics.from_vllm_config(base_vllm_config) |
| 557 | double_metrics = FfnMetrics.from_vllm_config(double_vllm_config) |
| 558 | triple_metrics = FfnMetrics.from_vllm_config(triple_vllm_config) |
| 559 | |
| 560 | ctx = ExecutionContext.from_single_request( |
| 561 | num_tokens=100, context_len=512, is_prefill=True |
| 562 | ) |
| 563 | |
| 564 | # Get total metrics - the key insight is that differences should be proportional |
| 565 | base_flops = base_metrics.get_num_flops(ctx) |
| 566 | double_flops = double_metrics.get_num_flops(ctx) |
| 567 | triple_flops = triple_metrics.get_num_flops(ctx) |
| 568 | |
| 569 | # The difference between double and base should equal one additional expert |
| 570 | one_expert_diff = double_flops - base_flops |
| 571 | |
| 572 | # The difference between triple and base should equal two additional experts |
| 573 | two_expert_diff = triple_flops - base_flops |
| 574 | |
| 575 | # Proportional scaling: 2 * (1 expert diff) should equal (2 expert diff) |
| 576 | assert two_expert_diff == 2 * one_expert_diff |
| 577 | |
| 578 | # Same logic applies to memory operations |
| 579 | base_read = base_metrics.get_read_bytes(ctx) |
| 580 | double_read = double_metrics.get_read_bytes(ctx) |
| 581 | triple_read = triple_metrics.get_read_bytes(ctx) |
| 582 | |
| 583 | one_expert_read_diff = double_read - base_read |
| 584 | two_expert_read_diff = triple_read - base_read |
| 585 | |
| 586 | assert two_expert_read_diff == 2 * one_expert_read_diff |
| 587 | |
| 588 | # Same for write bytes |
| 589 | base_write = base_metrics.get_write_bytes(ctx) |
| 590 | double_write = double_metrics.get_write_bytes(ctx) |
| 591 | triple_write = triple_metrics.get_write_bytes(ctx) |
| 592 | |
| 593 | one_expert_write_diff = double_write - base_write |
| 594 | two_expert_write_diff = triple_write - base_write |
| 595 | |
| 596 | assert two_expert_write_diff == 2 * one_expert_write_diff |
| 597 | |
| 598 | |
| 599 | def test_quantization_config_parser_fp8(): |
| 600 | """Test quantization parsers with fp8.""" |
| 601 | |
| 602 | class MockQuantConfig: |
| 603 | def get_name(self): |
| 604 | return "fp8" |
| 605 | |
| 606 | hf_config = Qwen3Config( |
| 607 | hidden_size=2048, num_attention_heads=16, num_hidden_layers=1 |
| 608 | ) |
| 609 | vllm_config = create_mock_vllm_config(hf_config, quant_config=MockQuantConfig()) |
| 610 | |
| 611 | attn_result = AttentionMetrics.get_parser().parse(vllm_config) |
| 612 | assert attn_result.weight_byte_size == 1 # fp8 |
| 613 | |
| 614 | ffn_result = FfnMetrics.get_parser().parse(vllm_config) |
| 615 | assert ffn_result.weight_byte_size == 1 # fp8 |
| 616 | |
| 617 | |
| 618 | def test_quantization_config_parser_mxfp4(): |
| 619 | """Test quantization parsers with mxfp4.""" |
| 620 | |
| 621 | class MockQuantConfig: |
| 622 | def get_name(self): |
| 623 | return "mxfp4" |
| 624 | |
| 625 | hf_config = Qwen3Config( |
| 626 | hidden_size=2048, intermediate_size=8192, num_hidden_layers=1 |
| 627 | ) |
| 628 | vllm_config = create_mock_vllm_config(hf_config, quant_config=MockQuantConfig()) |
| 629 | |
| 630 | ffn_result = FfnMetrics.get_parser().parse(vllm_config) |
| 631 | assert ffn_result.weight_byte_size == 0.5 # mxfp4 |
| 632 | |
| 633 | |
| 634 | #### Per-GPU Tests #### |
| 635 | |
| 636 | |
| 637 | def test_attention_per_gpu_with_tensor_parallelism(): |
| 638 | """Test attention metrics with tensor parallelism - per_gpu vs global.""" |
| 639 | hf_config = Qwen3Config( |
| 640 | hidden_size=4096, |
| 641 | num_attention_heads=32, |
| 642 | num_key_value_heads=8, |
| 643 | num_hidden_layers=24, |
| 644 | ) |
| 645 | |
| 646 | # Test with TP=4 |
| 647 | vllm_config = create_mock_vllm_config(hf_config, tensor_parallel_size=4) |
| 648 | metrics = AttentionMetrics.from_vllm_config(vllm_config) |
| 649 | |
| 650 | ctx = ExecutionContext.from_single_request( |
| 651 | num_tokens=128, context_len=1024, is_prefill=True |
| 652 | ) |
| 653 | |
| 654 | # Get global and per-gpu metrics |
| 655 | global_flops = metrics.get_num_flops(ctx, per_gpu=False) |
| 656 | per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True) |
| 657 | |
| 658 | # With TP=4, global flops should be 4x per-gpu flops (heads divided by 4) |
| 659 | assert global_flops == 4 * per_gpu_flops |
| 660 | |
| 661 | # Same for read/write bytes |
| 662 | global_read = metrics.get_read_bytes(ctx, per_gpu=False) |
| 663 | per_gpu_read = metrics.get_read_bytes(ctx, per_gpu=True) |
| 664 | # Reads should scale similarly (weight reads are divided by TP) |
| 665 | assert global_read > per_gpu_read |
| 666 | |
| 667 | global_write = metrics.get_write_bytes(ctx, per_gpu=False) |
| 668 | per_gpu_write = metrics.get_write_bytes(ctx, per_gpu=True) |
| 669 | assert global_write > per_gpu_write |
| 670 | |
| 671 | |
| 672 | def test_attention_per_gpu_with_pipeline_parallelism(): |
| 673 | """Test attention metrics with pipeline parallelism - per_gpu vs global.""" |
| 674 | hf_config = Qwen3Config( |
| 675 | hidden_size=2048, |
| 676 | num_attention_heads=16, |
| 677 | num_hidden_layers=32, |
| 678 | ) |
| 679 | |
| 680 | # Test with PP=4 |
| 681 | vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=4) |
| 682 | metrics = AttentionMetrics.from_vllm_config(vllm_config) |
| 683 | |
| 684 | ctx = ExecutionContext.from_single_request( |
| 685 | num_tokens=100, context_len=512, is_prefill=False |
| 686 | ) |
| 687 | |
| 688 | # Get global and per-gpu metrics |
| 689 | global_flops = metrics.get_num_flops(ctx, per_gpu=False) |
| 690 | per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True) |
| 691 | |
| 692 | # With PP=4, global flops should be 4x per-gpu flops (layers divided by 4) |
| 693 | assert global_flops == 4 * per_gpu_flops |
| 694 | |
| 695 | global_read = metrics.get_read_bytes(ctx, per_gpu=False) |
| 696 | per_gpu_read = metrics.get_read_bytes(ctx, per_gpu=True) |
| 697 | assert global_read == 4 * per_gpu_read |
| 698 | |
| 699 | |
| 700 | def test_ffn_per_gpu_with_tensor_parallelism(): |
| 701 | """Test FFN metrics with tensor parallelism - per_gpu vs global.""" |
| 702 | hf_config = Qwen3Config( |
| 703 | hidden_size=4096, |
| 704 | intermediate_size=14336, |
| 705 | num_hidden_layers=32, |
| 706 | ) |
| 707 | |
| 708 | # Test with DP=2, TP=4 (ffn_tp_size will be 8) |
| 709 | vllm_config = create_mock_vllm_config( |
| 710 | hf_config, |
| 711 | data_parallel_size=2, |
| 712 | tensor_parallel_size=4, |
| 713 | ) |
| 714 | metrics = FfnMetrics.from_vllm_config(vllm_config) |
| 715 | |
| 716 | # ffn_tp_size should be dp_size * tp_size = 8 (when EP not enabled) |
| 717 | assert metrics.ffn_tp_size == 8 |
| 718 | |
| 719 | ctx = ExecutionContext.from_single_request( |
| 720 | num_tokens=128, context_len=2048, is_prefill=True |
| 721 | ) |
| 722 | |
| 723 | # Get global and per-gpu metrics |
| 724 | global_flops = metrics.get_num_flops(ctx, per_gpu=False) |
| 725 | per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True) |
| 726 | |
| 727 | # With ffn_tp_size=8, global should be 8x per-gpu |
| 728 | assert global_flops == 8 * per_gpu_flops |
| 729 | |
| 730 | |
| 731 | def test_ffn_per_gpu_with_pipeline_parallelism(): |
| 732 | """Test FFN metrics with pipeline parallelism - per_gpu vs global.""" |
| 733 | hf_config = Qwen3Config( |
| 734 | hidden_size=2048, |
| 735 | intermediate_size=8192, |
| 736 | num_hidden_layers=24, |
| 737 | ) |
| 738 | |
| 739 | # Test with PP=6 |
| 740 | vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=6) |
| 741 | metrics = FfnMetrics.from_vllm_config(vllm_config) |
| 742 | |
| 743 | ctx = ExecutionContext.from_single_request( |
| 744 | num_tokens=100, context_len=512, is_prefill=True |
| 745 | ) |
| 746 | |
| 747 | # Get global and per-gpu metrics |
| 748 | global_flops = metrics.get_num_flops(ctx, per_gpu=False) |
| 749 | per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True) |
| 750 | |
| 751 | # With PP=6, global should be 6x per-gpu (layers divided by 6) |
| 752 | assert global_flops == 6 * per_gpu_flops |
| 753 | |
| 754 | |
| 755 | def test_moe_per_gpu_with_expert_parallelism(): |
| 756 | """Test MoE metrics with expert parallelism, covering num_activated_experts.""" |
| 757 | hf_config = Qwen3MoeConfig( |
| 758 | hidden_size=2048, |
| 759 | intermediate_size=8192, |
| 760 | num_hidden_layers=24, |
| 761 | num_experts=64, |
| 762 | num_experts_per_tok=8, |
| 763 | moe_intermediate_size=14336, |
| 764 | n_shared_experts=2, |
| 765 | ) |
| 766 | |
| 767 | # Test with DP=2, TP=4, EP enabled (ffn_ep_size will be 8) |
| 768 | vllm_config = create_mock_vllm_config( |
| 769 | hf_config, |
| 770 | data_parallel_size=2, |
| 771 | tensor_parallel_size=4, |
| 772 | enable_expert_parallel=True, |
| 773 | ) |
| 774 | metrics = FfnMetrics.from_vllm_config(vllm_config) |
| 775 | |
| 776 | # When EP enabled, ffn_ep_size = dp_size * tp_size = 8 |
| 777 | assert metrics.ffn_ep_size == 8 |
| 778 | assert metrics.ffn_tp_size == 1 |
| 779 | |
| 780 | ctx = ExecutionContext.from_single_request( |
| 781 | num_tokens=100, context_len=512, is_prefill=True |
| 782 | ) |
| 783 | |
| 784 | # Get per-gpu metrics |
| 785 | per_gpu_read_breakdown = metrics.get_read_bytes_breakdown(ctx, per_gpu=True) |
| 786 | global_read_breakdown = metrics.get_read_bytes_breakdown(ctx, per_gpu=False) |
| 787 | |
| 788 | # Verify that routed expert weight reads are reasonable |
| 789 | # With per_gpu=True, each GPU has 64/8 = 8 experts |
| 790 | # T=100, E_per_gpu=8/8=1, so T*E=100 expert activations |
| 791 | # num_activated_experts should be min(100, 8) = 8 |
| 792 | |
| 793 | # Check that weight reads scale appropriately |
| 794 | # Global has all 64 experts, per-gpu has 8 experts |
| 795 | # So weight reads should reflect this difference |
| 796 | if "routed_up_gate_weights" in per_gpu_read_breakdown: |
| 797 | per_gpu_weight_reads = per_gpu_read_breakdown["routed_up_gate_weights"] |
| 798 | global_weight_reads = global_read_breakdown["routed_up_gate_weights"] |
| 799 | |
| 800 | # The ratio should reflect the expert count difference |
| 801 | # This verifies the bug fix works correctly |
| 802 | assert per_gpu_weight_reads < global_weight_reads |
| 803 | |
| 804 | # Global should read more experts than per-gpu |
| 805 | # Exact ratio depends on num_activated_experts calculation |
| 806 | ratio = global_weight_reads / per_gpu_weight_reads |
| 807 | # Should be > 1 since global has more experts to read |
| 808 | assert ratio > 1 |
| 809 | |
| 810 | |
| 811 | def test_moe_per_gpu_expert_activation_accounting(): |
| 812 | """Test MoE expert-activation accounting with small batch sizes.""" |
| 813 | hf_config = Qwen3MoeConfig( |
| 814 | hidden_size=2048, |
| 815 | intermediate_size=8192, |
| 816 | num_hidden_layers=12, |
| 817 | num_experts=64, |
| 818 | num_experts_per_tok=8, |
| 819 | moe_intermediate_size=14336, |
| 820 | n_shared_experts=0, # No shared experts for this test |
| 821 | ) |
| 822 | |
| 823 | # Test with EP=8 |
| 824 | vllm_config = create_mock_vllm_config( |
| 825 | hf_config, |
| 826 | data_parallel_size=8, |
| 827 | enable_expert_parallel=True, |
| 828 | ) |
| 829 | metrics = FfnMetrics.from_vllm_config(vllm_config) |
| 830 | |
| 831 | # Small batch: T=10, E_per_gpu=8/8=1 |
| 832 | # Each GPU: T*E = 10*1 = 10 activations |
| 833 | # Experts per GPU: 64/8 = 8 |
| 834 | # So num_activated_experts should be min(10, 8) = 8 |
| 835 | small_ctx = ExecutionContext.from_single_request( |
| 836 | num_tokens=10, context_len=512, is_prefill=True |
| 837 | ) |
| 838 | small_read = metrics.get_read_bytes_breakdown(small_ctx, per_gpu=True) |
| 839 | |
| 840 | # Large batch: T=1000, E_per_gpu=1 |
| 841 | # Each GPU: T*E = 1000*1 = 1000 activations |
| 842 | # Experts per GPU: 8 |
| 843 | # So num_activated_experts should be min(1000, 8) = 8 (all experts activated) |
| 844 | large_ctx = ExecutionContext.from_single_request( |
| 845 | num_tokens=1000, context_len=512, is_prefill=True |
| 846 | ) |
| 847 | large_read = metrics.get_read_bytes_breakdown(large_ctx, per_gpu=True) |
| 848 | |
| 849 | # Weight reads should be similar (both activate all 8 experts per GPU) |
| 850 | # But activation reads should differ (proportional to T*E) |
| 851 | if "routed_up_gate_weights" in small_read: |
| 852 | small_weight = small_read["routed_up_gate_weights"] |
| 853 | large_weight = large_read["routed_up_gate_weights"] |
| 854 | |
| 855 | # Weight reads should be the same (both read all 8 experts) |
| 856 | assert small_weight == large_weight |
| 857 | |
| 858 | # But input activation reads should scale with T*E |
| 859 | small_input = small_read["routed_up_gate_input"] |
| 860 | large_input = large_read["routed_up_gate_input"] |
| 861 | assert large_input == 100 * small_input # 1000/10 = 100x |
| 862 | |
| 863 | |
| 864 | def test_unembed_per_gpu_with_tensor_parallelism(): |
| 865 | """Test unembed metrics with tensor parallelism - per_gpu vs global.""" |
| 866 | hf_config = Qwen3Config( |
| 867 | hidden_size=4096, |
| 868 | vocab_size=128000, |
| 869 | ) |
| 870 | |
| 871 | # Test with TP=8 |
| 872 | vllm_config = create_mock_vllm_config(hf_config, tensor_parallel_size=8) |
| 873 | metrics = UnembedMetrics.from_vllm_config(vllm_config) |
| 874 | |
| 875 | ctx = ExecutionContext.from_single_request( |
| 876 | num_tokens=100, context_len=512, is_prefill=True |
| 877 | ) |
| 878 | |
| 879 | # Get global and per-gpu metrics |
| 880 | global_flops = metrics.get_num_flops(ctx, per_gpu=False) |
| 881 | per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True) |
| 882 | |
| 883 | # With TP=8, vocab is divided by 8, so global should be 8x per-gpu |
| 884 | assert global_flops == 8 * per_gpu_flops |
| 885 | |
| 886 | # For read bytes, weight reads scale with TP but input reads don't (replicated) |
| 887 | global_read_breakdown = metrics.get_read_bytes_breakdown(ctx, per_gpu=False) |
| 888 | per_gpu_read_breakdown = metrics.get_read_bytes_breakdown(ctx, per_gpu=True) |
| 889 | |
| 890 | # Input reads should be the same (replicated across TP ranks) |
| 891 | assert global_read_breakdown["input"] == per_gpu_read_breakdown["input"] |
| 892 | |
| 893 | # Weight reads should scale 8x (divided by TP) |
| 894 | assert global_read_breakdown["weight"] == 8 * per_gpu_read_breakdown["weight"] |
| 895 | |
| 896 | |
| 897 | def test_model_metrics_per_gpu_aggregation(): |
| 898 | """Test ModelMetrics correctly aggregates per_gpu metrics across components.""" |
| 899 | hf_config = Qwen3Config( |
| 900 | hidden_size=2048, |
| 901 | num_attention_heads=16, |
| 902 | num_hidden_layers=12, |
| 903 | vocab_size=32000, |
| 904 | intermediate_size=8192, |
| 905 | ) |
| 906 | |
| 907 | # Test with mixed parallelism: TP=2, PP=2 |
| 908 | vllm_config = create_mock_vllm_config( |
| 909 | hf_config, |
| 910 | tensor_parallel_size=2, |
| 911 | pipeline_parallel_size=2, |
| 912 | ) |
| 913 | |
| 914 | model_metrics = ModelMetrics(vllm_config) |
| 915 | ctx = ExecutionContext.from_single_request( |
| 916 | num_tokens=100, context_len=512, is_prefill=True |
| 917 | ) |
| 918 | |
| 919 | # Get breakdowns for both modes |
| 920 | per_gpu_breakdown = model_metrics.get_num_flops_breakdown(ctx, per_gpu=True) |
| 921 | global_breakdown = model_metrics.get_num_flops_breakdown(ctx, per_gpu=False) |
| 922 | |
| 923 | # Verify breakdown sums match totals |
| 924 | per_gpu_total = model_metrics.get_num_flops(ctx, per_gpu=True) |
| 925 | global_total = model_metrics.get_num_flops(ctx, per_gpu=False) |
| 926 | |
| 927 | assert per_gpu_total == sum(per_gpu_breakdown.values()) |
| 928 | assert global_total == sum(global_breakdown.values()) |
| 929 | |
| 930 | # Global should be larger than per-gpu due to parallelism |
| 931 | assert global_total > per_gpu_total |
| 932 | |
| 933 | # With TP=2 and PP=2, the ratio depends on which parallelism applies to |
| 934 | # which component but we can verify that global is reasonably larger |
| 935 | ratio = global_total / per_gpu_total |
| 936 | assert ratio > 1 # Should be between PP and TP*PP depending on component mix |
| 937 | |
| 938 | |
| 939 | def test_attention_per_gpu_heads_not_evenly_divisible(): |
| 940 | """Test attention with heads not evenly divisible by TP.""" |
| 941 | hf_config = Qwen3Config( |
| 942 | hidden_size=2048, |
| 943 | num_attention_heads=17, # Not divisible by 4 |
| 944 | num_key_value_heads=5, # Not divisible by 4 |
| 945 | num_hidden_layers=8, |
| 946 | ) |
| 947 | |
| 948 | vllm_config = create_mock_vllm_config(hf_config, tensor_parallel_size=4) |
| 949 | metrics = AttentionMetrics.from_vllm_config(vllm_config) |
| 950 | |
| 951 | ctx = ExecutionContext.from_single_request( |
| 952 | num_tokens=64, context_len=256, is_prefill=True |
| 953 | ) |
| 954 | |
| 955 | # Should not crash and should handle max(1, ...) correctly |
| 956 | per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True) |
| 957 | global_flops = metrics.get_num_flops(ctx, per_gpu=False) |
| 958 | |
| 959 | # Both should be positive |
| 960 | assert per_gpu_flops > 0 |
| 961 | assert global_flops > 0 |
| 962 | assert global_flops > per_gpu_flops |
| 963 | |
| 964 | |
| 965 | # INT4 / FP4 quantization methods (weight_byte_size == 0.5) |
| 966 | _INT4_FP4_METHODS = [m for m, s in _QUANT_WEIGHT_BYTE_SIZE.items() if s == 0.5] |
| 967 | |
| 968 | |
| 969 | @pytest.mark.parametrize("quant_method", _INT4_FP4_METHODS) |
| 970 | def test_quantization_config_parser_int4_methods(quant_method): |
| 971 | """Test quantization parsers with INT4/FP4 methods (0.5 bytes).""" |
| 972 | |
| 973 | class MockQuantConfig: |
| 974 | def get_name(self): |
| 975 | return quant_method |
| 976 | |
| 977 | hf_config = Qwen3Config( |
| 978 | hidden_size=2048, |
| 979 | num_attention_heads=16, |
| 980 | intermediate_size=8192, |
| 981 | num_hidden_layers=1, |
| 982 | ) |
| 983 | vllm_config = create_mock_vllm_config(hf_config, quant_config=MockQuantConfig()) |
| 984 | |
| 985 | attn_result = AttentionMetrics.get_parser().parse(vllm_config) |
| 986 | assert attn_result.weight_byte_size == 0.5, ( |
| 987 | f"Expected 0.5 for {quant_method}, got {attn_result.weight_byte_size}" |
| 988 | ) |
| 989 | |
| 990 | ffn_result = FfnMetrics.get_parser().parse(vllm_config) |
| 991 | assert ffn_result.weight_byte_size == 0.5, ( |
| 992 | f"Expected 0.5 for {quant_method}, got {ffn_result.weight_byte_size}" |
| 993 | ) |
| 994 | |
| 995 | |
| 996 | # FP8 / INT8 quantization methods (weight_byte_size == 1) |
| 997 | _FP8_INT8_METHODS = [m for m, s in _QUANT_WEIGHT_BYTE_SIZE.items() if s == 1] |
| 998 | |
| 999 | |
| 1000 | @pytest.mark.parametrize("quant_method", _FP8_INT8_METHODS) |
| 1001 | def test_quantization_config_parser_fp8_methods(quant_method): |
| 1002 | """Test quantization parsers with FP8/INT8 methods (1 byte).""" |
| 1003 | |
| 1004 | class MockQuantConfig: |
| 1005 | def get_name(self): |
| 1006 | return quant_method |
| 1007 | |
| 1008 | hf_config = Qwen3Config( |
| 1009 | hidden_size=2048, |
| 1010 | num_attention_heads=16, |
| 1011 | intermediate_size=8192, |
| 1012 | num_hidden_layers=1, |
| 1013 | ) |
| 1014 | vllm_config = create_mock_vllm_config(hf_config, quant_config=MockQuantConfig()) |
| 1015 | |
| 1016 | attn_result = AttentionMetrics.get_parser().parse(vllm_config) |
| 1017 | assert attn_result.weight_byte_size == 1, ( |
| 1018 | f"Expected 1 for {quant_method}, got {attn_result.weight_byte_size}" |
| 1019 | ) |
| 1020 | |
| 1021 | ffn_result = FfnMetrics.get_parser().parse(vllm_config) |
| 1022 | assert ffn_result.weight_byte_size == 1, ( |
| 1023 | f"Expected 1 for {quant_method}, got {ffn_result.weight_byte_size}" |
| 1024 | ) |
| 1025 | |
| 1026 | |
| 1027 | def test_quantization_config_parser_unknown_method(): |
| 1028 | """Test that an unrecognized quant method raises InvalidComponent.""" |
| 1029 | |
| 1030 | class MockQuantConfig: |
| 1031 | def get_name(self): |
| 1032 | return "unknown_quant_method" |
| 1033 | |
| 1034 | hf_config = Qwen3Config( |
| 1035 | hidden_size=2048, |
| 1036 | num_attention_heads=16, |
| 1037 | intermediate_size=8192, |
| 1038 | num_hidden_layers=1, |
| 1039 | ) |
| 1040 | vllm_config = create_mock_vllm_config(hf_config, quant_config=MockQuantConfig()) |
| 1041 | |
| 1042 | with pytest.raises(InvalidComponent): |
| 1043 | AttentionMetrics.get_parser().parse(vllm_config) |
| 1044 | |
| 1045 | with pytest.raises(InvalidComponent): |
| 1046 | FfnMetrics.get_parser().parse(vllm_config) |
| 1047 | |
| 1048 | |
| 1049 | def test_quantized_model_metrics_aggregation(): |
| 1050 | """Test that ModelMetrics works end-to-end with a quantized model config.""" |
| 1051 | |
| 1052 | class MockQuantConfig: |
| 1053 | def get_name(self): |
| 1054 | return "gptq" |
| 1055 | |
| 1056 | hf_config = Qwen3Config( |
| 1057 | hidden_size=2048, |
| 1058 | num_attention_heads=16, |
| 1059 | num_hidden_layers=12, |
| 1060 | vocab_size=32000, |
| 1061 | intermediate_size=8192, |
| 1062 | ) |
| 1063 | vllm_config = create_mock_vllm_config(hf_config, quant_config=MockQuantConfig()) |
| 1064 | |
| 1065 | model_metrics = ModelMetrics(vllm_config) |
| 1066 | ctx = ExecutionContext.from_single_request( |
| 1067 | num_tokens=100, context_len=512, is_prefill=True |
| 1068 | ) |
| 1069 | |
| 1070 | # Should not crash and should produce valid metrics |
| 1071 | total_flops = model_metrics.get_num_flops(ctx) |
| 1072 | breakdown = model_metrics.get_num_flops_breakdown(ctx) |
| 1073 | |
| 1074 | assert total_flops > 0 |
| 1075 | assert total_flops == sum(breakdown.values()) |
| 1076 | |
| 1077 | |
| 1078 | #### MLA Attention Tests #### |
| 1079 | |
| 1080 | |
| 1081 | def test_mla_config_parser(): |
| 1082 | """Test MLAConfigParser extracts MLA-specific fields from DeepseekV3Config.""" |
| 1083 | hf_config = DeepseekV3Config( |
| 1084 | hidden_size=7168, |
| 1085 | num_attention_heads=128, |
| 1086 | num_hidden_layers=61, |
| 1087 | kv_lora_rank=512, |
| 1088 | qk_nope_head_dim=128, |
| 1089 | qk_rope_head_dim=64, |
| 1090 | v_head_dim=128, |
| 1091 | q_lora_rank=1536, |
| 1092 | ) |
| 1093 | vllm_config = create_mock_vllm_config(hf_config) |
| 1094 | |
| 1095 | parser_chain = MLAAttentionMetrics.get_parser() |
| 1096 | result = parser_chain.parse(vllm_config) |
| 1097 | |
| 1098 | assert result.kv_lora_rank == 512 |
| 1099 | assert result.qk_nope_head_dim == 128 |
| 1100 | assert result.qk_rope_head_dim == 64 |
| 1101 | assert result.v_head_dim == 128 |
| 1102 | assert result.q_lora_rank == 1536 |
| 1103 | assert result.num_attention_heads == 128 |
| 1104 | assert result.hidden_size == 7168 |
| 1105 | |
| 1106 | |
| 1107 | def test_mla_attention_metrics_decode(): |
| 1108 | """Test MLA decode metrics use compressed KV cache, not standard head_dim.""" |
| 1109 | hf_config = DeepseekV3Config( |
| 1110 | hidden_size=7168, |
| 1111 | num_attention_heads=128, |
| 1112 | num_hidden_layers=1, |
| 1113 | kv_lora_rank=512, |
| 1114 | qk_nope_head_dim=128, |
| 1115 | qk_rope_head_dim=64, |
| 1116 | v_head_dim=128, |
| 1117 | q_lora_rank=1536, |
| 1118 | ) |
| 1119 | vllm_config = create_mock_vllm_config(hf_config) |
| 1120 | metrics = MLAAttentionMetrics.from_vllm_config(vllm_config) |
| 1121 | |
| 1122 | # Single decode token with 1024 context |
| 1123 | ctx = ExecutionContext.from_single_request( |
| 1124 | num_tokens=1, context_len=1024, is_prefill=False |
| 1125 | ) |
| 1126 | |
| 1127 | write_breakdown = metrics.get_write_bytes_breakdown(ctx, per_gpu=False) |
| 1128 | |
| 1129 | # KV cache write should be 1 * (512 + 64) * cache_byte_size * 1 layer |
| 1130 | # = 576 * 2 = 1152 bytes (for bfloat16 cache) |
| 1131 | kv_compressed_dim = 512 + 64 # kv_lora_rank + qk_rope_head_dim |
| 1132 | expected_kv_cache_write = 1 * kv_compressed_dim * 2 * 1 # T * dim * bytes * L |
| 1133 | assert write_breakdown["kv_cache"] == expected_kv_cache_write |
| 1134 | |
| 1135 | # Verify read bytes include compressed KV cache reads for context |
| 1136 | read_breakdown = metrics.get_read_bytes_breakdown(ctx, per_gpu=False) |
| 1137 | assert "attn_input" in read_breakdown |
| 1138 | assert read_breakdown["attn_input"] > 0 |
| 1139 | |
| 1140 | |
| 1141 | def test_mla_attention_metrics_prefill(): |
| 1142 | """Test MLA prefill metrics account for low-rank Q and KV projections.""" |
| 1143 | hf_config = DeepseekV3Config( |
| 1144 | hidden_size=7168, |
| 1145 | num_attention_heads=128, |
| 1146 | num_hidden_layers=1, |
| 1147 | kv_lora_rank=512, |
| 1148 | qk_nope_head_dim=128, |
| 1149 | qk_rope_head_dim=64, |
| 1150 | v_head_dim=128, |
| 1151 | q_lora_rank=1536, |
| 1152 | ) |
| 1153 | vllm_config = create_mock_vllm_config(hf_config) |
| 1154 | metrics = MLAAttentionMetrics.from_vllm_config(vllm_config) |
| 1155 | |
| 1156 | ctx = ExecutionContext.from_single_request( |
| 1157 | num_tokens=2048, context_len=2048, is_prefill=True |
| 1158 | ) |
| 1159 | |
| 1160 | flops_breakdown = metrics.get_num_flops_breakdown(ctx, per_gpu=False) |
| 1161 | |
| 1162 | # Should have two-stage Q projection (q_a and q_b) |
| 1163 | assert "q_a_proj" in flops_breakdown |
| 1164 | assert "q_b_proj" in flops_breakdown |
| 1165 | assert "q_proj" not in flops_breakdown # Since q_lora_rank is not None |
| 1166 | |
| 1167 | # Should have KV projections |
| 1168 | assert "kv_a_proj" in flops_breakdown |
| 1169 | assert "kv_b_proj" in flops_breakdown |
| 1170 | |
| 1171 | # Should have attention and output |
| 1172 | assert "attn_qk" in flops_breakdown |
| 1173 | assert "attn_av" in flops_breakdown |
| 1174 | assert "out_proj" in flops_breakdown |
| 1175 | |
| 1176 | # Verify q_a_proj: 2 * T * D * q_lora_rank * L |
| 1177 | expected_q_a = 2 * 2048 * 7168 * 1536 * 1 |
| 1178 | assert flops_breakdown["q_a_proj"] == expected_q_a |
| 1179 | |
| 1180 | # Verify kv_a_proj: 2 * T * D * (kv_lora_rank + qk_rope_head_dim) * L |
| 1181 | expected_kv_a = 2 * 2048 * 7168 * (512 + 64) * 1 |
| 1182 | assert flops_breakdown["kv_a_proj"] == expected_kv_a |
| 1183 | |
| 1184 | |
| 1185 | def test_mla_kv_cache_vs_standard_attention(): |
| 1186 | """Test MLA KV cache writes are dramatically smaller than standard MHA.""" |
| 1187 | # MLA config (DeepSeek-V3 style) |
| 1188 | mla_config = DeepseekV3Config( |
| 1189 | hidden_size=7168, |
| 1190 | num_attention_heads=128, |
| 1191 | num_hidden_layers=1, |
| 1192 | kv_lora_rank=512, |
| 1193 | qk_nope_head_dim=128, |
| 1194 | qk_rope_head_dim=64, |
| 1195 | v_head_dim=128, |
| 1196 | ) |
| 1197 | mla_vllm_config = create_mock_vllm_config(mla_config) |
| 1198 | mla_metrics = MLAAttentionMetrics.from_vllm_config(mla_vllm_config) |
| 1199 | |
| 1200 | # Standard MHA config with same num_heads and head_dim |
| 1201 | standard_config = Qwen3Config( |
| 1202 | hidden_size=7168, |
| 1203 | num_attention_heads=128, |
| 1204 | num_key_value_heads=128, # MHA: same as num_heads |
| 1205 | num_hidden_layers=1, |
| 1206 | head_dim=128, |
| 1207 | ) |
| 1208 | standard_vllm_config = create_mock_vllm_config(standard_config) |
| 1209 | standard_metrics = AttentionMetrics.from_vllm_config(standard_vllm_config) |
| 1210 | |
| 1211 | # Compare KV cache write for 100 tokens |
| 1212 | ctx = ExecutionContext.from_single_request( |
| 1213 | num_tokens=100, context_len=100, is_prefill=True |
| 1214 | ) |
| 1215 | |
| 1216 | mla_write = mla_metrics.get_write_bytes_breakdown(ctx, per_gpu=False) |
| 1217 | standard_write = standard_metrics.get_write_bytes_breakdown(ctx, per_gpu=False) |
| 1218 | |
| 1219 | # MLA: T * (kv_lora_rank + qk_rope_head_dim) * cache_bytes * L |
| 1220 | # = 100 * 576 * 2 * 1 = 115,200 |
| 1221 | mla_kv_cache = mla_write["kv_cache"] |
| 1222 | |
| 1223 | # Standard: 2 * T * num_kv_heads * head_dim * cache_bytes * L |
| 1224 | # = 2 * 100 * 128 * 128 * 2 * 1 = 6,553,600 |
| 1225 | standard_kv_cache = standard_write["kv_cache"] |
| 1226 | |
| 1227 | # MLA KV cache should be dramatically smaller (about 57x) |
| 1228 | assert mla_kv_cache < standard_kv_cache |
| 1229 | ratio = standard_kv_cache / mla_kv_cache |
| 1230 | assert ratio > 50 # Should be ~56.9x |
| 1231 | |
| 1232 | |
| 1233 | def test_mla_per_gpu_with_tensor_parallelism(): |
| 1234 | """Test MLA metrics with tensor parallelism.""" |
| 1235 | hf_config = DeepseekV3Config( |
| 1236 | hidden_size=7168, |
| 1237 | num_attention_heads=128, |
| 1238 | num_hidden_layers=8, |
| 1239 | kv_lora_rank=512, |
| 1240 | qk_nope_head_dim=128, |
| 1241 | qk_rope_head_dim=64, |
| 1242 | v_head_dim=128, |
| 1243 | q_lora_rank=1536, |
| 1244 | ) |
| 1245 | |
| 1246 | # Test with TP=8 |
| 1247 | vllm_config = create_mock_vllm_config(hf_config, tensor_parallel_size=8) |
| 1248 | metrics = MLAAttentionMetrics.from_vllm_config(vllm_config) |
| 1249 | |
| 1250 | ctx = ExecutionContext.from_single_request( |
| 1251 | num_tokens=64, context_len=1024, is_prefill=True |
| 1252 | ) |
| 1253 | |
| 1254 | global_flops = metrics.get_num_flops(ctx, per_gpu=False) |
| 1255 | per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True) |
| 1256 | |
| 1257 | # Both should be positive |
| 1258 | assert global_flops > 0 |
| 1259 | assert per_gpu_flops > 0 |
| 1260 | # Global should exceed per-GPU |
| 1261 | assert global_flops > per_gpu_flops |
| 1262 | |
| 1263 | |
| 1264 | def test_mla_per_gpu_with_pipeline_parallelism(): |
| 1265 | """Test MLA metrics with pipeline parallelism.""" |
| 1266 | hf_config = DeepseekV3Config( |
| 1267 | hidden_size=7168, |
| 1268 | num_attention_heads=128, |
| 1269 | num_hidden_layers=16, # Divisible by PP |
| 1270 | kv_lora_rank=512, |
| 1271 | qk_nope_head_dim=128, |
| 1272 | qk_rope_head_dim=64, |
| 1273 | v_head_dim=128, |
| 1274 | q_lora_rank=1536, |
| 1275 | ) |
| 1276 | |
| 1277 | vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=4) |
| 1278 | metrics = MLAAttentionMetrics.from_vllm_config(vllm_config) |
| 1279 | |
| 1280 | ctx = ExecutionContext.from_single_request( |
| 1281 | num_tokens=1, context_len=512, is_prefill=False |
| 1282 | ) |
| 1283 | |
| 1284 | global_flops = metrics.get_num_flops(ctx, per_gpu=False) |
| 1285 | per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True) |
| 1286 | |
| 1287 | # With PP=4, layers are divided by 4 |
| 1288 | assert global_flops == 4 * per_gpu_flops |
| 1289 | |
| 1290 | |
| 1291 | def test_mla_model_metrics_excludes_standard_attention(): |
| 1292 | """Test that ModelMetrics uses MLAAttentionMetrics, not AttentionMetrics, |
| 1293 | for DeepSeek MLA models.""" |
| 1294 | hf_config = DeepseekV3Config( |
| 1295 | hidden_size=7168, |
| 1296 | num_attention_heads=128, |
| 1297 | num_hidden_layers=4, |
| 1298 | kv_lora_rank=512, |
| 1299 | qk_nope_head_dim=128, |
| 1300 | qk_rope_head_dim=64, |
| 1301 | v_head_dim=128, |
| 1302 | q_lora_rank=1536, |
| 1303 | ) |
| 1304 | vllm_config = create_mock_vllm_config(hf_config) |
| 1305 | model_metrics = ModelMetrics(vllm_config) |
| 1306 | |
| 1307 | # Should have MLAAttentionMetrics but NOT standard AttentionMetrics |
| 1308 | component_types = [m.component_type() for m in model_metrics.metrics] |
| 1309 | assert "mla_attn" in component_types |
| 1310 | assert "attn" not in component_types |
| 1311 | |
| 1312 | # Should still have FFN and unembed |
| 1313 | assert "ffn" in component_types |
| 1314 | assert "unembed" in component_types |
| 1315 | |
| 1316 | # Breakdowns should work end-to-end |
| 1317 | ctx = ExecutionContext.from_single_request( |
| 1318 | num_tokens=100, context_len=512, is_prefill=True |
| 1319 | ) |
| 1320 | total_flops = model_metrics.get_num_flops(ctx) |
| 1321 | breakdown = model_metrics.get_num_flops_breakdown(ctx) |
| 1322 | assert total_flops == sum(breakdown.values()) |
| 1323 | assert total_flops > 0 |
| 1324 | |
| 1325 | # Verify MLA-specific keys in breakdown |
| 1326 | assert any(k.startswith("mla_attn.") for k in breakdown) |
| 1327 | assert not any(k.startswith("attn.") for k in breakdown) |
| 1328 | |
| 1329 | |
| 1330 | def test_standard_attention_still_works_for_non_mla(): |
| 1331 | """Regression test: non-MLA models still use standard AttentionMetrics.""" |
| 1332 | hf_config = Qwen3Config( |
| 1333 | hidden_size=2048, |
| 1334 | num_attention_heads=16, |
| 1335 | num_hidden_layers=12, |
| 1336 | vocab_size=32000, |
| 1337 | intermediate_size=8192, |
| 1338 | ) |
| 1339 | vllm_config = create_mock_vllm_config(hf_config) |
| 1340 | model_metrics = ModelMetrics(vllm_config) |
| 1341 | |
| 1342 | component_types = [m.component_type() for m in model_metrics.metrics] |
| 1343 | assert "attn" in component_types |
| 1344 | assert "mla_attn" not in component_types |
| 1345 | |
| 1346 | ctx = ExecutionContext.from_single_request( |
| 1347 | num_tokens=100, context_len=512, is_prefill=True |
| 1348 | ) |
| 1349 | total_flops = model_metrics.get_num_flops(ctx) |
| 1350 | assert total_flops > 0 |
| 1351 | |
| 1352 | |
| 1353 | def test_mla_attention_scaling_with_layers(): |
| 1354 | """Test that MLA attention metrics scale proportionally with layers.""" |
| 1355 | base_config = DeepseekV3Config( |
| 1356 | hidden_size=7168, |
| 1357 | num_attention_heads=128, |
| 1358 | num_hidden_layers=8, |
| 1359 | kv_lora_rank=512, |
| 1360 | qk_nope_head_dim=128, |
| 1361 | qk_rope_head_dim=64, |
| 1362 | v_head_dim=128, |
| 1363 | q_lora_rank=1536, |
| 1364 | ) |
| 1365 | double_config = DeepseekV3Config( |
| 1366 | hidden_size=7168, |
| 1367 | num_attention_heads=128, |
| 1368 | num_hidden_layers=16, # Double layers |
| 1369 | kv_lora_rank=512, |
| 1370 | qk_nope_head_dim=128, |
| 1371 | qk_rope_head_dim=64, |
| 1372 | v_head_dim=128, |
| 1373 | q_lora_rank=1536, |
| 1374 | ) |
| 1375 | |
| 1376 | base_vllm = create_mock_vllm_config(base_config) |
| 1377 | double_vllm = create_mock_vllm_config(double_config) |
| 1378 | |
| 1379 | base_metrics = MLAAttentionMetrics.from_vllm_config(base_vllm) |
| 1380 | double_metrics = MLAAttentionMetrics.from_vllm_config(double_vllm) |
| 1381 | |
| 1382 | ctx = ExecutionContext.from_single_request( |
| 1383 | num_tokens=100, context_len=512, is_prefill=True |
| 1384 | ) |
| 1385 | |
| 1386 | # All metrics should double with double layers |
| 1387 | assert double_metrics.get_num_flops(ctx) == 2 * base_metrics.get_num_flops(ctx) |
| 1388 | assert double_metrics.get_read_bytes(ctx) == 2 * base_metrics.get_read_bytes(ctx) |
| 1389 | assert double_metrics.get_write_bytes(ctx) == 2 * base_metrics.get_write_bytes(ctx) |
| 1390 | |
| 1391 | |
| 1392 | #### Sliding Window Attention Tests #### |
| 1393 | |
| 1394 | |
| 1395 | def test_sliding_window_parser_no_swa(): |
| 1396 | """Test parser defaults to no SWA when sliding_window is unset.""" |
| 1397 | hf_config = Qwen3Config( |
| 1398 | hidden_size=2048, |
| 1399 | num_attention_heads=16, |
| 1400 | num_hidden_layers=12, |
| 1401 | ) |
| 1402 | vllm_config = create_mock_vllm_config(hf_config) |
| 1403 | parser = SlidingWindowAttentionParser() |
| 1404 | args = ParsedArgs() |
| 1405 | args.num_hidden_layers = 12 |
| 1406 | result = parser.parse(args, vllm_config) |
| 1407 | |
| 1408 | assert result.sliding_window is None |
| 1409 | assert result.num_swa_layers == 0 |
| 1410 | |
| 1411 | |
| 1412 | def test_sliding_window_parser_uniform_swa(): |
| 1413 | """Test SlidingWindowAttentionParser for full SWA model (e.g. Mistral).""" |
| 1414 | hf_config = MistralConfig( |
| 1415 | hidden_size=4096, |
| 1416 | num_attention_heads=32, |
| 1417 | num_key_value_heads=8, |
| 1418 | num_hidden_layers=32, |
| 1419 | sliding_window=4096, |
| 1420 | ) |
| 1421 | vllm_config = create_mock_vllm_config(hf_config) |
| 1422 | parser = SlidingWindowAttentionParser() |
| 1423 | args = ParsedArgs() |
| 1424 | args.num_hidden_layers = 32 |
| 1425 | result = parser.parse(args, vllm_config) |
| 1426 | |
| 1427 | assert result.sliding_window == 4096 |
| 1428 | assert result.num_swa_layers == 32 |
| 1429 | |
| 1430 | |
| 1431 | def test_sliding_window_parser_explicit_layer_types(): |
| 1432 | """Test SlidingWindowAttentionParser with explicit layer_types list.""" |
| 1433 | hf_config = Qwen3Config( |
| 1434 | hidden_size=2048, |
| 1435 | num_attention_heads=16, |
| 1436 | num_hidden_layers=6, |
| 1437 | ) |
| 1438 | hf_config.sliding_window = 2048 |
| 1439 | hf_config.layer_types = [ |
| 1440 | "sliding_attention", |
| 1441 | "full_attention", |
| 1442 | "sliding_attention", |
| 1443 | "full_attention", |
| 1444 | "sliding_attention", |
| 1445 | "full_attention", |
| 1446 | ] |
| 1447 | vllm_config = create_mock_vllm_config(hf_config) |
| 1448 | parser = SlidingWindowAttentionParser() |
| 1449 | args = ParsedArgs() |
| 1450 | args.num_hidden_layers = 6 |
| 1451 | result = parser.parse(args, vllm_config) |
| 1452 | |
| 1453 | assert result.sliding_window == 2048 |
| 1454 | assert result.num_swa_layers == 3 |
| 1455 | |
| 1456 | |
| 1457 | def test_sliding_window_parser_gemma2(): |
| 1458 | """Test SlidingWindowAttentionParser detects alternating SWA for Gemma 2.""" |
| 1459 | hf_config = Gemma2Config( |
| 1460 | hidden_size=2304, |
| 1461 | num_attention_heads=8, |
| 1462 | num_key_value_heads=4, |
| 1463 | num_hidden_layers=26, |
| 1464 | sliding_window=4096, |
| 1465 | ) |
| 1466 | vllm_config = create_mock_vllm_config(hf_config) |
| 1467 | parser = SlidingWindowAttentionParser() |
| 1468 | args = ParsedArgs() |
| 1469 | args.num_hidden_layers = 26 |
| 1470 | result = parser.parse(args, vllm_config) |
| 1471 | |
| 1472 | assert result.sliding_window == 4096 |
| 1473 | # Gemma 2 alternates: even layers are sliding window (26 layers -> 13 SWA) |
| 1474 | assert result.num_swa_layers == 13 |
| 1475 | |
| 1476 | |
| 1477 | def test_sliding_window_parser_disabled(): |
| 1478 | """Test SlidingWindowAttentionParser respects disable_sliding_window.""" |
| 1479 | hf_config = MistralConfig( |
| 1480 | hidden_size=4096, |
| 1481 | num_attention_heads=32, |
| 1482 | num_key_value_heads=8, |
| 1483 | num_hidden_layers=32, |
| 1484 | sliding_window=4096, |
| 1485 | ) |
| 1486 | vllm_config = create_mock_vllm_config(hf_config, disable_sliding_window=True) |
| 1487 | parser = SlidingWindowAttentionParser() |
| 1488 | args = ParsedArgs() |
| 1489 | args.num_hidden_layers = 32 |
| 1490 | result = parser.parse(args, vllm_config) |
| 1491 | |
| 1492 | assert result.sliding_window is None |
| 1493 | assert result.num_swa_layers == 0 |
| 1494 | |
| 1495 | |
| 1496 | def test_attention_metrics_sliding_window_flops_bounds(): |
| 1497 | """Test that SWA clamps attention core FLOPs when context_len > sliding_window.""" |
| 1498 | sw = 1024 |
| 1499 | dense_config = MistralConfig( |
| 1500 | hidden_size=2048, |
| 1501 | num_attention_heads=16, |
| 1502 | num_key_value_heads=16, |
| 1503 | num_hidden_layers=8, |
| 1504 | sliding_window=None, |
| 1505 | ) |
| 1506 | swa_config = MistralConfig( |
| 1507 | hidden_size=2048, |
| 1508 | num_attention_heads=16, |
| 1509 | num_key_value_heads=16, |
| 1510 | num_hidden_layers=8, |
| 1511 | sliding_window=sw, |
| 1512 | ) |
| 1513 | |
| 1514 | dense_vllm = create_mock_vllm_config(dense_config) |
| 1515 | swa_vllm = create_mock_vllm_config(swa_config) |
| 1516 | |
| 1517 | dense_metrics = AttentionMetrics.from_vllm_config(dense_vllm) |
| 1518 | swa_metrics = AttentionMetrics.from_vllm_config(swa_vllm) |
| 1519 | |
| 1520 | # 1. When context <= sliding_window: FLOPs should be identical |
| 1521 | short_ctx = ExecutionContext.from_single_request( |
| 1522 | num_tokens=1, context_len=512, is_prefill=False |
| 1523 | ) |
| 1524 | assert swa_metrics.get_num_flops(short_ctx) == dense_metrics.get_num_flops( |
| 1525 | short_ctx |
| 1526 | ) |
| 1527 | |
| 1528 | # 2. When context > sliding_window: FLOPs for SWA must be strictly smaller |
| 1529 | long_ctx = ExecutionContext.from_single_request( |
| 1530 | num_tokens=1, context_len=8192, is_prefill=False |
| 1531 | ) |
| 1532 | dense_flops = dense_metrics.get_num_flops_breakdown(long_ctx) |
| 1533 | swa_flops = swa_metrics.get_num_flops_breakdown(long_ctx) |
| 1534 | |
| 1535 | # Projections (qkv_proj, out_proj) are unaffected by sliding window |
| 1536 | assert swa_flops["qkv_proj"] == dense_flops["qkv_proj"] |
| 1537 | assert swa_flops["out_proj"] == dense_flops["out_proj"] |
| 1538 | |
| 1539 | # Attention core FLOPs are clamped: 8192 vs 1024 (8x difference) |
| 1540 | assert swa_flops["attn_qk"] < dense_flops["attn_qk"] |
| 1541 | assert swa_flops["attn_av"] < dense_flops["attn_av"] |
| 1542 | assert swa_flops["attn_qk"] * 8 == dense_flops["attn_qk"] |
| 1543 | assert swa_flops["attn_av"] * 8 == dense_flops["attn_av"] |
| 1544 | |
| 1545 | |
| 1546 | def test_attention_metrics_sliding_window_read_bytes(): |
| 1547 | """Test that SWA bounds decode KV-cache read bytes to sliding_window.""" |
| 1548 | sw = 2048 |
| 1549 | dense_config = MistralConfig( |
| 1550 | hidden_size=2048, |
| 1551 | num_attention_heads=16, |
| 1552 | num_key_value_heads=4, |
| 1553 | num_hidden_layers=4, |
| 1554 | sliding_window=None, |
| 1555 | ) |
| 1556 | swa_config = MistralConfig( |
| 1557 | hidden_size=2048, |
| 1558 | num_attention_heads=16, |
| 1559 | num_key_value_heads=4, |
| 1560 | num_hidden_layers=4, |
| 1561 | sliding_window=sw, |
| 1562 | ) |
| 1563 | |
| 1564 | dense_vllm = create_mock_vllm_config(dense_config) |
| 1565 | swa_vllm = create_mock_vllm_config(swa_config) |
| 1566 | |
| 1567 | dense_metrics = AttentionMetrics.from_vllm_config(dense_vllm) |
| 1568 | swa_metrics = AttentionMetrics.from_vllm_config(swa_vllm) |
| 1569 | |
| 1570 | # Long decode context: 8192 tokens |
| 1571 | ctx = ExecutionContext.from_single_request( |
| 1572 | num_tokens=1, context_len=8192, is_prefill=False |
| 1573 | ) |
| 1574 | dense_reads = dense_metrics.get_read_bytes_breakdown(ctx) |
| 1575 | swa_reads = swa_metrics.get_read_bytes_breakdown(ctx) |
| 1576 | |
| 1577 | # Attention input read traffic is lower for SWA |
| 1578 | assert swa_reads["attn_input"] < dense_reads["attn_input"] |
| 1579 | |
| 1580 | # Weight reads and projection inputs are identical |
| 1581 | assert swa_reads["qkv_weight"] == dense_reads["qkv_weight"] |
| 1582 | assert swa_reads["out_weight"] == dense_reads["out_weight"] |
| 1583 | assert swa_reads["qkv_input"] == dense_reads["qkv_input"] |
| 1584 | |
| 1585 | |
| 1586 | def test_attention_metrics_hybrid_alternating_layers(): |
| 1587 | """Test hybrid model where half layers are full attention and half are SWA.""" |
| 1588 | sw = 1024 |
| 1589 | hybrid_config = Qwen3Config( |
| 1590 | hidden_size=2048, |
| 1591 | num_attention_heads=16, |
| 1592 | num_key_value_heads=16, |
| 1593 | num_hidden_layers=4, |
| 1594 | ) |
| 1595 | hybrid_config.sliding_window = sw |
| 1596 | hybrid_config.layer_types = [ |
| 1597 | "sliding_attention", |
| 1598 | "full_attention", |
| 1599 | "sliding_attention", |
| 1600 | "full_attention", |
| 1601 | ] |
| 1602 | |
| 1603 | vllm_config = create_mock_vllm_config(hybrid_config) |
| 1604 | metrics = AttentionMetrics.from_vllm_config(vllm_config) |
| 1605 | |
| 1606 | assert metrics.num_swa_layers == 2 |
| 1607 | assert metrics.sliding_window == sw |
| 1608 | |
| 1609 | ctx = ExecutionContext.from_single_request( |
| 1610 | num_tokens=1, context_len=4096, is_prefill=False |
| 1611 | ) |
| 1612 | flops = metrics.get_num_flops_breakdown(ctx) |
| 1613 | |
| 1614 | # 2 full layers attend to 4096; 2 SWA layers attend to 1024 |
| 1615 | # Expected TC total = 2 * 4096 + 2 * 1024 = 8192 + 2048 = 10240 |
| 1616 | # Formula: 2 * q * d * (L_full * 4096 + L_swa * 1024) |
| 1617 | q = metrics.num_attention_heads |
| 1618 | d = metrics.head_dim |
| 1619 | expected_core_flops = 2 * q * d * (2 * 4096 + 2 * 1024) |
| 1620 | assert flops["attn_qk"] == expected_core_flops |
| 1621 | assert flops["attn_av"] == expected_core_flops |
| 1622 | |
| 1623 | |
| 1624 | def test_attention_metrics_sliding_window_per_gpu(): |
| 1625 | """Test that SWA metrics scale properly with tensor and pipeline parallelism.""" |
| 1626 | sw = 1024 |
| 1627 | hf_config = MistralConfig( |
| 1628 | hidden_size=4096, |
| 1629 | num_attention_heads=32, |
| 1630 | num_key_value_heads=8, |
| 1631 | num_hidden_layers=16, |
| 1632 | sliding_window=sw, |
| 1633 | ) |
| 1634 | single_vllm = create_mock_vllm_config(hf_config) |
| 1635 | tp_vllm = create_mock_vllm_config(hf_config, tensor_parallel_size=2) |
| 1636 | pp_vllm = create_mock_vllm_config(hf_config, pipeline_parallel_size=2) |
| 1637 | |
| 1638 | single_metrics = AttentionMetrics.from_vllm_config(single_vllm) |
| 1639 | tp_metrics = AttentionMetrics.from_vllm_config(tp_vllm) |
| 1640 | pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm) |
| 1641 | |
| 1642 | ctx = ExecutionContext.from_single_request( |
| 1643 | num_tokens=1, context_len=4096, is_prefill=False |
| 1644 | ) |
| 1645 | |
| 1646 | # Both TP=2 and PP=2 should halve the per-GPU attention core FLOPs |
| 1647 | single_flops = single_metrics.get_num_flops_breakdown(ctx, per_gpu=True) |
| 1648 | tp_flops = tp_metrics.get_num_flops_breakdown(ctx, per_gpu=True) |
| 1649 | pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True) |
| 1650 | |
| 1651 | assert tp_flops["attn_qk"] == single_flops["attn_qk"] // 2 |
| 1652 | assert pp_flops["attn_qk"] == single_flops["attn_qk"] // 2 |
| 1653 | |
| 1654 | |
| 1655 | def test_attention_metrics_hybrid_sliding_window_per_gpu(): |
| 1656 | """Test that hybrid SWA models preserve L_full + L_swa == L under PP.""" |
| 1657 | sw = 1024 |
| 1658 | hybrid_config = Gemma2Config( |
| 1659 | hidden_size=2048, |
| 1660 | num_attention_heads=16, |
| 1661 | num_key_value_heads=16, |
| 1662 | num_hidden_layers=30, |
| 1663 | sliding_window=sw, |
| 1664 | ) |
| 1665 | pp_vllm = create_mock_vllm_config(hybrid_config, pipeline_parallel_size=2) |
| 1666 | pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm) |
| 1667 | |
| 1668 | ctx = ExecutionContext.from_single_request( |
| 1669 | num_tokens=1, context_len=4096, is_prefill=False |
| 1670 | ) |
| 1671 | pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True) |
| 1672 | pp_read_bytes = pp_metrics.get_read_bytes_breakdown(ctx, per_gpu=True) |
| 1673 | |
| 1674 | # Under PP=2, each global count is halved: 30 layers -> 15 per GPU, and the |
| 1675 | # 15 SWA / 15 full split becomes 7.5 each. Gemma 2 alternates, so an even |
| 1676 | # split is the correct average stage; deriving it after reducing L instead |
| 1677 | # produced a lopsided 8/7. |
| 1678 | q = pp_metrics.num_attention_heads |
| 1679 | d = pp_metrics.head_dim |
| 1680 | expected_pp_attn = int(2 * q * d * (7.5 * 4096 + 7.5 * sw)) |
| 1681 | assert pp_flops["attn_qk"] == expected_pp_attn |
| 1682 | assert pp_flops["attn_av"] == expected_pp_attn |
| 1683 | assert "attn_input" in pp_read_bytes |
| 1684 | |
| 1685 | |
| 1686 | def test_attention_metrics_minority_swa_survives_pipeline_parallelism(): |
| 1687 | """A minority SWA sublayer must not round away under PP. |
| 1688 | |
| 1689 | Regression test for a 6-layer model with a single sliding-window layer at |
| 1690 | pp_size=2: deriving the split after reducing L gave |
| 1691 | round(3 * 1 / 6) == 0, billing all three per-GPU layers as full attention |
| 1692 | and dropping the sliding-window contribution entirely. |
| 1693 | """ |
| 1694 | sw = 512 |
| 1695 | context_len = 4096 |
| 1696 | hf_config = Qwen3Config( |
| 1697 | hidden_size=2048, |
| 1698 | num_attention_heads=16, |
| 1699 | num_key_value_heads=16, |
| 1700 | num_hidden_layers=6, |
| 1701 | ) |
| 1702 | hf_config.sliding_window = sw |
| 1703 | hf_config.layer_types = ["sliding_attention"] + ["full_attention"] * 5 |
| 1704 | |
| 1705 | pp_vllm = create_mock_vllm_config(hf_config, pipeline_parallel_size=2) |
| 1706 | pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm) |
| 1707 | assert pp_metrics.num_swa_layers == 1 |
| 1708 | |
| 1709 | ctx = ExecutionContext.from_single_request( |
| 1710 | num_tokens=1, context_len=context_len, is_prefill=False |
| 1711 | ) |
| 1712 | pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True) |
| 1713 | |
| 1714 | # 5 full and 1 SWA layer, each halved by PP=2: 2.5 full and 0.5 SWA. |
| 1715 | q = pp_metrics.num_attention_heads |
| 1716 | d = pp_metrics.head_dim |
| 1717 | expected = int(2 * q * d * (2.5 * context_len + 0.5 * sw)) |
| 1718 | assert pp_flops["attn_qk"] == expected |
| 1719 | assert pp_flops["attn_av"] == expected |
| 1720 | |
| 1721 | # The sliding-window layer must change the result, i.e. it is not being |
| 1722 | # billed at the full context length. |
| 1723 | all_full = int(2 * q * d * 3 * context_len) |
| 1724 | assert pp_flops["attn_qk"] < all_full |
| 1725 | |
| 1726 | |
| 1727 | def test_attention_metrics_per_gpu_layers_sum_to_whole_model(): |
| 1728 | """Per-GPU counts times pp_size must account for every layer. |
| 1729 | |
| 1730 | Flooring the layer count per stage silently drops the remainder, so a model |
| 1731 | whose layers do not divide evenly across stages under-reports its FLOPs. |
| 1732 | """ |
| 1733 | hf_config = Qwen3Config( |
| 1734 | hidden_size=2048, |
| 1735 | num_attention_heads=16, |
| 1736 | num_key_value_heads=16, |
| 1737 | num_hidden_layers=7, |
| 1738 | ) |
| 1739 | pp_size = 2 |
| 1740 | vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=pp_size) |
| 1741 | metrics = AttentionMetrics.from_vllm_config(vllm_config) |
| 1742 | |
| 1743 | ctx = ExecutionContext.from_single_request( |
| 1744 | num_tokens=8, context_len=1024, is_prefill=True |
| 1745 | ) |
| 1746 | per_gpu = metrics.get_num_flops_breakdown(ctx, per_gpu=True) |
| 1747 | whole = metrics.get_num_flops_breakdown(ctx, per_gpu=False) |
| 1748 | |
| 1749 | for key, value in whole.items(): |
| 1750 | assert per_gpu[key] * pp_size == pytest.approx(value, rel=1e-9) |