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

[Kernel][Perf] Add TP=2/4/8 per-rank shapes to the sm_120 batch-invariant matmul table (#58495)

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-09-29 23:51:50 +00:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent be255076d0
commit 4e0a414c44
1 changed files with 615 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/determinism/batch_invariant_configs.py
+615
View File
@@ -480,6 +480,621 @@ _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS: dict[
(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)),
),
),
},
}