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

[Kernel][Perf] Add sm_120 (RTX PRO 6000 / RTX 50) tuned configs for batch-invariant persistent matmul (#57456)

Browse Source 
Signed-off-by: LioEinaudi <zhao3024667639@gmail.com>
This commit is contained in:
Lio Einaudi
2026-09-19 09:30:55 -04:00
committed by GitHub
parent fbe8a157fb
commit 41c4a3ed4e
1 changed files with 229 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/determinism/batch_invariant_configs.py
+229
View File
@@ -254,6 +254,233 @@ _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS: dict[
),
),
},
"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 = (
@@ -265,6 +492,8 @@ _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: