Repository navigation
Commit e2d019f
authored
[Triton/Gluon] [JIT] [Feature] Add tuned native group32 A8W8 Triton GEMM for gfx950 (#5750)
* feat(gemm): add native group32 A8W8 Triton backend
Integrate unshuffled E4M3/E8M0 inputs into gemm_a8w8_blockscale. Keep
microscaling and packed-K kernels in Triton, with gfx950 tuning tables
loaded by get_gemm_config. Extend shared loading with JSON M bounds and
reduce config-copy overhead while retaining defensive nested copies.
Reuse the shared non-atomic split-K reduction and preserve legacy routes.
Validated with 98 AITER tests, 73,728 matching config selections, and 285
bitwise-equal BF16 migration cases; repeated 40 GPU cases six times.
Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com>
* Tune native group32 GEMM configurations for V4.1 projections
Sweep 18 gfx950 geometries across TP4, TP2, and attention-DP local
projections. Store tile, split-K, launch, and M-bound selections in the
existing JSON family; retain baseline intervals without stable gains.
Keep the existing large-M fallbacks and merge equal adjacent configs.
Exercise N-first masking independently of the selected traversal.
Validation: 49,304 candidates screened, 603 regression cases, 98 AITER
tests and 85 ATOM tests passed. Paired regression geometric-mean speedup
is 1.179x; changed configurations have no sampled slowdown above 1%.
Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com>
* Respect tuned backend selection before group32 GEMM fallback
Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com>
* Configure V4.1 group32 GEMMs to use the Triton backend
Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com>
* Preserve native split-K controls through the public A8W8 GEMM interface
Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com>
* fix(gemm): complete group32 dispatch and validation coverage
* perf(gemm): fuse group32 split-K reduction within one XCD
Small-M group32 GEMMs on narrow or deep projections are starved for CTAs:
N=512 gives the packed kernel 32 CTAs on 256 CUs, and splitting K costs a
second reduce launch. Both kernels now take FUSED_SPLITS, which sums the
split partials in the same launch.
A 1D grid places every split of a tile on one XCD (CTAs go to XCDs by
pid % 8, a multiple of every gfx950 XCD count). That XCD's L2 is then the
splits' coherence point, so the last CTA to arrive can read the other
partials after a workgroup-scope acq_rel counter add: no agent-scope
release, and no buffer_wbl2 of the whole L2. With agent scope the same
fused kernels are 2.3x slower, which is why a plain fused split-K lost to
the two-launch path. Partials are summed in split order whatever the
arrival order, so results are deterministic, and the last CTA re-zeroes
its counter, so graph replay is safe.
Counters follow the a16w16 split-K semaphores: one zeroed buffer per
(device, stream) from persistent_alloc, so concurrent streams never share
one. Where a stream's first use falls inside a capture that cannot
allocate outside the graph pool (torch < 2.10), the launch takes the
unfused path instead of failing the capture.
The fused path is selected only by the tuning table (FUSED_SPLITK on a
tile config, NUM_KSPLIT on a packed config). An explicit split_k keeps
the existing reduce-kernel path. With the original configs, all 567
sampled outputs across the 18 V4.1 geometries, tails, FP32/BF16 and
explicit split_k are bitwise identical to before.
Tests cover the reference match for both variants and output dtypes,
1x32 weight scales, 50-call determinism with counters re-zeroed, graph
replay on a warmed stream, first use inside a capture, and interleaved
launches on two streams.
* perf(gemm): tune fused split-K intervals and time every V4.1 group32 row
57 M intervals of 11 V4.1 geometries now use the fused split-K path, 42 at
M <= 64 and 15 above. Candidates were the best fused tile and packed
configs from full sweeps; an interval changed only if a candidate was at
least 3% faster at both of its bounds. Every changed interval was then
re-measured against its previous config in the same process (run_perftest,
rotated inputs, median of 5): at each CSV M inside it, or at its lower
bound, midpoint and upper bound where no CSV M falls inside. Three
intervals that did not hold up were left unchanged. No measured point is
slower than the previous config.
Speedup over the previous config, geometric mean of the rows re-measured:
o_b.dp M <= 64 1.26x draft_kvf M 129..320 1.22-1.26x
draft_kv M <= 64 1.22x sh_up.tp4 M 321..448 1.26x
sh_up.tp4 M <= 64 1.19x sh_up.tp2 M 129..192 1.20x
draft_kvf M <= 64 1.13x o_b.dp M 129..192 1.20x
o_b.tp2 M <= 64 1.12x draft_kv M 513..576 1.14x
The V4.1 model CSV previously left 188 us fields and every tflops, bw and
errRatio field empty. All 450 rows are now timed in one run with the final
configs, through gemm_a8w8_blockscale, with run_perftest and rotated
inputs, so the us column comes from a single method. tflops and bw use the
a8w8 blockscale tuner formula (bw excludes scales). errRatio is
checkAllclose at its default tolerance against an FP64 reference, as in
the tuner: at most 8.7e-4, and 0 rows above 0.05.
* perf(gemm): constexpr K loop, weight cache modifier and 68-bucket retune for group32 GEMM
Retune all 18 V4.1 geometries on the repo's 68 DSv4 M buckets (fused split-K, packed K_PACK=1, .cg weight loads); drop unused GROUP_SIZE_M; re-time the 1224-row CSV.
---------
Signed-off-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com>1 parent 37b1425 commit e2d019f
31 files changed
Lines changed: 23468 additions & 54 deletions
File tree
- aiter
- configs
- model_configs
- jit
- ops
- triton
- _triton_kernels/gemm/basic
- configs
- gfx950/triton/gemm/gemm_a8w8_blockscale_group32
- gemm/basic
- utils
- op_tests
- op_benchmarks/triton
- triton_tests/gemm/basic
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
| 1 | + | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
| 1 | + | |
Large diffs are not rendered by default.
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
135 | 135 | | |
136 | 136 | | |
137 | 137 | | |
| 138 | + | |
| 139 | + | |
| 140 | + | |
| 141 | + | |
| 142 | + | |
| 143 | + | |
| 144 | + | |
138 | 145 | | |
139 | 146 | | |
140 | 147 | | |
| |||
286 | 293 | | |
287 | 294 | | |
288 | 295 | | |
| 296 | + | |
| 297 | + | |
| 298 | + | |
| 299 | + | |
| 300 | + | |
| 301 | + | |
| 302 | + | |
| 303 | + | |
289 | 304 | | |
290 | 305 | | |
291 | 306 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
812 | 812 | | |
813 | 813 | | |
814 | 814 | | |
| 815 | + | |
| 816 | + | |
| 817 | + | |
815 | 818 | | |
816 | | - | |
| 819 | + | |
| 820 | + | |
| 821 | + | |
| 822 | + | |
| 823 | + | |
| 824 | + | |
| 825 | + | |
| 826 | + | |
| 827 | + | |
| 828 | + | |
| 829 | + | |
| 830 | + | |
| 831 | + | |
| 832 | + | |
| 833 | + | |
| 834 | + | |
| 835 | + | |
| 836 | + | |
817 | 837 | | |
818 | 838 | | |
819 | 839 | | |
| |||
830 | 850 | | |
831 | 851 | | |
832 | 852 | | |
| 853 | + | |
833 | 854 | | |
834 | 855 | | |
835 | 856 | | |
836 | 857 | | |
837 | 858 | | |
838 | 859 | | |
839 | 860 | | |
840 | | - | |
| 861 | + | |
841 | 862 | | |
842 | 863 | | |
843 | 864 | | |
844 | 865 | | |
845 | 866 | | |
846 | 867 | | |
847 | 868 | | |
| 869 | + | |
848 | 870 | | |
849 | | - | |
850 | | - | |
851 | | - | |
852 | | - | |
| 871 | + | |
| 872 | + | |
| 873 | + | |
| 874 | + | |
| 875 | + | |
| 876 | + | |
| 877 | + | |
| 878 | + | |
| 879 | + | |
| 880 | + | |
| 881 | + | |
| 882 | + | |
| 883 | + | |
| 884 | + | |
| 885 | + | |
| 886 | + | |
| 887 | + | |
| 888 | + | |
| 889 | + | |
| 890 | + | |
| 891 | + | |
| 892 | + | |
| 893 | + | |
| 894 | + | |
| 895 | + | |
| 896 | + | |
| 897 | + | |
| 898 | + | |
| 899 | + | |
| 900 | + | |
| 901 | + | |
| 902 | + | |
853 | 903 | | |
854 | 904 | | |
855 | 905 | | |
856 | | - | |
857 | 906 | | |
| 907 | + | |
858 | 908 | | |
| 909 | + | |
859 | 910 | | |
860 | 911 | | |
861 | 912 | | |
862 | | - | |
863 | | - | |
864 | | - | |
865 | | - | |
866 | | - | |
| 913 | + | |
| 914 | + | |
| 915 | + | |
| 916 | + | |
| 917 | + | |
| 918 | + | |
| 919 | + | |
| 920 | + | |
| 921 | + | |
| 922 | + | |
| 923 | + | |
| 924 | + | |
| 925 | + | |
| 926 | + | |
| 927 | + | |
| 928 | + | |
| 929 | + | |
| 930 | + | |
| 931 | + | |
| 932 | + | |
| 933 | + | |
| 934 | + | |
| 935 | + | |
| 936 | + | |
| 937 | + | |
| 938 | + | |
| 939 | + | |
| 940 | + | |
| 941 | + | |
| 942 | + | |
| 943 | + | |
| 944 | + | |
| 945 | + | |
| 946 | + | |
| 947 | + | |
| 948 | + | |
| 949 | + | |
| 950 | + | |
| 951 | + | |
| 952 | + | |
| 953 | + | |
| 954 | + | |
| 955 | + | |
| 956 | + | |
| 957 | + | |
867 | 958 | | |
868 | | - | |
869 | | - | |
870 | | - | |
871 | | - | |
872 | | - | |
873 | | - | |
874 | | - | |
875 | | - | |
876 | | - | |
877 | | - | |
878 | | - | |
879 | | - | |
880 | | - | |
881 | | - | |
882 | | - | |
883 | | - | |
884 | | - | |
885 | | - | |
886 | | - | |
887 | | - | |
888 | | - | |
889 | | - | |
890 | | - | |
891 | | - | |
892 | | - | |
893 | | - | |
894 | | - | |
895 | | - | |
896 | | - | |
897 | | - | |
898 | | - | |
899 | | - | |
900 | | - | |
901 | | - | |
902 | | - | |
903 | | - | |
| 959 | + | |
| 960 | + | |
| 961 | + | |
| 962 | + | |
| 963 | + | |
| 964 | + | |
| 965 | + | |
| 966 | + | |
| 967 | + | |
| 968 | + | |
| 969 | + | |
904 | 970 | | |
905 | 971 | | |
906 | 972 | | |
| |||
0 commit comments