diff --git a/vllm/model_executor/determinism/batch_invariant_configs.py b/vllm/model_executor/determinism/batch_invariant_configs.py index 39cda984f7..874447bb82 100644 --- a/vllm/model_executor/determinism/batch_invariant_configs.py +++ b/vllm/model_executor/determinism/batch_invariant_configs.py @@ -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: