diff --git a/src/contrib/common/gemm.py b/src/contrib/common/gemm.py index 21186a06..637fce60 100644 --- a/src/contrib/common/gemm.py +++ b/src/contrib/common/gemm.py @@ -134,7 +134,8 @@ def __init__(self, J, mfma_MN:int, wave_size:list, wave_cnt:list, - K, N): + K, N, + is_fp8=False): self.K = K self.N = N @@ -144,7 +145,11 @@ def __init__(self, J, assert wave_size_M % mfma_MN == 0 assert wave_size_N % mfma_MN == 0 # *2 for 8-bf16/fp16 so DWORDx4 lane-size can be used - self.mfma_K = (2*8 if mfma_MN == 32 else 2*16) + if is_fp8: + assert mfma_MN == 16, 'fp8 only support 16x16x16' + self.mfma_K = 2 * 32 + else: + self.mfma_K = (2*8 if mfma_MN == 32 else 2*16) # number of C/D regs per wave wave_nCM = wave_size_M // mfma_MN @@ -171,18 +176,25 @@ def __init__(self, J, self.wave_nCM = wave_nCM self.wave_nCN = wave_nCN self.wave_nCK = wave_nCK + self.is_fp8 = is_fp8 + if is_fp8: + self.sizeof_a = 1 + self.sizeof_b = 1 + else: + self.sizeof_a = J.sizeof_bf16 + self.sizeof_b = J.sizeof_bf16 def run(self, loaderA, loaderB, buff_c, M, debug_warp, skip_load): J = self.J - LDSA_size = self.wg_M * self.wg_K * J.sizeof_bf16 - LDSB_size = self.wg_N * self.wg_K * J.sizeof_bf16 + LDSA_size = self.wg_M * self.wg_K * self.sizeof_a + LDSB_size = self.wg_N * self.wg_K * self.sizeof_b ldsA = J.alloc_lds(LDSA_size) ldsB = J.alloc_lds(LDSB_size) # prefetch in memory-coalescing way # each lane prefetch DWORDx4 which is 8xhalf - num_lanes_per_row = J.div(J.sizeof_bf16 * self.wg_K, J.sizeof_DW4) + num_lanes_per_row = J.div(self.sizeof_a * self.wg_K, J.sizeof_DW4) dw4_prefetch_MN = self.wave_cnt * J.div(64, num_lanes_per_row) assert dw4_prefetch_MN >= 1 print(f"{self.wg_M=} {num_lanes_per_row=} {dw4_prefetch_MN=}") @@ -194,38 +206,38 @@ def swizzle(row, col): # each swizzle generates a new vaddr pattern, precompute all of them # each wave reads its own part - ds_readA_vaddr = J.gpr(self.wave_nCM, self.wave_nCK, "vu32") - ds_readB_vaddr = J.gpr(self.wave_nCN, self.wave_nCK, "vu32") + ds_readA_vaddr = J.gpr(1, self.wave_nCK, "vu32") + ds_readB_vaddr = J.gpr(1, self.wave_nCK, "vu32") # wave location warp_id_m = J.warp_id // self.wave_cnt_N warp_id_n = J.warp_id % self.wave_cnt_N warp_offset_m = warp_id_m * self.wave_size_M warp_offset_n = warp_id_n * self.wave_size_N - for m in range(self.wave_nCM): + for m in range(1): for k in range(self.wave_nCK): row = J.lane_id % self.mfma_MN + warp_offset_m + m*self.mfma_MN - col = J.lane_id // self.mfma_MN + (k * self.mfma_K * J.sizeof_bf16) // J.sizeof_DW4 + col = J.lane_id // self.mfma_MN + (k * self.mfma_K * self.sizeof_a) // J.sizeof_DW4 swizzle_col = swizzle(row, col) % (num_lanes_per_row) - ds_readA_vaddr[m, k] = J.gpr((row * (self.wg_K * J.sizeof_bf16)) + swizzle_col*(J.sizeof_DW4)) + ds_readA_vaddr[m, k] = J.gpr((row * (self.wg_K * self.sizeof_a)) + swizzle_col*(J.sizeof_DW4)) - for n in range(self.wave_nCN): + for n in range(1): for k in range(self.wave_nCK): row = J.lane_id % self.mfma_MN + warp_offset_n + n*self.mfma_MN - col = J.lane_id // self.mfma_MN + (k * self.mfma_K * J.sizeof_bf16) // J.sizeof_DW4 + col = J.lane_id // self.mfma_MN + (k * self.mfma_K * self.sizeof_b) // J.sizeof_DW4 swizzle_col = swizzle(row, col) % (num_lanes_per_row) - ds_readB_vaddr[n, k] = J.gpr((row * (self.wg_K * J.sizeof_bf16)) + swizzle_col*(J.sizeof_DW4)) + ds_readB_vaddr[n, k] = J.gpr((row * (self.wg_K * self.sizeof_b)) + swizzle_col*(J.sizeof_DW4)) Creg_size = (self.mfma_MN * self.mfma_MN)//64 mfma_C = self.J.gpr(self.wave_nCM, self.wave_nCN, Creg_size, f"af32") - ABReg_size = (self.mfma_MN * self.mfma_K * 2//4)//64 + ABReg_size = (self.mfma_MN * self.mfma_K * self.sizeof_a//4)//64 mfma_A = J.gpr(self.wave_nCM, self.wave_nCK, ABReg_size, "vbf16x2") mfma_B = J.gpr(self.wave_nCN, self.wave_nCK, ABReg_size, "vbf16x2") def ds_readA(m, k): - J.ds_read_b128(mfma_A[m,k], ds_readA_vaddr[m, k], mod=f"offset:{ldsA}") # vaddr, vdata offset gds + J.ds_read_b128(mfma_A[m,k], ds_readA_vaddr[0, k], mod=f"offset:{ldsA + m*self.mfma_MN*self.wg_K * self.sizeof_a}") # vaddr, vdata offset gds def ds_readB(n, k): - J.ds_read_b128(mfma_B[n,k], ds_readB_vaddr[n,k], mod=f"offset:{ldsB}") # vaddr, vdata offset gds + J.ds_read_b128(mfma_B[n,k], ds_readB_vaddr[0,k], mod=f"offset:{ldsB+n*self.mfma_MN*self.wg_K * self.sizeof_b}") # vaddr, vdata offset gds J.debug_setup((J.blockIdx.x[0] == 0) & (J.blockIdx.y[0] == 0) & (J.warp_id == debug_warp)) @@ -256,7 +268,7 @@ def ds_readB(n, k): mfma_C[:] = 0 # prelog 1: ds_write + prefetch - k_offset[0] = 0 if skip_load else (k_offset[0] + self.wg_K * J.sizeof_bf16) + k_offset[0] = 0 if skip_load else (k_offset[0] + self.wg_K * self.sizeof_a) loaderA.reset_offset(k_offset) loaderB.reset_offset(k_offset) @@ -297,6 +309,8 @@ def ds_readB(n, k): 16:("v_mfma_f32_16x16x32_bf16",16) if is_cdna4 else ("v_mfma_f32_16x16x16_bf16",16), 32:("v_mfma_f32_32x32x16_bf16",32) if is_cdna4 else ("v_mfma_f32_32x32x8_bf16",32) } + if self.is_fp8: + mfma_info[16] = ("v_mfma_f32_16x16x32_fp8_fp8", 16) mfma_name = mfma_info[self.mfma_MN][0] mfma_cycles = mfma_info[self.mfma_MN][1] @@ -327,11 +341,29 @@ def mfma_generator(k): mfma_C[m,n]) cur_k = J.gpr("su32", 0) k_loop_cnt = self.K//self.wg_K + mfma_cnt_dict = { + (256, 256) : { + 'ds_read': 2, + 'ds_write': 2, + 'prefetch': 8 + }, + (128, 256) : { + 'ds_read': 2, + 'ds_write': 3, + 'prefetch': 3 + }, + (64, 256) : { + 'ds_read': 1, + 'ds_write': 2, + 'prefetch': 1 + } + } + mfma_cnt = mfma_cnt_dict.get((self.wg_M, self.wg_N), mfma_cnt_dict[(256, 256)]) with J.While(cur_k[0] < k_loop_cnt): #for unroll in range(k_loop_cnt): - k_offset[0] = 0 if skip_load else (k_offset[0] + self.wg_K * J.sizeof_bf16) + k_offset[0] = 0 if skip_load else (k_offset[0] + self.wg_K * self.sizeof_a) loaderA.reset_offset(k_offset) loaderB.reset_offset(k_offset) @@ -344,12 +376,13 @@ def mfma_generator(k): for k in range(self.wave_nCK//2, self.wave_nCK): for m in range(self.wave_nCM): ds_readA(m, k) - J.emit(mfma0, 16*2) + J.emit(mfma0, 16*mfma_cnt['ds_read']) for n in range(self.wave_nCN): ds_readB(n, k) - J.emit(mfma0, 16*2) + J.emit(mfma0, 16*mfma_cnt['ds_read']) # ensure all waves has been finished reading LDS, so ds_write can overwrite it + J.emit(mfma0, 16*2) J.s_waitcnt(mod=f"lgkmcnt(0)") J.s_barrier() @@ -358,17 +391,17 @@ def mfma_generator(k): J.emit([mfma0, mfma1], 16) J.s_waitcnt(mod=f"vmcnt({num_prefetch_N + num_prefetch_M - 1})") loaderA.ds_write(r, ldsA) - J.emit([mfma0, mfma1], 16*2) + J.emit([mfma0, mfma1], 16*mfma_cnt['ds_write']) loaderA.prefetch(r) - J.emit([mfma0, mfma1], 16*8) + J.emit([mfma0, mfma1], 16*mfma_cnt['prefetch']) for r in range(num_prefetch_N): J.emit([mfma0, mfma1], 16) J.s_waitcnt(mod=f"vmcnt({num_prefetch_N + num_prefetch_M - 1})") loaderB.ds_write(r, ldsB) - J.emit([mfma0, mfma1], 16*2) + J.emit([mfma0, mfma1], 16*mfma_cnt['ds_write']) loaderB.prefetch(r) - J.emit([mfma0, mfma1], 16*8) + J.emit([mfma0, mfma1], 16*mfma_cnt['prefetch']) # enure mfma0 finished using part0 of mfma_A/mfma_B, before ds_read0 overwrites them # (most likely already empty and some part of mfma1 has been consumed) @@ -379,10 +412,10 @@ def mfma_generator(k): for k in range(0,self.wave_nCK//2): for m in range(self.wave_nCM): ds_readA(m, k) - J.emit([mfma1], 16*2) + J.emit([mfma1], 16*mfma_cnt['ds_read']) for n in range(self.wave_nCN): ds_readB(n, k) - J.emit([mfma1], 16*2) + J.emit([mfma1], 16*mfma_cnt['ds_read']) J.emit(mfma1) cur_k[0] += 1 @@ -456,7 +489,7 @@ def mfma_generator(k): for n in range(self.wave_nCN): row = J.lane_id % self.mfma_MN + warp_offset_m + m*self.mfma_MN col = J.lane_id // self.mfma_MN + n * (self.mfma_MN * J.sizeof_fp32 // J.sizeof_DW4) - voffset = J.gpr(row * (N*J.sizeof_fp32) + warp_offset_n*J.sizeof_fp32 + col*J.sizeof_DW4) + voffset = J.gpr(row * (self.N*J.sizeof_fp32) + warp_offset_n*J.sizeof_fp32 + col*J.sizeof_DW4) if self.mfma_MN == 16: buff_c.store_dwordx4(mfma_C[m,n], voffset, 0) elif self.mfma_MN == 32: @@ -503,7 +536,7 @@ def gemm_kernel(J, K, N, M01, GroupNum, gemm.wg_N, J.sizeof_bf16*gemm.wg_K, stride_bytes, total_wave_cnt, swizzle_row_div, skip_load) else: - loaderB = MFMA_DW4Loader_preshuffled(J, pB, actual_wg_M * K * J.sizeof_bf16, mfma_MN, + loaderB = MFMA_DW4Loader_preshuffled(J, pB, gemm.wg_N * K * J.sizeof_bf16, mfma_MN, gemm.wg_N, J.sizeof_bf16*gemm.wg_K, stride_bytes, total_wave_cnt, swizzle_row_div, skip_load) @@ -665,7 +698,7 @@ def test_gemm(mfma_MN, wave_size, wave_cnt, A_preshuffled = False, B_preshuffled #assert 0 #test_gemm(32, [128, 128], [2, 2], A_preshuffled = False, B_preshuffled = False) #test_gemm(16, [128, 128], [2, 2], A_preshuffled = False, B_preshuffled = False) - test_gemm(16, [64, 64], [2, 2], A_preshuffled = False, B_preshuffled = True) + test_gemm(16, [32, 128], [2, 2], A_preshuffled = False, B_preshuffled = True) #test_gemm(16, [128, 128], [2, 2], A_preshuffled = True, B_preshuffled = True) #test_gemm(32, [128, 128], [2, 2], A_preshuffled = True, B_preshuffled = True) assert 0 diff --git a/src/contrib/common/gemm_splitk.py b/src/contrib/common/gemm_splitk.py index 516ab458..d1da08fe 100644 --- a/src/contrib/common/gemm_splitk.py +++ b/src/contrib/common/gemm_splitk.py @@ -22,10 +22,13 @@ def gemm_splitk(J:JIT, BLOCK_TILE_SIZE_N = 32, BLOCK_TILE_SIZE_M = 16, USE_FP4_SHUFFLE_WEIGHT=False, - fp8_ptpc=True + quant_type_str='no', ): assert BLOCK_TILE_SIZE_M % 16 == 0, f'BLOCK_TILE_SIZE_M must be multiple of 16, current {BLOCK_TILE_SIZE_M=}' assert BLOCK_TILE_SIZE_N % 32 == 0, f'BLOCK_TILE_SIZE_N must be multiple of 32, current {BLOCK_TILE_SIZE_N=}' + fp8_ptpc = True if quant_type_str == 'per_Token' else False + fp8_per_tensor = True if quant_type_str == 'per_Tensor' else False + fp8_block128 = True if quant_type_str == 'per_1x128' else False sizeof_f32 = 4 sizeof_bf16 = 2 sizeof_w = sizeof_bf16 if weight_dtype == torch.bfloat16 else 1 @@ -77,11 +80,12 @@ def gemm_splitk(J:JIT, k_scale_n = div_up(div_up(K, 32), 8) // num_split_k v_w_scale = J.gpr(B_horz // 2, k_scale_n, 'vf32', align=4) k_scale_n_next_read_idx = 0 - elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and fp8_ptpc: - v_w_scale = J.gpr(B_horz, 4, 'vf32') - # for n in range(B_horz): - # J.global_load_dwordx4(v_w_scale[n], voffset_scale[n], p_w_scale) - elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and not fp8_ptpc: + elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz): + if fp8_per_tensor: + v_w_scale = J.gpr(2, 'vf32') + elif fp8_ptpc: + v_w_scale = J.gpr(B_horz, 4, 'vf32') + else: QUAN_BLK_SZ=128 #[ping-pong, ngroups] v_w_scale = J.gpr(2, 2, 'vf32') @@ -106,7 +110,7 @@ def load_gen(pp_reg_id, k=None): J.global_load_dword(v_w_scale[n, k_scale_n_next_read_idx], voffset_scale[n], p_w_scale, mod=f'offset:{k_scale_n_next_read_idx * 64 * sizeof_f32}') k_scale_n_next_read_idx += 1 k_scale_wip = B_horz // 2 - elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and not fp8_ptpc: + elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and fp8_block128: J.global_load_dword(v_w_scale[pp_reg_id, 0], voffset_scale[0], p_w_scale, mod=f'offset:{(k+1) * k_step_wg // QUAN_BLK_SZ * sizeof_f32}') J.global_load_dword(v_w_scale[pp_reg_id, 1], voffset_scale[1], p_w_scale, mod=f'offset:{(k+1) * k_step_wg // QUAN_BLK_SZ * sizeof_f32}') k_scale_wip = 2 @@ -165,7 +169,7 @@ def delayed_fma(): next(gen) next(gen) - elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and fp8_ptpc: + elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and (fp8_ptpc or fp8_per_tensor): v_w_f32 = J.gpr(2, 2, 2, 'vf32', align=4) v_w_bf16 = J.gpr(B_horz, 2, 2, 'vf32', align=4) # kl = 16 would be divided into 2 steps. Each accumulate 8 in K dimension. @@ -198,7 +202,7 @@ def delayed_fma(): yield J.v_mfma_f32_16x16x16_bf16(C_reg[n, m], v_w_bf16[n, 0], A_reg[pp_reg_id, m, i, 0], C_reg[n, m]) yield J.v_mfma_f32_16x16x16_bf16(C_reg[n, m], v_w_bf16[n, 1], A_reg[pp_reg_id, m, i, 1], C_reg[n, m]) - elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and not fp8_ptpc: + elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and fp8_block128: # cdna3 path: if is_cdna4 == False: v_w_f32 = J.gpr(2, 2, 2, 'vf32', align=4) @@ -289,6 +293,9 @@ def tail(pp_reg_id, k=None): for n in range(B_horz): J.global_load_dwordx4(v_w_scale[n], voffset_scale[n], p_w_scale) J.s_waitcnt(mod=f"vmcnt({B_horz})") + elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and fp8_per_tensor: + J.global_load_dword(v_w_scale[0], voffset_scale[0], p_w_scale) + J.s_waitcnt(mod=f"vmcnt(1)") else: J.s_waitcnt(mod=f"vmcnt(0)") @@ -344,3 +351,10 @@ def tail(pp_reg_id, k=None): for m in range(A_vert): J.v_pk_mul_f32(C_reg[n, m, 0:1], C_reg[n, m, 0:1], v_w_scale[n, 0:1]) J.v_pk_mul_f32(C_reg[n, m, 2:3], C_reg[n, m, 2:3], v_w_scale[n, 2:3]) + if (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and fp8_per_tensor: + J.s_waitcnt(mod=f"vmcnt(0)") + v_w_scale[1] = v_w_scale[0] + for n in range(B_horz): + for m in range(A_vert): + J.v_pk_mul_f32(C_reg[n, m, 0:1], C_reg[n, m, 0:1], v_w_scale) + J.v_pk_mul_f32(C_reg[n, m, 2:3], C_reg[n, m, 2:3], v_w_scale) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index 73aff4d4..cd4ba509 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -14,7 +14,9 @@ "moe_gemm_batch", "moe_2stage_splitk", "moe_2stage_gateup", + "moe_2stage_gateup_ref", "moe_2stage_down", + "moe_2stage_down_ref", "moe_1stage_splitk", "moe_2stage_down_loopn", ] @@ -53,6 +55,7 @@ def moe_gemm_batch1(J:JIT, M:"int", N:"int", K:"int", + quant_type_str, ): BLOCK_TILE_SIZE_M = 16 # nBM * 16 @@ -103,7 +106,7 @@ def moe_gemm_batch1(J:JIT, p_input[:] = p_input[:] + e_idx * (K * sizeof_bf16) voffset_b[1] = voffset_b[0] + K * (16 * sizeof_w) p_output[:] = p_output[:] + output_offset - if weight_dtype != torch.bfloat16: + if weight_dtype != torch.bfloat16 and quant_type_str == 'per_Token': p_w_scale[:] += (16 if with_silu else 32) * sizeof_f32 * J.blockIdx.x # 16 elements for a if fp8 @@ -119,12 +122,15 @@ def moe_gemm_batch1(J:JIT, voffset_scale = J.gpr(2, 'vu32') if weight_dtype != torch.bfloat16: - voffset_scale[0] = J.gpr(s_e_id * (N * sizeof_f32)) + lane_div_16 * (4 * sizeof_f32) - voffset_scale[1] = voffset_scale[0] + (N // 2 * sizeof_f32 if with_silu else 16 * sizeof_f32) + if quant_type_str == 'per_Tensor': + voffset_scale[0] = J.gpr(s_e_id * sizeof_f32) + else: + voffset_scale[0] = J.gpr(s_e_id * (N * sizeof_f32)) + lane_div_16 * (4 * sizeof_f32) + voffset_scale[1] = voffset_scale[0] + (N // 2 * sizeof_f32 if with_silu else 16 * sizeof_f32) gemm_splitk(J, weight_dtype, K, N, num_split_k, buff_a, buff_b, p_w_scale, - voffset_a, voffset_b, voffset_scale, C_reg) + voffset_a, voffset_b, voffset_scale, C_reg, quant_type_str=quant_type_str) s_cvt_bf16_bias = J.gpr(1, "su32") s_cvt_bf16_bias[0] = 0x00008000 @@ -206,7 +212,8 @@ def moe_gemm_batch(J:JIT, M:"int", N:"int", K:"int", - TOPK:"int"): + TOPK:"int", + quant_type_str,): BLOCK_TILE_SIZE_M = 16 # nBM * 16 BLOCK_TILE_SIZE_N = 32 # nBN * 32 BLOCK_SIZE_K = 32 @@ -284,13 +291,16 @@ def moe_gemm_batch(J:JIT, voffset_scale = J.gpr(2, 'vu32') if weight_dtype != torch.bfloat16: - p_w_scale[:] += (16 if with_silu else 32) * sizeof_f32 * J.blockIdx.x - voffset_scale[0] = J.gpr(s_e_id * (N * sizeof_f32)) + lane_div_16 * (4 * sizeof_f32) - voffset_scale[1] = voffset_scale[0] + (N // 2 * sizeof_f32 if with_silu else 16 * sizeof_f32) + if quant_type_str == 'per_Tensor': + voffset_scale[0] = J.gpr(s_e_id * sizeof_f32) + else: + p_w_scale[:] += (16 if with_silu else 32) * sizeof_f32 * J.blockIdx.x + voffset_scale[0] = J.gpr(s_e_id * (N * sizeof_f32)) + lane_div_16 * (4 * sizeof_f32) + voffset_scale[1] = voffset_scale[0] + (N // 2 * sizeof_f32 if with_silu else 16 * sizeof_f32) gemm_splitk(J, weight_dtype, K, N, num_split_k, buff_a, buff_b, p_w_scale, - voffset_a, voffset_b, voffset_scale, C_reg) + voffset_a, voffset_b, voffset_scale, C_reg, quant_type_str=quant_type_str) s_cvt_bf16_bias = J.gpr(1, "su32") s_cvt_bf16_bias[0] = 0x00008000 @@ -378,7 +388,9 @@ def moe_2stage_splitk(J:JIT, p_num_valid_ids:"void*", p_w_scale:"float*", M:"int", - fp8_ptpc,): + quant_type_str,): + fp8_ptpc = True if quant_type_str == 'per_Token' else False + fp8_per_tensor = True if quant_type_str == 'per_Tensor' else False assert weight_dtype == torch.bfloat16 or weight_dtype == torch.float4_e2m1fn_x2 or weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz if weight_dtype == torch.float4_e2m1fn_x2: if with_silu: @@ -549,38 +561,42 @@ def get_k_bytes(k_in_elements): voffset_scale[0] = J.gpr(s_e_id * (N * k_scale_stride)) + J.threadIdx.x * sizeof_f32 for m in range(1, B_horz // 2): voffset_scale[m] = voffset_scale[0] + 32 * k_scale_stride * m - elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and fp8_ptpc: - voffset_scale = J.gpr(B_horz, 'vu32') - p_w_scale[:] += (BLOCK_TILE_SIZE_N_HALF if with_silu else BLOCK_TILE_SIZE_N) * sizeof_f32 * J.blockIdx.x - voffset_scale[0] = J.gpr(s_e_id * (N * sizeof_f32)) + lane_div_16 * (4 * sizeof_f32) - if with_silu: - voffset_scale[B_horz // 2] = voffset_scale[0] + N // 2 * sizeof_f32 - for m in range(1, B_horz // 2): - voffset_scale[m] = voffset_scale[0] + 16 * sizeof_f32 * m - voffset_scale[B_horz // 2 + m] = voffset_scale[B_horz // 2] + 16 * sizeof_f32 * m - else: - for m in range(1, B_horz): - voffset_scale[m] = voffset_scale[0] + 16 * sizeof_f32 * m - elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and not fp8_ptpc: - assert K %(num_split_k*128) == 0 ,f' K %(num_split_k*128) must be 0' - # the scale layout is [E, N//128, K//128] - k_scale_stride = K // 128 * sizeof_f32 - # would need BLOCK_TILE_SIZE_N//128 *2 scale offset register for 1st stage. - voffset_scale = J.gpr(2, 'vu32') - # N_wg scale offset on expert wg and N wg. - if with_silu: - p_w_scale[:] += BLOCK_TILE_SIZE_N_HALF* J.blockIdx.x // 128 * k_scale_stride + s_e_id * (N//128 * k_scale_stride) - else: - p_w_scale[:] += BLOCK_TILE_SIZE_N * J.blockIdx.x //128 * k_scale_stride + s_e_id * (N//128 * k_scale_stride ) - if with_silu: - #[E, INTER_SIZE_TP*2//128, HIDDEN_SIZE//128] - voffset_scale[0] = J.threadIdx.x //128 * sizeof_f32 - # N tile offset within wave offset - voffset_scale[1] = voffset_scale[0] + (N//2//128 * k_scale_stride ) + elif (weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz): + if fp8_per_tensor: + voffset_scale = J.gpr(B_horz, 'vu32') + voffset_scale[0] = J.gpr(s_e_id * sizeof_f32) + elif fp8_ptpc: + voffset_scale = J.gpr(B_horz, 'vu32') + p_w_scale[:] += (BLOCK_TILE_SIZE_N_HALF if with_silu else BLOCK_TILE_SIZE_N) * sizeof_f32 * J.blockIdx.x + voffset_scale[0] = J.gpr(s_e_id * (N * sizeof_f32)) + lane_div_16 * (4 * sizeof_f32) + if with_silu: + voffset_scale[B_horz // 2] = voffset_scale[0] + N // 2 * sizeof_f32 + for m in range(1, B_horz // 2): + voffset_scale[m] = voffset_scale[0] + 16 * sizeof_f32 * m + voffset_scale[B_horz // 2 + m] = voffset_scale[B_horz // 2] + 16 * sizeof_f32 * m + else: + for m in range(1, B_horz): + voffset_scale[m] = voffset_scale[0] + 16 * sizeof_f32 * m else: - # 64 lane offset - voffset_scale[0] = 0 - voffset_scale[1] = 0 + assert K %(num_split_k*128) == 0 ,f' K %(num_split_k*128) must be 0' + # the scale layout is [E, N//128, K//128] + k_scale_stride = K // 128 * sizeof_f32 + # would need BLOCK_TILE_SIZE_N//128 *2 scale offset register for 1st stage. + voffset_scale = J.gpr(2, 'vu32') + # N_wg scale offset on expert wg and N wg. + if with_silu: + p_w_scale[:] += BLOCK_TILE_SIZE_N_HALF* J.blockIdx.x // 128 * k_scale_stride + s_e_id * (N//128 * k_scale_stride) + else: + p_w_scale[:] += BLOCK_TILE_SIZE_N * J.blockIdx.x //128 * k_scale_stride + s_e_id * (N//128 * k_scale_stride ) + if with_silu: + #[E, INTER_SIZE_TP*2//128, HIDDEN_SIZE//128] + voffset_scale[0] = J.threadIdx.x //128 * sizeof_f32 + # N tile offset within wave offset + voffset_scale[1] = voffset_scale[0] + (N//2//128 * k_scale_stride ) + else: + # 64 lane offset + voffset_scale[0] = 0 + voffset_scale[1] = 0 if (weight_dtype == torch.bfloat16 or weight_dtype == torch.float8_e4m3fn or weight_dtype == torch.float8_e4m3fnuz) and fp8_ptpc and not with_silu: for m in range(A_vert): v_weight[m, 1] = v_weight[m, 0] @@ -594,7 +610,7 @@ def get_k_bytes(k_in_elements): gemm_splitk(J, weight_dtype, K, N, num_split_k, buff_a, buff_b, p_w_scale, - voffset_a, voffset_b, voffset_scale, C_reg, BLOCK_TILE_SIZE_N=BLOCK_TILE_SIZE_N, BLOCK_TILE_SIZE_M=BLOCK_TILE_SIZE_M, USE_FP4_SHUFFLE_WEIGHT=USE_FP4_SHUFFLE_WEIGHT, fp8_ptpc=fp8_ptpc) + voffset_a, voffset_b, voffset_scale, C_reg, BLOCK_TILE_SIZE_N=BLOCK_TILE_SIZE_N, BLOCK_TILE_SIZE_M=BLOCK_TILE_SIZE_M, USE_FP4_SHUFFLE_WEIGHT=USE_FP4_SHUFFLE_WEIGHT, quant_type_str=quant_type_str) s_cvt_bf16_bias = J.gpr(1, "su32") s_cvt_bf16_bias[0] = 0x00008000 @@ -1138,6 +1154,46 @@ def loop_body(): loop_body() +def xcd_swizzle(J, blk1d, num_blocks, num_oc_blocks, NUM_XCD, NUM_CU_PER_XCD): + NUM_CU = NUM_XCD * NUM_CU_PER_XCD + num_groupped_blocks = num_blocks // NUM_CU * NUM_CU + blk_m = J.gpr("su32") + blk_n = J.gpr("su32") + if 0 and num_oc_blocks == 16: + # in unit of 4x8 [256x256] blocks + with J.If(blk1d < num_groupped_blocks) as If: + blk_base = (blk1d // NUM_CU) * NUM_CU + cu_id = blk1d % NUM_CU + xcd_id = cu_id % 8 # 0~8 + xcd_cu = cu_id // 8 # 0~31 + coord_n = (xcd_id % 2)*8 + (xcd_cu % 8) + coord_m = (xcd_id // 2)*4 + (xcd_cu // 8) + task_id = coord_m * num_oc_blocks + coord_n + new_blk1d = blk_base + task_id + blk_m[0] = new_blk1d // num_oc_blocks + blk_n[0] = new_blk1d - blk_m * num_oc_blocks + If.Else() + blk_m[0] = blk1d // num_oc_blocks + blk_n[0] = blk1d - blk_m * num_oc_blocks + elif 0 and num_oc_blocks <= 4: + with J.If(blk1d < num_groupped_blocks) as If: + blk_base = (blk1d // NUM_CU) * NUM_CU + cu_id = blk1d - blk1d // NUM_CU * NUM_CU + xcd_id = cu_id % NUM_XCD + xcd_cu = cu_id // NUM_XCD + task_id = xcd_id * NUM_CU_PER_XCD + xcd_cu + new_blk1d = blk_base + task_id + blk_m[0] = new_blk1d // num_oc_blocks + blk_n[0] = new_blk1d - blk_m * num_oc_blocks + + If.Else() + blk_m[0] = blk1d // num_oc_blocks + blk_n[0] = blk1d - blk_m * num_oc_blocks + else: + blk_m[0] = blk1d // num_oc_blocks + blk_n[0] = blk1d - blk_m * num_oc_blocks + return blk_m, blk_n + @jit() def moe_2stage_gateup(J:JIT, weight_dtype, @@ -1146,6 +1202,8 @@ def moe_2stage_gateup(J:JIT, N, # compile-time args BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + quant_type_w, # weight-quantization methods, activation is always per-token + p_id:"void*", p_input:"void*", p_weight:"void*", p_output:"void*", @@ -1153,8 +1211,11 @@ def moe_2stage_gateup(J:JIT, p_sorted_ids:"void*", p_sorted_expert_ids:"void*", p_num_valid_ids:"void*", - p_w_scale:"float*", - M:"int",): + pt_scale: "float*", + pc_scale:"float*", + M:"int", + num_blocks:"int", + dyn,): class MFMA_DW4Loader: def __init__(self, J:JIT, ptr, buff_size, row_index, wg_M:int, row_bytes:int, stride_bytes:int, @@ -1213,7 +1274,7 @@ def ds_write(self, index, lds_base): class MFMA_DW4Loader_preshuffled: def __init__(self, J:JIT, ptr, buff_size, mfma_MN, wg_M:int, row_bytes:int, stride_bytes:int, - wave_cnt:int, swizzle_row_div:int, + wave_cnt:int, swizzle_row_div:int, blockn_idx, skip_load:bool = False, N=0): self.buff = J.Buffer(ptr, buff_size) sizeof_DWORDX4 = 16 @@ -1236,7 +1297,9 @@ def __init__(self, J:JIT, ptr, buff_size, mfma_MN, self.prefetch_reg = J.gpr(wave_prefetches, 4, "vu32") self.prefetch_sbase = J.gpr("su32") self.prefetch_offsets = J.gpr(wave_prefetches, "su32") - self.prefetch_voffset = lane_id * sizeof_DWORDX4 + BLOCK_TILE_SIZE_N // 2 * stride_bytes * J.blockIdx.x + blockn_idx_v = J.gpr(1, "vu32") + blockn_idx_v[0] = blockn_idx + self.prefetch_voffset = lane_id * sizeof_DWORDX4 + BLOCK_TILE_SIZE_N // 2 * stride_bytes * blockn_idx_v prefetch_id = 0 for wave in range(wave_cnt): with J.If(warp_id[0] == wave): @@ -1295,139 +1358,327 @@ def ds_write(self, index, lds_base): BLOCK_TILE_SIZE_N_HALF = BLOCK_TILE_SIZE_N // 2 + is_fp8 = str(weight_dtype).startswith("torch.float8_e4m3") + if is_fp8: + assert quant_type_w == "QuantType.per_Token" or quant_type_w == "QuantType.per_Tensor" + sizeof_bf16 = 2 sizeof_f32 = 4 sizeof_DWORDX4 = 16 + sizeof_a = sizeof_bf16 if weight_dtype == torch.bfloat16 else 1 sizeof_w = sizeof_bf16 if weight_dtype == torch.bfloat16 else 1 - stride_A = K * sizeof_bf16 + stride_A = K * sizeof_a stride_B = K * sizeof_w # 4(mfma columns) * 16(dwordx4) * 2(unroll 2 times) BLOCK_TILE_SIZE_K = 4 * 16 * 2 / sizeof_w read_col_lanes_per_wave = 4 * 16 * 2 // 16 read_row_lanes_per_wave = 64 // read_col_lanes_per_wave - assert BLOCK_TILE_SIZE_M % (4 * read_row_lanes_per_wave) == 0 - A_vert = BLOCK_TILE_SIZE_M // 4 // read_row_lanes_per_wave - # grid: N // 32, sorted_expert_ids.shape[0] - # expert index in p_sorted_expert_ids - e_idx = J.blockIdx.y - s_e_id = J.gpr(1, 'su32') - J.s_load_dword(s_e_id, p_sorted_expert_ids, e_idx[0] * 4) max_id = J.gpr(1, 'su32') J.s_load_dword(max_id, p_num_valid_ids, 0) - J.s_waitcnt(mod=f"lgkmcnt(0)") - # invalid padding section - J.Jump("continue_following", e_idx * BLOCK_TILE_SIZE_M < max_id) - J.s_endpgm() - J.Label("continue_following") - J.debug_setup((J.warp_id[0] == 0) & (J.blockIdx.x[0] == 0) & (e_idx[0] == 0)) - - # hide following initialization into s_waitcnt - p_sorted_ids[:] += e_idx * (BLOCK_TILE_SIZE_M * 4) - # one WG per CU, 4 waves split on K lane_mod_16 = get_lane_id_mod(J, 16) lane_div_16 = get_lane_id_div(J, 16) warp_id = J.gpr("su32") J.v_readfirstlane_b32(warp_id, J.threadIdx.x[0] // 64) + J.s_waitcnt(mod=f"lgkmcnt(0)") - v_sorted_id = J.gpr(A_vert, 'vu32') - v_sorted_id_off = J.gpr((J.lane_id // read_col_lanes_per_wave + J.warp_id * read_row_lanes_per_wave) * 4) - J.global_load_dword(v_sorted_id[0], v_sorted_id_off, p_sorted_ids) - p_sorted_ids_tmp = J.gpr(2, 'su32') - p_sorted_ids_tmp[0] = p_sorted_ids[0] - p_sorted_ids_tmp[1] = p_sorted_ids[1] - for n in range(1, A_vert): - p_sorted_ids_tmp[:] += read_row_lanes_per_wave * 4 * sizeof_f32 - J.global_load_dword(v_sorted_id[n], v_sorted_id_off, p_sorted_ids_tmp) - # for write, thread layout [64 // 4(=thread lanes of N), 4(=TILE_N // 2wave // 2(gate+up) * sizeof_bf16) // 16] - write_col_lanes_per_wave = BLOCK_TILE_SIZE_N // 2 // 2 * sizeof_bf16 // 16 - write_row_lanes_per_wave = 64 // write_col_lanes_per_wave - assert BLOCK_TILE_SIZE_M % (4 * write_row_lanes_per_wave) == 0 - A_vert_write = BLOCK_TILE_SIZE_M // 2 // write_row_lanes_per_wave - v_sorted_id_write = J.gpr(A_vert_write, 'vu32') - v_sorted_id_write_off = J.gpr((J.lane_id // write_col_lanes_per_wave + J.warp_id // 2 * (A_vert_write * write_row_lanes_per_wave)) * 4) - J.global_load_dword(v_sorted_id_write[0], v_sorted_id_write_off, p_sorted_ids) - for n in range(1, A_vert_write): - p_sorted_ids[:] += write_row_lanes_per_wave * sizeof_f32 - J.global_load_dword(v_sorted_id_write[n], v_sorted_id_write_off, p_sorted_ids) - - p_output[:] = p_output[:] + BLOCK_TILE_SIZE_N_HALF * sizeof_bf16 * J.blockIdx.x + def loop_body(idx): + # grid: N // 32, sorted_expert_ids.shape[0] + # expert index in p_sorted_expert_ids + e_idx, blockn_idx = xcd_swizzle(J, idx, num_blocks, N // BLOCK_TILE_SIZE_N, 4, 20) + # invalid padding section + J.Jump("continue_following", e_idx * BLOCK_TILE_SIZE_M < max_id) + J.s_endpgm() + J.Label("continue_following") + s_e_id = J.gpr(1, 'su32') + J.s_load_dword(s_e_id, p_sorted_expert_ids, e_idx[0] * 4) - # wait for v_sorted_id - J.s_waitcnt(mod=f"vmcnt(0)") - v_token_id = J.gpr(A_vert, 'vu32') - for m in range(A_vert): - v_token_id[m] = v_sorted_id[m] & 0xffffff + J.s_waitcnt(mod=f"lgkmcnt(0)") + J.debug_setup((J.warp_id[0] == 0) & (blockn_idx[0] == 0) & (e_idx[0] == 0)) - p_weight[:] = p_weight[:] + s_e_id * (N * K * sizeof_w) + # hide following initialization into s_waitcnt + p_cur_sorted_ids = J.gpr(2, 'su32') + p_cur_sorted_ids[:] = p_sorted_ids[:] + e_idx * (BLOCK_TILE_SIZE_M * 4) + VALID_TILE_SIZES = [s for s in [64, 128] if s < BLOCK_TILE_SIZE_M] + s_sorted_id_tile = J.gpr(len(VALID_TILE_SIZES), 'su32') + for n in range(len(VALID_TILE_SIZES)): + J.s_load_dword(s_sorted_id_tile[n], p_cur_sorted_ids, VALID_TILE_SIZES[n] * J.sizeof_DW) + + if len(VALID_TILE_SIZES): + J.s_waitcnt(mod=f"lgkmcnt(0)") + # one WG per CU, 4 waves split on K + + p_cur_output = J.gpr(2, 'su32') + p_cur_output[:] = p_output[:] + BLOCK_TILE_SIZE_N_HALF * sizeof_bf16 * blockn_idx + + def kernel(CUR_BLOCK_TILE_SIZE_M): + assert CUR_BLOCK_TILE_SIZE_M % (4 * read_row_lanes_per_wave) == 0 + A_vert = CUR_BLOCK_TILE_SIZE_M // 4 // read_row_lanes_per_wave + v_sorted_id = J.gpr(A_vert, 'vu32') + v_sorted_id_off = J.gpr((J.lane_id // read_col_lanes_per_wave + J.warp_id * read_row_lanes_per_wave) * 4) + for n in range(A_vert): + J.global_load_dword(v_sorted_id[n], v_sorted_id_off, p_cur_sorted_ids, mod=f'offset:{n * read_row_lanes_per_wave * 4 * sizeof_f32}') + # for write, thread layout [64 // 4(=thread lanes of N), 4(=TILE_N // 2wave // 2(gate+up) * sizeof_bf16) // 16] + write_col_lanes_per_wave = BLOCK_TILE_SIZE_N // 2 // 2 * sizeof_bf16 // 16 + write_row_lanes_per_wave = 64 // write_col_lanes_per_wave + assert CUR_BLOCK_TILE_SIZE_M % (4 * write_row_lanes_per_wave) == 0 + A_vert_write = CUR_BLOCK_TILE_SIZE_M // 2 // write_row_lanes_per_wave + v_sorted_id_write = J.gpr(A_vert_write, 'vu32') + v_sorted_id_write_off = J.gpr((J.lane_id // write_col_lanes_per_wave + J.warp_id // 2 * (A_vert_write * write_row_lanes_per_wave)) * 4) + for n in range(A_vert_write): + J.global_load_dword(v_sorted_id_write[n], v_sorted_id_write_off, p_cur_sorted_ids, mod=f'offset:{n * write_row_lanes_per_wave * sizeof_f32}') + + if is_fp8: + v_sorted_id_scale = J.gpr(CUR_BLOCK_TILE_SIZE_M // 2 // 16, 'vu32') + v_sorted_id_scale_off = J.gpr((lane_mod_16 + J.warp_id // 2 * (CUR_BLOCK_TILE_SIZE_M // 2)) * 4) + for n in range(CUR_BLOCK_TILE_SIZE_M // 2 // 16): + J.global_load_dword(v_sorted_id_scale[n], v_sorted_id_scale_off, p_cur_sorted_ids, mod=f'offset:{n * 16 * sizeof_f32}') + + # wait for v_sorted_id + J.s_waitcnt(mod=f"vmcnt(0)") + v_token_id = J.gpr(A_vert, 'vu32') + for m in range(A_vert): + v_token_id[m] = v_sorted_id[m] & 0xffffff + + p_cur_weight = J.gpr(2, 'su32') + p_cur_weight[:] = p_weight[:] + s_e_id * (N * K * sizeof_w) + + gemm = UGEMM(J, 16, [CUR_BLOCK_TILE_SIZE_M // 2, BLOCK_TILE_SIZE_N // 2], [2, 2], K, N, is_fp8) + swizzle_row_div = 1 + total_wave_cnt = 4 + loaderA = MFMA_DW4Loader(J, p_input, M * (K * sizeof_a), v_token_id, + gemm.wg_M, sizeof_a*gemm.wg_K, stride_A, + total_wave_cnt, swizzle_row_div, False) + loaderB = MFMA_DW4Loader_preshuffled(J, p_cur_weight, N * K * sizeof_w, 16, + gemm.wg_N, sizeof_w*gemm.wg_K, stride_B, + total_wave_cnt, swizzle_row_div, blockn_idx, False, N) + + # C_reg: [wave_nCM, wave_nCN] + C_reg = gemm.run(loaderA, loaderB, None, M, 0, 0) + + if is_fp8: + buff_sa = J.Buffer(pt_scale, M * J.sizeof("fp32")) + v_pt_scales = J.gpr(gemm.wave_nCM, 'vf32') + for m in range(gemm.wave_nCM): + voffset_sa = J.gpr("vu32", (v_sorted_id_scale[m] & 0xffffff) * J.sizeof_DW) + buff_sa.load_dword(v_pt_scales[m], voffset_sa, 0) + + if quant_type_w == "QuantType.per_Token": + cur_pc_scale = J.gpr(2, 'su32') + cur_pc_scale[:] = pc_scale[:] + (s_e_id * (N * J.sizeof_DW) + BLOCK_TILE_SIZE_N_HALF * J.sizeof_DW * blockn_idx + J.warp_id % 2 * BLOCK_TILE_SIZE_N_HALF // 2 * J.sizeof_DW) # per-channel scales for weights + vaddr_pc_scale = J.gpr("vu32", (J.lane_id // 16) * J.sizeof_DW4) + scales_pc = J.gpr(gemm.wave_nCN, 4, "vf32") + for n in range(0, gemm.wave_nCN, 2): + J.global_load_dwordx4(scales_pc[n], vaddr_pc_scale, cur_pc_scale, mod=f"offset:{n // 2 * 16 * J.sizeof_DW}") + vaddr_pc_scale += N // 2 * J.sizeof_DW + for n in range(1, gemm.wave_nCN, 2): + J.global_load_dwordx4(scales_pc[n], vaddr_pc_scale, cur_pc_scale, mod=f"offset:{n // 2 * 16 * J.sizeof_DW}") + + if quant_type_w == "QuantType.per_Tensor": + scales_pc = J.gpr("vf32") + voffset_sa = J.gpr(1, 'vu32') + voffset_sa[0] = s_e_id * J.sizeof_DW # each expert has a per-tensor scale + J.global_load_dword(scales_pc[0], voffset_sa, pc_scale) + + # vaddr_pc_scale[0] += num_mfma_n * 16 * J.sizeof_DW + J.s_waitcnt(mod=f"vmcnt(0)") + + s_cvt_bf16_bias = J.gpr(1, "su32") + s_cvt_bf16_bias[0] = 0x00008000 + + # swizzle-LDS to form better memory-coelascing VMEM + # each warp has its own [mfma_M x wave_size_N] output buffer + wave_size_N = BLOCK_TILE_SIZE_N // 2 // 2 + lds_out = J.alloc_lds(4 * 16 * wave_size_N * sizeof_bf16) + lds_warp_offset = warp_id * (16 * wave_size_N * sizeof_bf16) + num_lanes_per_row = write_col_lanes_per_wave + assert 64 % num_lanes_per_row == 0 + rows_per_read = 64 // num_lanes_per_row + assert rows_per_read > 0 + assert 16 % rows_per_read == 0 + + vdata = J.gpr(4, "vu32", align=4) + + assert A_vert_write % gemm.wave_nCM == 0 + write_times16 = A_vert_write // gemm.wave_nCM + for m in range(gemm.wave_nCM): + for n in range(gemm.wave_nCN // 2): + for i in range(4): + tmp = J.gpr(2, "vf32", align=2) + J.v_accvgpr_read_b32(tmp[0], C_reg[m, 2 * n + 0, i]) + J.v_accvgpr_read_b32(tmp[1], C_reg[m, 2 * n + 1, i]) + if is_fp8: + tmp[0] *= v_pt_scales[m] + tmp[1] *= v_pt_scales[m] + if quant_type_w == "QuantType.per_Token": + tmp[0] *= scales_pc[2 * n + 0, i] + tmp[1] *= scales_pc[2 * n + 1, i] + if quant_type_w == "QuantType.per_Tensor": + tmp[0] *= scales_pc[0] + tmp[1] *= scales_pc[0] + tmp[0] = J.gpr(tmp[1] * J.silu(tmp[0])) + J.v_add_u32(vdata[i], tmp[0], s_cvt_bf16_bias) + vouts = J.gpr(2, "vf32") + J.pk_f32_to_bf16(vouts[0], vdata[0], vdata[1]) + J.pk_f32_to_bf16(vouts[1], vdata[2], vdata[3]) + + row = lane_mod_16 + col = lane_div_16 + n * (16 * sizeof_bf16 // 8) + # writing unit is 8 bytes while reading unit is 16 bytes + swizzle_col = (col // 2 ^ row) % (num_lanes_per_row) + vaddr_w = J.gpr((row) * (wave_size_N * sizeof_bf16) + \ + lds_warp_offset + \ + (swizzle_col * 16 + (col & 1) * 8)) + J.ds_write_b64(vaddr_w, vouts, mod=f"offset:{lds_out}") + + for r in range(0, 16, rows_per_read): + cur_m = v_sorted_id_write[m * write_times16 + r // rows_per_read] & 0xffffff + voffset = J.gpr(cur_m * (N // 2 * sizeof_bf16 * TOPK) + + (v_sorted_id_write[m * write_times16 + r // rows_per_read] >> 24) * (N // 2 * sizeof_bf16) + + (J.lane_id % write_col_lanes_per_wave) * sizeof_DWORDX4 + J.warp_id % 2 * (write_col_lanes_per_wave * sizeof_DWORDX4)) + row = J.lane_id // num_lanes_per_row + r + col = J.lane_id % num_lanes_per_row + swizzle_col = (row ^ col) % num_lanes_per_row + vaddr_r = J.gpr((swizzle_col) * sizeof_DWORDX4 + \ + (row) * (wave_size_N * sizeof_bf16) + \ + lds_warp_offset) + J.ds_read_b128(vdata, vaddr_r, mod=f"offset:{lds_out}") + J.s_waitcnt(mod=f"lgkmcnt(0)") + with J.ExecMask(cur_m < M[0], early_skip=False): + J.global_store_dwordx4(voffset, vdata, p_cur_output) - gemm = UGEMM(J, 16, [BLOCK_TILE_SIZE_M // 2, BLOCK_TILE_SIZE_N // 2], [2, 2], K, N) - swizzle_row_div = 1 - total_wave_cnt = 4 - loaderA = MFMA_DW4Loader(J, p_input, M * (K * sizeof_bf16), v_token_id, - gemm.wg_M, sizeof_bf16*gemm.wg_K, stride_A, - total_wave_cnt, swizzle_row_div, False) - loaderB = MFMA_DW4Loader_preshuffled(J, p_weight, N * K * sizeof_bf16, 16, - gemm.wg_N, sizeof_bf16*gemm.wg_K, stride_B, - total_wave_cnt, swizzle_row_div, False, N) - # C_reg: [wave_nCM, wave_nCN] - C_reg = gemm.run(loaderA, loaderB, None, M, 0, 0) + J.free_lds(lds_out) - s_cvt_bf16_bias = J.gpr(1, "su32") - s_cvt_bf16_bias[0] = 0x00008000 + J.get_sgpr_const(0x8000) + J.get_sgpr_const(0x3020706) - # swizzle-LDS to form better memory-coelascing VMEM - # each warp has its own [mfma_M x wave_size_N] output buffer - wave_size_N = BLOCK_TILE_SIZE_N // 2 // 2 - lds_out = J.alloc_lds(4 * 16 * wave_size_N * sizeof_bf16) - lds_warp_offset = warp_id * (16 * wave_size_N * sizeof_bf16) - num_lanes_per_row = write_col_lanes_per_wave - assert 64 % num_lanes_per_row == 0 - rows_per_read = 64 // num_lanes_per_row - assert rows_per_read > 0 - assert 16 % rows_per_read == 0 - - vdata = J.gpr(4, "vu32", align=4) - - assert A_vert_write % gemm.wave_nCM == 0 - write_times16 = A_vert_write // gemm.wave_nCM - for m in range(gemm.wave_nCM): - for n in range(gemm.wave_nCN // 2): - for i in range(4): - tmp = J.gpr(2, "vf32", align=2) - J.v_accvgpr_read_b32(tmp[0], C_reg[m, 2 * n + 0, i]) - J.v_accvgpr_read_b32(tmp[1], C_reg[m, 2 * n + 1, i]) - tmp[0] = J.gpr(tmp[1] * J.silu(tmp[0])) - J.v_add_u32(vdata[i], tmp[0], s_cvt_bf16_bias) - vouts = J.gpr(2, "vf32") - J.pk_f32_to_bf16(vouts[0], vdata[0], vdata[1]) - J.pk_f32_to_bf16(vouts[1], vdata[2], vdata[3]) - - row = lane_mod_16 - col = lane_div_16 + n * (16 * sizeof_bf16 // 8) - # writing unit is 8 bytes while reading unit is 16 bytes - swizzle_col = (col // 2 ^ row) % (num_lanes_per_row) - vaddr_w = J.gpr((row) * (wave_size_N * sizeof_bf16) + \ - lds_warp_offset + \ - (swizzle_col * 16 + (col & 1) * 8)) - J.ds_write_b64(vaddr_w, vouts, mod=f"offset:{lds_out}") - - for r in range(0, 16, rows_per_read): - cur_m = v_sorted_id_write[m * write_times16 + r // rows_per_read] & 0xffffff - voffset = J.gpr(cur_m * (N // 2 * sizeof_bf16 * TOPK) + - (v_sorted_id_write[m * write_times16 + r // rows_per_read] >> 24) * (N // 2 * sizeof_bf16) + - (J.lane_id % write_col_lanes_per_wave) * sizeof_DWORDX4 + J.warp_id % 2 * (write_col_lanes_per_wave * sizeof_DWORDX4)) - row = J.lane_id // num_lanes_per_row + r - col = J.lane_id % num_lanes_per_row - swizzle_col = (row ^ col) % num_lanes_per_row - vaddr_r = J.gpr((swizzle_col) * sizeof_DWORDX4 + \ - (row) * (wave_size_N * sizeof_bf16) + \ - lds_warp_offset) - J.ds_read_b128(vdata, vaddr_r, mod=f"offset:{lds_out}") + if len(VALID_TILE_SIZES) == 1: + with J.If((s_sorted_id_tile[0] >> 24) != TOPK) as If: + kernel(BLOCK_TILE_SIZE_M) + If.Else() + kernel(VALID_TILE_SIZES[0]) + else: + assert len(VALID_TILE_SIZES) == 0 + kernel(BLOCK_TILE_SIZE_M) + + if dyn: + with J.While(): + idx = J.gpr(1, 'su32') + idx[0] = 0xffffffff + v_idx = J.gpr(1, 'vu32') + J.s_barrier() + with J.If(warp_id[0] == 0): + J.s_atomic_inc(idx, p_id, 0, mod='glc') + J.s_waitcnt(mod=f"lgkmcnt(0)") + v_idx[0] = idx[0] + J.ds_write_b32(0 * lane_mod_16, v_idx, mod=f"offset:{0}") + J.s_waitcnt(mod=f"lgkmcnt(0)") + J.s_barrier() + J.ds_read_b32(v_idx, 0 * lane_mod_16, mod=f"offset:{0}") J.s_waitcnt(mod=f"lgkmcnt(0)") - with J.ExecMask(cur_m < M[0]): - J.global_store_dwordx4(voffset, vdata, p_output) + J.v_readfirstlane_b32(idx, v_idx) + J.s_barrier() + loop_body(idx) + else: + loop_body(J.blockIdx.x) + +def de_shuffle_weight(weight_, mfma_MN = 16): + org_dtype = weight_.dtype + if org_dtype == torch.float4_e2m1fn_x2: + weight_ = weight_.view(torch.int8) + + K = weight_.shape[-1] + M = weight_.numel() // K + + weight = weight_.view(M, K) + K_bytes = K * weight.itemsize + sizeof_DW4 = 16 + mfma_K_lanes = 64 // mfma_MN + mfma_K_L = sizeof_DW4//weight.itemsize + mfma_K = mfma_K_lanes * mfma_K_L + + assert M % mfma_MN == 0 + mfma_K_bytes = mfma_K_lanes * sizeof_DW4 + assert K_bytes % mfma_K_bytes == 0 + #x = x.reshape(M//mfma_MN, mfma_MN, K//mfma_K, mfma_K_lanes, mfma_K_L) + #x = x.permute(0,2,3,1,4) + + assert K % mfma_K == 0 + weight = weight.reshape(M//mfma_MN, K//mfma_K, mfma_K_lanes, mfma_MN, mfma_K_L) + weight = weight.permute(0,3,1,2,4) + weight = weight.reshape(M, K).contiguous().view(weight_.shape) + return weight.view(org_dtype) + +def moe_2stage_gateup_ref(gemm1_in_q, gemm1_in_scale, + w1, w1_scale, + cur_out, + TILE_M, + sorted_ids, sorted_expert_ids, sorted_weights, num_valid_ids, TOPK): + w2_deshuffle = de_shuffle_weight(w1) + max_id = num_valid_ids[0] + for e_idx in range(sorted_expert_ids.shape[0]): + if e_idx * TILE_M >= max_id: + break + i0 = e_idx*TILE_M + i1 = i0 + TILE_M + s_e_id = sorted_expert_ids[e_idx] + + ids = sorted_ids[i0:i1].clone() + tok_ids = ids & 0xFFFFFF + tok_topk = ids >> 24 + valid_mask = tok_topk < torch.tensor(TOPK) + + scales_pt = gemm1_in_scale[tok_ids[valid_mask]] + scales_pc = w1_scale[s_e_id] + + src = gemm1_in_q[tok_ids[valid_mask], ...] + + src = src.to(torch.float) + wei = w2_deshuffle[s_e_id, ...].to(torch.float) + + act = src.to(torch.float) @ wei.t().to(torch.float) + + act *= (scales_pt[:, None] * scales_pc[None, :]).squeeze(dim=-1) + gate, up = act.chunk(2, dim=-1) + act = torch.nn.functional.silu(gate) * up + + cur_out[tok_ids[valid_mask], tok_topk[valid_mask], ...] = act.to(cur_out.dtype) + +def moe_2stage_down_ref(gemm1_out_q, gemm1_out_scale, + w2, w2_scale, + cur_out, + TILE_M, + sorted_ids, sorted_expert_ids, sorted_weights, num_valid_ids): + w2_deshuffle = de_shuffle_weight(w2) + B, TOPK, _ = gemm1_out_q.shape + gemm1_out_scale = gemm1_out_scale.view(B, TOPK) + max_id = num_valid_ids[0] + for e_idx in range(sorted_expert_ids.shape[0]): + if e_idx * TILE_M >= max_id: + break + i0 = e_idx*TILE_M + i1 = i0 + TILE_M + s_e_id = sorted_expert_ids[e_idx] + + ids = sorted_ids[i0:i1].clone() + tok_ids = ids & 0xFFFFFF + tok_topk = ids >> 24 + valid_mask = tok_topk < torch.tensor(TOPK) + + scales_pt = gemm1_out_scale[tok_ids[valid_mask], tok_topk[valid_mask]].flatten() + scales_pc = w2_scale[s_e_id].flatten() + + src = gemm1_out_q[tok_ids[valid_mask], tok_topk[valid_mask], ...] + + src = src.to(torch.float) + wei = w2_deshuffle[s_e_id, ...].to(torch.float) + + act = src.to(torch.float) @ wei.t().to(torch.float) + + act *= scales_pt[:, None] * scales_pc[None, :] + + cur_out[tok_ids[valid_mask], ...] += act * sorted_weights[i0:i1][valid_mask, None] @jit() def moe_2stage_down(J:JIT, @@ -1436,6 +1687,8 @@ def moe_2stage_down(J:JIT, with_silu, # BLOCK_TILE_SIZE_M, # 32 BLOCK_TILE_SIZE_N, # 64 + quant_type_w, + p_id:"void*", p_input:"void*", # [8192, 8, 128] p_weight:"void*", # [128, 2048, 128] p_output:"void*", # [8192, 2048] @@ -1443,29 +1696,32 @@ def moe_2stage_down(J:JIT, p_sorted_weights:"float*", # [69624] p_sorted_expert_ids:"void*", # [2176] p_num_valid_ids:"void*", # [2] value: [65536, 8192] - p_w_scale:"float*", - M:"int",): + pt_scale: "float*", + pc_scale:"float*", + M:"int", + num_blocks:"int", + dyn,): + if dyn: + SUB_M = 64 + assert BLOCK_TILE_SIZE_M % SUB_M == 0 + else: + SUB_M = BLOCK_TILE_SIZE_M sizeof_w = J.sizeof(weight_dtype) dtype_A = "bf16" + is_fp8 = str(weight_dtype).startswith("torch.float8_e4m3") + + if is_fp8: + fp8_ptpc = {"quant_type_w":quant_type_w} + dtype_A = "fp8" + else: + fp8_ptpc = None - e_idx = J.blockIdx.y - s_e_id = J.gpr(1, 'su32') - J.s_load_dword(s_e_id, p_sorted_expert_ids, e_idx[0] * 4) max_id = J.gpr(1, 'su32') J.s_load_dword(max_id, p_num_valid_ids, 0) - J.s_waitcnt(mod=f"lgkmcnt(0)") - # invalid padding section - J.Jump("continue_following", e_idx * BLOCK_TILE_SIZE_M < max_id) - J.s_endpgm() - J.Label("continue_following") - - p_sorted_ids[:] += e_idx * (BLOCK_TILE_SIZE_M * 4) - p_sorted_weights[:] += e_idx * (BLOCK_TILE_SIZE_M * 4) - p_weight[:] += s_e_id * (N * K * sizeof_w) mfma_MN = 16 mfma_K = (64//mfma_MN) * (J.sizeof_DW4//J.sizeof(dtype_A)) - num_mfma_m = J.div(BLOCK_TILE_SIZE_M, mfma_MN) + num_mfma_m = J.div(SUB_M, mfma_MN) num_mfma_k = J.div(K, mfma_K) A = J.gpr(num_mfma_m, num_mfma_k, 4, "abf16x2") @@ -1473,45 +1729,136 @@ def moe_2stage_down(J:JIT, # collect token id row = J.lane_id % mfma_MN col = J.lane_id // mfma_MN - v_sorted_id = J.gpr(num_mfma_m, 'vu32') - for m in range(num_mfma_m): - J.global_load_dword(v_sorted_id[m], row*J.sizeof_DW + m*mfma_MN*J.sizeof_DW, p_sorted_ids) + M_TOPK = J.gpr(M[0] * TOPK) - J.s_waitcnt(mod=f"vmcnt(0)") + J.get_sgpr_const(0x8000) + J.get_sgpr_const(0x3020706) + J.s_waitcnt(mod=f"lgkmcnt(0)") - vaddr = J.gpr(num_mfma_m, "vu32") - for m in range(num_mfma_m): - vaddr[m] = (v_sorted_id[m] & 0xFFFFFF) * (TOPK * K * sizeof_w) + (v_sorted_id[m]>>24) *(K * sizeof_w) + col * J.sizeof_DW4 + def loop_body(idx, run_in_m_sub_tiles): + e_idx, sub_m = xcd_swizzle(J, idx, num_blocks, BLOCK_TILE_SIZE_M // SUB_M, 4, 20) + if not run_in_m_sub_tiles: + sub_m = 0 + # invalid padding section + J.Jump("continue_following", e_idx * BLOCK_TILE_SIZE_M < max_id) + J.s_endpgm() + J.Label("continue_following") + # sub_m = J.blockIdx.x + s_e_id = J.gpr(1, 'su32') + J.s_load_dword(s_e_id, p_sorted_expert_ids, e_idx[0] * 4) + J.s_waitcnt(mod=f"lgkmcnt(0)") - buff_a = J.Buffer(p_input, M * TOPK * K * J.sizeof(dtype_A)) - for m in range(num_mfma_m): - for k in range(num_mfma_k): - buff_a.load_dwordx4(A[m,k], vaddr[m], 0, offset12=k*mfma_K*J.sizeof(dtype_A)) - - v_sorted_weights = J.gpr('vf32') - assert BLOCK_TILE_SIZE_M <= 256 - with J.ExecMask(J.threadIdx.x < BLOCK_TILE_SIZE_M): - J.global_load_dword(v_sorted_weights, J.threadIdx.x * 4, p_sorted_weights) + p_cur_sorted_ids = J.gpr(2, 'su32') + p_cur_sorted_ids[:] = p_sorted_ids[:] + (e_idx * (BLOCK_TILE_SIZE_M * 4) + sub_m * (SUB_M * 4)) + p_cur_sorted_weights = J.gpr(2, 'su32') + p_cur_sorted_weights[:] = p_sorted_weights[:] + (e_idx * (BLOCK_TILE_SIZE_M * 4) + sub_m * (SUB_M * 4)) + p_cur_weight = J.gpr(2, 'su32') + p_cur_weight[:] = p_weight[:] + s_e_id * (N * K * sizeof_w) + cur_pc_scale = J.gpr(2, 'su32') + if quant_type_w == "QuantType.per_Tensor": + cur_pc_scale[:] = pc_scale[:] + s_e_id * (1 * J.sizeof_DW) # per-tensor scales for weights + else: + cur_pc_scale[:] = pc_scale[:] + s_e_id * (N * J.sizeof_DW) # per-channel scales for weights - J.s_waitcnt(mod=f"vmcnt(0)") + v_sorted_id = J.gpr(num_mfma_m, 'vu32') + for m in range(num_mfma_m): + J.global_load_dword(v_sorted_id[m], row*J.sizeof_DW + m*mfma_MN*J.sizeof_DW, p_cur_sorted_ids) - lds_weights = J.alloc_lds(256*4) - lds_token_ids = J.alloc_lds(256*4) + J.s_waitcnt(mod=f"vmcnt(0)") - with J.ExecMask(J.threadIdx.x < BLOCK_TILE_SIZE_M): - J.ds_write_b32(J.threadIdx.x * 4, v_sorted_weights, mod=f"offset:{lds_weights}") - with J.ExecMask(J.threadIdx.x < 16): + vaddr = J.gpr(num_mfma_m, "vu32") for m in range(num_mfma_m): - J.ds_write_b32(J.lane_id * 4, v_sorted_id[m] & 0xffffff, mod=f"offset:{lds_token_ids + m*16*4}") + vaddr[m] = (v_sorted_id[m] & 0xFFFFFF) * (TOPK * K * sizeof_w) + (v_sorted_id[m]>>24) *(K * sizeof_w) + col * J.sizeof_DW4 - J.s_waitcnt(mod=f"lgkmcnt(0)") - J.s_barrier() + if is_fp8: + buff_sa = J.Buffer(pt_scale, M * TOPK * J.sizeof("fp32")) + v_pt_scales = J.gpr(num_mfma_m, 'vf32') + for m in range(num_mfma_m): + voffset_sa = J.gpr("vu32", (v_sorted_id[m] & 0xFFFFFF) * (TOPK * J.sizeof_DW) + (v_sorted_id[m]>>24) * J.sizeof_DW) + buff_sa.load_dword(v_pt_scales[m], voffset_sa, 0) + fp8_ptpc["v_pt_scales"] = v_pt_scales - num_mfma_n = 1 if BLOCK_TILE_SIZE_M > 64 else 2 - - down_kernel(J, mfma_MN, num_mfma_n, BLOCK_TILE_SIZE_M, N, K, - A, lds_token_ids, lds_weights, - p_weight, p_output, M) + if quant_type_w == "QuantType.per_Token": + fp8_ptpc["pc_scale"] = cur_pc_scale + + if quant_type_w == "QuantType.per_Tensor": + v_pc_scales = J.gpr('vf32') + v_offset_zero = J.gpr('vu32', 0) + J.global_load_dword(v_pc_scales[0], v_offset_zero, cur_pc_scale) + fp8_ptpc["v_pc_scales"] = v_pc_scales + + buff_a = J.Buffer(p_input, M * TOPK * K * J.sizeof(dtype_A)) + for m in range(num_mfma_m): + for k in range(num_mfma_k): + buff_a.load_dwordx4(A[m,k], vaddr[m], 0, offset12=k*mfma_K*J.sizeof(dtype_A)) + + v_sorted_weights = J.gpr('vf32') + assert SUB_M <= 256 + with J.ExecMask(J.threadIdx.x < SUB_M): + J.global_load_dword(v_sorted_weights, J.threadIdx.x * 4, p_cur_sorted_weights) + + J.s_waitcnt(mod=f"vmcnt(0)") + + lds_weights = J.alloc_lds(256*4) + lds_token_ids = J.alloc_lds(256*4) + + with J.ExecMask(J.threadIdx.x < SUB_M): + J.ds_write_b32(J.threadIdx.x * 4, v_sorted_weights, mod=f"offset:{lds_weights}") + with J.ExecMask(J.threadIdx.x < 16): + for m in range(num_mfma_m): + J.ds_write_b32(J.lane_id * 4, (v_sorted_id[m] & 0xffffff) * TOPK + (v_sorted_id[m]>>24), mod=f"offset:{lds_token_ids + m*16*4}") + + J.s_waitcnt(mod=f"lgkmcnt(0)") + J.s_barrier() + + if run_in_m_sub_tiles: + s_sorted_ids = J.gpr(1, 'su32') + J.v_readfirstlane_b32(s_sorted_ids[0], v_sorted_id[0]) + with J.If((s_sorted_ids[0] >> 24) != TOPK): + num_mfma_n = 1 if SUB_M > 64 else 2 + down_kernel(J, mfma_MN, num_mfma_n, SUB_M, N, K, + A, lds_token_ids, lds_weights, + p_cur_weight, p_output, M_TOPK, fp8_ptpc) + else: + # 0,1,2,3,4,5,6,7, num_mfma_m + VALID_TILE_SIZES = [s for s in [64,96,128] if s < BLOCK_TILE_SIZE_M] + s_sorted_ids = J.gpr(len(VALID_TILE_SIZES), 'su32') + for i, TILE_SIZE in enumerate(VALID_TILE_SIZES): + J.v_readfirstlane_b32(s_sorted_ids[i], v_sorted_id[TILE_SIZE//mfma_MN]) + + for i, TILE_SIZE in enumerate(VALID_TILE_SIZES): + with J.If((s_sorted_ids[i] >> 24) == TOPK): + num_mfma_n = 1 if TILE_SIZE > 64 else 2 + down_kernel(J, mfma_MN, num_mfma_n, TILE_SIZE, N, K, + A, lds_token_ids, lds_weights, + p_cur_weight, p_output, M_TOPK, fp8_ptpc) + J.s_endpgm() + + num_mfma_n = 1 if BLOCK_TILE_SIZE_M > 64 else 2 + down_kernel(J, mfma_MN, num_mfma_n, BLOCK_TILE_SIZE_M, N, K, + A, lds_token_ids, lds_weights, + p_cur_weight, p_output, M_TOPK, fp8_ptpc) + if dyn: + with J.While(): + idx = J.gpr(1, 'su32') + idx[0] = 0xffffffff + v_idx = J.gpr(1, 'vu32') + J.s_barrier() + with J.If(J.warp_id[0] == 0): + J.s_atomic_inc(idx, p_id, 0, mod='glc') + J.s_waitcnt(mod=f"lgkmcnt(0)") + v_idx[0] = idx[0] + J.ds_write_b32(0 * row, v_idx, mod=f"offset:{0}") + J.s_waitcnt(mod=f"lgkmcnt(0)") + J.s_barrier() + J.ds_read_b32(v_idx, 0 * row, mod=f"offset:{0}") + J.s_waitcnt(mod=f"lgkmcnt(0)") + J.v_readfirstlane_b32(idx, v_idx) + J.s_barrier() + + loop_body(idx, True) + else: + loop_body(J.blockIdx.x, False) def down_kernel(J, mfma_MN, num_mfma_n, BM, N, K, A, # A = J.gpr(num_mfma_m, num_mfma_k, 4, "abf16x2") @@ -1519,33 +1866,51 @@ def down_kernel(J, mfma_MN, num_mfma_n, BM, N, K, lds_weights, # 256 fp32 weights pB:"void*", pC:"void*", - M): + M, + fp8_ptpc): + num_warps = 4 + sizeof_w = J.sizeof_bf16 if fp8_ptpc is None else J.sizeof("fp8") # given DW4 lane size, how many bf16 items along K direction - mfma_K = (64//mfma_MN) * (J.sizeof_DW4//J.sizeof_bf16) # mfma_K = 32 + mfma_K = (64//mfma_MN) * (J.sizeof_DW4//sizeof_w) # each DW4 vgpr holds a mfma_K=32(bf16) or 64(fp8) # load A [BM x K] bf16 into AccGPRs num_mfma_m = J.div(BM, mfma_MN) num_mfma_k = J.div(K, mfma_K) # K=96, num_mfma_k=3 # 4 warps work in parallel along N dimension - buff_b = J.Buffer(pB, N * K * J.sizeof_bf16) + buff_b = J.Buffer(pB, N * K * sizeof_w) # ping-pong buffer - B = J.gpr(2, num_mfma_n, num_mfma_k, 4, "vbf16x2") + B = J.gpr(2, num_mfma_n, num_mfma_k, 4, "abf16x2") # 4 x (8,bf16) or (16,fp8) C = J.gpr(2, num_mfma_m, num_mfma_n, 4, "vf32") # prelog0, load Bn0 # prelog1, load Bn1, compute Cn0 # loop: load Bn2, compute Cn1, store Cn0 to LDS & load Cn0 & store to HBM - voff_b = J.gpr(J.lane_id * J.sizeof_DW4 + J.gpr(J.warp_id * (mfma_MN * K * J.sizeof_bf16))) + voff_b = J.gpr(J.lane_id * J.sizeof_DW4 + J.gpr(J.warp_id * (mfma_MN * K * sizeof_w))) soff_b = J.gpr("su32") soff_b[0] = 0 + + if fp8_ptpc is not None: + if fp8_ptpc['quant_type_w'] == "QuantType.per_Token": + vaddr_pc_scale = J.gpr("vu32", (J.lane_id // 16) * J.sizeof_DW4 + J.warp_id * (mfma_MN * J.sizeof_DW)) # point to the start of pc scales for this WG + scales_pc = J.gpr(2, num_mfma_n, 4, "vf32") + buff_sb = J.Buffer(fp8_ptpc["pc_scale"], N * J.sizeof_DW) + def loadB_generator(index): for n in range(num_mfma_n): for k in range(num_mfma_k): yield 1 buff_b.load_dwordx4(B[index,n,k], voff_b, soff_b) - soff_b[0] = soff_b[0] + mfma_MN * mfma_K * J.sizeof_bf16 - soff_b[0] = soff_b[0] + (3*num_mfma_k * mfma_MN * mfma_K * J.sizeof_bf16) + soff_b[0] = soff_b[0] + mfma_MN * mfma_K * sizeof_w + soff_b[0] = soff_b[0] + (3*num_mfma_k * mfma_MN * mfma_K * sizeof_w) + if fp8_ptpc is not None and fp8_ptpc['quant_type_w'] == "QuantType.per_Token": + # load scales_pc for next block, each MFMA 16x16 block-B needs 16 scales + # fp8_ptpc["v_pt_scales"] = v_pt_scales + # fp8_ptpc["pc_scale"] = pc_scale + for n in range(num_mfma_n): + yield 1 + buff_sb.load_dwordx4(scales_pc[index, n], vaddr_pc_scale, 0, offset12=n*num_warps*mfma_MN*J.sizeof_DW) + vaddr_pc_scale[0] += num_warps * num_mfma_n * mfma_MN * J.sizeof_DW def mfma_generator(index): for k in range(num_mfma_k): @@ -1553,12 +1918,39 @@ def mfma_generator(index): for n in range(num_mfma_n): Ci = 0 if k == 0 else C[index,m,n] yield 16 - J.v_mfma_f32_16x16x16_bf16(C[index,m,n], B[index,n,k,0:1], A[m,k,0:1], Ci) + if fp8_ptpc is not None: + J.v_mfma_f32_16x16x32_fp8_fp8(C[index,m,n], B[index,n,k,0:1], A[m,k,0:1], Ci) + else: + J.v_mfma_f32_16x16x16_bf16(C[index,m,n], B[index,n,k,0:1], A[m,k,0:1], Ci) for k in range(num_mfma_k): for m in range(num_mfma_m): for n in range(num_mfma_n): yield 16 - J.v_mfma_f32_16x16x16_bf16(C[index,m,n], B[index,n,k,2:3], A[m,k,2:3], C[index,m,n]) + if fp8_ptpc is not None: + J.v_mfma_f32_16x16x32_fp8_fp8(C[index,m,n], B[index,n,k,2:3], A[m,k,2:3], C[index,m,n]) + else: + J.v_mfma_f32_16x16x16_bf16(C[index,m,n], B[index,n,k,2:3], A[m,k,2:3], C[index,m,n]) + if fp8_ptpc is not None: + # dequantize C here since all computations over K have finished + # scales_pt are resident, scales_pc are + for m in range(num_mfma_m): + for n in range(num_mfma_n): + yield 16 + C[index,m,n,0] *= fp8_ptpc["v_pt_scales"][m] + C[index,m,n,1] *= fp8_ptpc["v_pt_scales"][m] + C[index,m,n,2] *= fp8_ptpc["v_pt_scales"][m] + C[index,m,n,3] *= fp8_ptpc["v_pt_scales"][m] + yield 16 + if fp8_ptpc['quant_type_w'] == "QuantType.per_Token": + C[index,m,n,0] *= scales_pc[index, n, 0] + C[index,m,n,1] *= scales_pc[index, n, 1] + C[index,m,n,2] *= scales_pc[index, n, 2] + C[index,m,n,3] *= scales_pc[index, n, 3] + if fp8_ptpc['quant_type_w'] == "QuantType.per_Tensor": + C[index,m,n,0] *= fp8_ptpc["v_pc_scales"][0] + C[index,m,n,1] *= fp8_ptpc["v_pc_scales"][0] + C[index,m,n,2] *= fp8_ptpc["v_pc_scales"][0] + C[index,m,n,3] *= fp8_ptpc["v_pc_scales"][0] # prelog0, load Bn0 J.emit(loadB_generator(0)) @@ -1573,7 +1965,7 @@ def mfma_generator(index): # loop: load Bn2, compute Cn1, store Cn0 to LDS & load Cn0 & store to HBM s_cvt_bf16_bias = J.get_sgpr_const(0x00008000) - vmem_lane_size = J.sizeof_DW + vmem_lane_size = J.sizeof_DW4 lds_padding = (4 if vmem_lane_size == J.sizeof_DW else 8) * J.sizeof_bf16 # to avoid bank-conflict lds_width = num_mfma_n * 4 * mfma_MN * J.sizeof_bf16 @@ -1594,7 +1986,7 @@ def mfma_generator(index): voff_c_lds_r = J.gpr(row * lds_stride + col * (vmem_lane_size)) vmem_stride = N * J.sizeof_bf16 v_weights = J.gpr(num_mfma_m, 2, "vf32") # pkmul - voff_vmem = J.gpr(num_loads, "vu32") + voff_vmem = J.gpr(num_loads, 2, "vu32") for m in range(num_mfma_m): J.ds_read_b32(v_weights[m,0], (m*mfma_MN + (J.lane_id % mfma_MN))*4, mod=f"offset:{lds_weights}") @@ -1609,11 +2001,17 @@ def mfma_generator(index): for m in range(num_mfma_m): v_weights[m,1] = v_weights[m,0] + saddr_dummy = J.gpr(2, "su32", 0) + voff_vmem_base = J.gpr(2, "vu32", pC[0], pC[1]) + J.v_lshl_add_u64(voff_vmem_base, J.gpr(2, "vu32", col * (vmem_lane_size), 0), 0, voff_vmem_base) for i in range(num_loads): - voff_vmem[i] = voff_vmem_row[i] * vmem_stride + col * (vmem_lane_size) + #voff_vmem[i] = voff_vmem_row[i] * vmem_stride + col * (vmem_lane_size) + J.v_mad_u64_u32(voff_vmem[i], saddr_dummy, voff_vmem_row[i], J.gpr("su32", vmem_stride), voff_vmem_base) temp_c = J.gpr(num_loads, vmem_lane_size//J.sizeof_DW, "vbf16x2") + voff_vmem_step = J.gpr(2, "vu32", 4 * num_mfma_n * mfma_MN * J.sizeof_bf16, 0) + def loop_body(ni): J.s_waitcnt(mod=f"vmcnt({num_loads})") @@ -1663,16 +2061,20 @@ def loop_body(ni): for i in range(num_loads): J.s_waitcnt(mod=f"lgkmcnt({min(15,num_loads - i - 1)})") if vmem_lane_size == J.sizeof_DW4: - J.global_store_dwordx4(voff_vmem[i], temp_c[i], pC) # this is fast: (48us) + with J.ExecMask(voff_vmem_row[i] < M[0], early_skip=False): + J.global_store_dwordx4(voff_vmem[i], temp_c[i], "off", mod="nt sc1") # this is fast: (48us) else: assert vmem_lane_size == J.sizeof_DW # the bigger the M is, the bigger the perf-diff is with J.ExecMask(voff_vmem_row[i] < M[0], early_skip=False): J.global_atomic_pk_add_bf16(voff_vmem[i], temp_c[i], pC) # this is much slower than directly store (60us) - J.emit(mfma1, 32) + J.emit(mfma1, 128) + if vmem_lane_size == J.sizeof_DW4: + J.v_lshl_add_u64(voff_vmem[i], voff_vmem_step, 0, voff_vmem[i]) J.emit(mfma1) - pC[:] += (4 * num_mfma_n * mfma_MN * J.sizeof_bf16) + if vmem_lane_size != J.sizeof_DW4: + pC[:] += (4 * num_mfma_n * mfma_MN * J.sizeof_bf16) loop_i = J.gpr("su32") loop_i[0] = 0 @@ -1685,6 +2087,8 @@ def loop_body(ni): if loop_cnt % 2: loop_body(0) + J.free_lds(lds) + @jit(with_debug_log=False) def moe_1stage_splitk(J:JIT, weight_dtype, diff --git a/src/contrib/moe_gemm_mxfp4.py b/src/contrib/moe_gemm_mxfp4.py index 393f0da5..89abb332 100644 --- a/src/contrib/moe_gemm_mxfp4.py +++ b/src/contrib/moe_gemm_mxfp4.py @@ -216,8 +216,14 @@ def moe_gemm_final_reduce_bf16(J, TOPK, OC, input:"void*", output:"void*", num_t J.s_min_u32(tok1, tok1[0], num_tokens_total[0]) - input[:] += tok0[0] * (TOPK * OC * J.sizeof_bf16) - output[:] += tok0[0] * (OC * J.sizeof_bf16) + offset_64bit = J.gpr(2,"su32") + J.s_mul_hi_u32(offset_64bit[1], tok0[0], (TOPK * OC * J.sizeof_bf16)) + J.s_mul_i32(offset_64bit[0], tok0[0], (TOPK * OC * J.sizeof_bf16)) + input[:] += offset_64bit + + J.s_mul_hi_u32(offset_64bit[1], tok0[0], (OC * J.sizeof_bf16)) + J.s_mul_i32(offset_64bit[0], tok0[0], (OC * J.sizeof_bf16)) + output[:] += offset_64bit buff = J.Buffer(input, (tok1[0] - tok0[0]) * (TOPK * OC * J.sizeof_bf16)) buff_out = J.Buffer(output, (tok1[0] - tok0[0]) * (OC * J.sizeof_bf16)) diff --git a/src/core/asmjit.py b/src/core/asmjit.py index 06b5872f..a1e9b970 100644 --- a/src/core/asmjit.py +++ b/src/core/asmjit.py @@ -3281,6 +3281,15 @@ def emit(generators:list, cycles:int=99999999): @cache def get_sgpr_const(self, value): + """ + if it's cached at jit-generation time inside a conditional branch, if we also initialize it here, + then the other references outside the conditional branch may got a sgpr w/o proper initialization. + a standard compiler can guarantee the initialization happens at proper places, but jit is not that smart + what we can do here is check if we are inside a branch, if so, give an error + """ + YELLOW = "\033[1;33m" + END = "\033[0m" + print(f"{YELLOW} Caution: get_sgpr_const({hex(value)}) inside branch could lead to uninitialized sgpr {END}") sgpr = self.gpr("su32", name=f"sgpr_const_{value}") sgpr[0] = value return sgpr diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index 04cf9bf9..2b9a2916 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -14,6 +14,30 @@ USE_FP4_SHUFFLE_WEIGHT = 1 +DDD = int(os.getenv("DDD", "0")) + +def is_arch_type(arch): + props = torch.cuda.get_device_properties() + return arch in props.gcnArchName + +def get_fp8type(): + return torch.float8_e4m3fn if is_arch_type('950') else torch.float8_e4m3fnuz + +def get_fp4type_if_valid(): + return torch.float4_e2m1fn_x2 if is_arch_type('950') else None + +prec_bf16 = (torch.bfloat16, aiter.QuantType.No) +prec_fp8_ptpc = (get_fp8type(), aiter.QuantType.per_Token) +prec_fp8_b = (get_fp8type(), aiter.QuantType.per_1x128) +prec_fp8_t = (get_fp8type(), aiter.QuantType.per_Tensor) +prec_mxfp4 = (get_fp4type_if_valid(), aiter.QuantType.per_1x32) + +quant2str_dict = { + aiter.QuantType.per_Token: 'per_Token', + aiter.QuantType.per_1x128: 'per_1x128', + aiter.QuantType.per_Tensor: 'per_Tensor', + aiter.QuantType.per_1x32: 'per_1x32', +} def _run_aiter(hidden_states, w1, # [expert(local_expert:EP), inter_dim*2, dim] N,K @@ -22,8 +46,9 @@ def _run_aiter(hidden_states, topk_ids, w1_scale: Optional[torch.tensor] = None, # [expert(local_expert:EP), inter_dim, 1] w2_scale: Optional[torch.tensor] = None, # [expert(local_expert:EP), model_dim, 1] - fp8_ptpc=True): + quant_type = aiter.QuantType.No): from aiter.fused_moe import fused_moe + """ from aiter import QuantType if w1.dtype == torch.float4_e2m1fn_x2: quant_type = QuantType.per_1x32 @@ -33,6 +58,7 @@ def _run_aiter(hidden_states, quant_type = QuantType.per_128x128 else: quant_type = QuantType.per_Token + """ return fused_moe( hidden_states, w1, @@ -80,618 +106,786 @@ def expert_forward(n, x): def wei_is_fp8(weight_type): return weight_type == torch.float8_e4m3fn or weight_type == torch.float8_e4m3fnuz -def _run_batch(kernel_type, B=1, weight_type=torch.bfloat16, TILE_M=16, TILE_N=32, run_count=10, HIDDEN_SIZE=2048, INTER_SIZE=1024, TOPK=8, E=128, TP=8, fp8_ptpc=True): - INTER_SIZE_TP = INTER_SIZE // TP - BUF_COPY = 32 - hidden_states = (torch.randn([BUF_COPY, B, HIDDEN_SIZE], dtype=torch.bfloat16) + 1)*0.001 - if weight_type == torch.bfloat16: - w_ = torch.randn([E, INTER_SIZE_TP * 2, HIDDEN_SIZE], dtype=weight_type) - w1_ref = w_ - w1 = [w_.clone() for _ in range(BUF_COPY)] - w_ = torch.randn([E, HIDDEN_SIZE, INTER_SIZE_TP], dtype=weight_type) - w2_ref = w_ - w2 = [w_.clone() for _ in range(BUF_COPY)] - w1_scale = [None] * BUF_COPY - w2_scale = [None] * BUF_COPY - elif weight_type == torch.float4_e2m1fn_x2: - import aiter - from aiter.utility import fp4_utils - from aiter.ops.shuffle import shuffle_weight - # w_ = torch.randn([E, INTER_SIZE_TP * 2, HIDDEN_SIZE], dtype=torch.bfloat16) - w_ = torch.randn([E, INTER_SIZE_TP * 2, HIDDEN_SIZE], dtype=torch.bfloat16) - w1_qt, w1_qt_scale_ = aiter.get_torch_quant(aiter.QuantType.per_1x32)(w_, quant_dtype=weight_type) - - #w1_qt_scale_[...] = 1.0 - w1_f32 = fp4_utils.mxfp4_to_f32(w1_qt).to(dtype=torch.bfloat16).reshape(E, INTER_SIZE_TP * 2, HIDDEN_SIZE // 32, 32) - w1_scale_f32 = fp4_utils.e8m0_to_f32(w1_qt_scale_).to(dtype=torch.bfloat16).reshape(E, INTER_SIZE_TP * 2, HIDDEN_SIZE // 32, 1) - w1_ref = (w1_f32 * w1_scale_f32).reshape(E, INTER_SIZE_TP * 2, HIDDEN_SIZE) - w1_qt_scale = fp4_utils.e8m0_shuffle(w1_qt_scale_) - if USE_FP4_SHUFFLE_WEIGHT: - w1 = [shuffle_weight(w1_qt) for _ in range(BUF_COPY)] - else: +from dataclasses import dataclass +import aiter +from aiter.ops.shuffle import shuffle_weight +from aiter.fused_moe import moe_sorting +from aiter.utility import fp4_utils +from aiter.ops.quant import pertoken_quant + +ext_topk_ids = None + +def quant_expert_weights(w1, quant_type, dtype): + if quant_type == aiter.QuantType.per_Token: + torch_quant = aiter.get_torch_quant(aiter.QuantType.per_Token) + w1_qt, w1s = torch_quant(w1, quant_dtype=dtype) + w1_ref = (w1_qt.to(dtype=w1.dtype) * w1s).to(dtype=w1.dtype) + return w1_qt, w1s, w1_ref + + if quant_type == aiter.QuantType.per_Tensor: + fmax = torch.finfo(dtype).max + w1s = w1.float().abs().amax(dim=(1,2)) / fmax + w1_qt = ( + (w1.float() / w1s.view(E, 1, 1)).clamp(-fmax, fmax).to(dtype) + ) + w1_ref = (w1_qt.to(dtype=w1.dtype) * w1s.view(E, 1, 1)).to(dtype=w1.dtype) + return w1_qt, w1s, w1_ref + assert 0, quant_type + +@dataclass +class TestCase: + # all MOE problem specification is data member + TILE_M:int + TILE_N:int + HIDDEN_SIZE:int + INTER_SIZE_TP:int + E:int + TOPK:int + DYN_SCHEDULE:bool = False # valid for 'mxn_2s' + STAGE2_TILE_N:int = 0 + run_count:int = 10 + INTER_SIZE_TP_ADJ:int = 0 + + perf:list = None + + def __call__(self, kernel_type, weight_type, quant_type, B=1, run_count=0): + run_count = self.run_count if run_count <= 0 else run_count + fp8_quant_type = quant_type + TILE_M=self.TILE_M + TILE_N=self.TILE_N + HIDDEN_SIZE = self.HIDDEN_SIZE + INTER_SIZE_TP = self.INTER_SIZE_TP + STAGE2_TILE_N=self.STAGE2_TILE_N + if STAGE2_TILE_N == 0: + STAGE2_TILE_N = TILE_N + + # adjust INTER_SIZE_TP + if kernel_type == 'aiter' and weight_type == torch.float4_e2m1fn_x2 and INTER_SIZE_TP % 128 != 0: + INTER_SIZE_TP = div_up(INTER_SIZE_TP, 128) * 128 + + if (weight_type == torch.float8_e4m3fn or weight_type == torch.float8_e4m3fnuz) and INTER_SIZE_TP % 128 != 0: + #INTER_SIZE_TP = div_up(INTER_SIZE_TP, 128) * 128 + pass + + self.INTER_SIZE_TP_ADJ = INTER_SIZE_TP + + E=self.E + TOPK=self.TOPK + global ext_topk_ids + BUF_COPY = 32 + hidden_states = (torch.randn([BUF_COPY, B, HIDDEN_SIZE], dtype=torch.bfloat16) + 1)*0.001 + if weight_type == torch.bfloat16: + w_ = torch.randn([E, INTER_SIZE_TP * 2, HIDDEN_SIZE], dtype=weight_type) + w1_ref = w_ + w1 = [w_.clone() for _ in range(BUF_COPY)] + w_ = torch.randn([E, HIDDEN_SIZE, INTER_SIZE_TP], dtype=weight_type) + w2_ref = w_ + w2 = [w_.clone() for _ in range(BUF_COPY)] + w1_scale = [None] * BUF_COPY + w2_scale = [None] * BUF_COPY + elif weight_type == torch.float4_e2m1fn_x2: + # w_ = torch.randn([E, INTER_SIZE_TP * 2, HIDDEN_SIZE], dtype=torch.bfloat16) + w_ = torch.randn([E, INTER_SIZE_TP * 2, HIDDEN_SIZE], dtype=torch.bfloat16) + w1_qt, w1_qt_scale_ = aiter.get_torch_quant(aiter.QuantType.per_1x32)(w_, quant_dtype=weight_type) + + #w1_qt_scale_[...] = 1.0 + w1_f32 = fp4_utils.mxfp4_to_f32(w1_qt).to(dtype=torch.bfloat16).reshape(E, INTER_SIZE_TP * 2, HIDDEN_SIZE // 32, 32) + w1_scale_f32 = fp4_utils.e8m0_to_f32(w1_qt_scale_).to(dtype=torch.bfloat16).reshape(E, INTER_SIZE_TP * 2, HIDDEN_SIZE // 32, 1) + w1_ref = (w1_f32 * w1_scale_f32).reshape(E, INTER_SIZE_TP * 2, HIDDEN_SIZE) + w1_qt_scale = fp4_utils.e8m0_shuffle(w1_qt_scale_) + if USE_FP4_SHUFFLE_WEIGHT: + w1 = [shuffle_weight(w1_qt) for _ in range(BUF_COPY)] + else: + w1 = [w1_qt.clone() for _ in range(BUF_COPY)] + w1_scale = [w1_qt_scale.clone() for _ in range(BUF_COPY)] + # w_ = torch.randn([E, HIDDEN_SIZE, INTER_SIZE_TP], dtype=torch.bfloat16) + w_ = torch.randn([E, HIDDEN_SIZE, INTER_SIZE_TP], dtype=torch.bfloat16) + w2_qt, w2_qt_scale_ = aiter.get_torch_quant(aiter.QuantType.per_1x32)(w_, quant_dtype=weight_type) + #w2_qt_scale_[...] = 1.0 + + w2_f32 = fp4_utils.mxfp4_to_f32(w2_qt).to(dtype=torch.bfloat16).reshape(E, HIDDEN_SIZE, INTER_SIZE_TP // 32, 32) + w2_scale_f32 = fp4_utils.e8m0_to_f32(w2_qt_scale_).to(dtype=torch.bfloat16).reshape(E, HIDDEN_SIZE, INTER_SIZE_TP // 32, 1) + w2_ref = (w2_f32 * w2_scale_f32).reshape(E, HIDDEN_SIZE, INTER_SIZE_TP) + # pad scale + w2_qt_scale_pad = torch.zeros(w2_qt_scale_.shape[0], div_up(w2_qt_scale_.shape[1], 8) * 8, dtype=w2_qt_scale_.dtype) + w2_qt_scale_pad[:, :w2_qt_scale_.shape[1]] = w2_qt_scale_ + w2_qt_scale = fp4_utils.e8m0_shuffle(w2_qt_scale_pad) + if USE_FP4_SHUFFLE_WEIGHT: + w2 = [shuffle_weight(w2_qt) for _ in range(BUF_COPY)] + else: + w2 = [w2_qt.clone() for _ in range(BUF_COPY)] + w2_scale = [w2_qt_scale.clone() for _ in range(BUF_COPY)] + elif (fp8_quant_type == aiter.QuantType.per_Token or fp8_quant_type == aiter.QuantType.per_Tensor) and wei_is_fp8(weight_type): + w_ = torch.randn([E, INTER_SIZE_TP * 2, HIDDEN_SIZE], dtype=torch.bfloat16) + w1_qt, w1_qt_scale, w1_ref = quant_expert_weights(w_, fp8_quant_type, weight_type) w1 = [w1_qt.clone() for _ in range(BUF_COPY)] - w1_scale = [w1_qt_scale.clone() for _ in range(BUF_COPY)] - # w_ = torch.randn([E, HIDDEN_SIZE, INTER_SIZE_TP], dtype=torch.bfloat16) - w_ = torch.randn([E, HIDDEN_SIZE, INTER_SIZE_TP], dtype=torch.bfloat16) - w2_qt, w2_qt_scale_ = aiter.get_torch_quant(aiter.QuantType.per_1x32)(w_, quant_dtype=weight_type) - #w2_qt_scale_[...] = 1.0 - - w2_f32 = fp4_utils.mxfp4_to_f32(w2_qt).to(dtype=torch.bfloat16).reshape(E, HIDDEN_SIZE, INTER_SIZE_TP // 32, 32) - w2_scale_f32 = fp4_utils.e8m0_to_f32(w2_qt_scale_).to(dtype=torch.bfloat16).reshape(E, HIDDEN_SIZE, INTER_SIZE_TP // 32, 1) - w2_ref = (w2_f32 * w2_scale_f32).reshape(E, HIDDEN_SIZE, INTER_SIZE_TP) - # pad scale - w2_qt_scale_pad = torch.zeros(w2_qt_scale_.shape[0], div_up(w2_qt_scale_.shape[1], 8) * 8, dtype=w2_qt_scale_.dtype) - w2_qt_scale_pad[:, :w2_qt_scale_.shape[1]] = w2_qt_scale_ - w2_qt_scale = fp4_utils.e8m0_shuffle(w2_qt_scale_pad) - if USE_FP4_SHUFFLE_WEIGHT: - w2 = [shuffle_weight(w2_qt) for _ in range(BUF_COPY)] - else: + for e in w1: e.is_shuffled = True + w1_scale = [w1_qt_scale.clone() for _ in range(BUF_COPY)] + w_ = torch.randn([E, HIDDEN_SIZE, INTER_SIZE_TP], dtype=torch.bfloat16) + w2_qt, w2_qt_scale, w2_ref = quant_expert_weights(w_, fp8_quant_type, weight_type) w2 = [w2_qt.clone() for _ in range(BUF_COPY)] - w2_scale = [w2_qt_scale.clone() for _ in range(BUF_COPY)] - elif fp8_ptpc and wei_is_fp8(weight_type): - import aiter - torch_quant = aiter.get_torch_quant(aiter.QuantType.per_Token) - w_ = torch.randn([E, INTER_SIZE_TP * 2, HIDDEN_SIZE], dtype=torch.bfloat16) - w1_qt, w1_qt_scale = torch_quant(w_, quant_dtype=weight_type) - w1_ref = (w1_qt.to(dtype=torch.bfloat16) * w1_qt_scale).to(dtype=torch.bfloat16) - w1 = [w1_qt.clone() for _ in range(BUF_COPY)] - w1_scale = [w1_qt_scale.clone() for _ in range(BUF_COPY)] - w_ = torch.randn([E, HIDDEN_SIZE, INTER_SIZE_TP], dtype=torch.bfloat16) - w2_qt, w2_qt_scale = torch_quant(w_, quant_dtype=weight_type) - w2_ref = (w2_qt.to(dtype=torch.bfloat16) * w2_qt_scale).to(dtype=torch.bfloat16) - w2 = [w2_qt.clone() for _ in range(BUF_COPY)] - w2_scale = [w2_qt_scale.clone() for _ in range(BUF_COPY)] - elif not fp8_ptpc and wei_is_fp8(weight_type): - def weight_per_128x128_quant(weight, quant_dtype): - from aiter.ops.quant import pertoken_quant - E, dim1, dim2 = weight.shape - assert dim1 % 128 == 0 and dim2 % 128 == 0, f"weight shape {weight.shape} is not aligned to 128 for per 128x128 quantization" - - weight_blocks = weight.view( - E, dim1 // 128, 128, dim2 // 128, 128 - ) # [E, num_blocks_dim1, 128, num_blocks_dim2, 128] - weight_blocks = weight_blocks.permute( - 0, 1, 3, 2, 4 - ).contiguous() # [E, num_blocks_dim1, num_blocks_dim2, 128, 128] - weight_blocks = weight_blocks.view( - E, -1, 128 * 128 - ) # [E, num_blocks, 128*128] - weight_qt, weight_scale = pertoken_quant( - weight_blocks, quant_dtype=quant_dtype - ) - weight_qt = weight_qt.view( - E, dim1 // 128, dim2 // 128, 128, 128 - ) # [E, num_blocks_dim1, num_blocks_dim2, 128, 128] - weight_qt = weight_qt.permute( - 0, 1, 3, 2, 4 - ).contiguous() # [E, num_blocks_dim1, 128, num_blocks_dim2, 128] - weight_qt = weight_qt.view(E, dim1, dim2) # [E, dim1, dim2] - weight_scale = weight_scale.view( - E, dim1 // 128, dim2 // 128 - ) # [E, num_blocks_dim1, num_blocks_dim2] - return weight_qt, weight_scale - - QUAN_BLOCK_SZ=128 - assert HIDDEN_SIZE%QUAN_BLOCK_SZ == 0 and INTER_SIZE_TP%QUAN_BLOCK_SZ==0, f"HIDDEN_SIZE and INTER_SIZE/TP must be multiples of {QUAN_BLOCK_SZ} for per block quantization" - assert QUAN_BLOCK_SZ%TILE_N == 0, f"{QUAN_BLOCK_SZ=} must be multiples of {TILE_N=} for per block quantization" - - import aiter - w_ = torch.randn([E*INTER_SIZE_TP * 2 * HIDDEN_SIZE // 128, 128], dtype=torch.bfloat16) / 2.0 - - w1_qt, w1_qt_scale = weight_per_128x128_quant(w_.view(E, INTER_SIZE_TP * 2, HIDDEN_SIZE), quant_dtype=weight_type) - # w1_qt_scale[...] = 1.0 - # print(w1_qt_scale) - - # print(f'==========={w1_qt.shape=}, {w1_qt_scale.shape=}') - w1_ref = (w1_qt.to(dtype=torch.bfloat16).view(E, INTER_SIZE_TP * 2//128, 128, HIDDEN_SIZE//128, 128) * w1_qt_scale.view((E, INTER_SIZE_TP * 2//128, 1, HIDDEN_SIZE//128, 1))).to(dtype=torch.bfloat16) - #w1_ref = w_ - w1_ref = w1_ref.view(E, INTER_SIZE_TP * 2, HIDDEN_SIZE) - w1_qt = w1_qt.view(E, INTER_SIZE_TP * 2, HIDDEN_SIZE) - w1_qt_scale = w1_qt_scale.view(E, INTER_SIZE_TP * 2//128, HIDDEN_SIZE//128) - w1 = [w1_qt.clone() for _ in range(BUF_COPY)] - w1_scale = [w1_qt_scale.clone() for _ in range(BUF_COPY)] - - w_ = torch.randn([E*HIDDEN_SIZE, INTER_SIZE_TP], dtype=torch.bfloat16) / 2.0 - w2_qt, w2_qt_scale = weight_per_128x128_quant(w_.view(E, HIDDEN_SIZE, INTER_SIZE_TP), quant_dtype=weight_type) - # w2_qt_scale[...] = 1.0 - # print(f'==========={w2_qt.shape=}, {w2_qt_scale.shape=}') - w2_ref = (w2_qt.to(dtype=torch.bfloat16).view(E, HIDDEN_SIZE//128, 128, INTER_SIZE_TP//128, 128) * w2_qt_scale.view(E, HIDDEN_SIZE//128, 1, INTER_SIZE_TP//128, 1)).to(dtype=torch.bfloat16) - w2_ref = w2_ref.view(E, HIDDEN_SIZE, INTER_SIZE_TP) - w2_qt = w2_qt.view(E, HIDDEN_SIZE, INTER_SIZE_TP) - w2_qt_scale = w2_qt_scale.view(E, HIDDEN_SIZE//128, INTER_SIZE_TP//128) - w2 = [w2_qt.clone() for _ in range(BUF_COPY)] - w2_scale = [w2_qt_scale.clone() for _ in range(BUF_COPY)] - else: - assert 0, f'not support weight type "{weight_type}"' - - topk_weight = torch.randn([BUF_COPY, B, TOPK], dtype=torch.float32) - topk_ids = torch.ones([BUF_COPY, B, TOPK], dtype=torch.int32) - # make a B*TOPK seq, which contains 0...E-1 0...E-1 - rep_e = div_up(B * TOPK, E) - topk_ids_1d = torch.ones([rep_e, E], dtype=torch.int32) - topk_ids_1d[:, ] = torch.randperm(E, dtype=torch.int32) - topk_ids[:, ] = topk_ids_1d.reshape(-1)[ : B * TOPK].reshape(B, TOPK) - access_expert = torch.unique(topk_ids[0]) - access_expert = access_expert.shape[0] - - flops = 2 * B * TOPK * (HIDDEN_SIZE * INTER_SIZE_TP * 2 + HIDDEN_SIZE * INTER_SIZE_TP) - if weight_type == torch.bfloat16: - ele_size = 2 - elif weight_type == torch.float4_e2m1fn_x2: - ele_size = 0.5 - else: - ele_size = 1 - mem_size = B * HIDDEN_SIZE * 2 + (HIDDEN_SIZE * INTER_SIZE_TP * 2 + HIDDEN_SIZE * INTER_SIZE_TP) * access_expert * ele_size - - import aiter - from aiter.ops.shuffle import shuffle_weight - from aiter.fused_moe import moe_sorting - from aiter.utility import fp4_utils - - - def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_ptpc): - B = hidden_states.shape[0] - E, N1, K1 = w1.shape - N2, K2 = w2.shape[1], w2.shape[2] - gemm1_out = torch.empty([B, TOPK, N1 // 2], dtype=hidden_states.dtype, device=hidden_states.device) - #print(topk_weight.shape, topk_weight.dtype) - #assert 0 - if kernel_type == '16x32_2s_b1': - # test moe_gemm_batch: 2 stages, BLOCK_TILE_M=16, BLOCK_TILE_N=32, batch == 1 - cur_out = torch.zeros([1, N2], dtype=hidden_states.dtype, device=hidden_states.device) - moe_gemm_batch1([N1 // 32, TOPK],[256], w1.dtype, True, hidden_states.data_ptr(), w1.data_ptr(), gemm1_out.data_ptr(), topk_ids.data_ptr(), topk_weight.data_ptr(), w1_scale.data_ptr() if w1_scale is not None else 0, 1, N1, K1) - moe_gemm_batch1([N2 // 32, TOPK],[64], w1.dtype, False, gemm1_out.data_ptr(), w2.data_ptr(), cur_out.data_ptr(), topk_ids.data_ptr(), topk_weight.data_ptr(), w2_scale.data_ptr() if w2_scale is not None else 0, 1, N2, K2) - elif kernel_type == '16x32_2s_b': - # test moe_gemm_batch: 2 stages, BLOCK_TILE_M=16, BLOCK_TILE_N=32 - sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, cur_out = moe_sorting( - topk_ids, - topk_weight, - E, - K1, # reduce dim is same with output dim - hidden_states.dtype, - 16, - None, - None, - 0, - ) - grid = sorted_expert_ids.shape[0] - if B * TOPK <= E: - grid = B * TOPK - moe_gemm_batch([N1 // 32, grid], [256], - w1.dtype, True, - hidden_states.data_ptr(), w1.data_ptr(), gemm1_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), w1_scale.data_ptr() if w1_scale is not None else 0, B, N1, K1, TOPK) - # moe_gemm_batch([N2 // 32, grid], [64], - # w1.dtype, False, - # gemm1_out.data_ptr(), w2.data_ptr(), cur_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), w2_scale.data_ptr() if w2_scale is not None else 0, B, N2, K2, TOPK) - num_CU = torch.cuda.get_device_properties().multi_processor_count - BLOCK_N = 1024 - if (w1.dtype == torch.float8_e4m3fn or w1.dtype == torch.float8_e4m3fnuz) and fp8_ptpc and N2 // BLOCK_N * grid >= num_CU: - BLOCK_TILE_SIZE_M = 16 - BLOCK_TILE_SIZE_N = 16 - assert N2 % BLOCK_N == 0 - use_atomic_write = B < 8 - NUM_STAGES = 3 - gemm2_out = cur_out - if not use_atomic_write: - gemm2_out = torch.empty([B, TOPK, HIDDEN_SIZE], dtype=torch.bfloat16) - moe_2stage_down_loopn([N2 // BLOCK_N, grid], [256], - w1.dtype, TOPK, K2, N2, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - gemm1_out.data_ptr(), w2.data_ptr(), gemm2_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), - w2_scale.data_ptr() if w2_scale is not None else 0, B, fp8_ptpc, BLOCK_N, use_atomic_write, NUM_STAGES) - if not use_atomic_write: - cur_out = torch.sum(gemm2_out, dim=1) + for e in w2: e.is_shuffled = True + w2_scale = [w2_qt_scale.clone() for _ in range(BUF_COPY)] + elif wei_is_fp8(weight_type): + def weight_per_128x128_quant(weight, quant_dtype): + E, dim1, dim2 = weight.shape + assert dim1 % 128 == 0 and dim2 % 128 == 0, f"weight shape {weight.shape} is not aligned to 128 for per 128x128 quantization" + + weight_blocks = weight.view( + E, dim1 // 128, 128, dim2 // 128, 128 + ) # [E, num_blocks_dim1, 128, num_blocks_dim2, 128] + weight_blocks = weight_blocks.permute( + 0, 1, 3, 2, 4 + ).contiguous() # [E, num_blocks_dim1, num_blocks_dim2, 128, 128] + weight_blocks = weight_blocks.view( + E, -1, 128 * 128 + ) # [E, num_blocks, 128*128] + weight_qt, weight_scale = pertoken_quant( + weight_blocks, quant_dtype=quant_dtype + ) + weight_qt = weight_qt.view( + E, dim1 // 128, dim2 // 128, 128, 128 + ) # [E, num_blocks_dim1, num_blocks_dim2, 128, 128] + weight_qt = weight_qt.permute( + 0, 1, 3, 2, 4 + ).contiguous() # [E, num_blocks_dim1, 128, num_blocks_dim2, 128] + weight_qt = weight_qt.view(E, dim1, dim2) # [E, dim1, dim2] + weight_scale = weight_scale.view( + E, dim1 // 128, dim2 // 128 + ) # [E, num_blocks_dim1, num_blocks_dim2] + return weight_qt, weight_scale + + QUAN_BLOCK_SZ=128 + assert HIDDEN_SIZE%QUAN_BLOCK_SZ == 0 and INTER_SIZE_TP%QUAN_BLOCK_SZ==0, f"HIDDEN_SIZE and INTER_SIZE/TP must be multiples of {QUAN_BLOCK_SZ} for per block quantization" + assert QUAN_BLOCK_SZ%TILE_N == 0, f"{QUAN_BLOCK_SZ=} must be multiples of {TILE_N=} for per block quantization" + + w_ = torch.randn([E*INTER_SIZE_TP * 2 * HIDDEN_SIZE // 128, 128], dtype=torch.bfloat16) / 2.0 + + w1_qt, w1_qt_scale = weight_per_128x128_quant(w_.view(E, INTER_SIZE_TP * 2, HIDDEN_SIZE), quant_dtype=weight_type) + # w1_qt_scale[...] = 1.0 + # print(w1_qt_scale) + + # print(f'==========={w1_qt.shape=}, {w1_qt_scale.shape=}') + w1_ref = (w1_qt.to(dtype=torch.bfloat16).view(E, INTER_SIZE_TP * 2//128, 128, HIDDEN_SIZE//128, 128) * w1_qt_scale.view((E, INTER_SIZE_TP * 2//128, 1, HIDDEN_SIZE//128, 1))).to(dtype=torch.bfloat16) + #w1_ref = w_ + w1_ref = w1_ref.view(E, INTER_SIZE_TP * 2, HIDDEN_SIZE) + w1_qt = w1_qt.view(E, INTER_SIZE_TP * 2, HIDDEN_SIZE) + w1_qt_scale = w1_qt_scale.view(E, INTER_SIZE_TP * 2//128, HIDDEN_SIZE//128) + w1 = [w1_qt.clone() for _ in range(BUF_COPY)] + w1_scale = [w1_qt_scale.clone() for _ in range(BUF_COPY)] + + w_ = torch.randn([E*HIDDEN_SIZE, INTER_SIZE_TP], dtype=torch.bfloat16) / 2.0 + w2_qt, w2_qt_scale = weight_per_128x128_quant(w_.view(E, HIDDEN_SIZE, INTER_SIZE_TP), quant_dtype=weight_type) + # w2_qt_scale[...] = 1.0 + # print(f'==========={w2_qt.shape=}, {w2_qt_scale.shape=}') + w2_ref = (w2_qt.to(dtype=torch.bfloat16).view(E, HIDDEN_SIZE//128, 128, INTER_SIZE_TP//128, 128) * w2_qt_scale.view(E, HIDDEN_SIZE//128, 1, INTER_SIZE_TP//128, 1)).to(dtype=torch.bfloat16) + w2_ref = w2_ref.view(E, HIDDEN_SIZE, INTER_SIZE_TP) + w2_qt = w2_qt.view(E, HIDDEN_SIZE, INTER_SIZE_TP) + w2_qt_scale = w2_qt_scale.view(E, HIDDEN_SIZE//128, INTER_SIZE_TP//128) + w2 = [w2_qt.clone() for _ in range(BUF_COPY)] + w2_scale = [w2_qt_scale.clone() for _ in range(BUF_COPY)] + else: + assert 0, f'not support weight type "{weight_type}"' + + topk_weight = torch.randn([BUF_COPY, B, TOPK], dtype=torch.float32) + topk_ids = torch.ones([BUF_COPY, B, TOPK], dtype=torch.int32) + if ext_topk_ids is not None: + topk_ids[...] = ext_topk_ids[None, :, :] + # make a B*TOPK seq, which contains 0...E-1 0...E-1 + rep_e = div_up(B * TOPK, E) + topk_ids_1d = torch.ones([rep_e, E], dtype=torch.int32) + topk_ids_1d[:, ] = torch.randperm(E, dtype=torch.int32) + topk_ids[:, ] = topk_ids_1d.reshape(-1)[ : B * TOPK].reshape(B, TOPK) + access_expert = torch.unique(topk_ids[0]) + access_expert = access_expert.shape[0] + + flops = 2 * B * TOPK * (HIDDEN_SIZE * INTER_SIZE_TP * 2 + HIDDEN_SIZE * INTER_SIZE_TP) + if weight_type == torch.bfloat16: + ele_size = 2 + elif weight_type == torch.float4_e2m1fn_x2: + ele_size = 0.5 + else: + ele_size = 1 + mem_size = B * HIDDEN_SIZE * 2 + (HIDDEN_SIZE * INTER_SIZE_TP * 2 + HIDDEN_SIZE * INTER_SIZE_TP) * access_expert * ele_size + + def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_quant_type): + fp8_ptpc = (fp8_quant_type == aiter.QuantType.per_Token) + quant_type_str = quant2str_dict.get(fp8_quant_type, 'no') + B = hidden_states.shape[0] + E, N1, K1 = w1.shape + N2, K2 = w2.shape[1], w2.shape[2] + gemm1_out = torch.empty([B, TOPK, N1 // 2], dtype=hidden_states.dtype, device=hidden_states.device) + #print(topk_weight.shape, topk_weight.dtype) + #assert 0 + if kernel_type == '16x32_2s_b1': + # test moe_gemm_batch: 2 stages, BLOCK_TILE_M=16, BLOCK_TILE_N=32, batch == 1 + cur_out = torch.zeros([1, N2], dtype=hidden_states.dtype, device=hidden_states.device) + moe_gemm_batch1([N1 // 32, TOPK],[256], w1.dtype, True, hidden_states.data_ptr(), w1.data_ptr(), gemm1_out.data_ptr(), topk_ids.data_ptr(), topk_weight.data_ptr(), w1_scale.data_ptr() if w1_scale is not None else 0, 1, N1, K1, quant_type_str) + moe_gemm_batch1([N2 // 32, TOPK],[64], w1.dtype, False, gemm1_out.data_ptr(), w2.data_ptr(), cur_out.data_ptr(), topk_ids.data_ptr(), topk_weight.data_ptr(), w2_scale.data_ptr() if w2_scale is not None else 0, 1, N2, K2, quant_type_str) + elif kernel_type == '16x32_2s_b': + # test moe_gemm_batch: 2 stages, BLOCK_TILE_M=16, BLOCK_TILE_N=32 + sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, cur_out = moe_sorting( + topk_ids, + topk_weight, + E, + K1, # reduce dim is same with output dim + hidden_states.dtype, + 16, + None, + None, + 0, + ) + grid = sorted_expert_ids.shape[0] + if B * TOPK <= E: + grid = B * TOPK + moe_gemm_batch([N1 // 32, grid], [256], + w1.dtype, True, + hidden_states.data_ptr(), w1.data_ptr(), gemm1_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), w1_scale.data_ptr() if w1_scale is not None else 0, B, N1, K1, TOPK, quant_type_str) + # moe_gemm_batch([N2 // 32, grid], [64], + # w1.dtype, False, + # gemm1_out.data_ptr(), w2.data_ptr(), cur_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), w2_scale.data_ptr() if w2_scale is not None else 0, B, N2, K2, TOPK) + num_CU = torch.cuda.get_device_properties().multi_processor_count + BLOCK_N = 1024 + if (w1.dtype == torch.float8_e4m3fn or w1.dtype == torch.float8_e4m3fnuz) and fp8_ptpc and N2 // BLOCK_N * grid >= num_CU and 32 >= B >= 16: + BLOCK_TILE_SIZE_M = 16 + BLOCK_TILE_SIZE_N = 16 + assert N2 % BLOCK_N == 0 + use_atomic_write = B < 8 + NUM_STAGES = 3 + gemm2_out = cur_out + if not use_atomic_write: + gemm2_out = torch.empty([B, TOPK, HIDDEN_SIZE], dtype=torch.bfloat16) + moe_2stage_down_loopn([N2 // BLOCK_N, grid], [256], + w1.dtype, TOPK, K2, N2, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + gemm1_out.data_ptr(), w2.data_ptr(), gemm2_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), + w2_scale.data_ptr() if w2_scale is not None else 0, B, fp8_ptpc, BLOCK_N, use_atomic_write, NUM_STAGES) + if not use_atomic_write: + cur_out = torch.sum(gemm2_out, dim=1) - else: - BLOCK_TILE_SIZE_M = 16 - BLOCK_TILE_SIZE_N = 64 + else: + BLOCK_TILE_SIZE_M = 16 + BLOCK_TILE_SIZE_N = 64 + moe_2stage_splitk([N2 // BLOCK_TILE_SIZE_N, grid], [64], + w1.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + gemm1_out.data_ptr(), w2.data_ptr(), cur_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), w2_scale.data_ptr() if w2_scale is not None else 0, B, quant_type_str) + elif kernel_type == 'mxn_splitk_2s': + # test moe_gemm_batch_vmn: 2 stages, m/n can be set + if weight_type == torch.float4_e2m1fn_x2: + K1 *= 2 + K2 *= 2 + sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, cur_out = moe_sorting( + topk_ids, + topk_weight, + E, + K1, # reduce dim is same with output dim + hidden_states.dtype, + TILE_M, + None, + None, + 0, + ) + BLOCK_TILE_SIZE_M = TILE_M + BLOCK_TILE_SIZE_N = TILE_N + grid = sorted_expert_ids.shape[0] + if B * TOPK <= E: + grid = B * TOPK + moe_2stage_splitk([N1 // BLOCK_TILE_SIZE_N, grid], [256], + w1.dtype, TOPK, K1, N1, True, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + hidden_states.data_ptr(), w1.data_ptr(), gemm1_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), w1_scale.data_ptr() if w1_scale is not None else 0, B, fp8_ptpc) moe_2stage_splitk([N2 // BLOCK_TILE_SIZE_N, grid], [64], w1.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, gemm1_out.data_ptr(), w2.data_ptr(), cur_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), w2_scale.data_ptr() if w2_scale is not None else 0, B, fp8_ptpc) - elif kernel_type == 'mxn_splitk_2s': - # test moe_gemm_batch_vmn: 2 stages, m/n can be set - if weight_type == torch.float4_e2m1fn_x2: - K1 *= 2 - K2 *= 2 - sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, cur_out = moe_sorting( - topk_ids, - topk_weight, - E, - K1, # reduce dim is same with output dim - hidden_states.dtype, - TILE_M, - None, - None, - 0, - ) - BLOCK_TILE_SIZE_M = TILE_M - BLOCK_TILE_SIZE_N = TILE_N - grid = sorted_expert_ids.shape[0] - if B * TOPK <= E: - grid = B * TOPK - moe_2stage_splitk([N1 // BLOCK_TILE_SIZE_N, grid], [256], - w1.dtype, TOPK, K1, N1, True, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - hidden_states.data_ptr(), w1.data_ptr(), gemm1_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), w1_scale.data_ptr() if w1_scale is not None else 0, B, fp8_ptpc) - moe_2stage_splitk([N2 // BLOCK_TILE_SIZE_N, grid], [64], - w1.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - gemm1_out.data_ptr(), w2.data_ptr(), cur_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), w2_scale.data_ptr() if w2_scale is not None else 0, B, fp8_ptpc) - elif kernel_type == 'mxn_splitk_1s': - # test moe_gemm_stage1 - sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, cur_out = moe_sorting( - topk_ids, - topk_weight, - E, - K1, # reduce dim is same with output dim - hidden_states.dtype, - TILE_M, - None, - None, - 0, - ) - BLOCK_TILE_SIZE_M = TILE_M - BLOCK_TILE_SIZE_N = TILE_N - moe_1stage_splitk([1, sorted_expert_ids.shape[0]], [256], - w1.dtype, TOPK, K1, N1, N2, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - hidden_states.data_ptr(), w1.data_ptr(), w1_scale.data_ptr() if w1_scale is not None else 0, w2.data_ptr(), w2_scale.data_ptr() if w2_scale is not None else 0, - cur_out.data_ptr(), - sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), B) - elif kernel_type == 'mxn_2s': - #assert weight_type == torch.bfloat16, f'mxn_2s only support bfloat16, but got {weight_type}' - # test moe_gemm_batch_vmn: 2 stages, m/n can be set - sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, cur_out = moe_sorting( - topk_ids, - topk_weight, - E, - N2, # reduce dim is same with output dim - hidden_states.dtype, - TILE_M, - None, - None, - 0, - ) - #print(f"================ {hidden_states.shape=} {hidden_states.dtype} {topk_ids.shape} {topk_weight.shape} {E} {K1}-{N2} {TILE_M} {cur_out.shape=}") - if weight_type == torch.float4_e2m1fn_x2: - # if B <= 1024: - # a1, a1_scale = fused_dynamic_mxfp4_quant_moe_sort( - # hidden_states, - # sorted_ids=sorted_ids, - # num_valid_ids=num_valid_ids, - # token_num=token_num, - # topk=1, - # block_size=block_size_M, - # ) - # else: - from aiter.utility.fp4_utils import moe_mxfp4_sort - quant_func = aiter.get_hip_quant(aiter.QuantType.per_1x32) - hidden_states_q, hidden_states_scale = quant_func( - hidden_states, - scale=None, - quant_dtype=torch.float4_e2m1fn_x2, - num_rows=None, + elif kernel_type == 'mxn_splitk_1s': + # test moe_gemm_stage1 + sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, cur_out = moe_sorting( + topk_ids, + topk_weight, + E, + K1, # reduce dim is same with output dim + hidden_states.dtype, + TILE_M, + None, + None, + 0, ) - # TODO: it seems assume using 8(x32)blocks - hidden_states_scale = moe_mxfp4_sort( - hidden_states_scale, - sorted_ids=sorted_ids, - num_valid_ids=num_valid_ids, - token_num=B, - block_size=TILE_M, + BLOCK_TILE_SIZE_M = TILE_M + BLOCK_TILE_SIZE_N = TILE_N + moe_1stage_splitk([1, sorted_expert_ids.shape[0]], [256], + w1.dtype, TOPK, K1, N1, N2, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + hidden_states.data_ptr(), w1.data_ptr(), w1_scale.data_ptr() if w1_scale is not None else 0, w2.data_ptr(), w2_scale.data_ptr() if w2_scale is not None else 0, + cur_out.data_ptr(), + sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), B) + elif kernel_type == 'mxn_2s': + #assert weight_type == torch.bfloat16, f'mxn_2s only support bfloat16, but got {weight_type}' + # test moe_gemm_batch_vmn: 2 stages, m/n can be set + sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, cur_out = moe_sorting( + topk_ids, + topk_weight, + E, + N2, # reduce dim is same with output dim + hidden_states.dtype, + TILE_M, + None, + None, + 0, ) - - # TODO: call kernel - # gemm1_out : torch.empty([B, TOPK, N1 // 2], dtype=hidden_states.dtype, device=hidden_states.device) - if 0: - if 0: - torch.save((sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, w1, w1_scale, - hidden_states_q, hidden_states_scale, gemm1_out), 'tensors_tuple2.pt') - assert 0 - moe_gemm_ref(TILE_M, TILE_N, True, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, - w1, w1_scale, - #hidden_states, None, - hidden_states_q, hidden_states_scale, - gemm1_out) - gemm1_out_q, gemm1_out_scale = quant_func( - gemm1_out.view(B*TOPK, -1), + #print(f"================ {hidden_states.shape=} {hidden_states.dtype} {topk_ids.shape} {topk_weight.shape} {E} {K1}-{N2} {TILE_M} {cur_out.shape=}") + if weight_type == torch.float4_e2m1fn_x2: + # if B <= 1024: + # a1, a1_scale = fused_dynamic_mxfp4_quant_moe_sort( + # hidden_states, + # sorted_ids=sorted_ids, + # num_valid_ids=num_valid_ids, + # token_num=token_num, + # topk=1, + # block_size=block_size_M, + # ) + # else: + from aiter.utility.fp4_utils import moe_mxfp4_sort + quant_func = aiter.get_hip_quant(aiter.QuantType.per_1x32) + hidden_states_q, hidden_states_scale = quant_func( + hidden_states, scale=None, quant_dtype=torch.float4_e2m1fn_x2, num_rows=None, ) - gemm1_out_scale = moe_mxfp4_sort( - gemm1_out_scale[: B * TOPK, :].view(B, TOPK, -1), + # TODO: it seems assume using 8(x32)blocks + hidden_states_scale = moe_mxfp4_sort( + hidden_states_scale, sorted_ids=sorted_ids, num_valid_ids=num_valid_ids, token_num=B, block_size=TILE_M, ) + + # TODO: call kernel + # gemm1_out : torch.empty([B, TOPK, N1 // 2], dtype=hidden_states.dtype, device=hidden_states.device) if 0: - torch.save((sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, w2, w2_scale, - gemm1_out_q, gemm1_out_scale, cur_out), 'tensors_tuple.pt') + if 0: + torch.save((sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, w1, w1_scale, + hidden_states_q, hidden_states_scale, gemm1_out), 'tensors_tuple2.pt') + assert 0 + moe_gemm_ref(TILE_M, TILE_N, True, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, + w1, w1_scale, + #hidden_states, None, + hidden_states_q, hidden_states_scale, + gemm1_out) + gemm1_out_q, gemm1_out_scale = quant_func( + gemm1_out.view(B*TOPK, -1), + scale=None, + quant_dtype=torch.float4_e2m1fn_x2, + num_rows=None, + ) + gemm1_out_scale = moe_mxfp4_sort( + gemm1_out_scale[: B * TOPK, :].view(B, TOPK, -1), + sorted_ids=sorted_ids, + num_valid_ids=num_valid_ids, + token_num=B, + block_size=TILE_M, + ) + if 0: + torch.save((sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, w2, w2_scale, + gemm1_out_q, gemm1_out_scale, cur_out), 'tensors_tuple.pt') + assert 0 + moe_gemm_ref(TILE_M, TILE_N, False, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, + w2, w2_scale, + #gemm1_out, None, + gemm1_out_q.view(B, TOPK, -1), gemm1_out_scale, + cur_out) + else: + gateup_OC = w1.shape[1] + assert gateup_OC % TILE_N == 0 + num_oc_blocks = gateup_OC // TILE_N + num_e_blocks = sorted_expert_ids.shape[0] + moe_gemm_mxfp4([num_oc_blocks, num_e_blocks],[256], + TILE_M, TILE_N, + w1.shape[0], w1.shape[1], w1.shape[2], + True, TOPK, # gate_up, + sorted_ids.data_ptr(), + sorted_weights.data_ptr(), + sorted_expert_ids.data_ptr(), + num_valid_ids.data_ptr(), + w1.data_ptr(), w1_scale.data_ptr(), + hidden_states_q.data_ptr(), hidden_states_scale.data_ptr(), + gemm1_out.data_ptr(), B) + + gemm1_out_q, gemm1_out_scale = quant_func( + gemm1_out.view(B*TOPK, -1), + scale=None, + quant_dtype=torch.float4_e2m1fn_x2, + num_rows=None, + ) + gemm1_out_scale = moe_mxfp4_sort( + gemm1_out_scale[: B * TOPK, :].view(B, TOPK, -1), + sorted_ids=sorted_ids, + num_valid_ids=num_valid_ids, + token_num=B, + block_size=TILE_M, + ) + + down_OC = w2.shape[1] + assert down_OC % TILE_N == 0 + num_oc_blocks = down_OC // TILE_N + num_e_blocks = sorted_expert_ids.shape[0] + gemm2_out = torch.empty(B, TOPK, N2, dtype=torch.bfloat16) + moe_gemm_mxfp4([num_oc_blocks, num_e_blocks],[256], + TILE_M, TILE_N, + w2.shape[0], w2.shape[1], w2.shape[2], + False, TOPK, # gate_up, + sorted_ids.data_ptr(), + sorted_weights.data_ptr(), + sorted_expert_ids.data_ptr(), + num_valid_ids.data_ptr(), + w2.data_ptr(), w2_scale.data_ptr(), + gemm1_out_q.data_ptr(), gemm1_out_scale.data_ptr(), + #cur_out.data_ptr(), + gemm2_out.data_ptr(), + B) + if 1: + num_WG = 256 * 2 + num_tokens_wg = B // num_WG + num_extra_tokens = B % num_WG + moe_gemm_final_reduce_bf16([num_WG], [64], TOPK, N2, + gemm2_out.data_ptr(), + cur_out.data_ptr(), + num_tokens_wg, num_extra_tokens, B) + ''' + for i_tok in range(B): + cur_out[i_tok,:] = 0 + for topk in range(TOPK): + cur_out[i_tok,:] += gemm2_out[i_tok, topk] * topk_weight[i_tok, topk] + ''' + elif weight_type == torch.bfloat16: + # moe_gemm_ref(TILE_M, TILE_N, True, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, + # w1, w1_scale, hidden_states, None, gemm1_out) + # moe_gemm_ref(TILE_M, TILE_N, False, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, + # w2, w2_scale, gemm1_out, None, cur_out) + + BLOCK_TILE_SIZE_M = TILE_M + BLOCK_TILE_SIZE_N = TILE_N + dyn_schedule = self.DYN_SCHEDULE + if dyn_schedule: + grid_gate_up = torch.cuda.get_device_properties().multi_processor_count + grid_down = torch.cuda.get_device_properties().multi_processor_count * 2 # occupancy is 2 + else: + grid_gate_up = N1 // BLOCK_TILE_SIZE_N * sorted_expert_ids.shape[0] + grid_down = sorted_expert_ids.shape[0] + + id_buf = torch.zeros(64, dtype=torch.int32) + with cudaPerf(2 * B * TOPK * HIDDEN_SIZE * INTER_SIZE_TP * 2, HIDDEN_SIZE * INTER_SIZE_TP * access_expert * ele_size * 2, name=f"up") as p: + moe_2stage_gateup([grid_gate_up], [256], + w1.dtype, TOPK, K1, N1, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, str(fp8_quant_type), + id_buf.data_ptr(), hidden_states.data_ptr(), w1.data_ptr(), gemm1_out.data_ptr(), sorted_ids.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), + None, w1_scale, B, N1 // BLOCK_TILE_SIZE_N * sorted_expert_ids.shape[0], dyn_schedule) + gemm2_out = torch.empty(B, TOPK, N2, dtype=torch.bfloat16, device=hidden_states.device) + id_buf2 = torch.zeros(64, dtype=torch.int32) + moe_2stage_down([grid_down], [256], + w1.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, STAGE2_TILE_N, str(fp8_quant_type), + id_buf2, gemm1_out, w2, gemm2_out, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, None, w2_scale, B, sorted_expert_ids.shape[0], dyn_schedule) + + num_WG = 80 * 4 + num_tokens_wg = B // num_WG + num_extra_tokens = B % num_WG + moe_gemm_final_reduce_bf16([num_WG], [64], TOPK, N2, + gemm2_out.data_ptr(), + cur_out.data_ptr(), + num_tokens_wg, num_extra_tokens, B) + elif fp8_quant_type == aiter.QuantType.per_Token or fp8_quant_type == aiter.QuantType.per_Tensor: + BLOCK_TILE_SIZE_M = TILE_M + BLOCK_TILE_SIZE_N = TILE_N + dyn_schedule = self.DYN_SCHEDULE + if dyn_schedule: + grid_gate_up = torch.cuda.get_device_properties().multi_processor_count + grid_down = torch.cuda.get_device_properties().multi_processor_count * 2 # occupancy is 2 + else: + grid_gate_up = N1 // BLOCK_TILE_SIZE_N * sorted_expert_ids.shape[0] + grid_down = sorted_expert_ids.shape[0] + # always quantize activation with aiter.QuantType.per_Token + # fp8_quant_type is only for weights + quant_func = aiter.get_hip_quant(aiter.QuantType.per_Token) + + with cudaPerf(0, 0, name=f"quant_up") as p: + hidden_states_q, hidden_states_scale = quant_func( + hidden_states, + scale=None, + quant_dtype=weight_type, + num_rows=None, + ) + if 0: + hs_s = hidden_states_scale + w1_s = w1_scale + if fp8_quant_type == aiter.QuantType.per_Tensor: + w1_s = w1_s[:,None,None].expand(-1,N1,1) + moe_2stage_gateup_ref(hidden_states_q, hs_s, + w1, w1_s, + gemm1_out, + TILE_M, + sorted_ids, sorted_expert_ids, sorted_weights, num_valid_ids, TOPK) + else: + id_buf = torch.zeros(64, dtype=torch.int32) + with cudaPerf(2 * B * TOPK * HIDDEN_SIZE * INTER_SIZE_TP * 2, HIDDEN_SIZE * INTER_SIZE_TP * access_expert * ele_size * 2, name=f"up") as p: + moe_2stage_gateup([grid_gate_up], [256], + w1.dtype, TOPK, K1, N1, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, str(fp8_quant_type), + id_buf, + hidden_states_q, w1, + gemm1_out, + sorted_ids, + sorted_expert_ids, + num_valid_ids, + hidden_states_scale, + w1_scale, B, N1 // BLOCK_TILE_SIZE_N * sorted_expert_ids.shape[0], + dyn_schedule) + # down + with cudaPerf(0, 0, name=f"quant_down") as p: + gemm1_out_q, gemm1_out_scale = quant_func( + gemm1_out.view(B * TOPK, -1), + scale=None, + quant_dtype=w2.dtype, + num_rows=None, + ) + + if DDD: + print(sorted_expert_ids) + total_wasted = 0 + for k in range(0, sorted_expert_ids.numel()): + if k*BLOCK_TILE_SIZE_M > num_valid_ids[0]: break + n_valid = ((sorted_ids[k*BLOCK_TILE_SIZE_M:(k+1)*BLOCK_TILE_SIZE_M] >> 24) < TOPK).sum().item() + if n_valid < BLOCK_TILE_SIZE_M: + print(f" expert {sorted_expert_ids[k]}: {n_valid}/{BLOCK_TILE_SIZE_M}") + total_wasted += BLOCK_TILE_SIZE_M - n_valid + else: + print(f" expert {sorted_expert_ids[k]}: ______") + print(f"Total wasted: {total_wasted/sorted_ids.numel()*100:.2f} %") assert 0 - moe_gemm_ref(TILE_M, TILE_N, False, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, - w2, w2_scale, - #gemm1_out, None, - gemm1_out_q.view(B, TOPK, -1), gemm1_out_scale, - cur_out) - else: - gateup_OC = w1.shape[1] - assert gateup_OC % TILE_N == 0 - num_oc_blocks = gateup_OC // TILE_N - num_e_blocks = sorted_expert_ids.shape[0] - moe_gemm_mxfp4([num_oc_blocks, num_e_blocks],[256], - TILE_M, TILE_N, - w1.shape[0], w1.shape[1], w1.shape[2], - True, TOPK, # gate_up, - sorted_ids.data_ptr(), - sorted_weights.data_ptr(), - sorted_expert_ids.data_ptr(), - num_valid_ids.data_ptr(), - w1.data_ptr(), w1_scale.data_ptr(), - hidden_states_q.data_ptr(), hidden_states_scale.data_ptr(), - gemm1_out.data_ptr(), B) - - gemm1_out_q, gemm1_out_scale = quant_func( - gemm1_out.view(B*TOPK, -1), - scale=None, - quant_dtype=torch.float4_e2m1fn_x2, - num_rows=None, - ) - gemm1_out_scale = moe_mxfp4_sort( - gemm1_out_scale[: B * TOPK, :].view(B, TOPK, -1), - sorted_ids=sorted_ids, - num_valid_ids=num_valid_ids, - token_num=B, - block_size=TILE_M, - ) - down_OC = w2.shape[1] - assert down_OC % TILE_N == 0 - num_oc_blocks = down_OC // TILE_N - num_e_blocks = sorted_expert_ids.shape[0] - gemm2_out = torch.empty(B, TOPK, N2, dtype=torch.bfloat16) - moe_gemm_mxfp4([num_oc_blocks, num_e_blocks],[256], - TILE_M, TILE_N, - w2.shape[0], w2.shape[1], w2.shape[2], - False, TOPK, # gate_up, - sorted_ids.data_ptr(), - sorted_weights.data_ptr(), - sorted_expert_ids.data_ptr(), - num_valid_ids.data_ptr(), - w2.data_ptr(), w2_scale.data_ptr(), - gemm1_out_q.data_ptr(), gemm1_out_scale.data_ptr(), - #cur_out.data_ptr(), - gemm2_out.data_ptr(), - B) - if 1: - num_WG = 256 * 2 - num_tokens_wg = B // num_WG - num_extra_tokens = B % num_WG - moe_gemm_final_reduce_bf16([num_WG], [64], TOPK, N2, - gemm2_out.data_ptr(), - cur_out.data_ptr(), - num_tokens_wg, num_extra_tokens, B) - ''' - for i_tok in range(B): - cur_out[i_tok,:] = 0 - for topk in range(TOPK): - cur_out[i_tok,:] += gemm2_out[i_tok, topk] * topk_weight[i_tok, topk] - ''' - else: - moe_gemm_ref(TILE_M, TILE_N, True, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, - w1, w1_scale, hidden_states, None, gemm1_out) - moe_gemm_ref(TILE_M, TILE_N, False, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, - w2, w2_scale, gemm1_out, None, cur_out) + if 0: + g1_s = gemm1_out_scale + w2_s = w2_scale + if fp8_quant_type == aiter.QuantType.per_Tensor: + w2_s = w2_s[:,None,None].expand(-1,N2,-1) + moe_2stage_down_ref(gemm1_out_q.view(B, TOPK, -1), g1_s, + w2, w2_s, + cur_out, + TILE_M, + sorted_ids, sorted_expert_ids, sorted_weights, num_valid_ids) + else: + gemm2_out = torch.empty(B, TOPK, N2, dtype=torch.bfloat16, device=gemm1_out_q.device) + down_mem_size = HIDDEN_SIZE * INTER_SIZE_TP * access_expert * ele_size + B * TOPK * INTER_SIZE_TP * 2 + B * TOPK * HIDDEN_SIZE * 2 + id_buf2 = torch.zeros(64, dtype=torch.int32) + with cudaPerf(2 * B * TOPK * HIDDEN_SIZE * INTER_SIZE_TP, down_mem_size, name=f"down") as p: + moe_2stage_down([grid_down], [256], + w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, STAGE2_TILE_N, str(fp8_quant_type), + id_buf2, gemm1_out_q, w2, + gemm2_out, #cur_out, + sorted_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + gemm1_out_scale, + w2_scale, + B, + sorted_expert_ids.shape[0], + dyn_schedule) + with cudaPerf(rw_bytes=B*(TOPK+1)*N2*2, name="reduce") as p: + if 0: + cur_out = gemm2_out.sum(dim=1) + else: + num_WG = 80 * 4 + num_tokens_wg = B // num_WG + num_extra_tokens = B % num_WG + moe_gemm_final_reduce_bf16([num_WG], [64], TOPK, N2, + gemm2_out.data_ptr(), + cur_out.data_ptr(), + num_tokens_wg, num_extra_tokens, B) - BLOCK_TILE_SIZE_M = TILE_M - BLOCK_TILE_SIZE_N = TILE_N - #moe_2stage_gateup([N1 // BLOCK_TILE_SIZE_N, sorted_expert_ids.shape[0]], [256], - # w1.dtype, TOPK, K1, N1, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - # hidden_states.data_ptr(), w1.data_ptr(), gemm1_out.data_ptr(), sorted_ids.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), w1_scale.data_ptr() if w1_scale is not None else 0, B) - #moe_2stage_down([1, sorted_expert_ids.shape[0]], [256], - # w1.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - # gemm1_out.data_ptr(), w2.data_ptr(), cur_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), w2_scale.data_ptr() if w2_scale is not None else 0, B) - - # BLOCK_TILE_SIZE_N = 64 - # moe_2stage_splitk([N2 // BLOCK_TILE_SIZE_N, sorted_expert_ids.shape[0]], [64], - # w1.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - # gemm1_out.data_ptr(), w2.data_ptr(), cur_out.data_ptr(), sorted_ids.data_ptr(), sorted_weights.data_ptr(), sorted_expert_ids.data_ptr(), num_valid_ids.data_ptr(), w2_scale.data_ptr() if w2_scale is not None else 0, B) - else: - assert 0, f'not support kernel type "{kernel_type}"' - return cur_out - - tflops_res = [] - latencies = [] - bw = [] - if kernel_type == 'aiter': - # aiter preshuffle seems doesn't help much for fp8 - # if wei_is_fp8(w1[0].dtype): - # j = 0 - # for _ in range(run_count): - # w1[j] = shuffle_weight(w1[j], layout=(16, 16)) - # w2[j] = shuffle_weight(w2[j], layout=(16, 16)) - # j = (j + 1) % BUF_COPY - i = 0 - for _ in range(run_count): - with cudaPerf(flops, mem_size, name=f"{kernel_type}[{B=},{str(weight_type).split('.')[1]}]") as p: - _run_aiter(hidden_states=hidden_states[i], w1=w1[i], w2=w2[i], topk_weight=topk_weight[i], topk_ids=topk_ids[i], w1_scale=w1_scale[i], w2_scale=w2_scale[i], fp8_ptpc=fp8_ptpc) - i = (i + 1) % BUF_COPY - tflops_res.append(p.tflops()) - latencies.append(p.dt()) - bw.append(p.bw()) - else: - if weight_type == torch.float4_e2m1fn_x2: - # fp4 no shuffle - w1_qt_aiter = w1[0] - w2_qt_aiter = w2[0] - else: - w1_qt_aiter = shuffle_weight(w1[0], layout=(16, 16)) - w2_qt_aiter = shuffle_weight(w2[0], layout=(16, 16)) - ref_out = get_torch_ref(hidden_states=hidden_states[0], w1=w1_ref, w2=w2_ref, topk_weight=topk_weight[0], topk_ids=topk_ids[0]) - # aiter_out = _run_aiter(hidden_states=hidden_states[0], w1=w1[0], w2=w2[0], topk_weight=topk_weight[0], topk_ids=topk_ids[0], w1_scale=w1_scale[0], w2_scale=w2_scale[0]) - cur_out = run(hidden_states=hidden_states[0], w1=w1_qt_aiter, w2=w2_qt_aiter, topk_weight=topk_weight[0], topk_ids=topk_ids[0], w1_scale=w1_scale[0], w2_scale=w2_scale[0], fp8_ptpc=fp8_ptpc) - #print(f">>>>>>>>>>>>>>> {calc_diff(aiter_out, ref_out)=} ") - #print(f">>>>>>>>>>>>>>> {calc_diff(cur_out, ref_out)=} ") - #print(f">>>>>>>>>>>>>>> {calc_diff(aiter_out, cur_out)=} ") - - i = 0 - for _ in range(run_count): - with cudaPerf(flops, mem_size, name=f"{kernel_type}[{B=},{str(weight_type).split('.')[1]}]") as p: - run(hidden_states=hidden_states[i], w1=w1[i], w2=w2[i], topk_weight=topk_weight[i], topk_ids=topk_ids[i], w1_scale=w1_scale[i], w2_scale=w2_scale[i], fp8_ptpc=fp8_ptpc) - i = (i + 1) % BUF_COPY - tflops_res.append(p.tflops()) - latencies.append(p.dt()) - bw.append(p.bw()) - - #print(f">>>>>>>>>>>>>>> {calc_diff(aiter_out, ref_out)=} ") - #print(f">>>>>>>>>>>>>>> {calc_diff(cur_out, ref_out)=} ") - #print(f">>>>>>>>>>>>>>> {calc_diff(aiter_out, cur_out)=} ") - - diff = calc_diff(ref_out, cur_out) - if 1 and diff > 0.02: - #if not torch.allclose(ref_out, cur_out, rtol=0.02, atol=0.02): - print(ref_out) - print(cur_out) - idx = torch.where(torch.abs(ref_out - cur_out) > 0.01) - if len(idx[0]): - print(f'idx = {idx}\nref={ref_out[idx]}\ncur={cur_out[idx]}\n{len(idx[0])}') - assert 0, f"{kernel_type=}, {B=}, {weight_type=}, {TILE_M=}, {TILE_N=}, {run_count=}" + else: + assert 0, f'not support kernel type "{kernel_type}"' + return cur_out + + tflops_res = [] + latencies = [] + bw = [] + if kernel_type == 'aiter': + # aiter preshuffle seems doesn't help much for fp8 + # if wei_is_fp8(w1[0].dtype): + # j = 0 + # for _ in range(run_count): + # w1[j] = shuffle_weight(w1[j], layout=(16, 16)) + # w2[j] = shuffle_weight(w2[j], layout=(16, 16)) + # j = (j + 1) % BUF_COPY + i = 0 + for _ in range(run_count): + with cudaPerf(flops, mem_size, name=f"{kernel_type}[{B=},{str(weight_type).split('.')[1]}]") as p: + _run_aiter(hidden_states=hidden_states[i], w1=w1[i], w2=w2[i], topk_weight=topk_weight[i], topk_ids=topk_ids[i], w1_scale=w1_scale[i], w2_scale=w2_scale[i], quant_type=quant_type) + i = (i + 1) % BUF_COPY + tflops_res.append(p.tflops()) + latencies.append(p.dt()) + bw.append(p.bw()) + diff = 0 else: - quantype="" - if wei_is_fp8(weight_type): - if fp8_ptpc: - quantype = " @ PTPC" - else: - quantype = " @ blockwise" - print(f"{kernel_type}[{B=} {weight_type=}{quantype}] acc OK") - if run_count > 0: - return {'flops': sum(tflops_res[1:])/len(tflops_res[1:]), # tflops - 'latency': sum(latencies[1:])/len(latencies[1:]) * 1e6, # us - 'bw': sum(bw[1:]) / len(bw[1:])} # GB/s - -def is_arch_type(arch): - props = torch.cuda.get_device_properties() - return arch in props.gcnArchName + if weight_type == torch.float4_e2m1fn_x2: + # fp4 no shuffle + w1_qt_aiter = w1[0] + w2_qt_aiter = w2[0] + else: + w1_qt_aiter = shuffle_weight(w1[0], layout=(16, 16)) + w2_qt_aiter = shuffle_weight(w2[0], layout=(16, 16)) + ref_out = get_torch_ref(hidden_states=hidden_states[0], w1=w1_ref, w2=w2_ref, topk_weight=topk_weight[0], topk_ids=topk_ids[0]) + # aiter_out = _run_aiter(hidden_states=hidden_states[0], w1=w1[0], w2=w2[0], topk_weight=topk_weight[0], topk_ids=topk_ids[0], w1_scale=w1_scale[0], w2_scale=w2_scale[0]) + cur_out = run(hidden_states=hidden_states[0], w1=w1_qt_aiter, w2=w2_qt_aiter, topk_weight=topk_weight[0], topk_ids=topk_ids[0], w1_scale=w1_scale[0], w2_scale=w2_scale[0], fp8_quant_type=fp8_quant_type) + #print(f">>>>>>>>>>>>>>> {calc_diff(aiter_out, ref_out)=} ") + #print(f">>>>>>>>>>>>>>> {calc_diff(cur_out, ref_out)=} ") + #print(f">>>>>>>>>>>>>>> {calc_diff(aiter_out, cur_out)=} ") + + i = 0 + for _ in range(run_count): + with cudaPerf(flops, mem_size, name=f"{kernel_type}[{B=},{str(weight_type).split('.')[1]}]") as p: + run(hidden_states=hidden_states[i], w1=w1[i], w2=w2[i], topk_weight=topk_weight[i], topk_ids=topk_ids[i], w1_scale=w1_scale[i], w2_scale=w2_scale[i], fp8_quant_type=fp8_quant_type) + i = (i + 1) % BUF_COPY + tflops_res.append(p.tflops()) + latencies.append(p.dt()) + bw.append(p.bw()) + + #print(f">>>>>>>>>>>>>>> {calc_diff(aiter_out, ref_out)=} ") + #print(f">>>>>>>>>>>>>>> {calc_diff(cur_out, ref_out)=} ") + #print(f">>>>>>>>>>>>>>> {calc_diff(aiter_out, cur_out)=} ") + + diff = calc_diff(ref_out, cur_out)#, diff_thr=0.01) + if 1 and diff > 0.02: + #if not torch.allclose(ref_out, cur_out, rtol=0.02, atol=0.02): + print(ref_out) + print(cur_out) + idx = torch.where(torch.abs(ref_out - cur_out) > 0.01) + if len(idx[0]): + print(f'idx = {idx}\nref={ref_out[idx]}\ncur={cur_out[idx]}\n{len(idx[0])}') + assert 0, f"{kernel_type=}, {B=}, {weight_type=}, {TILE_M=}, {TILE_N=}, {run_count=}" + else: + quantype="" + if wei_is_fp8(weight_type): + if fp8_quant_type == aiter.QuantType.per_Token: + quantype = "@PTPC" + elif fp8_quant_type == aiter.QuantType.per_Tensor: + quantype = "@Tensor" + else: + quantype = "@blockwise" + print(f"{kernel_type}[{B=} {weight_type=}{quantype}] acc OK err {diff=:.6f}") + if run_count > 0: + return {'flops': sum(tflops_res[1:])/len(tflops_res[1:]), # tflops + 'latency': sum(latencies[1:])/len(latencies[1:]) * 1e6, # us + 'bw': sum(bw[1:]) / len(bw[1:]), + "diff" : diff} # GB/s + + # special path for batch1 + def entry_b1(self, prec): + if self.perf is None: + self.perf = [] + kernel_type = '16x32_2s_b1' + perf = {} + perf[kernel_type] = {} + perf_prec = {} + for weight_type, quant_type in prec: + if weight_type is None: continue + perf_prec[1] = self(kernel_type, weight_type, quant_type, B=1) + perf[kernel_type][str(weight_type)] = perf_prec + self.perf.append(perf) + + def entry_common(self, kernel_type, batch, prec): + if self.perf is None: + self.perf = [] + perf = {} + perf[kernel_type] = {} + for weight_type, quant_type in prec: + if weight_type is None: continue + perf_prec = {} + for i in batch: + ret = self(kernel_type, weight_type, quant_type, i) + key = f'{i}' + if self.INTER_SIZE_TP_ADJ != self.INTER_SIZE_TP: + key += f" (adjusted INTER_SIZE_TP={self.INTER_SIZE_TP}=>{self.INTER_SIZE_TP_ADJ})" + perf_prec[key] = ret -def get_fp8type(): - return torch.float8_e4m3fn if is_arch_type('950') else torch.float8_e4m3fnuz + perf[kernel_type][str(weight_type)+"@"+str(quant_type)] = perf_prec + self.perf.append(perf) -def get_fp4type_if_valid(): - return torch.float4_e2m1fn_x2 if is_arch_type('950') else None + def show_perf(self, desc=''): + print(f'\nsummary[{desc}]:') + for perf in self.perf: + for kernel, vals in perf.items(): + for prec, vals_ in vals.items(): + for b, data in vals_.items(): + print(f'{kernel}[{prec:<4} B={b:<4}]: {data["latency"]:5.0f} us, {data["bw"]:6.1f} GB/s, {data["flops"]:4.1f} tflops, diff : {data["diff"]:.6f}') -# special path for batch1 -def entry_b1(prec=[torch.bfloat16], HIDDEN_SIZE=2048, INTER_SIZE=1024, TOPK=8, E=128, TP=8, run_count=10): - kernel_type = '16x32_2s_b1' - perf = {} - perf[kernel_type] = {} - perf_prec = {} - - for weight_type in prec: - if weight_type is None: continue - perf_prec[1] = _run_batch(kernel_type, B=1, weight_type=weight_type, run_count=run_count, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TOPK=TOPK, E=E, TP=TP) - perf[kernel_type][str(weight_type)] = perf_prec - return perf - -def entry_common(kernel_type, batch, prec=[torch.bfloat16], TILE_M=32, TILE_N=64, HIDDEN_SIZE=2048, INTER_SIZE=1024, TOPK=10, E=512, TP=8, run_count=10, fp8_ptpc = True): - perf = {} - perf[kernel_type] = {} - for weight_type in prec: - if weight_type is None: continue - perf_prec = {} - org_INTER_SIZE = INTER_SIZE - - if (weight_type == torch.float8_e4m3fn or weight_type == torch.float8_e4m3fnuz) and INTER_SIZE // TP % 128 != 0: - INTER_SIZE = div_up(INTER_SIZE // TP, 128) * 128 * TP - if kernel_type == 'aiter' and weight_type == torch.float4_e2m1fn_x2 and INTER_SIZE // TP % 128 != 0: - INTER_SIZE = div_up(INTER_SIZE // TP, 128) * 128 * TP - for i in batch: - if org_INTER_SIZE != INTER_SIZE: - key = f'{i} (adjusted INTER_SIZE={INTER_SIZE})' - else: - key = f'{i}' - - perf_prec[key] = _run_batch(kernel_type, B=i, weight_type=weight_type, TILE_M=TILE_M, TILE_N=TILE_N, run_count=run_count, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TOPK=TOPK, E=E, TP=TP, fp8_ptpc=fp8_ptpc) - quan_type="" - if wei_is_fp8(weight_type): - if fp8_ptpc: - quan_type = " @ PTPC" - else: - quan_type = " @ blockwise" - perf[kernel_type][str(weight_type)+quan_type] = perf_prec - INTER_SIZE = org_INTER_SIZE - - return perf def init_env(): torch.set_printoptions(linewidth=3000, sci_mode=False, edgeitems=8, ) torch.set_default_device('cuda') torch.manual_seed(0) + #pyhip.set_device() -def test_acc(TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE=2048, TP=8): +def test_acc(test): init_env() #entry_common('aiter', batch=[8192], prec=[torch.float4_e2m1fn_x2], TILE_M=128, TILE_N=128, run_count=2, HIDDEN_SIZE=4096, INTER_SIZE=2048, TP=8) # entry_common('mxn_splitk_2s', batch=[16], prec=[torch.float4_e2m1fn_x2], TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=0) #entry_common('mxn_2s', batch=[8192], test_fp8=False, TILE_M=128, TILE_N=128, run_count=0) #assert 0,"========================" + batch = list(range(2, 64)) # fix TILE_M=16, TILE_N=32 - entry_b1(run_count=0, prec=[torch.bfloat16, get_fp8type()], HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP) # batch 1 - entry_common('16x32_2s_b', batch=batch, prec=[torch.bfloat16, get_fp8type()], TILE_M=16, TILE_N=32, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=0) + test.entry_b1(run_count=0, prec=[prec_bf16, prec_fp8_ptpc]) # batch 1 + test.entry_common('16x32_2s_b', batch=batch, prec=[prec_bf16, prec_fp8_ptpc]) batch += list(range(128, 256)) batch += [i * 256 for i in range(1, 4)] batch += [i * 2048 for i in range(1, 5)] batch += list(range(2048 * 3, 2048 * 3 + 256)) - entry_common('mxn_splitk_2s', batch=batch, prec=[torch.bfloat16, get_fp4type_if_valid()], TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=0) + test.entry_common('mxn_splitk_2s', batch=batch, prec=[prec_bf16, prec_mxfp4]) # TILE_M/N is configurable - entry_common('mxn_splitk_2s', batch=batch, prec=[get_fp8type()], TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=0) - entry_common('mxn_splitk_2s', batch=batch, prec=[get_fp8type()], TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=0, fp8_ptpc=False) + test.entry_common('mxn_splitk_2s', batch=batch, prec=[prec_fp8_ptpc]) + test.entry_common('mxn_splitk_2s', batch=batch, prec=[prec_fp8_b]) # TODO: support fp8 - entry_common('mxn_splitk_1s', batch=batch, prec=[torch.bfloat16], TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=0) + test.entry_common('mxn_splitk_1s', batch=batch, prec=[prec_bf16]) # entry_common('mxn_2s', batch=batch, prec=[torch.bfloat16], TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=0) -def show_perf(perflist): - print('\nsummary:') - for perf in perflist: - for kernel, vals in perf.items(): - for prec, vals_ in vals.items(): - for b, data in vals_.items(): - print(f'{kernel}[{prec:<4} B={b:<4}]: {data["latency"]:5.0f} us, {data["bw"]:6.1f} GB/s, {data["flops"]:4.1f} tflops') @pytest.mark.parametrize("batch", [[1, 2, 4, 8, 12, 16, 32, 64]]) -def test_small_batch_perf(batch, HIDDEN_SIZE=4096, INTER_SIZE=2048, TP=8): +def test_small_batch_perf(batch, TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK): init_env() - perf = [] + test = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK) if is_arch_type('942'): - perf.append(entry_common('aiter', batch, prec=[torch.bfloat16, get_fp8type()], HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP)) + test.entry_common('aiter', batch, prec=[prec_bf16, prec_fp8_ptpc]) else: - perf.append(entry_common('aiter', batch, prec=[torch.bfloat16], HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP)) + test.entry_common('aiter', batch, prec=[prec_bf16]) # fix TILE_M=16, TILE_N=32 - perf.append(entry_b1(prec=[torch.bfloat16, get_fp8type()], HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP)) # batch 1 - perf.append(entry_common('16x32_2s_b', batch=batch, prec=[get_fp8type()], TILE_M=16, TILE_N=32, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP)) - show_perf(perf) + test.entry_b1(prec=[prec_bf16, prec_fp8_ptpc]) + test.entry_common('16x32_2s_b', batch=batch, prec=[prec_fp8_ptpc]) + test.show_perf() @pytest.mark.parametrize("batch", [[16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]]) -def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE=2048, TP=8): +def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE_TP=2048, E=512, TOPK=10, test_sets=['aiter', 'mxn_2s', 'mxn_splitk_2s']): init_env() - perf = [] - # perf.append(entry_common('aiter', batch, prec=[torch.bfloat16, get_fp8type(), get_fp4type_if_valid()], HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, TILE_M=TILE_M, TILE_N=TILE_N)) - # perf.append(entry_common('aiter', batch, prec=[torch.bfloat16, get_fp8type()], HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, TILE_M=TILE_M, TILE_N=TILE_N)) - # perf.append(entry_common('aiter', batch, prec=[get_fp8type()], HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, TILE_M=TILE_M, TILE_N=TILE_N, TOPK=10, E=512, fp8_ptpc=False)) + test = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK) + if 'aiter' in test_sets: + test.entry_common('aiter', batch, prec=[prec_bf16, prec_fp8_ptpc, prec_mxfp4]) + # perf.append(entry_common('aiter', batch, prec=[torch.bfloat16, get_fp8type()], HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, TILE_M=TILE_M, TILE_N=TILE_N)) + # perf.append(entry_common('aiter', batch, prec=[get_fp8type()], HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, TILE_M=TILE_M, TILE_N=TILE_N, TOPK=10, E=512, fp8_ptpc=False)) # TODO: support fp8 # perf.append(entry_common('mxn_splitk_1s', batch=batch, prec=[torch.bfloat16], TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP)) - # perf.append(entry_common('mxn_2s', batch=batch, prec=[torch.bfloat16], TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP)) - # TILE_M/N is configurable - perf.append(entry_common('mxn_splitk_2s', batch=batch, prec=[torch.bfloat16, get_fp4type_if_valid()], TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, TOPK=10, E=512, INTER_SIZE=INTER_SIZE, TP=TP)) - perf.append(entry_common('mxn_splitk_2s', batch=batch, prec=[get_fp8type()], TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TOPK=10, E=512, TP=TP, fp8_ptpc=True)) - perf.append(entry_common('mxn_splitk_2s', batch=batch, prec=[ get_fp8type()], TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TOPK=10, E=512, TP=TP, fp8_ptpc=False)) - show_perf(perf) + if 'mxn_2s' in test_sets: + test.entry_common('mxn_2s', batch=batch, prec=[ prec_fp8_ptpc]) + if 'mxn_splitk_2s' in test_sets: + # TILE_M/N is configurable + test.entry_common('mxn_splitk_2s', batch=batch, prec=[torch.bfloat16, prec_mxfp4]) + test.entry_common('mxn_splitk_2s', batch=batch, prec=[prec_fp8_ptpc]) + test.entry_common('mxn_splitk_2s', batch=batch, prec=[prec_fp8_b]) + test.show_perf() if __name__ == '__main__': TILE_M = 16 TILE_N = 64 - HIDDEN_SIZE = 4096 - INTER_SIZE = 1024 - TP = 8 + HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK = 4096, 128, 512, 10 + HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK = 4096, 192, 192, 8 init_env() @@ -702,11 +896,58 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE=2048, TP #entry_common('mxn_2s', batch=batch, prec=[torch.float4_e2m1fn_x2], TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP) #with torchPerf(): # entry_common('aiter', batch, prec=[get_fp4type_if_valid()], HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, TILE_M=TILE_M, TILE_N=TILE_N) - if 1: - test_acc(TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP) + if 0: + test_acc(TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE_TP=INTER_SIZE_TP) batch = [1, 2, 4, 8, 12, 16, 32] #batch = [1, 2, 4] - test_small_batch_perf(batch, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP) + test_small_batch_perf(batch, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE_TP=INTER_SIZE_TP) # batch = [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] # batch = [1, 2, 4, 8, 16, 32,64,128, 256] # test_perf(batch, TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP) + else: + batch = [512, 1024, 2048] + #batch = [8192,] + if 0: + ext_topk_ids = torch.load("/root/tingqli/topk_ids/topk_ids_79.pt") + print(ext_topk_ids.shape) + batch[0] = ext_topk_ids.shape[0] + + # Hunyuan + for prec in [prec_fp8_t]: + TILE_M, TILE_N = 16, 64 + batch = [2, 4, 8, 16, 32, 64, 128, 256] + test_dec = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK) + test_dec.entry_common('aiter', [1] + batch, prec=[prec]) + test_dec.entry_common('16x32_2s_b1', [1], prec=[prec]) + test_dec.entry_common('16x32_2s_b', batch, prec=[prec]) + + for TILE_M in [64, 128]: + TILE_N = 128 + batch = [512,1024,2048,4096,8192, 16384, 32768, 65536, 131072] + test_prefill = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK) + test_prefill.entry_common('aiter', batch, prec=[prec]) + test_prefill.entry_common('mxn_2s', batch, prec=[prec]) + test_dec.show_perf('hunyuan dec') + test_prefill.show_perf('hunyuan prefill') + + # Qwen3.5 + HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK = 4096, 128, 512, 10 + for prec in [prec_fp8_ptpc]: + TILE_M, TILE_N = 16, 64 + batch = [2, 4, 8, 16, 32, 64, 128, 256] + test_dec = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK) + test_dec.entry_common('aiter', [1] + batch, prec=[prec]) + test_dec.entry_common('16x32_2s_b1', [1], prec=[prec]) + test_dec.entry_common('16x32_2s_b', batch, prec=[prec]) + + for DYN in [True, False]: + GATE_TILE_N, DOWN_TILE_N = 128, 128 + if DYN: + GATE_TILE_N = 256 + for TILE_M in [64, 128]: + batch = [512,1024,2048,4096,8192, 16384, 32768, 65536, 131072] + test_prefill = TestCase(TILE_M, GATE_TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK, DYN_SCHEDULE=DYN, STAGE2_TILE_N=DOWN_TILE_N) + test_prefill.entry_common('aiter', batch, prec=[prec]) + test_prefill.entry_common('mxn_2s', batch, prec=[prec]) + test_dec.show_perf(f'qwen dec {TILE_M=} {DYN=}') + test_prefill.show_perf(f'qwen prefill {TILE_M=} {DYN=}')