diff --git a/tests/v1/determinism/test_matmul_batch_invariant.py b/tests/v1/determinism/test_matmul_batch_invariant.py index 8838b6be27..531c5370df 100644 --- a/tests/v1/determinism/test_matmul_batch_invariant.py +++ b/tests/v1/determinism/test_matmul_batch_invariant.py @@ -12,6 +12,10 @@ import torch from utils import skip_unsupported from vllm.model_executor.layers.batch_invariant import matmul_batch_invariant +from vllm.model_executor.layers.batch_invariant_configs import ( + _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS, + _get_tuned_matmul_arch_family, +) from vllm.platforms import current_platform DEVICE_TYPE = current_platform.device_type @@ -103,3 +107,26 @@ def test_matmul_batch_invariance(dtype): batch_output_a = batch_output[3] assert torch.equal(standard_output[0], batch_output_a) + + +@skip_unsupported +@pytest.mark.parametrize("m", [8, 32, 256, 2048]) +def test_matmul_batch_invariance_across_tuned_m_buckets(m): + # Tuned M buckets must preserve each row's K-reduction order. + capability = ( + current_platform.get_device_capability() if current_platform.is_cuda() else None + ) + arch_family = _get_tuned_matmul_arch_family(capability) + if arch_family not in _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS: + pytest.skip("No tuned persistent matmul config for this architecture") + + device = torch.device(DEVICE_TYPE) + n = k = 2048 + torch.manual_seed(42) + a = torch.rand((m, k), dtype=torch.bfloat16, device=device) + b = torch.rand((k, n), dtype=torch.bfloat16, device=device) + + single_output = matmul_batch_invariant(a[:1], b) + batch_output = matmul_batch_invariant(a, b) + + assert torch.equal(single_output[0], batch_output[0]) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index f0521d954c..7db629ebad 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -8,6 +8,10 @@ from typing import Any import torch import vllm.envs as envs +from vllm.model_executor.layers.batch_invariant_configs import ( + _get_matmul_config, + resolve_tuned_matmul_configs, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.mem_utils import get_max_shared_memory_bytes @@ -200,7 +204,7 @@ def matmul_persistent( B_LARGE=b.numel() > 2**31, C_LARGE=c.numel() > 2**31, HAS_BIAS=bias is not None, - **configs[dtype], + **_get_matmul_config(M, N, K, dtype, configs[dtype]), ) return c @@ -997,6 +1001,7 @@ def override_envs_for_invariance(): def init_batch_invariance(): # this will hit all the csrc overrides as well if envs.VLLM_BATCH_INVARIANT: + resolve_tuned_matmul_configs() override_envs_for_invariance() enable_batch_invariant_mode() diff --git a/vllm/model_executor/layers/batch_invariant_configs.py b/vllm/model_executor/layers/batch_invariant_configs.py new file mode 100644 index 0000000000..4570a491ce --- /dev/null +++ b/vllm/model_executor/layers/batch_invariant_configs.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import NamedTuple + +import torch + +from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability + + +class _MatmulMConfig(NamedTuple): + block_m: int + block_n: int + num_warps: int + num_stages: int + + +class _MatmulShapeConfig(NamedTuple): + block_k: int + m_buckets: tuple[tuple[int, _MatmulMConfig], ...] + + +_BATCH_INVARIANT_MATMUL_TUNED_CONFIGS: dict[ + str, dict[tuple[int, int], _MatmulShapeConfig] +] = { + "ada": { + (12288, 2048): _MatmulShapeConfig( + block_k=64, + m_buckets=( + (1, _MatmulMConfig(16, 128, 8, 2)), + (4, _MatmulMConfig(16, 64, 8, 5)), + (8, _MatmulMConfig(16, 128, 8, 4)), + (16, _MatmulMConfig(16, 128, 4, 5)), + (32, _MatmulMConfig(32, 128, 4, 5)), + (64, _MatmulMConfig(32, 256, 4, 3)), + (256, _MatmulMConfig(128, 128, 8, 3)), + (512, _MatmulMConfig(128, 64, 4, 4)), + (1024, _MatmulMConfig(128, 128, 8, 3)), + (2048, _MatmulMConfig(128, 128, 4, 3)), + ), + ), + (2048, 6144): _MatmulShapeConfig( + block_k=128, + m_buckets=( + (1, _MatmulMConfig(16, 64, 4, 4)), + (4, _MatmulMConfig(16, 64, 8, 4)), + (8, _MatmulMConfig(16, 32, 8, 4)), + (16, _MatmulMConfig(16, 32, 8, 4)), + (32, _MatmulMConfig(16, 64, 8, 3)), + (64, _MatmulMConfig(32, 64, 4, 4)), + (256, _MatmulMConfig(32, 64, 4, 4)), + (512, _MatmulMConfig(64, 64, 4, 3)), + (1024, _MatmulMConfig(64, 64, 4, 3)), + (2048, _MatmulMConfig(64, 64, 4, 3)), + ), + ), + (4096, 2048): _MatmulShapeConfig( + block_k=64, + m_buckets=( + (1, _MatmulMConfig(16, 256, 8, 3)), + (4, _MatmulMConfig(16, 64, 4, 3)), + (8, _MatmulMConfig(16, 128, 8, 2)), + (16, _MatmulMConfig(16, 32, 4, 5)), + (32, _MatmulMConfig(16, 64, 4, 4)), + (64, _MatmulMConfig(64, 64, 8, 3)), + (256, _MatmulMConfig(64, 64, 4, 5)), + (512, _MatmulMConfig(64, 64, 4, 5)), + (1024, _MatmulMConfig(128, 64, 4, 4)), + (2048, _MatmulMConfig(128, 64, 8, 4)), + ), + ), + (151936, 2048): _MatmulShapeConfig( + block_k=64, + m_buckets=( + (1, _MatmulMConfig(16, 32, 8, 5)), + (4, _MatmulMConfig(16, 32, 8, 5)), + (8, _MatmulMConfig(16, 32, 8, 5)), + (16, _MatmulMConfig(16, 32, 4, 4)), + (32, _MatmulMConfig(16, 64, 4, 4)), + (64, _MatmulMConfig(32, 64, 4, 4)), + (256, _MatmulMConfig(128, 128, 4, 3)), + (512, _MatmulMConfig(128, 128, 4, 3)), + (1024, _MatmulMConfig(128, 128, 4, 3)), + (2048, _MatmulMConfig(128, 128, 4, 3)), + ), + ), + (2048, 2048): _MatmulShapeConfig( + block_k=64, + m_buckets=( + (1, _MatmulMConfig(32, 128, 8, 5)), + (4, _MatmulMConfig(16, 256, 8, 3)), + (8, _MatmulMConfig(32, 128, 8, 3)), + (16, _MatmulMConfig(32, 128, 8, 3)), + (32, _MatmulMConfig(16, 128, 8, 5)), + (64, _MatmulMConfig(16, 256, 4, 3)), + (256, _MatmulMConfig(64, 128, 8, 4)), + (512, _MatmulMConfig(64, 64, 4, 5)), + (1024, _MatmulMConfig(64, 64, 4, 5)), + (2048, _MatmulMConfig(64, 64, 4, 5)), + ), + ), + }, + "hopper": { + (12288, 2048): _MatmulShapeConfig( + block_k=128, + m_buckets=( + (1, _MatmulMConfig(16, 256, 4, 3)), + (4, _MatmulMConfig(16, 256, 4, 4)), + (8, _MatmulMConfig(16, 64, 4, 4)), + (16, _MatmulMConfig(16, 64, 4, 4)), + (32, _MatmulMConfig(32, 64, 8, 5)), + (64, _MatmulMConfig(64, 64, 4, 4)), + (256, _MatmulMConfig(64, 128, 4, 3)), + (512, _MatmulMConfig(128, 128, 8, 3)), + (1024, _MatmulMConfig(128, 256, 8, 2)), + (2048, _MatmulMConfig(128, 256, 8, 2)), + ), + ), + (2048, 6144): _MatmulShapeConfig( + block_k=128, + m_buckets=( + (1, _MatmulMConfig(16, 32, 4, 4)), + (4, _MatmulMConfig(16, 64, 4, 4)), + (8, _MatmulMConfig(16, 64, 4, 4)), + (16, _MatmulMConfig(16, 32, 4, 3)), + (32, _MatmulMConfig(32, 32, 8, 3)), + (64, _MatmulMConfig(64, 32, 4, 5)), + (256, _MatmulMConfig(64, 128, 4, 3)), + (512, _MatmulMConfig(128, 128, 4, 3)), + (1024, _MatmulMConfig(64, 64, 4, 5)), + (2048, _MatmulMConfig(64, 128, 4, 4)), + ), + ), + (4096, 2048): _MatmulShapeConfig( + block_k=128, + m_buckets=( + (1, _MatmulMConfig(16, 64, 8, 5)), + (4, _MatmulMConfig(16, 64, 4, 5)), + (8, _MatmulMConfig(16, 128, 4, 3)), + (16, _MatmulMConfig(16, 64, 4, 5)), + (32, _MatmulMConfig(16, 64, 4, 4)), + (64, _MatmulMConfig(64, 32, 4, 3)), + (256, _MatmulMConfig(64, 256, 4, 2)), + (512, _MatmulMConfig(64, 64, 4, 3)), + (1024, _MatmulMConfig(64, 128, 4, 3)), + (2048, _MatmulMConfig(64, 64, 4, 4)), + ), + ), + (151936, 2048): _MatmulShapeConfig( + block_k=128, + m_buckets=( + (1, _MatmulMConfig(16, 256, 4, 4)), + (4, _MatmulMConfig(16, 256, 4, 4)), + (8, _MatmulMConfig(16, 256, 4, 3)), + (16, _MatmulMConfig(16, 256, 4, 3)), + (32, _MatmulMConfig(32, 256, 8, 3)), + (64, _MatmulMConfig(64, 128, 4, 3)), + (256, _MatmulMConfig(64, 128, 4, 3)), + (512, _MatmulMConfig(128, 128, 8, 3)), + (1024, _MatmulMConfig(128, 256, 8, 2)), + (2048, _MatmulMConfig(128, 256, 8, 2)), + ), + ), + (2048, 2048): _MatmulShapeConfig( + block_k=128, + m_buckets=( + (1, _MatmulMConfig(16, 64, 4, 5)), + (4, _MatmulMConfig(16, 128, 8, 3)), + (8, _MatmulMConfig(16, 32, 4, 3)), + (16, _MatmulMConfig(32, 64, 8, 5)), + (32, _MatmulMConfig(16, 128, 4, 5)), + (64, _MatmulMConfig(32, 64, 4, 2)), + (256, _MatmulMConfig(128, 64, 4, 4)), + (512, _MatmulMConfig(64, 256, 4, 2)), + (1024, _MatmulMConfig(64, 64, 4, 3)), + (2048, _MatmulMConfig(64, 128, 4, 3)), + ), + ), + }, +} + +_TUNED_MATMUL_CONFIGS_FOR_DEVICE: dict[tuple[int, int], _MatmulShapeConfig] | None = ( + None +) +_TUNED_MATMUL_CONFIGS_RESOLVED = False + + +def _get_tuned_matmul_arch_family(capability: DeviceCapability | None) -> str | None: + if capability is None: + return None + if capability.major == 9: + return "hopper" + if capability.major == 8 and capability.minor == 9: + return "ada" + return None + + +def resolve_tuned_matmul_configs() -> None: + global _TUNED_MATMUL_CONFIGS_FOR_DEVICE + global _TUNED_MATMUL_CONFIGS_RESOLVED + + if _TUNED_MATMUL_CONFIGS_RESOLVED: + return + + # Normal vLLM initialization resolves this before torch.compile tracing. + capability = ( + current_platform.get_device_capability() if current_platform.is_cuda() else None + ) + arch_family = _get_tuned_matmul_arch_family(capability) + if arch_family is None: + _TUNED_MATMUL_CONFIGS_FOR_DEVICE = None + else: + _TUNED_MATMUL_CONFIGS_FOR_DEVICE = _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS.get( + arch_family + ) + _TUNED_MATMUL_CONFIGS_RESOLVED = True + + +def _get_matmul_config( + M: int, + N: int, + K: int, + dtype: torch.dtype, + default: dict[str, int], +) -> dict[str, int]: + if not _TUNED_MATMUL_CONFIGS_RESOLVED: + resolve_tuned_matmul_configs() + if dtype != torch.bfloat16: + return default + device_configs = _TUNED_MATMUL_CONFIGS_FOR_DEVICE + if device_configs is None: + return default + shape_config = device_configs.get((N, K)) + if shape_config is None: + return default + # Values above the tuned range reuse the largest bucket; + # shape-wide BLOCK_K keeps this batch-invariant. + m_config = shape_config.m_buckets[-1][1] + for max_m, bucket_config in shape_config.m_buckets: + if max_m >= M: + m_config = bucket_config + break + return { + "BLOCK_SIZE_M": m_config.block_m, + "BLOCK_SIZE_N": m_config.block_n, + "BLOCK_SIZE_K": shape_config.block_k, + "GROUP_SIZE_M": 8, + "num_warps": m_config.num_warps, + "num_stages": m_config.num_stages, + }