609 lines
24 KiB
Python
609 lines
24 KiB
Python
# 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)),
|
||
),
|
||
),
|
||
},
|
||
"blackwell": {
|
||
(12288, 2048): _MatmulShapeConfig(
|
||
block_k=64,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 128, 8, 5)),
|
||
(4, _MatmulMConfig(16, 128, 4, 5)),
|
||
(8, _MatmulMConfig(16, 128, 4, 5)),
|
||
(16, _MatmulMConfig(16, 128, 4, 5)),
|
||
(32, _MatmulMConfig(32, 128, 4, 5)),
|
||
(64, _MatmulMConfig(64, 128, 8, 5)),
|
||
(256, _MatmulMConfig(128, 256, 4, 4)),
|
||
(512, _MatmulMConfig(128, 256, 4, 4)),
|
||
(1024, _MatmulMConfig(128, 256, 4, 4)),
|
||
(2048, _MatmulMConfig(128, 256, 4, 4)),
|
||
),
|
||
),
|
||
(2048, 6144): _MatmulShapeConfig(
|
||
block_k=128,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 32, 4, 5)),
|
||
(4, _MatmulMConfig(16, 32, 4, 5)),
|
||
(8, _MatmulMConfig(16, 32, 4, 5)),
|
||
(16, _MatmulMConfig(16, 32, 4, 5)),
|
||
(32, _MatmulMConfig(16, 32, 4, 5)),
|
||
(64, _MatmulMConfig(32, 32, 4, 5)),
|
||
(256, _MatmulMConfig(64, 64, 4, 5)),
|
||
(512, _MatmulMConfig(128, 64, 4, 4)),
|
||
(1024, _MatmulMConfig(128, 128, 8, 3)),
|
||
(2048, _MatmulMConfig(128, 128, 8, 3)),
|
||
),
|
||
),
|
||
(4096, 2048): _MatmulShapeConfig(
|
||
block_k=128,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 64, 4, 5)),
|
||
(4, _MatmulMConfig(16, 64, 4, 5)),
|
||
(8, _MatmulMConfig(16, 32, 4, 5)),
|
||
(16, _MatmulMConfig(16, 32, 4, 5)),
|
||
(32, _MatmulMConfig(32, 32, 4, 5)),
|
||
(64, _MatmulMConfig(32, 64, 4, 4)),
|
||
(256, _MatmulMConfig(128, 64, 4, 4)),
|
||
(512, _MatmulMConfig(128, 128, 8, 3)),
|
||
(1024, _MatmulMConfig(128, 128, 8, 3)),
|
||
(2048, _MatmulMConfig(128, 128, 8, 3)),
|
||
),
|
||
),
|
||
(151936, 2048): _MatmulShapeConfig(
|
||
block_k=64,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 256, 4, 5)),
|
||
(4, _MatmulMConfig(16, 256, 4, 5)),
|
||
(8, _MatmulMConfig(16, 256, 4, 5)),
|
||
(16, _MatmulMConfig(16, 256, 4, 5)),
|
||
(32, _MatmulMConfig(64, 256, 4, 5)),
|
||
(64, _MatmulMConfig(64, 256, 4, 5)),
|
||
(256, _MatmulMConfig(128, 256, 4, 4)),
|
||
(512, _MatmulMConfig(128, 256, 4, 4)),
|
||
(1024, _MatmulMConfig(128, 256, 4, 4)),
|
||
(2048, _MatmulMConfig(128, 256, 4, 4)),
|
||
),
|
||
),
|
||
(2048, 2048): _MatmulShapeConfig(
|
||
block_k=128,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 32, 4, 4)),
|
||
(4, _MatmulMConfig(16, 32, 4, 4)),
|
||
(8, _MatmulMConfig(16, 32, 4, 5)),
|
||
(16, _MatmulMConfig(16, 32, 8, 5)),
|
||
(32, _MatmulMConfig(16, 32, 4, 4)),
|
||
(64, _MatmulMConfig(32, 32, 4, 5)),
|
||
(256, _MatmulMConfig(64, 64, 4, 4)),
|
||
(512, _MatmulMConfig(128, 64, 4, 4)),
|
||
(1024, _MatmulMConfig(128, 128, 4, 3)),
|
||
(2048, _MatmulMConfig(128, 128, 8, 3)),
|
||
),
|
||
),
|
||
},
|
||
"sm120": {
|
||
(12288, 2048): _MatmulShapeConfig( # Qwen3-1.7B gate_up_proj
|
||
block_k=64,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 128, 8, 5)),
|
||
(4, _MatmulMConfig(16, 128, 8, 4)),
|
||
(8, _MatmulMConfig(16, 128, 8, 4)),
|
||
(16, _MatmulMConfig(16, 128, 4, 4)),
|
||
(32, _MatmulMConfig(32, 128, 8, 5)),
|
||
(64, _MatmulMConfig(64, 128, 4, 3)),
|
||
(256, _MatmulMConfig(64, 64, 4, 4)),
|
||
(512, _MatmulMConfig(64, 128, 4, 5)),
|
||
(1024, _MatmulMConfig(128, 128, 4, 3)),
|
||
(2048, _MatmulMConfig(128, 128, 4, 3)),
|
||
),
|
||
),
|
||
(2048, 6144): _MatmulShapeConfig( # Qwen3-1.7B down_proj
|
||
block_k=64,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 32, 8, 5)),
|
||
(4, _MatmulMConfig(16, 32, 8, 3)),
|
||
(8, _MatmulMConfig(16, 32, 8, 4)),
|
||
(16, _MatmulMConfig(16, 32, 4, 5)),
|
||
(32, _MatmulMConfig(32, 32, 4, 5)),
|
||
(64, _MatmulMConfig(32, 32, 4, 4)),
|
||
(256, _MatmulMConfig(64, 64, 4, 5)),
|
||
(512, _MatmulMConfig(128, 64, 4, 3)),
|
||
(1024, _MatmulMConfig(128, 128, 4, 3)),
|
||
(2048, _MatmulMConfig(128, 64, 8, 4)),
|
||
),
|
||
),
|
||
(4096, 2048): _MatmulShapeConfig( # Qwen3-1.7B qkv_proj
|
||
block_k=128,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 64, 4, 4)),
|
||
(4, _MatmulMConfig(16, 64, 8, 5)),
|
||
(8, _MatmulMConfig(16, 64, 8, 5)),
|
||
(16, _MatmulMConfig(16, 64, 8, 3)),
|
||
(32, _MatmulMConfig(16, 64, 8, 5)),
|
||
(64, _MatmulMConfig(32, 64, 4, 2)),
|
||
(256, _MatmulMConfig(64, 128, 4, 3)),
|
||
(512, _MatmulMConfig(64, 64, 4, 3)),
|
||
(1024, _MatmulMConfig(64, 128, 4, 3)),
|
||
(2048, _MatmulMConfig(64, 128, 4, 3)),
|
||
),
|
||
),
|
||
(151936, 2048): _MatmulShapeConfig( # Qwen3-1.7B lm_head
|
||
block_k=64,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 64, 8, 5)),
|
||
(4, _MatmulMConfig(16, 128, 4, 4)),
|
||
(8, _MatmulMConfig(16, 128, 4, 5)),
|
||
(16, _MatmulMConfig(16, 128, 8, 3)),
|
||
(32, _MatmulMConfig(32, 128, 8, 3)),
|
||
(64, _MatmulMConfig(32, 128, 8, 5)),
|
||
(256, _MatmulMConfig(128, 128, 4, 3)),
|
||
(512, _MatmulMConfig(128, 128, 8, 3)),
|
||
(1024, _MatmulMConfig(128, 128, 4, 3)),
|
||
(2048, _MatmulMConfig(128, 128, 4, 3)),
|
||
),
|
||
),
|
||
(2048, 2048): _MatmulShapeConfig( # Qwen3-1.7B o_proj
|
||
block_k=32,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 128, 4, 3)),
|
||
(4, _MatmulMConfig(16, 128, 4, 4)),
|
||
(8, _MatmulMConfig(16, 64, 4, 5)),
|
||
(16, _MatmulMConfig(32, 32, 8, 4)),
|
||
(32, _MatmulMConfig(16, 64, 4, 5)),
|
||
(64, _MatmulMConfig(32, 32, 4, 5)),
|
||
(256, _MatmulMConfig(128, 32, 4, 5)),
|
||
(512, _MatmulMConfig(128, 64, 4, 4)),
|
||
(1024, _MatmulMConfig(128, 128, 8, 5)),
|
||
(2048, _MatmulMConfig(128, 256, 8, 4)),
|
||
),
|
||
),
|
||
(19456, 2560): _MatmulShapeConfig( # Qwen3-4B gate_up_proj
|
||
block_k=64,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 128, 8, 4)),
|
||
(4, _MatmulMConfig(16, 128, 8, 3)),
|
||
(8, _MatmulMConfig(16, 128, 8, 5)),
|
||
(16, _MatmulMConfig(16, 128, 8, 5)),
|
||
(32, _MatmulMConfig(32, 128, 4, 4)),
|
||
(64, _MatmulMConfig(64, 128, 8, 4)),
|
||
(256, _MatmulMConfig(128, 128, 4, 3)),
|
||
(512, _MatmulMConfig(64, 128, 4, 5)),
|
||
(1024, _MatmulMConfig(128, 128, 4, 3)),
|
||
(2048, _MatmulMConfig(128, 128, 4, 3)),
|
||
),
|
||
),
|
||
(2560, 9728): _MatmulShapeConfig( # Qwen3-4B down_proj
|
||
block_k=128,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 64, 4, 5)),
|
||
(4, _MatmulMConfig(16, 64, 4, 4)),
|
||
(8, _MatmulMConfig(16, 64, 4, 4)),
|
||
(16, _MatmulMConfig(16, 32, 4, 5)),
|
||
(32, _MatmulMConfig(32, 32, 4, 3)),
|
||
(64, _MatmulMConfig(32, 64, 8, 4)),
|
||
(256, _MatmulMConfig(64, 64, 4, 3)),
|
||
(512, _MatmulMConfig(64, 128, 4, 3)),
|
||
(1024, _MatmulMConfig(64, 128, 4, 3)),
|
||
(2048, _MatmulMConfig(64, 128, 4, 3)),
|
||
),
|
||
),
|
||
(6144, 2560): _MatmulShapeConfig( # Qwen3-4B qkv_proj
|
||
block_k=64,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 128, 8, 4)),
|
||
(4, _MatmulMConfig(16, 256, 4, 3)),
|
||
(8, _MatmulMConfig(16, 128, 4, 5)),
|
||
(16, _MatmulMConfig(16, 64, 8, 3)),
|
||
(32, _MatmulMConfig(16, 128, 8, 4)),
|
||
(64, _MatmulMConfig(64, 64, 4, 5)),
|
||
(256, _MatmulMConfig(128, 128, 8, 3)),
|
||
(512, _MatmulMConfig(64, 64, 4, 5)),
|
||
(1024, _MatmulMConfig(64, 128, 8, 4)),
|
||
(2048, _MatmulMConfig(128, 128, 4, 3)),
|
||
),
|
||
),
|
||
(151936, 2560): _MatmulShapeConfig( # Qwen3-4B lm_head
|
||
block_k=64,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 64, 8, 5)),
|
||
(4, _MatmulMConfig(16, 64, 8, 5)),
|
||
(8, _MatmulMConfig(16, 64, 4, 5)),
|
||
(16, _MatmulMConfig(16, 64, 4, 5)),
|
||
(32, _MatmulMConfig(32, 64, 4, 5)),
|
||
(64, _MatmulMConfig(64, 64, 4, 5)),
|
||
(256, _MatmulMConfig(128, 128, 4, 3)),
|
||
(512, _MatmulMConfig(128, 128, 4, 3)),
|
||
(1024, _MatmulMConfig(128, 128, 8, 3)),
|
||
(2048, _MatmulMConfig(128, 128, 4, 3)),
|
||
),
|
||
),
|
||
(2560, 4096): _MatmulShapeConfig( # Qwen3-4B o_proj
|
||
block_k=128,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 64, 8, 3)),
|
||
(4, _MatmulMConfig(16, 64, 8, 4)),
|
||
(8, _MatmulMConfig(16, 64, 8, 5)),
|
||
(16, _MatmulMConfig(16, 64, 4, 4)),
|
||
(32, _MatmulMConfig(16, 32, 4, 4)),
|
||
(64, _MatmulMConfig(16, 64, 8, 3)),
|
||
(256, _MatmulMConfig(64, 64, 8, 3)),
|
||
(512, _MatmulMConfig(64, 128, 4, 3)),
|
||
(1024, _MatmulMConfig(64, 128, 4, 3)),
|
||
(2048, _MatmulMConfig(64, 128, 4, 3)),
|
||
),
|
||
),
|
||
(24576, 4096): _MatmulShapeConfig( # Qwen3-8B gate_up_proj
|
||
block_k=64,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 256, 8, 3)),
|
||
(4, _MatmulMConfig(16, 256, 8, 3)),
|
||
(8, _MatmulMConfig(16, 256, 8, 3)),
|
||
(16, _MatmulMConfig(16, 256, 8, 3)),
|
||
(32, _MatmulMConfig(32, 256, 8, 3)),
|
||
(64, _MatmulMConfig(64, 256, 4, 3)),
|
||
(256, _MatmulMConfig(64, 128, 4, 5)),
|
||
(512, _MatmulMConfig(128, 128, 4, 3)),
|
||
(1024, _MatmulMConfig(128, 128, 4, 3)),
|
||
(2048, _MatmulMConfig(128, 128, 4, 3)),
|
||
),
|
||
),
|
||
(4096, 12288): _MatmulShapeConfig( # Qwen3-8B down_proj
|
||
block_k=128,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 64, 4, 5)),
|
||
(4, _MatmulMConfig(16, 64, 4, 5)),
|
||
(8, _MatmulMConfig(16, 64, 4, 4)),
|
||
(16, _MatmulMConfig(16, 32, 4, 4)),
|
||
(32, _MatmulMConfig(32, 32, 8, 5)),
|
||
(64, _MatmulMConfig(32, 64, 4, 5)),
|
||
(256, _MatmulMConfig(64, 128, 4, 3)),
|
||
(512, _MatmulMConfig(64, 64, 4, 3)),
|
||
(1024, _MatmulMConfig(64, 128, 4, 3)),
|
||
(2048, _MatmulMConfig(64, 128, 4, 3)),
|
||
),
|
||
),
|
||
(6144, 4096): _MatmulShapeConfig( # Qwen3-8B qkv_proj
|
||
block_k=64,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 64, 8, 4)),
|
||
(4, _MatmulMConfig(16, 128, 8, 4)),
|
||
(8, _MatmulMConfig(16, 128, 4, 4)),
|
||
(16, _MatmulMConfig(16, 64, 4, 5)),
|
||
(32, _MatmulMConfig(16, 128, 8, 5)),
|
||
(64, _MatmulMConfig(64, 64, 4, 4)),
|
||
(256, _MatmulMConfig(128, 128, 8, 3)),
|
||
(512, _MatmulMConfig(64, 64, 4, 5)),
|
||
(1024, _MatmulMConfig(128, 64, 4, 3)),
|
||
(2048, _MatmulMConfig(128, 128, 4, 3)),
|
||
),
|
||
),
|
||
(151936, 4096): _MatmulShapeConfig( # Qwen3-8B lm_head
|
||
block_k=64,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 64, 8, 5)),
|
||
(4, _MatmulMConfig(16, 64, 8, 5)),
|
||
(8, _MatmulMConfig(16, 64, 4, 5)),
|
||
(16, _MatmulMConfig(16, 64, 4, 5)),
|
||
(32, _MatmulMConfig(32, 64, 4, 5)),
|
||
(64, _MatmulMConfig(32, 128, 8, 5)),
|
||
(256, _MatmulMConfig(128, 128, 4, 3)),
|
||
(512, _MatmulMConfig(128, 128, 8, 3)),
|
||
(1024, _MatmulMConfig(128, 128, 4, 3)),
|
||
(2048, _MatmulMConfig(128, 128, 4, 3)),
|
||
),
|
||
),
|
||
(4096, 4096): _MatmulShapeConfig( # Qwen3-8B o_proj
|
||
block_k=128,
|
||
m_buckets=(
|
||
(1, _MatmulMConfig(16, 64, 8, 4)),
|
||
(4, _MatmulMConfig(16, 64, 4, 3)),
|
||
(8, _MatmulMConfig(16, 32, 4, 4)),
|
||
(16, _MatmulMConfig(16, 32, 4, 3)),
|
||
(32, _MatmulMConfig(16, 64, 4, 5)),
|
||
(64, _MatmulMConfig(64, 32, 8, 4)),
|
||
(256, _MatmulMConfig(64, 128, 4, 3)),
|
||
(512, _MatmulMConfig(64, 64, 4, 3)),
|
||
(1024, _MatmulMConfig(64, 128, 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 == 12:
|
||
return "sm120"
|
||
if capability.major == 10:
|
||
return "blackwell"
|
||
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,
|
||
}
|
||
|
||
|
||
def _get_descriptor_matmul_config(
|
||
M: int, N: int, K: int, dtype: torch.dtype
|
||
) -> dict[str, int]:
|
||
"""Select a config for the XPU descriptor matmul from shape and dtype.
|
||
|
||
The default 128×128×64 config is optimized for large square-ish GEMMs but
|
||
wastes resources on skinny shapes (small M during decode, or very tall/thin
|
||
matrices during prefill). Shape-dependent tuning closes the gap with oneDNN.
|
||
|
||
Only BLOCK_SIZE_M/N and the launch parameters vary with M; dtype alone
|
||
decides BLOCK_SIZE_K, so the K-reduction order never depends on batch size.
|
||
"""
|
||
# fp32 uses smaller BLOCK_SIZE_K due to register pressure
|
||
block_k = 32 if dtype == torch.float32 else 64
|
||
|
||
if M <= 16:
|
||
# Decode: M=1-16. Tiny M means most of a 128-row tile is wasted.
|
||
# Use small M-block, wide N-block to maximize useful work per tile.
|
||
return {
|
||
"BLOCK_SIZE_M": 16,
|
||
"BLOCK_SIZE_N": 256,
|
||
"BLOCK_SIZE_K": block_k,
|
||
"GROUP_SIZE_M": 1,
|
||
"num_stages": 4,
|
||
"num_warps": 8,
|
||
}
|
||
elif M <= 64:
|
||
# Small batch decode or very short prefill.
|
||
return {
|
||
"BLOCK_SIZE_M": 32,
|
||
"BLOCK_SIZE_N": 128,
|
||
"BLOCK_SIZE_K": block_k,
|
||
"GROUP_SIZE_M": 4,
|
||
"num_stages": 4,
|
||
"num_warps": 8,
|
||
}
|
||
else:
|
||
# Medium and large prefill (M > 64). 64×128 tiles provide the best
|
||
# balance of register pressure vs parallelism on Intel XPU.
|
||
# M=2048, N=4096 → 32×32 = 1024 tiles, well above ~160 compute units.
|
||
return {
|
||
"BLOCK_SIZE_M": 64,
|
||
"BLOCK_SIZE_N": 128,
|
||
"BLOCK_SIZE_K": block_k,
|
||
"GROUP_SIZE_M": 8,
|
||
"num_stages": 3,
|
||
"num_warps": 8,
|
||
}
|