From 4e0a414c440e489d472a797ff1028cfc41565693 Mon Sep 17 00:00:00 2001 From: Lio Einaudi Date: Wed, 30 Sep 2026 07:51:50 +0800 Subject: [PATCH] [Kernel][Perf] Add TP=2/4/8 per-rank shapes to the sm_120 batch-invariant matmul table (#58495) Signed-off-by: LioEinaudi Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../determinism/batch_invariant_configs.py | 615 ++++++++++++++++++ 1 file changed, 615 insertions(+) diff --git a/vllm/model_executor/determinism/batch_invariant_configs.py b/vllm/model_executor/determinism/batch_invariant_configs.py index 874447bb82..4fff768c1f 100644 --- a/vllm/model_executor/determinism/batch_invariant_configs.py +++ b/vllm/model_executor/determinism/batch_invariant_configs.py @@ -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)), + ), + ), }, }