import time
import torch
import triton
import triton.language as tl
from triton import Config

from task import input_t, output_t

# best version 3

# fp8_gemm_configs = [
#     Config({'BLOCK_SIZE_M': block_m, 'BLOCK_SIZE_N': block_n, 'BLOCK_SIZE_K': 128, "GROUP_SIZE_M": 4}, num_stages=num_stages, num_warps=8)
#     for block_m in [16, 32, 64] for block_n in [32, 64, 128] for num_stages in [2, 3, 4, 5]
# ]

# @triton.autotune(configs=fp8_gemm_configs, key=['N', 'K'])
@triton.jit
def fp8_gemm_kernel(a_ptr, b_ptr, c_ptr,
                    a_s_ptr, b_s_ptr,
                    M, N: tl.constexpr, K: tl.constexpr,
                    BLOCK_SIZE_M: tl.constexpr,
                    BLOCK_SIZE_N: tl.constexpr,
                    BLOCK_SIZE_K: tl.constexpr,
                    GROUP_SIZE_M: tl.constexpr):
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

    # pid_m = pid // num_pid_n
    # pid_n = pid % num_pid_n

    NUM_XCDS: tl.constexpr = 8
    GRID_MN = tl.cdiv(M, BLOCK_SIZE_M) * tl.cdiv(N, BLOCK_SIZE_N)
    ## pid remapping on xcds
    # Number of pids per XCD in the new arrangement
    pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
    # When GRID_MN cannot divide NUM_XCDS, some xcds will have
    # pids_per_xcd pids, the other will have pids_per_xcd - 1 pids.
    # We calculate the number of xcds that have pids_per_xcd pids as
    # tall_xcds
    tall_xcds = GRID_MN % NUM_XCDS
    tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
    # Compute current XCD and local pid within the XCD
    xcd = pid % NUM_XCDS
    local_pid = pid // NUM_XCDS
    # Calculate new pid based on the new grouping
    # Note that we need to consider the following two cases:
    # 1. the current pid is on a tall xcd
    # 2. the current pid is on a short xcd
    if xcd < tall_xcds:
        pid = xcd * pids_per_xcd + local_pid
    else:
        pid = (
            tall_xcds * pids_per_xcd
            + (xcd - tall_xcds) * (pids_per_xcd - 1)
            + local_pid
        )

    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    k = tl.cdiv(K, BLOCK_SIZE_K)
    offs_m = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
    b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
    a_s_ptrs = a_s_ptr + offs_m
    b_s_ptrs = b_s_ptr + (offs_n // BLOCK_SIZE_K) * k

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for i in range(k):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - i * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - i * BLOCK_SIZE_K, other=0.0)
        a_s = tl.load(a_s_ptrs)
        b_s = tl.load(b_s_ptrs)
        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K
        a_s_ptrs += M
        b_s_ptrs += 1
    c = accumulator.to(c_ptr.dtype.element_ty)
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=mask)


def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    a = a.contiguous()
    b = b.contiguous()
    a_scale = a_scale.t().contiguous()
    b_scale = b_scale.contiguous()
    K = a.size(-1)
    M = a.numel() // K
    N = b.size(0)

    if M < 2048:
        BLOCK_SIZE_M = 32
    else:
        BLOCK_SIZE_M = 64
    BLOCK_SIZE_N = 64
    BLOCK_SIZE_K = 128
    GROUP_SIZE_M = 16
    num_stages = 2
    num_warps = 8
    waves_per_eu = 2

    grid = (triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N),)
    fp8_gemm_kernel_begin_time = time.perf_counter_ns()
    torch.cuda.synchronize()
    fp8_gemm_kernel[grid](a, b, c, a_scale, b_scale, M, N, K,
                        BLOCK_SIZE_M=BLOCK_SIZE_M,
                        BLOCK_SIZE_N=BLOCK_SIZE_N,
                        BLOCK_SIZE_K=BLOCK_SIZE_K,
                        GROUP_SIZE_M=GROUP_SIZE_M,
                        num_stages=num_stages,
                        num_warps=num_warps,
                        waves_per_eu=waves_per_eu,
                        kpack=2,
                        matrix_instr_nonkdim=16,)
    torch.cuda.synchronize()
    fp8_gemm_kernel_end_time = time.perf_counter_ns()
    print(fp8_gemm_kernel_end_time - fp8_gemm_kernel_begin_time)
    # bm = fp8_gemm_kernel.best_config.kwargs["BLOCK_SIZE_M"]
    # bn = fp8_gemm_kernel.best_config.kwargs["BLOCK_SIZE_N"]
    # ns = fp8_gemm_kernel.best_config.num_stages
    # m = M / 64
    # n = N / 64
    # k = K / 64
    # if (m != 16):
    #     print(f"{N},{K},{bm},{bn},{ns}")
    return c

python eval.py benchmark /tmp/tmp_phbnh9i
755147183
325911
265167
262257
265029
261060
258115
260138
260271
260189
256665
260158
263316
260673
260513
257181
257538
257748
258373
256568
257157
257858
261356
260278
260841
258436
256105
263061
259603
258957
260813
262488
260927
258169
258557
257557
260237
257290
263337
259742
257842
257532
255874
256633
262430
260028
258646
258607
258838
256769
258455
257772
254954
259045
260662
255181
258674
259505
256421
257432
259921
255701
258419
256893
259788
259770
264277
266001
258060
257173
256972
258036
259703
264358
259886
256610
258886
256369
272427
261171
258987
255731
256830
257853
257288
262758
257000
256799
256379
260670
262078
264497
257825
255245
256033
257435
265121
261416
256755
259507
260405
294862
282381
265318
264915
264563
265169
266644
267163
266443
265731
262962
260477
261429
265778
260769
262397
261955
260548
265444
263155
263796
264556
260492
265971
261880
265181
269179
266911
268172
270180
260673
262103
263321
262860
263354
263763
265294
262214
268278
266079
260134
263493
259580
262807
262288
263197
262336
265295
261209
261573
264952
261362
260585
259393
261762
262356
261894
265694
267052
263303
261829
264157
264087
265216
264372
264164
263359
262631
260577
264539
260116
262097
267942
261637
264161
262093
258018
261000
262422
258228
257329
260628
260950
263095
262213
260471
259349
260891
262782
262162
262138
260117
262325
260217
261228
257027
263678
259120
260972
266124
262753
254746945
145171
125882
123743
126022
129332
124125
122997
119761
122097
120467
121826
124210
125074
120225
121744
113591
123586
119417
152685
117599
118341
116598
117913
116281
148406
115412
115989
115757
117514
116524
115484
113401
114414
114258
114900
115610
115149
113617
116312
114748
113966
115565
113647
115928
114132
116075
114685
115157
118374
115907
115039
114273
113284
114987
117808
115117
115636
116845
115078
115050
114587
114372
118386
116756
116258
113982
115553
113323
114549
113933
115651
114120
114414
113348
113202
113982
118939
116734
114799
115463
114615
118640
114392
115865
113740
116210
125756
114970
115900
113353
115829
113842
112365
115132
115139
111740
114128
114253
115199
115186
265416732
203446
182174
176620
178439
172554
173502
174618
174462
174510
179081
172262
174090
160921
159565
167214
161738
176232
151029
153565
151369
154012
151703
151434
153421
156082
156348
154474
156285
156212
155504
151358
158415
156521
156356
153316
154454
152835
155663
152068
154946
152773
153617
153219
153148
154797
155233
150979
154665
153836
156142
152250
153315
153326
159936
154715
155401
154271
155235
158654
154872
154127
153305
151869
155722
151143
153159
155228
157842
153594
154442
150511
153918
153598
154170
155471
153355
151485
158385
153397
152818
154559
152315
152233
156037
156460
153859
154461
157823
156393
151454
153442
156168
152782
156498
151858
153271
154544
151288
153895
155069
225116567
114021
82624
112691
111728
111706
79420
76552
76796
108601
76348
76016
108003
75539
77488
110360
112072
90067
110430
74952
73987
74191
73295
106539
74857
71703
105735
72943
71766
72720
75021
75253
77102
109332
73194
73212
73643
72454
74822
74938
72085
72182
73493
72136
72224
73707
73670
72732
73599
72552
72369
72344
73525
71944
106567
74213
73075
72475
72072
71309
73121
72331
73559
73615
72341
74117
74886
75227
77598
72371
72955
72786
71626
73527
72451
71663
74658
106231
72550
72904
72575
73048
71699
72422
80365
72416
73686
73467
72873
72035
72344
105780
74741
71342
71577
72829
72142
71934
73162
75350
73187
257202539
255896
236512
235857
239648
238416
230261
231219
236597
244531
231645
248687
231733
246592
231066
234557
232943
246404
243943
232337
243111
228893
229116
227413
226376
238800
238078
227842
227712
227209
226266
228851
232612
228630
229645
230310
230328
242658
230197
234593
249787
233266
230669
228736
229248
242308
225838
227078
226540
226969
228218
235179
225750
229366
235383
229485
241599
227737
243729
228844
226202
224599
228023
228337
226401
226618
225245
227791
240411
227510
241114
226933
228475
226717
225946
242553
228872
225161
227348
239253
227701
226355
226055
225920
227387
240678
226275
224733
227462
243629
226455
229578
228105
228435
229281
227359
241302
227147
227242
237912
227688
225565709
575700
557138
526331
540371
529836
535455
542323
531159
533544
527141
521540
528502
534503
535648
525719
529122
520139
522917
523924
525563
521565
533304
518486
529599
521611
521549
526228
524516
517669
523553
522044
519667
522321
523723
526083
518487
523434
522458
522027
522707
519578
526632
523754
519926
523037
529566
528429
516275
533982
520092
523921
523607
520675
531231
521003
531191
527662
521026
523187
522971
526140
518388
525835
522804
528409
523325
523378
518521
524945
522090
519234
525626
529091
522217
523025
531562
531171
523258
529721
523923
523903
519571
520395
522708
519572
521010
525608
519705
527611
519336
519699
518858
532331
522726
528508
528790
523728
525326
523133
520254
231194426
293290
252829
260005
266076
257994
270460
276610
269748
283567
245126
261671
246830
297525
241827
260062
277473
242356
264552
277775
256855
258696
263367
260082
263660
257002
240115
260738
245863
241887
248032
266090
277353
261630
249413
280515
268174
278739
267584
261887
268062
265171
266811
262334
279214
259761
268223
265320
242277
260914
265491
255347
261758
264026
275653
276655
259239
277696
285671
270162
263434
272212
270747
283433
268985
247235
272107
264029
245731
280031
261031
284076
268933
264891
266449
246793
279624
270153
272423
242679
266514
268894
267870
244673
264919
243515
261976
264067
277529
278206
262485
266860
246293
282087
243912
262649
253206
244682
258407
275835
260765
226155080
206607
185058
177612
180605
172927
176211
180967
179985
177129
174284
176254
176340
179016
198384
164955
164148
162477
181466
152702
154163
153359
154872
153002
154861
157151
156811
158296
156024
157251
161108
159818
155274
154715
157446
152770
155157
156994
160846
156972
159767
162467
153657
157347
153765
159265
153535
154150
155642
156214
155226
156226
154921
157110
152938
157030
152950
154069
157438
153871
153966
154675
152895
155534
158294
154140
153639
157647
153382
154928
155785
158887
155088
152308
158752
153309
155572
154230
152831
155252
155731
155789
154326
154265
154558
155660
156275
154526
152596
156218
154877
152979
154571
153632
159898
156991
152658
152439
157270
153410
158113
221274810
110708
80982
76898
79214
76427
77185
77070
74330
73405
76681
74542
75207
73513
74254
74679
75241
74136
76377
74910
76434
74764
74937
71157
75051
78379
77706
74696
74531
74769
76013
76506
75493
75157
76767
76829
73742
75827
73852
75632
72697
76371
76875
74748
74197
72721
75174
74057
73593
75674
74944
73609
73936
75112
73874
73151
76553
74631
74504
75293
75599
74295
75325
75496
74965
75552
77696
74807
73173
80890
73590
73393
74488
74271
74379
73273
75240
75039
74048
74129
74828
73804
74988
74717
73156
75550
73211
73222
74205
74029
73208
74656
73820
71966
74536
75650
74007
79039
73891
74714
72452
245918174
1334382
1302466
1168027
1053997
1049264
1037455
1049083
1043909
1039343
1042611
1047191
1039992
1070811
1044006
1042998
1049311
1049633
1045831
1040296
1044371
1056115
1063233
1050977
1044074
1061692
1049220
1044799
1045029
1044524
1052985
1038520
1056125
1057250
1040895
1048123
1038211
1043421
1051488
1042318
1041911
1045626
1062129
1065847
1049046
1054906
1038009
1046364
1040944
1040780
1046103
1049092
1059460
1064505
1065177
1057030
1060499
1044027
1047485
1037919
1043801
1044999
1081898
1067137
1056034
1050146
1046257
1043383
1044118
1044043
1041185
1044391
1054120
1061926
1056106
1046336
1036697
1066336
1057973
1048622
1049227
1057512
1070938
1068713
1041813
1052539
1045921
1052005
1045510
1037702
1040946
1051232
1052731
1045185
1047062
1043606
1057755
1045216
1041629
1035268
1044973
248627168
330129
303940
297636
299218
290500
330999
287079
292867
327485
289943
280892
280113
283813
281150
281276
286148
279438
280121
279586
279408
281634
280606
281809
280938
280297
280406
280180
281627
277578
277357
279494
281406
278706
280693
277829
282035
280548
284387
278261
282272
286260
288810
289866
288646
284175
287492
287839
286210
289794
290112
284462
289722
289491
286118
286088
283284
285271
290452
276649
278954
277022
279568
281643
281321
280855
278485
284980
281153
279890
281286
278984
280814
281904
283106
281570
279687
281243
281751
280272
281515
281166
280139
278204
281920
280451
280349
281320
320724
284923
279380
283384
279272
278684
278909
285191
278662
280211
278474
278754
279836
247533324
1057443
928577
792839
794456
792351
789432
789280
793352
785970
783139
786041
783026
786516
794059
785022
784789
796307
789429
792614
794203
792938
793383
794590
789604
795215
797484
797469
796835
795682
789355
789845
796490
786377
793304
788777
789442
787546
806459
794742
795791
791829
789775
792560
788340
789911
792351
789014
789173
793754
788045
792894
787055
797665
792416
791350
794463
793882
789698
801951
792984
793815
786241
789598
796993
787943
794552
792154
795876
791680
795040
794991
790635
790108
790750
785789
800310
789790
798154
797764
790923
784040
796065
798974
788996
792602
788927
799649
794221
797036
799015
793680
787982
784460
791256
803449
799993
795409
792056
787583
787513
227996232
173506
152149
147524
145893
146944
144698
147614
147545
145977
146279
146383
173830
144534
145742
147909
141591
144528
143913
144057
143074
144090
145090
144529
147505
141824
144078
143677
143705
143605
152121
147104
144854
145069
143604
143805
145782
144427
143382
142963
145450
146298
144181
142379
143135
144021
141735
146692
143933
143611
141644
143275
145118
143195
145136
142691
141047
143709
142630
144441
144419
147755
145374
141121
144472
143081
142761
146625
143974
144003
142322
143628
144186
142890
142829
141456
141136
143697
142210
144434
141683
142635
178561
145047
144510
143113
143138
142604
143696
143203
142688
141801
149726
141995
143463
141943
178898
141975
142994
141578
141980
248921189
686030
595737
660697
579743
587642
584680
577154
580948
574605
590776
569912
600503
581613
566743
582856
577306
571725
605525
560540
591727
586660
587739
615067
601542
591385
590179
588281
587076
593603
584519
584469
589853
590717
604502
588479
575255
601671
609608
590832
588756
591510
593374
592243
594017
592072
588832
593742
582965
603518
591378
605584
582511
575493
573912
572233
569557
566915
572352
565050
581082
567795
570244
573252
571375
598047
579572
580805
580810
580855
600012
587767
601730
582948
600554
582203
600595
579458
585954
587551
580827
576861
576116
577311
574298
577646
572164
578748
573190
574028
587964
574060
576433
590695
573256
579490
597562
579184
587595
578302
583440
250573730
1836514
1738368
1701197
1717830
1676941
1708547
1761262
1737525
1700729
1694446
1714191
1739775
1755558
1747211
1749979
1768595
1726780
1716075
1739260
1743290
1759421
1741087
1732019
1752648
1740245
1755166
1726604
1748494
1760688
1767884
1736269
1719752
1709736
1754270
1767335
1771093
1746883
1739880
1754001
1772562
1761007
1732080
1730597
1757483
1765903
1753809
1749017
1734286
1744563
1771741
1741970
1711818
1723761
1763450
1765533
1721495
1688497
1731234
1746153
1736865
1718337
1709272
1718872
1751465
1742524
1740142
1734826
1767119
1751040
1716839
1714375
1719161
1744741
1728844
1761478
1742955
1733101
1763378
1752768
1743593
1717483
1740933
1776912
1725136
1736022
1734744
1756746
1779628
1753606
1715034
1742935
1759748
1763057
1735259
1766416
1746537
1759693
1771073
1756975
1751698
250673744
763397
758961
715086
700425
719617
700030
702028
704294
702717
696938
712803
703082
710431
708874
747106
716952
712835
715568
724306
724108
721081
722037
716479
716541
715561
719409
725283
720032
715806
719237
720101
717937
714528
707651
713393
719381
720688
728679
722415
720091
719176
717336
718386
714739
782868
703686
704227
718283
698921
697497
696220
694696
693090
695737
690804
691673
709122
706311
693502
698795
711518
698290
708552
704613
720508
706991
708079
723412
704151
692298
699163
696804
697217
698899
695123
710794
693665
693264
693141
692167
693957
694367
691785
709273
693658
708999
699501
693579
694401
708363
690796
709921
691911
696807
699941
698610
696727
695347
695300
694203
245479119
1032730
932555
825612
793417
781795
791082
786412
787506
782455
780743
784942
792468
784688
789644
787312
783146
782968
775392
783289
796933
789428
785105
782873
789189
786830
784947
789114
780865
783073
785138
786170
781956
785304
786782
786890
788135
783230
786970
784052
785992
790673
786652
782216
791478
787201
790107
783182
782505
791525
783117
789176
793204
786950
788093
785585
784380
789510
785604
792960
786321
794421
786197
786040
788863
778396
790678
788986
794804
786634
785333
791490
782375
780886
782043
783258
791281
779228
790690
785033
788552
774186
791357
786970
792143
790418
789416
787329
782509
791278
796842
795111
787841
784510
795981
790495
789208
785341
784407
782987
790029
248613041
179733
155804
154415
152363
152192
150970
151145
150321
151181
149366
149779
169702
149435
153707
149743
148730
180740
148119
151061
150382
147089
146346
148080
148492
146819
149134
148899
154595
147028
155695
149017
149525
145601
146593
150849
151284
149833
150553
149951
148998
150207
147186
148759
147994
146775
147348
148098
147531
148211
148729
145732
146464
150134
145589
148352
145875
148790
150807
152971
148311
149132
146772
147264
146884
148812
148400
149157
148440
147784
149261
146183
149998
150769
149305
148758
150104
147988
147993
146176
148040
146798
148688
149034
149908
150247
149514
149890
147423
148339
148516
148401
147455
146908
148584
148011
146647
148316
148059
147937
146325
