@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 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>

Lio Einaudi committed Sep 30, 2026 at 07:51 UTC 4e0a414c440e489d472a797ff1028cfc41565693
1 file changed +615
vllm/model_executor/determinism/batch_invariant_configs.py
+615
@@ -480,6 +480,621 @@ _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS: dict[
480 (2048, _MatmulMConfig(64, 128, 4, 3)),
481 ),
482 ),
483 + (6144, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp2_gate_up_proj
484 + block_k=64,
485 + m_buckets=(
486 + (1, _MatmulMConfig(16, 32, 8, 3)),
487 + (4, _MatmulMConfig(16, 128, 8, 5)),
488 + (8, _MatmulMConfig(16, 128, 4, 5)),
489 + (16, _MatmulMConfig(16, 32, 4, 5)),
490 + (32, _MatmulMConfig(32, 128, 8, 5)),
491 + (64, _MatmulMConfig(32, 128, 4, 4)),
492 + (256, _MatmulMConfig(128, 128, 4, 3)),
493 + (512, _MatmulMConfig(64, 64, 4, 5)),
494 + (1024, _MatmulMConfig(128, 64, 4, 3)),
495 + (2048, _MatmulMConfig(128, 128, 4, 3)),
496 + ),
497 + ),
498 + (2048, 3072): _MatmulShapeConfig( # Qwen3-1.7B tp2_down_proj
499 + block_k=128,
500 + m_buckets=(
501 + (1, _MatmulMConfig(16, 32, 8, 5)),
502 + (4, _MatmulMConfig(16, 64, 4, 3)),
503 + (8, _MatmulMConfig(16, 64, 8, 4)),
504 + (16, _MatmulMConfig(16, 64, 8, 3)),
505 + (32, _MatmulMConfig(16, 64, 8, 3)),
506 + (64, _MatmulMConfig(16, 64, 4, 5)),
507 + (256, _MatmulMConfig(64, 64, 8, 3)),
508 + (512, _MatmulMConfig(64, 128, 4, 3)),
509 + (1024, _MatmulMConfig(64, 64, 4, 3)),
510 + (2048, _MatmulMConfig(64, 128, 4, 3)),
511 + ),
512 + ),
513 + (2048, 1024): _MatmulShapeConfig( # Qwen3-1.7B tp2_o_proj
514 + block_k=64,
515 + m_buckets=(
516 + (1, _MatmulMConfig(16, 256, 8, 2)),
517 + (4, _MatmulMConfig(16, 32, 8, 4)),
518 + (8, _MatmulMConfig(32, 128, 8, 5)),
519 + (16, _MatmulMConfig(32, 128, 4, 4)),
520 + (32, _MatmulMConfig(64, 128, 4, 3)),
521 + (64, _MatmulMConfig(16, 256, 4, 3)),
522 + (256, _MatmulMConfig(32, 128, 4, 3)),
523 + (512, _MatmulMConfig(64, 128, 4, 4)),
524 + (1024, _MatmulMConfig(128, 128, 8, 3)),
525 + (2048, _MatmulMConfig(64, 128, 4, 5)),
526 + ),
527 + ),
528 + (75968, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp2_lm_head
529 + block_k=64,
530 + m_buckets=(
531 + (1, _MatmulMConfig(16, 256, 8, 3)),
532 + (4, _MatmulMConfig(16, 256, 8, 3)),
533 + (8, _MatmulMConfig(16, 256, 8, 3)),
534 + (16, _MatmulMConfig(16, 256, 4, 3)),
535 + (32, _MatmulMConfig(32, 256, 8, 3)),
536 + (64, _MatmulMConfig(64, 256, 8, 3)),
537 + (256, _MatmulMConfig(64, 128, 4, 5)),
538 + (512, _MatmulMConfig(128, 128, 4, 3)),
539 + (1024, _MatmulMConfig(128, 128, 4, 3)),
540 + (2048, _MatmulMConfig(128, 128, 4, 3)),
541 + ),
542 + ),
543 + (3072, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp4_gate_up_proj
544 + block_k=64,
545 + m_buckets=(
546 + (1, _MatmulMConfig(16, 256, 4, 3)),
547 + (4, _MatmulMConfig(32, 128, 4, 5)),
548 + (8, _MatmulMConfig(16, 128, 8, 3)),
549 + (16, _MatmulMConfig(32, 128, 8, 4)),
550 + (32, _MatmulMConfig(32, 128, 8, 5)),
551 + (64, _MatmulMConfig(64, 32, 8, 3)),
552 + (256, _MatmulMConfig(128, 64, 8, 3)),
553 + (512, _MatmulMConfig(128, 128, 4, 3)),
554 + (1024, _MatmulMConfig(64, 64, 4, 5)),
555 + (2048, _MatmulMConfig(64, 128, 8, 4)),
556 + ),
557 + ),
558 + (2048, 1536): _MatmulShapeConfig( # Qwen3-1.7B tp4_down_proj
559 + block_k=64,
560 + m_buckets=(
561 + (1, _MatmulMConfig(16, 128, 8, 2)),
562 + (4, _MatmulMConfig(16, 256, 4, 3)),
563 + (8, _MatmulMConfig(32, 128, 4, 2)),
564 + (16, _MatmulMConfig(32, 32, 8, 5)),
565 + (32, _MatmulMConfig(32, 32, 8, 4)),
566 + (64, _MatmulMConfig(32, 128, 4, 4)),
567 + (256, _MatmulMConfig(64, 64, 4, 3)),
568 + (512, _MatmulMConfig(128, 64, 8, 4)),
569 + (1024, _MatmulMConfig(128, 128, 4, 3)),
570 + (2048, _MatmulMConfig(128, 64, 8, 4)),
571 + ),
572 + ),
573 + (1024, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp4_qkv_proj
574 + block_k=64,
575 + m_buckets=(
576 + (1, _MatmulMConfig(32, 128, 8, 4)),
577 + (4, _MatmulMConfig(16, 32, 8, 4)),
578 + (8, _MatmulMConfig(16, 32, 4, 4)),
579 + (16, _MatmulMConfig(32, 128, 8, 3)),
580 + (32, _MatmulMConfig(32, 128, 8, 4)),
581 + (64, _MatmulMConfig(32, 64, 4, 5)),
582 + (256, _MatmulMConfig(64, 32, 4, 3)),
583 + (512, _MatmulMConfig(64, 64, 4, 3)),
584 + (1024, _MatmulMConfig(64, 128, 4, 5)),
585 + (2048, _MatmulMConfig(128, 128, 8, 3)),
586 + ),
587 + ),
588 + (2048, 512): _MatmulShapeConfig( # Qwen3-1.7B tp4_o_proj
589 + block_k=64,
590 + m_buckets=(
591 + (1, _MatmulMConfig(32, 32, 4, 5)),
592 + (4, _MatmulMConfig(64, 128, 4, 2)),
593 + (8, _MatmulMConfig(16, 128, 4, 2)),
594 + (16, _MatmulMConfig(16, 256, 8, 3)),
595 + (32, _MatmulMConfig(64, 32, 8, 4)),
596 + (64, _MatmulMConfig(16, 128, 8, 4)),
597 + (256, _MatmulMConfig(64, 64, 8, 5)),
598 + (512, _MatmulMConfig(32, 128, 8, 4)),
599 + (1024, _MatmulMConfig(64, 64, 8, 3)),
600 + (2048, _MatmulMConfig(64, 128, 8, 3)),
601 + ),
602 + ),
603 + (37984, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp4_lm_head
604 + block_k=64,
605 + m_buckets=(
606 + (1, _MatmulMConfig(16, 128, 8, 4)),
607 + (4, _MatmulMConfig(16, 128, 8, 5)),
608 + (8, _MatmulMConfig(16, 128, 4, 5)),
609 + (16, _MatmulMConfig(16, 128, 8, 4)),
610 + (32, _MatmulMConfig(32, 128, 4, 5)),
611 + (64, _MatmulMConfig(64, 128, 4, 5)),
612 + (256, _MatmulMConfig(64, 128, 4, 5)),
613 + (512, _MatmulMConfig(128, 128, 4, 3)),
614 + (1024, _MatmulMConfig(128, 128, 4, 3)),
615 + (2048, _MatmulMConfig(128, 128, 4, 3)),
616 + ),
617 + ),
618 + (1536, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp8_gate_up_proj
619 + block_k=64,
620 + m_buckets=(
621 + (1, _MatmulMConfig(32, 128, 8, 3)),
622 + (4, _MatmulMConfig(16, 256, 4, 3)),
623 + (8, _MatmulMConfig(16, 32, 8, 5)),
624 + (16, _MatmulMConfig(32, 32, 8, 3)),
625 + (32, _MatmulMConfig(32, 128, 8, 4)),
626 + (64, _MatmulMConfig(16, 128, 8, 4)),
627 + (256, _MatmulMConfig(32, 128, 8, 3)),
628 + (512, _MatmulMConfig(64, 128, 8, 4)),
629 + (1024, _MatmulMConfig(128, 128, 4, 3)),
630 + (2048, _MatmulMConfig(64, 64, 4, 4)),
631 + ),
632 + ),
633 + (2048, 768): _MatmulShapeConfig( # Qwen3-1.7B tp8_down_proj
634 + block_k=128,
635 + m_buckets=(
636 + (1, _MatmulMConfig(32, 64, 4, 3)),
637 + (4, _MatmulMConfig(32, 64, 8, 3)),
638 + (8, _MatmulMConfig(16, 32, 4, 5)),
639 + (16, _MatmulMConfig(32, 64, 4, 2)),
640 + (32, _MatmulMConfig(16, 64, 8, 3)),
641 + (64, _MatmulMConfig(16, 32, 8, 2)),
642 + (256, _MatmulMConfig(32, 32, 4, 4)),
643 + (512, _MatmulMConfig(32, 64, 4, 5)),
644 + (1024, _MatmulMConfig(64, 128, 4, 3)),
645 + (2048, _MatmulMConfig(64, 128, 4, 3)),
646 + ),
647 + ),
648 + (512, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp8_qkv_proj
649 + block_k=64,
650 + m_buckets=(
651 + (1, _MatmulMConfig(64, 64, 8, 5)),
652 + (4, _MatmulMConfig(16, 32, 4, 4)),
653 + (8, _MatmulMConfig(16, 32, 4, 5)),
654 + (16, _MatmulMConfig(32, 128, 8, 3)),
655 + (32, _MatmulMConfig(64, 64, 8, 5)),
656 + (64, _MatmulMConfig(32, 128, 8, 4)),
657 + (256, _MatmulMConfig(16, 128, 4, 4)),
658 + (512, _MatmulMConfig(32, 64, 8, 3)),
659 + (1024, _MatmulMConfig(64, 64, 4, 4)),
660 + (2048, _MatmulMConfig(64, 128, 4, 4)),
661 + ),
662 + ),
663 + (2048, 256): _MatmulShapeConfig( # Qwen3-1.7B tp8_o_proj
664 + block_k=64,
665 + m_buckets=(
666 + (1, _MatmulMConfig(32, 128, 8, 2)),
667 + (4, _MatmulMConfig(16, 128, 4, 2)),
668 + (8, _MatmulMConfig(32, 256, 4, 2)),
669 + (16, _MatmulMConfig(32, 32, 8, 4)),
670 + (32, _MatmulMConfig(16, 256, 4, 3)),
671 + (64, _MatmulMConfig(128, 32, 8, 3)),
672 + (256, _MatmulMConfig(16, 32, 8, 2)),
673 + (512, _MatmulMConfig(16, 32, 4, 3)),
674 + (1024, _MatmulMConfig(64, 256, 4, 3)),
675 + (2048, _MatmulMConfig(64, 128, 8, 4)),
676 + ),
677 + ),
678 + (18992, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp8_lm_head
679 + block_k=64,
680 + m_buckets=(
681 + (1, _MatmulMConfig(16, 64, 8, 5)),
682 + (4, _MatmulMConfig(16, 128, 8, 5)),
683 + (8, _MatmulMConfig(16, 128, 8, 3)),
684 + (16, _MatmulMConfig(16, 128, 4, 5)),
685 + (32, _MatmulMConfig(32, 128, 8, 5)),
686 + (64, _MatmulMConfig(64, 128, 8, 4)),
687 + (256, _MatmulMConfig(128, 128, 4, 3)),
688 + (512, _MatmulMConfig(128, 128, 4, 3)),
689 + (1024, _MatmulMConfig(128, 128, 8, 3)),
690 + (2048, _MatmulMConfig(128, 128, 4, 3)),
691 + ),
692 + ),
693 + (9728, 2560): _MatmulShapeConfig( # Qwen3-4B tp2_gate_up_proj
694 + block_k=64,
695 + m_buckets=(
696 + (1, _MatmulMConfig(16, 32, 8, 4)),
697 + (4, _MatmulMConfig(16, 64, 8, 3)),
698 + (8, _MatmulMConfig(16, 128, 4, 5)),
699 + (16, _MatmulMConfig(16, 32, 4, 5)),
700 + (32, _MatmulMConfig(32, 64, 8, 4)),
701 + (64, _MatmulMConfig(32, 128, 4, 3)),
702 + (256, _MatmulMConfig(128, 128, 4, 3)),
703 + (512, _MatmulMConfig(128, 128, 4, 3)),
704 + (1024, _MatmulMConfig(64, 128, 4, 5)),
705 + (2048, _MatmulMConfig(128, 128, 4, 3)),
706 + ),
707 + ),
708 + (2560, 4864): _MatmulShapeConfig( # Qwen3-4B tp2_down_proj
709 + block_k=128,
710 + m_buckets=(
711 + (1, _MatmulMConfig(16, 32, 4, 3)),
712 + (4, _MatmulMConfig(16, 128, 8, 3)),
713 + (8, _MatmulMConfig(16, 64, 8, 4)),
714 + (16, _MatmulMConfig(32, 64, 8, 4)),
715 + (32, _MatmulMConfig(16, 64, 4, 3)),
716 + (64, _MatmulMConfig(32, 32, 8, 4)),
717 + (256, _MatmulMConfig(64, 64, 4, 3)),
718 + (512, _MatmulMConfig(64, 128, 4, 3)),
719 + (1024, _MatmulMConfig(64, 128, 4, 3)),
720 + (2048, _MatmulMConfig(64, 128, 4, 3)),
721 + ),
722 + ),
723 + (3072, 2560): _MatmulShapeConfig( # Qwen3-4B tp2_qkv_proj
724 + block_k=128,
725 + m_buckets=(
726 + (1, _MatmulMConfig(16, 128, 8, 3)),
727 + (4, _MatmulMConfig(16, 32, 8, 2)),
728 + (8, _MatmulMConfig(16, 32, 8, 4)),
729 + (16, _MatmulMConfig(16, 64, 4, 3)),
730 + (32, _MatmulMConfig(32, 64, 4, 4)),
731 + (64, _MatmulMConfig(16, 128, 4, 3)),
732 + (256, _MatmulMConfig(64, 128, 4, 3)),
733 + (512, _MatmulMConfig(64, 64, 4, 3)),
734 + (1024, _MatmulMConfig(64, 128, 4, 3)),
735 + (2048, _MatmulMConfig(64, 128, 4, 3)),
736 + ),
737 + ),
738 + (2560, 2048): _MatmulShapeConfig( # Qwen3-4B tp2_o_proj
739 + block_k=64,
740 + m_buckets=(
741 + (1, _MatmulMConfig(32, 128, 8, 4)),
742 + (4, _MatmulMConfig(32, 128, 4, 4)),
743 + (8, _MatmulMConfig(16, 128, 4, 3)),
744 + (16, _MatmulMConfig(32, 128, 4, 5)),
745 + (32, _MatmulMConfig(16, 128, 8, 4)),
746 + (64, _MatmulMConfig(16, 128, 8, 4)),
747 + (256, _MatmulMConfig(128, 32, 8, 5)),
748 + (512, _MatmulMConfig(128, 64, 8, 3)),
749 + (1024, _MatmulMConfig(128, 128, 4, 3)),
750 + (2048, _MatmulMConfig(128, 128, 8, 3)),
751 + ),
752 + ),
753 + (75968, 2560): _MatmulShapeConfig( # Qwen3-4B tp2_lm_head
754 + block_k=64,
755 + m_buckets=(
756 + (1, _MatmulMConfig(16, 64, 8, 5)),
757 + (4, _MatmulMConfig(16, 64, 8, 5)),
758 + (8, _MatmulMConfig(16, 64, 4, 5)),
759 + (16, _MatmulMConfig(16, 64, 4, 5)),
760 + (32, _MatmulMConfig(32, 64, 4, 5)),
761 + (64, _MatmulMConfig(64, 256, 4, 3)),
762 + (256, _MatmulMConfig(64, 128, 4, 5)),
763 + (512, _MatmulMConfig(128, 128, 4, 3)),
764 + (1024, _MatmulMConfig(128, 128, 4, 3)),
765 + (2048, _MatmulMConfig(128, 128, 8, 3)),
766 + ),
767 + ),
768 + (4864, 2560): _MatmulShapeConfig( # Qwen3-4B tp4_gate_up_proj
769 + block_k=64,
770 + m_buckets=(
771 + (1, _MatmulMConfig(16, 32, 8, 5)),
772 + (4, _MatmulMConfig(16, 128, 4, 4)),
773 + (8, _MatmulMConfig(16, 32, 8, 4)),
774 + (16, _MatmulMConfig(16, 128, 8, 4)),
775 + (32, _MatmulMConfig(16, 128, 4, 5)),
776 + (64, _MatmulMConfig(16, 128, 4, 3)),
777 + (256, _MatmulMConfig(64, 128, 4, 5)),
778 + (512, _MatmulMConfig(128, 128, 4, 3)),
779 + (1024, _MatmulMConfig(128, 128, 4, 3)),
780 + (2048, _MatmulMConfig(128, 128, 4, 3)),
781 + ),
782 + ),
783 + (2560, 2432): _MatmulShapeConfig( # Qwen3-4B tp4_down_proj
784 + block_k=64,
785 + m_buckets=(
786 + (1, _MatmulMConfig(16, 32, 8, 4)),
787 + (4, _MatmulMConfig(16, 32, 4, 5)),
788 + (8, _MatmulMConfig(32, 64, 8, 4)),
789 + (16, _MatmulMConfig(32, 128, 8, 5)),
790 + (32, _MatmulMConfig(16, 128, 4, 4)),
791 + (64, _MatmulMConfig(64, 32, 4, 3)),
792 + (256, _MatmulMConfig(32, 128, 8, 4)),
793 + (512, _MatmulMConfig(64, 128, 4, 5)),
794 + (1024, _MatmulMConfig(128, 128, 4, 3)),
795 + (2048, _MatmulMConfig(128, 128, 4, 3)),
796 + ),
797 + ),
798 + (1536, 2560): _MatmulShapeConfig( # Qwen3-4B tp4_qkv_proj
799 + block_k=64,
800 + m_buckets=(
801 + (1, _MatmulMConfig(16, 32, 8, 4)),
802 + (4, _MatmulMConfig(16, 32, 8, 3)),
803 + (8, _MatmulMConfig(16, 32, 4, 4)),
804 + (16, _MatmulMConfig(32, 128, 8, 3)),
805 + (32, _MatmulMConfig(16, 128, 4, 3)),
806 + (64, _MatmulMConfig(16, 64, 8, 4)),
807 + (256, _MatmulMConfig(128, 32, 4, 5)),
808 + (512, _MatmulMConfig(64, 128, 4, 5)),
809 + (1024, _MatmulMConfig(64, 256, 4, 3)),
810 + (2048, _MatmulMConfig(64, 128, 4, 4)),
811 + ),
812 + ),
813 + (2560, 1024): _MatmulShapeConfig( # Qwen3-4B tp4_o_proj
814 + block_k=64,
815 + m_buckets=(
816 + (1, _MatmulMConfig(32, 128, 8, 5)),
817 + (4, _MatmulMConfig(16, 32, 8, 3)),
818 + (8, _MatmulMConfig(16, 128, 4, 3)),
819 + (16, _MatmulMConfig(32, 256, 4, 3)),
820 + (32, _MatmulMConfig(16, 128, 4, 2)),
821 + (64, _MatmulMConfig(32, 256, 8, 3)),
822 + (256, _MatmulMConfig(32, 128, 4, 5)),
823 + (512, _MatmulMConfig(64, 128, 8, 3)),
824 + (1024, _MatmulMConfig(128, 128, 4, 3)),
825 + (2048, _MatmulMConfig(128, 128, 4, 3)),
826 + ),
827 + ),
828 + (37984, 2560): _MatmulShapeConfig( # Qwen3-4B tp4_lm_head
829 + block_k=128,
830 + m_buckets=(
831 + (1, _MatmulMConfig(16, 128, 8, 3)),
832 + (4, _MatmulMConfig(16, 128, 8, 3)),
833 + (8, _MatmulMConfig(16, 128, 4, 3)),
834 + (16, _MatmulMConfig(16, 32, 4, 5)),
835 + (32, _MatmulMConfig(16, 32, 4, 5)),
836 + (64, _MatmulMConfig(32, 32, 4, 5)),
837 + (256, _MatmulMConfig(64, 128, 4, 3)),
838 + (512, _MatmulMConfig(64, 128, 4, 3)),
839 + (1024, _MatmulMConfig(64, 128, 4, 3)),
840 + (2048, _MatmulMConfig(64, 128, 4, 3)),
841 + ),
842 + ),
843 + (2432, 2560): _MatmulShapeConfig( # Qwen3-4B tp8_gate_up_proj
844 + block_k=64,
845 + m_buckets=(
846 + (1, _MatmulMConfig(16, 32, 4, 5)),
847 + (4, _MatmulMConfig(32, 128, 4, 3)),
848 + (8, _MatmulMConfig(16, 128, 8, 3)),
849 + (16, _MatmulMConfig(16, 128, 4, 4)),
850 + (32, _MatmulMConfig(16, 128, 8, 4)),
851 + (64, _MatmulMConfig(16, 128, 4, 3)),
852 + (256, _MatmulMConfig(64, 64, 4, 5)),
853 + (512, _MatmulMConfig(64, 128, 4, 4)),
854 + (1024, _MatmulMConfig(128, 128, 8, 3)),
855 + (2048, _MatmulMConfig(128, 128, 4, 3)),
856 + ),
857 + ),
858 + (2560, 1216): _MatmulShapeConfig( # Qwen3-4B tp8_down_proj
859 + block_k=64,
860 + m_buckets=(
861 + (1, _MatmulMConfig(16, 256, 8, 2)),
862 + (4, _MatmulMConfig(16, 128, 8, 4)),
863 + (8, _MatmulMConfig(16, 32, 4, 3)),
864 + (16, _MatmulMConfig(32, 32, 4, 3)),
865 + (32, _MatmulMConfig(16, 128, 8, 4)),
866 + (64, _MatmulMConfig(32, 128, 8, 4)),
867 + (256, _MatmulMConfig(32, 128, 4, 5)),
868 + (512, _MatmulMConfig(64, 128, 8, 4)),
869 + (1024, _MatmulMConfig(64, 256, 4, 3)),
870 + (2048, _MatmulMConfig(128, 128, 4, 3)),
871 + ),
872 + ),
873 + (768, 2560): _MatmulShapeConfig( # Qwen3-4B tp8_qkv_proj
874 + block_k=128,
875 + m_buckets=(
876 + (1, _MatmulMConfig(32, 32, 4, 2)),
877 + (4, _MatmulMConfig(16, 64, 8, 3)),
878 + (8, _MatmulMConfig(16, 64, 8, 4)),
879 + (16, _MatmulMConfig(16, 64, 4, 5)),
880 + (32, _MatmulMConfig(32, 64, 8, 4)),
881 + (64, _MatmulMConfig(16, 64, 8, 4)),
882 + (256, _MatmulMConfig(64, 32, 4, 5)),
883 + (512, _MatmulMConfig(64, 64, 8, 3)),
884 + (1024, _MatmulMConfig(64, 128, 4, 3)),
885 + (2048, _MatmulMConfig(64, 64, 4, 3)),
886 + ),
887 + ),
888 + (2560, 512): _MatmulShapeConfig( # Qwen3-4B tp8_o_proj
889 + block_k=64,
890 + m_buckets=(
891 + (1, _MatmulMConfig(16, 128, 8, 4)),
892 + (4, _MatmulMConfig(16, 32, 4, 2)),
893 + (8, _MatmulMConfig(16, 128, 4, 5)),
894 + (16, _MatmulMConfig(32, 32, 4, 5)),
895 + (32, _MatmulMConfig(32, 256, 4, 3)),
896 + (64, _MatmulMConfig(32, 128, 8, 5)),
897 + (256, _MatmulMConfig(16, 256, 4, 2)),
898 + (512, _MatmulMConfig(32, 128, 8, 4)),
899 + (1024, _MatmulMConfig(128, 128, 8, 3)),
900 + (2048, _MatmulMConfig(128, 128, 4, 3)),
901 + ),
902 + ),
903 + (18992, 2560): _MatmulShapeConfig( # Qwen3-4B tp8_lm_head
904 + block_k=64,
905 + m_buckets=(
906 + (1, _MatmulMConfig(16, 64, 8, 4)),
907 + (4, _MatmulMConfig(16, 64, 8, 5)),
908 + (8, _MatmulMConfig(16, 128, 8, 3)),
909 + (16, _MatmulMConfig(16, 128, 8, 5)),
910 + (32, _MatmulMConfig(32, 128, 8, 5)),
911 + (64, _MatmulMConfig(64, 128, 8, 4)),
912 + (256, _MatmulMConfig(128, 128, 4, 3)),
913 + (512, _MatmulMConfig(128, 128, 4, 3)),
914 + (1024, _MatmulMConfig(128, 128, 4, 3)),
915 + (2048, _MatmulMConfig(128, 128, 4, 3)),
916 + ),
917 + ),
918 + (12288, 4096): _MatmulShapeConfig( # Qwen3-8B tp2_gate_up_proj
919 + block_k=64,
920 + m_buckets=(
921 + (1, _MatmulMConfig(16, 128, 8, 3)),
922 + (4, _MatmulMConfig(16, 128, 8, 5)),
923 + (8, _MatmulMConfig(16, 128, 8, 4)),
924 + (16, _MatmulMConfig(16, 128, 4, 4)),
925 + (32, _MatmulMConfig(32, 128, 4, 4)),
926 + (64, _MatmulMConfig(64, 128, 4, 5)),
927 + (256, _MatmulMConfig(64, 64, 4, 5)),
928 + (512, _MatmulMConfig(128, 64, 4, 3)),
929 + (1024, _MatmulMConfig(128, 128, 4, 3)),
930 + (2048, _MatmulMConfig(128, 128, 4, 3)),
931 + ),
932 + ),
933 + (4096, 6144): _MatmulShapeConfig( # Qwen3-8B tp2_down_proj
934 + block_k=64,
935 + m_buckets=(
936 + (1, _MatmulMConfig(16, 32, 8, 5)),
937 + (4, _MatmulMConfig(16, 32, 8, 5)),
938 + (8, _MatmulMConfig(16, 32, 8, 3)),
939 + (16, _MatmulMConfig(16, 32, 4, 5)),
940 + (32, _MatmulMConfig(16, 64, 8, 4)),
941 + (64, _MatmulMConfig(32, 64, 4, 5)),
942 + (256, _MatmulMConfig(64, 128, 4, 5)),
943 + (512, _MatmulMConfig(128, 128, 4, 3)),
944 + (1024, _MatmulMConfig(128, 64, 8, 4)),
945 + (2048, _MatmulMConfig(128, 128, 4, 3)),
946 + ),
947 + ),
948 + (3072, 4096): _MatmulShapeConfig( # Qwen3-8B tp2_qkv_proj
949 + block_k=64,
950 + m_buckets=(
951 + (1, _MatmulMConfig(16, 32, 8, 4)),
952 + (4, _MatmulMConfig(16, 32, 8, 3)),
953 + (8, _MatmulMConfig(16, 128, 4, 3)),
954 + (16, _MatmulMConfig(16, 32, 4, 4)),
955 + (32, _MatmulMConfig(32, 32, 8, 5)),
956 + (64, _MatmulMConfig(16, 128, 8, 4)),
957 + (256, _MatmulMConfig(64, 128, 4, 5)),
958 + (512, _MatmulMConfig(128, 128, 8, 3)),
959 + (1024, _MatmulMConfig(64, 64, 4, 5)),
960 + (2048, _MatmulMConfig(128, 64, 4, 3)),
961 + ),
962 + ),
963 + (75968, 4096): _MatmulShapeConfig( # Qwen3-8B tp2_lm_head
964 + block_k=64,
965 + m_buckets=(
966 + (1, _MatmulMConfig(16, 256, 8, 3)),
967 + (4, _MatmulMConfig(16, 256, 8, 3)),
968 + (8, _MatmulMConfig(16, 256, 8, 3)),
969 + (16, _MatmulMConfig(16, 256, 8, 3)),
970 + (32, _MatmulMConfig(32, 256, 8, 3)),
971 + (64, _MatmulMConfig(64, 256, 8, 3)),
972 + (256, _MatmulMConfig(128, 128, 4, 3)),
973 + (512, _MatmulMConfig(128, 128, 8, 3)),
974 + (1024, _MatmulMConfig(128, 128, 4, 3)),
975 + (2048, _MatmulMConfig(128, 128, 4, 3)),
976 + ),
977 + ),
978 + (4096, 3072): _MatmulShapeConfig( # Qwen3-8B tp4_down_proj
979 + block_k=64,
980 + m_buckets=(
981 + (1, _MatmulMConfig(16, 256, 8, 3)),
982 + (4, _MatmulMConfig(16, 128, 4, 5)),
983 + (8, _MatmulMConfig(16, 128, 4, 4)),
984 + (16, _MatmulMConfig(16, 64, 4, 5)),
985 + (32, _MatmulMConfig(32, 128, 4, 3)),
986 + (64, _MatmulMConfig(16, 128, 4, 4)),
987 + (256, _MatmulMConfig(64, 128, 4, 5)),
988 + (512, _MatmulMConfig(128, 128, 4, 3)),
989 + (1024, _MatmulMConfig(128, 64, 8, 4)),
990 + (2048, _MatmulMConfig(128, 128, 4, 3)),
991 + ),
992 + ),
993 + (1536, 4096): _MatmulShapeConfig( # Qwen3-8B tp4_qkv_proj
994 + block_k=128,
995 + m_buckets=(
996 + (1, _MatmulMConfig(16, 64, 4, 5)),
997 + (4, _MatmulMConfig(16, 64, 8, 3)),
998 + (8, _MatmulMConfig(32, 64, 4, 5)),
999 + (16, _MatmulMConfig(16, 64, 8, 4)),
1000 + (32, _MatmulMConfig(16, 32, 8, 3)),
1001 + (64, _MatmulMConfig(16, 64, 8, 5)),
1002 + (256, _MatmulMConfig(64, 64, 4, 3)),
1003 + (512, _MatmulMConfig(64, 128, 4, 3)),
1004 + (1024, _MatmulMConfig(64, 64, 4, 3)),
1005 + (2048, _MatmulMConfig(64, 128, 4, 3)),
1006 + ),
1007 + ),
1008 + (4096, 1024): _MatmulShapeConfig( # Qwen3-8B tp4_o_proj
1009 + block_k=64,
1010 + m_buckets=(
1011 + (1, _MatmulMConfig(16, 128, 8, 4)),
1012 + (4, _MatmulMConfig(32, 128, 8, 4)),
1013 + (8, _MatmulMConfig(16, 256, 8, 3)),
1014 + (16, _MatmulMConfig(16, 128, 8, 3)),
1015 + (32, _MatmulMConfig(16, 128, 4, 4)),
1016 + (64, _MatmulMConfig(64, 64, 4, 3)),
1017 + (256, _MatmulMConfig(128, 64, 4, 3)),
1018 + (512, _MatmulMConfig(64, 64, 4, 5)),
1019 + (1024, _MatmulMConfig(128, 64, 4, 4)),
1020 + (2048, _MatmulMConfig(128, 128, 4, 3)),
1021 + ),
1022 + ),
1023 + (37984, 4096): _MatmulShapeConfig( # Qwen3-8B tp4_lm_head
1024 + block_k=64,
1025 + m_buckets=(
1026 + (1, _MatmulMConfig(16, 128, 8, 3)),
1027 + (4, _MatmulMConfig(16, 128, 4, 5)),
1028 + (8, _MatmulMConfig(16, 128, 4, 5)),
1029 + (16, _MatmulMConfig(16, 128, 8, 3)),
1030 + (32, _MatmulMConfig(32, 128, 8, 3)),
1031 + (64, _MatmulMConfig(64, 128, 8, 3)),
1032 + (256, _MatmulMConfig(64, 128, 4, 5)),
1033 + (512, _MatmulMConfig(128, 128, 4, 3)),
1034 + (1024, _MatmulMConfig(128, 128, 4, 3)),
1035 + (2048, _MatmulMConfig(128, 128, 4, 3)),
1036 + ),
1037 + ),
1038 + (4096, 1536): _MatmulShapeConfig( # Qwen3-8B tp8_down_proj
1039 + block_k=32,
1040 + m_buckets=(
1041 + (1, _MatmulMConfig(32, 128, 8, 5)),
1042 + (4, _MatmulMConfig(16, 256, 4, 5)),
1043 + (8, _MatmulMConfig(16, 64, 4, 4)),
1044 + (16, _MatmulMConfig(16, 128, 8, 3)),
1045 + (32, _MatmulMConfig(16, 64, 4, 3)),
1046 + (64, _MatmulMConfig(16, 128, 4, 4)),
1047 + (256, _MatmulMConfig(64, 128, 4, 4)),
1048 + (512, _MatmulMConfig(128, 128, 8, 5)),
1049 + (1024, _MatmulMConfig(128, 256, 8, 4)),
1050 + (2048, _MatmulMConfig(128, 128, 8, 4)),
1051 + ),
1052 + ),
1053 + (768, 4096): _MatmulShapeConfig( # Qwen3-8B tp8_qkv_proj
1054 + block_k=128,
1055 + m_buckets=(
1056 + (1, _MatmulMConfig(16, 32, 8, 5)),
1057 + (4, _MatmulMConfig(16, 32, 8, 3)),
1058 + (8, _MatmulMConfig(32, 64, 4, 5)),
1059 + (16, _MatmulMConfig(32, 64, 8, 3)),
1060 + (32, _MatmulMConfig(32, 32, 8, 5)),
1061 + (64, _MatmulMConfig(64, 32, 8, 3)),
1062 + (256, _MatmulMConfig(32, 64, 4, 4)),
1063 + (512, _MatmulMConfig(64, 64, 8, 3)),
1064 + (1024, _MatmulMConfig(64, 128, 4, 3)),
1065 + (2048, _MatmulMConfig(64, 64, 4, 3)),
1066 + ),
1067 + ),
1068 + (4096, 512): _MatmulShapeConfig( # Qwen3-8B tp8_o_proj
1069 + block_k=64,
1070 + m_buckets=(
1071 + (1, _MatmulMConfig(16, 128, 8, 2)),
1072 + (4, _MatmulMConfig(32, 128, 4, 2)),
1073 + (8, _MatmulMConfig(16, 256, 4, 2)),
1074 + (16, _MatmulMConfig(32, 256, 4, 3)),
1075 + (32, _MatmulMConfig(32, 256, 8, 3)),
1076 + (64, _MatmulMConfig(32, 128, 4, 5)),
1077 + (256, _MatmulMConfig(64, 32, 8, 4)),
1078 + (512, _MatmulMConfig(128, 128, 8, 3)),
1079 + (1024, _MatmulMConfig(128, 64, 4, 4)),
1080 + (2048, _MatmulMConfig(128, 128, 8, 3)),
1081 + ),
1082 + ),
1083 + (18992, 4096): _MatmulShapeConfig( # Qwen3-8B tp8_lm_head
1084 + block_k=64,
1085 + m_buckets=(
1086 + (1, _MatmulMConfig(16, 64, 8, 5)),
1087 + (4, _MatmulMConfig(16, 64, 8, 4)),
1088 + (8, _MatmulMConfig(16, 64, 8, 5)),
1089 + (16, _MatmulMConfig(16, 128, 4, 3)),
1090 + (32, _MatmulMConfig(32, 64, 8, 5)),
1091 + (64, _MatmulMConfig(64, 128, 8, 3)),
1092 + (256, _MatmulMConfig(128, 128, 4, 3)),
1093 + (512, _MatmulMConfig(128, 128, 4, 3)),
1094 + (1024, _MatmulMConfig(128, 128, 4, 3)),
1095 + (2048, _MatmulMConfig(128, 128, 4, 3)),
1096 + ),
1097 + ),
1098 },
1099 }
1100