iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
c19c263ac8ea3fd6e1d2d689f0e4e725855d89cf
vLLM/vllm/model_executor/determinism/batch_invariant_configs.py
T
Lio Einaudiandmergify[bot] <37929162+mergify[bot]@users.noreply.github.com> 4e0a414c44 [Kernel][Perf] Add TP=2/4/8 per-rank shapes to the sm_120 batch-invariant matmul table (#58495) 
Signed-off-by: LioEinaudi <zhao3024667639@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-09-29 23:51:50 +00:00

1224 lines
51 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# 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)),
),
),
(6144, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp2_gate_up_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 32, 8, 3)),
(4, _MatmulMConfig(16, 128, 8, 5)),
(8, _MatmulMConfig(16, 128, 4, 5)),
(16, _MatmulMConfig(16, 32, 4, 5)),
(32, _MatmulMConfig(32, 128, 8, 5)),
(64, _MatmulMConfig(32, 128, 4, 4)),
(256, _MatmulMConfig(128, 128, 4, 3)),
(512, _MatmulMConfig(64, 64, 4, 5)),
(1024, _MatmulMConfig(128, 64, 4, 3)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(2048, 3072): _MatmulShapeConfig( # Qwen3-1.7B tp2_down_proj
block_k=128,
m_buckets=(
(1, _MatmulMConfig(16, 32, 8, 5)),
(4, _MatmulMConfig(16, 64, 4, 3)),
(8, _MatmulMConfig(16, 64, 8, 4)),
(16, _MatmulMConfig(16, 64, 8, 3)),
(32, _MatmulMConfig(16, 64, 8, 3)),
(64, _MatmulMConfig(16, 64, 4, 5)),
(256, _MatmulMConfig(64, 64, 8, 3)),
(512, _MatmulMConfig(64, 128, 4, 3)),
(1024, _MatmulMConfig(64, 64, 4, 3)),
(2048, _MatmulMConfig(64, 128, 4, 3)),
),
),
(2048, 1024): _MatmulShapeConfig( # Qwen3-1.7B tp2_o_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 256, 8, 2)),
(4, _MatmulMConfig(16, 32, 8, 4)),
(8, _MatmulMConfig(32, 128, 8, 5)),
(16, _MatmulMConfig(32, 128, 4, 4)),
(32, _MatmulMConfig(64, 128, 4, 3)),
(64, _MatmulMConfig(16, 256, 4, 3)),
(256, _MatmulMConfig(32, 128, 4, 3)),
(512, _MatmulMConfig(64, 128, 4, 4)),
(1024, _MatmulMConfig(128, 128, 8, 3)),
(2048, _MatmulMConfig(64, 128, 4, 5)),
),
),
(75968, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp2_lm_head
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, 4, 3)),
(32, _MatmulMConfig(32, 256, 8, 3)),
(64, _MatmulMConfig(64, 256, 8, 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)),
),
),
(3072, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp4_gate_up_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 256, 4, 3)),
(4, _MatmulMConfig(32, 128, 4, 5)),
(8, _MatmulMConfig(16, 128, 8, 3)),
(16, _MatmulMConfig(32, 128, 8, 4)),
(32, _MatmulMConfig(32, 128, 8, 5)),
(64, _MatmulMConfig(64, 32, 8, 3)),
(256, _MatmulMConfig(128, 64, 8, 3)),
(512, _MatmulMConfig(128, 128, 4, 3)),
(1024, _MatmulMConfig(64, 64, 4, 5)),
(2048, _MatmulMConfig(64, 128, 8, 4)),
),
),
(2048, 1536): _MatmulShapeConfig( # Qwen3-1.7B tp4_down_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 128, 8, 2)),
(4, _MatmulMConfig(16, 256, 4, 3)),
(8, _MatmulMConfig(32, 128, 4, 2)),
(16, _MatmulMConfig(32, 32, 8, 5)),
(32, _MatmulMConfig(32, 32, 8, 4)),
(64, _MatmulMConfig(32, 128, 4, 4)),
(256, _MatmulMConfig(64, 64, 4, 3)),
(512, _MatmulMConfig(128, 64, 8, 4)),
(1024, _MatmulMConfig(128, 128, 4, 3)),
(2048, _MatmulMConfig(128, 64, 8, 4)),
),
),
(1024, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp4_qkv_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(32, 128, 8, 4)),
(4, _MatmulMConfig(16, 32, 8, 4)),
(8, _MatmulMConfig(16, 32, 4, 4)),
(16, _MatmulMConfig(32, 128, 8, 3)),
(32, _MatmulMConfig(32, 128, 8, 4)),
(64, _MatmulMConfig(32, 64, 4, 5)),
(256, _MatmulMConfig(64, 32, 4, 3)),
(512, _MatmulMConfig(64, 64, 4, 3)),
(1024, _MatmulMConfig(64, 128, 4, 5)),
(2048, _MatmulMConfig(128, 128, 8, 3)),
),
),
(2048, 512): _MatmulShapeConfig( # Qwen3-1.7B tp4_o_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(32, 32, 4, 5)),
(4, _MatmulMConfig(64, 128, 4, 2)),
(8, _MatmulMConfig(16, 128, 4, 2)),
(16, _MatmulMConfig(16, 256, 8, 3)),
(32, _MatmulMConfig(64, 32, 8, 4)),
(64, _MatmulMConfig(16, 128, 8, 4)),
(256, _MatmulMConfig(64, 64, 8, 5)),
(512, _MatmulMConfig(32, 128, 8, 4)),
(1024, _MatmulMConfig(64, 64, 8, 3)),
(2048, _MatmulMConfig(64, 128, 8, 3)),
),
),
(37984, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp4_lm_head
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 128, 8, 4)),
(4, _MatmulMConfig(16, 128, 8, 5)),
(8, _MatmulMConfig(16, 128, 4, 5)),
(16, _MatmulMConfig(16, 128, 8, 4)),
(32, _MatmulMConfig(32, 128, 4, 5)),
(64, _MatmulMConfig(64, 128, 4, 5)),
(256, _MatmulMConfig(64, 128, 4, 5)),
(512, _MatmulMConfig(128, 128, 4, 3)),
(1024, _MatmulMConfig(128, 128, 4, 3)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(1536, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp8_gate_up_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(32, 128, 8, 3)),
(4, _MatmulMConfig(16, 256, 4, 3)),
(8, _MatmulMConfig(16, 32, 8, 5)),
(16, _MatmulMConfig(32, 32, 8, 3)),
(32, _MatmulMConfig(32, 128, 8, 4)),
(64, _MatmulMConfig(16, 128, 8, 4)),
(256, _MatmulMConfig(32, 128, 8, 3)),
(512, _MatmulMConfig(64, 128, 8, 4)),
(1024, _MatmulMConfig(128, 128, 4, 3)),
(2048, _MatmulMConfig(64, 64, 4, 4)),
),
),
(2048, 768): _MatmulShapeConfig( # Qwen3-1.7B tp8_down_proj
block_k=128,
m_buckets=(
(1, _MatmulMConfig(32, 64, 4, 3)),
(4, _MatmulMConfig(32, 64, 8, 3)),
(8, _MatmulMConfig(16, 32, 4, 5)),
(16, _MatmulMConfig(32, 64, 4, 2)),
(32, _MatmulMConfig(16, 64, 8, 3)),
(64, _MatmulMConfig(16, 32, 8, 2)),
(256, _MatmulMConfig(32, 32, 4, 4)),
(512, _MatmulMConfig(32, 64, 4, 5)),
(1024, _MatmulMConfig(64, 128, 4, 3)),
(2048, _MatmulMConfig(64, 128, 4, 3)),
),
),
(512, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp8_qkv_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(64, 64, 8, 5)),
(4, _MatmulMConfig(16, 32, 4, 4)),
(8, _MatmulMConfig(16, 32, 4, 5)),
(16, _MatmulMConfig(32, 128, 8, 3)),
(32, _MatmulMConfig(64, 64, 8, 5)),
(64, _MatmulMConfig(32, 128, 8, 4)),
(256, _MatmulMConfig(16, 128, 4, 4)),
(512, _MatmulMConfig(32, 64, 8, 3)),
(1024, _MatmulMConfig(64, 64, 4, 4)),
(2048, _MatmulMConfig(64, 128, 4, 4)),
),
),
(2048, 256): _MatmulShapeConfig( # Qwen3-1.7B tp8_o_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(32, 128, 8, 2)),
(4, _MatmulMConfig(16, 128, 4, 2)),
(8, _MatmulMConfig(32, 256, 4, 2)),
(16, _MatmulMConfig(32, 32, 8, 4)),
(32, _MatmulMConfig(16, 256, 4, 3)),
(64, _MatmulMConfig(128, 32, 8, 3)),
(256, _MatmulMConfig(16, 32, 8, 2)),
(512, _MatmulMConfig(16, 32, 4, 3)),
(1024, _MatmulMConfig(64, 256, 4, 3)),
(2048, _MatmulMConfig(64, 128, 8, 4)),
),
),
(18992, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp8_lm_head
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 64, 8, 5)),
(4, _MatmulMConfig(16, 128, 8, 5)),
(8, _MatmulMConfig(16, 128, 8, 3)),
(16, _MatmulMConfig(16, 128, 4, 5)),
(32, _MatmulMConfig(32, 128, 8, 5)),
(64, _MatmulMConfig(64, 128, 8, 4)),
(256, _MatmulMConfig(128, 128, 4, 3)),
(512, _MatmulMConfig(128, 128, 4, 3)),
(1024, _MatmulMConfig(128, 128, 8, 3)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(9728, 2560): _MatmulShapeConfig( # Qwen3-4B tp2_gate_up_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 32, 8, 4)),
(4, _MatmulMConfig(16, 64, 8, 3)),
(8, _MatmulMConfig(16, 128, 4, 5)),
(16, _MatmulMConfig(16, 32, 4, 5)),
(32, _MatmulMConfig(32, 64, 8, 4)),
(64, _MatmulMConfig(32, 128, 4, 3)),
(256, _MatmulMConfig(128, 128, 4, 3)),
(512, _MatmulMConfig(128, 128, 4, 3)),
(1024, _MatmulMConfig(64, 128, 4, 5)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(2560, 4864): _MatmulShapeConfig( # Qwen3-4B tp2_down_proj
block_k=128,
m_buckets=(
(1, _MatmulMConfig(16, 32, 4, 3)),
(4, _MatmulMConfig(16, 128, 8, 3)),
(8, _MatmulMConfig(16, 64, 8, 4)),
(16, _MatmulMConfig(32, 64, 8, 4)),
(32, _MatmulMConfig(16, 64, 4, 3)),
(64, _MatmulMConfig(32, 32, 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)),
),
),
(3072, 2560): _MatmulShapeConfig( # Qwen3-4B tp2_qkv_proj
block_k=128,
m_buckets=(
(1, _MatmulMConfig(16, 128, 8, 3)),
(4, _MatmulMConfig(16, 32, 8, 2)),
(8, _MatmulMConfig(16, 32, 8, 4)),
(16, _MatmulMConfig(16, 64, 4, 3)),
(32, _MatmulMConfig(32, 64, 4, 4)),
(64, _MatmulMConfig(16, 128, 4, 3)),
(256, _MatmulMConfig(64, 128, 4, 3)),
(512, _MatmulMConfig(64, 64, 4, 3)),
(1024, _MatmulMConfig(64, 128, 4, 3)),
(2048, _MatmulMConfig(64, 128, 4, 3)),
),
),
(2560, 2048): _MatmulShapeConfig( # Qwen3-4B tp2_o_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(32, 128, 8, 4)),
(4, _MatmulMConfig(32, 128, 4, 4)),
(8, _MatmulMConfig(16, 128, 4, 3)),
(16, _MatmulMConfig(32, 128, 4, 5)),
(32, _MatmulMConfig(16, 128, 8, 4)),
(64, _MatmulMConfig(16, 128, 8, 4)),
(256, _MatmulMConfig(128, 32, 8, 5)),
(512, _MatmulMConfig(128, 64, 8, 3)),
(1024, _MatmulMConfig(128, 128, 4, 3)),
(2048, _MatmulMConfig(128, 128, 8, 3)),
),
),
(75968, 2560): _MatmulShapeConfig( # Qwen3-4B tp2_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, 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, 8, 3)),
),
),
(4864, 2560): _MatmulShapeConfig( # Qwen3-4B tp4_gate_up_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 32, 8, 5)),
(4, _MatmulMConfig(16, 128, 4, 4)),
(8, _MatmulMConfig(16, 32, 8, 4)),
(16, _MatmulMConfig(16, 128, 8, 4)),
(32, _MatmulMConfig(16, 128, 4, 5)),
(64, _MatmulMConfig(16, 128, 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)),
),
),
(2560, 2432): _MatmulShapeConfig( # Qwen3-4B tp4_down_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 32, 8, 4)),
(4, _MatmulMConfig(16, 32, 4, 5)),
(8, _MatmulMConfig(32, 64, 8, 4)),
(16, _MatmulMConfig(32, 128, 8, 5)),
(32, _MatmulMConfig(16, 128, 4, 4)),
(64, _MatmulMConfig(64, 32, 4, 3)),
(256, _MatmulMConfig(32, 128, 8, 4)),
(512, _MatmulMConfig(64, 128, 4, 5)),
(1024, _MatmulMConfig(128, 128, 4, 3)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(1536, 2560): _MatmulShapeConfig( # Qwen3-4B tp4_qkv_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 32, 8, 4)),
(4, _MatmulMConfig(16, 32, 8, 3)),
(8, _MatmulMConfig(16, 32, 4, 4)),
(16, _MatmulMConfig(32, 128, 8, 3)),
(32, _MatmulMConfig(16, 128, 4, 3)),
(64, _MatmulMConfig(16, 64, 8, 4)),
(256, _MatmulMConfig(128, 32, 4, 5)),
(512, _MatmulMConfig(64, 128, 4, 5)),
(1024, _MatmulMConfig(64, 256, 4, 3)),
(2048, _MatmulMConfig(64, 128, 4, 4)),
),
),
(2560, 1024): _MatmulShapeConfig( # Qwen3-4B tp4_o_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(32, 128, 8, 5)),
(4, _MatmulMConfig(16, 32, 8, 3)),
(8, _MatmulMConfig(16, 128, 4, 3)),
(16, _MatmulMConfig(32, 256, 4, 3)),
(32, _MatmulMConfig(16, 128, 4, 2)),
(64, _MatmulMConfig(32, 256, 8, 3)),
(256, _MatmulMConfig(32, 128, 4, 5)),
(512, _MatmulMConfig(64, 128, 8, 3)),
(1024, _MatmulMConfig(128, 128, 4, 3)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(37984, 2560): _MatmulShapeConfig( # Qwen3-4B tp4_lm_head
block_k=128,
m_buckets=(
(1, _MatmulMConfig(16, 128, 8, 3)),
(4, _MatmulMConfig(16, 128, 8, 3)),
(8, _MatmulMConfig(16, 128, 4, 3)),
(16, _MatmulMConfig(16, 32, 4, 5)),
(32, _MatmulMConfig(16, 32, 4, 5)),
(64, _MatmulMConfig(32, 32, 4, 5)),
(256, _MatmulMConfig(64, 128, 4, 3)),
(512, _MatmulMConfig(64, 128, 4, 3)),
(1024, _MatmulMConfig(64, 128, 4, 3)),
(2048, _MatmulMConfig(64, 128, 4, 3)),
),
),
(2432, 2560): _MatmulShapeConfig( # Qwen3-4B tp8_gate_up_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 32, 4, 5)),
(4, _MatmulMConfig(32, 128, 4, 3)),
(8, _MatmulMConfig(16, 128, 8, 3)),
(16, _MatmulMConfig(16, 128, 4, 4)),
(32, _MatmulMConfig(16, 128, 8, 4)),
(64, _MatmulMConfig(16, 128, 4, 3)),
(256, _MatmulMConfig(64, 64, 4, 5)),
(512, _MatmulMConfig(64, 128, 4, 4)),
(1024, _MatmulMConfig(128, 128, 8, 3)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(2560, 1216): _MatmulShapeConfig( # Qwen3-4B tp8_down_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 256, 8, 2)),
(4, _MatmulMConfig(16, 128, 8, 4)),
(8, _MatmulMConfig(16, 32, 4, 3)),
(16, _MatmulMConfig(32, 32, 4, 3)),
(32, _MatmulMConfig(16, 128, 8, 4)),
(64, _MatmulMConfig(32, 128, 8, 4)),
(256, _MatmulMConfig(32, 128, 4, 5)),
(512, _MatmulMConfig(64, 128, 8, 4)),
(1024, _MatmulMConfig(64, 256, 4, 3)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(768, 2560): _MatmulShapeConfig( # Qwen3-4B tp8_qkv_proj
block_k=128,
m_buckets=(
(1, _MatmulMConfig(32, 32, 4, 2)),
(4, _MatmulMConfig(16, 64, 8, 3)),
(8, _MatmulMConfig(16, 64, 8, 4)),
(16, _MatmulMConfig(16, 64, 4, 5)),
(32, _MatmulMConfig(32, 64, 8, 4)),
(64, _MatmulMConfig(16, 64, 8, 4)),
(256, _MatmulMConfig(64, 32, 4, 5)),
(512, _MatmulMConfig(64, 64, 8, 3)),
(1024, _MatmulMConfig(64, 128, 4, 3)),
(2048, _MatmulMConfig(64, 64, 4, 3)),
),
),
(2560, 512): _MatmulShapeConfig( # Qwen3-4B tp8_o_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 128, 8, 4)),
(4, _MatmulMConfig(16, 32, 4, 2)),
(8, _MatmulMConfig(16, 128, 4, 5)),
(16, _MatmulMConfig(32, 32, 4, 5)),
(32, _MatmulMConfig(32, 256, 4, 3)),
(64, _MatmulMConfig(32, 128, 8, 5)),
(256, _MatmulMConfig(16, 256, 4, 2)),
(512, _MatmulMConfig(32, 128, 8, 4)),
(1024, _MatmulMConfig(128, 128, 8, 3)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(18992, 2560): _MatmulShapeConfig( # Qwen3-4B tp8_lm_head
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 64, 8, 4)),
(4, _MatmulMConfig(16, 64, 8, 5)),
(8, _MatmulMConfig(16, 128, 8, 3)),
(16, _MatmulMConfig(16, 128, 8, 5)),
(32, _MatmulMConfig(32, 128, 8, 5)),
(64, _MatmulMConfig(64, 128, 8, 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)),
),
),
(12288, 4096): _MatmulShapeConfig( # Qwen3-8B tp2_gate_up_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 128, 8, 3)),
(4, _MatmulMConfig(16, 128, 8, 5)),
(8, _MatmulMConfig(16, 128, 8, 4)),
(16, _MatmulMConfig(16, 128, 4, 4)),
(32, _MatmulMConfig(32, 128, 4, 4)),
(64, _MatmulMConfig(64, 128, 4, 5)),
(256, _MatmulMConfig(64, 64, 4, 5)),
(512, _MatmulMConfig(128, 64, 4, 3)),
(1024, _MatmulMConfig(128, 128, 4, 3)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(4096, 6144): _MatmulShapeConfig( # Qwen3-8B tp2_down_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 32, 8, 5)),
(4, _MatmulMConfig(16, 32, 8, 5)),
(8, _MatmulMConfig(16, 32, 8, 3)),
(16, _MatmulMConfig(16, 32, 4, 5)),
(32, _MatmulMConfig(16, 64, 8, 4)),
(64, _MatmulMConfig(32, 64, 4, 5)),
(256, _MatmulMConfig(64, 128, 4, 5)),
(512, _MatmulMConfig(128, 128, 4, 3)),
(1024, _MatmulMConfig(128, 64, 8, 4)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(3072, 4096): _MatmulShapeConfig( # Qwen3-8B tp2_qkv_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 32, 8, 4)),
(4, _MatmulMConfig(16, 32, 8, 3)),
(8, _MatmulMConfig(16, 128, 4, 3)),
(16, _MatmulMConfig(16, 32, 4, 4)),
(32, _MatmulMConfig(32, 32, 8, 5)),
(64, _MatmulMConfig(16, 128, 8, 4)),
(256, _MatmulMConfig(64, 128, 4, 5)),
(512, _MatmulMConfig(128, 128, 8, 3)),
(1024, _MatmulMConfig(64, 64, 4, 5)),
(2048, _MatmulMConfig(128, 64, 4, 3)),
),
),
(75968, 4096): _MatmulShapeConfig( # Qwen3-8B tp2_lm_head
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, 8, 3)),
(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, 3072): _MatmulShapeConfig( # Qwen3-8B tp4_down_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 256, 8, 3)),
(4, _MatmulMConfig(16, 128, 4, 5)),
(8, _MatmulMConfig(16, 128, 4, 4)),
(16, _MatmulMConfig(16, 64, 4, 5)),
(32, _MatmulMConfig(32, 128, 4, 3)),
(64, _MatmulMConfig(16, 128, 4, 4)),
(256, _MatmulMConfig(64, 128, 4, 5)),
(512, _MatmulMConfig(128, 128, 4, 3)),
(1024, _MatmulMConfig(128, 64, 8, 4)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(1536, 4096): _MatmulShapeConfig( # Qwen3-8B tp4_qkv_proj
block_k=128,
m_buckets=(
(1, _MatmulMConfig(16, 64, 4, 5)),
(4, _MatmulMConfig(16, 64, 8, 3)),
(8, _MatmulMConfig(32, 64, 4, 5)),
(16, _MatmulMConfig(16, 64, 8, 4)),
(32, _MatmulMConfig(16, 32, 8, 3)),
(64, _MatmulMConfig(16, 64, 8, 5)),
(256, _MatmulMConfig(64, 64, 4, 3)),
(512, _MatmulMConfig(64, 128, 4, 3)),
(1024, _MatmulMConfig(64, 64, 4, 3)),
(2048, _MatmulMConfig(64, 128, 4, 3)),
),
),
(4096, 1024): _MatmulShapeConfig( # Qwen3-8B tp4_o_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 128, 8, 4)),
(4, _MatmulMConfig(32, 128, 8, 4)),
(8, _MatmulMConfig(16, 256, 8, 3)),
(16, _MatmulMConfig(16, 128, 8, 3)),
(32, _MatmulMConfig(16, 128, 4, 4)),
(64, _MatmulMConfig(64, 64, 4, 3)),
(256, _MatmulMConfig(128, 64, 4, 3)),
(512, _MatmulMConfig(64, 64, 4, 5)),
(1024, _MatmulMConfig(128, 64, 4, 4)),
(2048, _MatmulMConfig(128, 128, 4, 3)),
),
),
(37984, 4096): _MatmulShapeConfig( # Qwen3-8B tp4_lm_head
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 128, 8, 3)),
(4, _MatmulMConfig(16, 128, 4, 5)),
(8, _MatmulMConfig(16, 128, 4, 5)),
(16, _MatmulMConfig(16, 128, 8, 3)),
(32, _MatmulMConfig(32, 128, 8, 3)),
(64, _MatmulMConfig(64, 128, 8, 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, 1536): _MatmulShapeConfig( # Qwen3-8B tp8_down_proj
block_k=32,
m_buckets=(
(1, _MatmulMConfig(32, 128, 8, 5)),
(4, _MatmulMConfig(16, 256, 4, 5)),
(8, _MatmulMConfig(16, 64, 4, 4)),
(16, _MatmulMConfig(16, 128, 8, 3)),
(32, _MatmulMConfig(16, 64, 4, 3)),
(64, _MatmulMConfig(16, 128, 4, 4)),
(256, _MatmulMConfig(64, 128, 4, 4)),
(512, _MatmulMConfig(128, 128, 8, 5)),
(1024, _MatmulMConfig(128, 256, 8, 4)),
(2048, _MatmulMConfig(128, 128, 8, 4)),
),
),
(768, 4096): _MatmulShapeConfig( # Qwen3-8B tp8_qkv_proj
block_k=128,
m_buckets=(
(1, _MatmulMConfig(16, 32, 8, 5)),
(4, _MatmulMConfig(16, 32, 8, 3)),
(8, _MatmulMConfig(32, 64, 4, 5)),
(16, _MatmulMConfig(32, 64, 8, 3)),
(32, _MatmulMConfig(32, 32, 8, 5)),
(64, _MatmulMConfig(64, 32, 8, 3)),
(256, _MatmulMConfig(32, 64, 4, 4)),
(512, _MatmulMConfig(64, 64, 8, 3)),
(1024, _MatmulMConfig(64, 128, 4, 3)),
(2048, _MatmulMConfig(64, 64, 4, 3)),
),
),
(4096, 512): _MatmulShapeConfig( # Qwen3-8B tp8_o_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 128, 8, 2)),
(4, _MatmulMConfig(32, 128, 4, 2)),
(8, _MatmulMConfig(16, 256, 4, 2)),
(16, _MatmulMConfig(32, 256, 4, 3)),
(32, _MatmulMConfig(32, 256, 8, 3)),
(64, _MatmulMConfig(32, 128, 4, 5)),
(256, _MatmulMConfig(64, 32, 8, 4)),
(512, _MatmulMConfig(128, 128, 8, 3)),
(1024, _MatmulMConfig(128, 64, 4, 4)),
(2048, _MatmulMConfig(128, 128, 8, 3)),
),
),
(18992, 4096): _MatmulShapeConfig( # Qwen3-8B tp8_lm_head
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 64, 8, 5)),
(4, _MatmulMConfig(16, 64, 8, 4)),
(8, _MatmulMConfig(16, 64, 8, 5)),
(16, _MatmulMConfig(16, 128, 4, 3)),
(32, _MatmulMConfig(32, 64, 8, 5)),
(64, _MatmulMConfig(64, 128, 8, 3)),
(256, _MatmulMConfig(128, 128, 4, 3)),
(512, _MatmulMConfig(128, 128, 4, 3)),
(1024, _MatmulMConfig(128, 128, 4, 3)),
(2048, _MatmulMConfig(128, 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,
}
Reference in New Issue View Git Blame Copy Permalink