diff --git a/tests/v1/metrics/test_perf_metrics.py b/tests/v1/metrics/test_perf_metrics.py index 3868cc772e..c0997ec459 100644 --- a/tests/v1/metrics/test_perf_metrics.py +++ b/tests/v1/metrics/test_perf_metrics.py @@ -133,7 +133,56 @@ def test_base_config_parser(): assert result.num_attention_heads == 16 assert result.num_hidden_layers == 24 assert result.weight_byte_size == 2 # float16 is 2 bytes - assert result.activation_byte_size == 2 # default activation size + assert result.activation_byte_size == 2 # float16 activations are 2 bytes + + +@pytest.mark.parametrize( + "model_dtype,expected_byte_size", + [("float16", 2), ("bfloat16", 2), ("float32", 4)], +) +def test_base_config_parser_activation_byte_size_follows_dtype( + model_dtype: str, expected_byte_size: int +): + """Activation traffic is sized from the model's compute dtype. + + Activations are produced in the model dtype, so pinning them to 2 bytes + under-reported memory traffic by half for a float32 model while that same + model's weight traffic was sized correctly. + """ + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_hidden_layers=24, + ) + vllm_config = create_mock_vllm_config(hf_config, model_dtype=model_dtype) + + result = BaseConfigParser().parse(ParsedArgs(), vllm_config) + + assert result.activation_byte_size == expected_byte_size + assert result.weight_byte_size == expected_byte_size + + +def test_attention_read_bytes_scale_with_activation_dtype(): + """A float32 model reads twice the activation bytes of a bfloat16 one.""" + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=4, + ) + ctx = ExecutionContext.from_single_request( + num_tokens=128, context_len=128, is_prefill=True + ) + + bf16 = AttentionMetrics.from_vllm_config( + create_mock_vllm_config(hf_config, model_dtype="bfloat16") + ).get_read_bytes_breakdown(ctx) + fp32 = AttentionMetrics.from_vllm_config( + create_mock_vllm_config(hf_config, model_dtype="float32") + ).get_read_bytes_breakdown(ctx) + + # qkv_input is pure activation traffic, so it doubles exactly. + assert fp32["qkv_input"] == 2 * bf16["qkv_input"] def test_base_attention_config_parser_with_gqa(): diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index dd046a8250..1cd8fa79f1 100644 --- a/vllm/v1/metrics/perf.py +++ b/vllm/v1/metrics/perf.py @@ -369,9 +369,10 @@ class BaseConfigParser(Parser): args.weight_byte_size = get_dtype_size(torch_dtype) - # FIXME: handle this better by parsing whether activations use - # bf16, fp32, etc... - args.activation_byte_size = 2 + # Activations are produced in the model's compute dtype. Quantization + # overrides weight_byte_size below but leaves activations alone, so + # this stays keyed on the model dtype. + args.activation_byte_size = get_dtype_size(torch_dtype) args.dp_size = vllm_config.parallel_config.data_parallel_size args.tp_size = vllm_config.parallel_config.tensor_parallel_size