iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Kernel][Perf] Per-architecture tuned configs for batch-invariant persistent matmul (~3x decode kernels on RTX 4090D/H20) (#53247)

Browse Source 
Signed-off-by: LioEinaudi <zhao3024667639@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Lio Einaudi
2026-08-24 11:28:51 -04:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 29c9af5211
commit 7797b6022c
3 changed files with 283 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/determinism/test_matmul_batch_invariant.py
+27
View File
@@ -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])
vllm/model_executor/layers/batch_invariant.py
+6 -1
View File
@@ -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()
vllm/model_executor/layers/batch_invariant_configs.py
+250
View File
@@ -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,
}