Wokaround for SWDEV-470361

[gshtras]

Calling the version of setProblem that does not cause integer overflow on large gemm shapes

[gshtras]

Tested on

import torch
from vllm._C import ops as vllm_ops

m_vals = [16, 128, 8192, 10240, 57344]
n_vals = [1,2,4,6,8,16,24,32,40,48,56,64,72,80,88,96,104,112,120,128,136,144,152,160,168,176,184,192,200,208,216,224,232,240,248,256,257,259,262,263,266,268,270,273,275,277,282,283,286,289,290,292,293,295,298,304,310,312,315,318,321,324,328,329,333,340,341,343,347,351,354,356,362,365,367,369,371,375,378,379,380,382,385,387,390,395,396,398,402,404,407,411,416,417,424,425,430,432,439,442,449,452,454,456,459,462,466,471,473,475,481,486,488,491,499,502,506,513,517,518,523,525,530,531,534,540,548,551,557,562,568,575,576,580,584,586,590,595,598,602,605,606,612,615,618,624,628,632,635,639,642,644,649,653,659,661,666,669,671,675,679,684,687,692,698,701,705,712,720,727,731,733,736,740,743,744,747,750,756,758,761,764,767,768,774,775,781,783,789,790,796,802,808,812,814,819,821,828,830,837,841,846,847,850,859,861,866,869,878,884,889,892,897,900,904,905,908,912,917,919,923,929,936,937,938,940,941,942,944,945,946,947,949,950,951,952,954,955,956,957,958,959,960,961,962,963,964,965,966,967,969,971,972,974,975,977,978,980,981,982,983,984,985,986,987,988,989,990,991,992,993,994,995,996,997,998,999,1000,1001,1002,1003,1004,1005,1006,1007,1008,1009,1010,1011,1012,1013,1015,1016,1017,1019,1020,1021,1022,1023,1024,1049,1055,1064,1071,1072,1073,1075,1085,1087,1089,1094,1095,1100,1102,1104,1105,1107,1108,1110,1112,1114,1115,1116,1117,1118,1119,1120,1123,1124,1125,1126,1127,1128,1129,1133,1134,1135,1136,1138,1139,1140,1141,1142,1150,1152,1154,1156,1159,1160,1161,1162,1166,1167,1169,1170,1171,1175,1180,1184,1186,1189,1190,1197,1198,1199,1200,1202,1203,1208,1209,1210,1211,1220,1222,1223,1224,1226,1228,1232,1234,1236,1242,1243,1246,1248,1249,1250,1251,1254,1255,1256,1261,1263,1266,1267,1269,1271,1274,1285,1289,1290,1292,1293,1296,1300,1302,1305,1306,1311,1317,1324,1327,1328,1330,1333,1334,1336,1339,1343,1344,1346,1348,1353,1355,1356,1357,1358,1359,1363,1365,1366,1368,1369,1372,1379,1380,1381,1383,1385,1386,1389,1390,1391,1392,1393,1397,1398,1399,1401,1402,1403,1405,1408,1409,1410,1413,1414,1417,1425,1427,1429,1430,1434,1443,1444,1446,1447,1448,1453,1456,1457,1458,1462,1464,1469,1471,1473,1478,1480,1481,1487,1490,1502,1504,1507,1509,1511,1514,1515,1518,1531,1536,1538,1541,1547,1548,1563,1566,1573,1579,1582,1585,1587,1588,1592,1595,1601,1602,1605,1608,1610,1612,1615,1620,1623,1636,1642,1646,1652,1657,1675,1682,1683,1684,1685,1687,1692,1695,1698,1700,1714,1716,1717,1728,1730,1734,1739,1744,1745,1752,1754,1755,1762,1773,1774,1795,1800,1801,1807,1816,1829,1831,1844,1846,1868,1887,1897,1898,1900,1915,1933,1937,1940,1960,1966,1967,1969,1970,1971,1987,1991,2032,2035,2037,2082,2091,2113,2116,2120,2129,2142,2156,2199,2201,2234,2235,2250,2266,2282,2287,2288,2302,2304,2312,2339,2357,2386,2400,2401,2418,2443,2444,2493,2497,2503,2510,2546,2549,2552,2559,2572,2575,2599,2625,2636,2637,2687,2747,2750,2751,2765,2885,2957,2982,3002,3011,3058,3067,3105,3155,3212,3272,3301,3306,3360,3530,3579,32768,33238,56662,65536,]
k_vals = [16, 128, 8192, 28672]

def test(m, n, k):
    print(f"Testing m={m}, n={n}, k={k}")
    a = (torch.rand(n, k, device='cuda')*480. - 240.).to(torch.float8_e4m3fnuz)
    b = (torch.rand(m, k, device='cuda')*480. - 240.).to(torch.float8_e4m3fnuz)
    i = torch.tensor(1, dtype=torch.float32, device='cuda')
    i1 = torch.tensor(0.01, dtype=torch.float32, device='cuda')
    res1 = vllm_ops.fp8_gemm_16(a, b.t(), i1, i1, 36526)
    res2, _ = torch._scaled_mm(a, b.t(), out_dtype=torch.float16, scale_a=i1, scale_b=i1,scale_result=i,bias=None)

    assert torch.allclose(res1, res2, atol=0.01, rtol=1e-3)

for m in m_vals:
    for n in n_vals:
        for k in k_vals:
            test(m, n, k)