From 9ca026e246ae5f1f69fc58a23b2a63d6c1c21daa Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Tue, 12 May 2026 04:08:33 +0000 Subject: [PATCH 01/24] [jit][moe] prefill fp8 placeholder version --- tests/contrib/moe/test_moe.py | 83 ++++++++++++++++++++++++----------- 1 file changed, 57 insertions(+), 26 deletions(-) diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index 04cf9bf9..20abf9a2 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -474,25 +474,49 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt 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) + 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 - #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) + 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) + elif fp8_ptpc: + BLOCK_TILE_SIZE_M = TILE_M + BLOCK_TILE_SIZE_N = TILE_N + quant_func = aiter.get_hip_quant(aiter.QuantType.per_Token) + if 1: + moe_2stage_gateup([N1 // BLOCK_TILE_SIZE_N, sorted_expert_ids.shape[0]], [256], + w1_ref.dtype, TOPK, K1, N1, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + hidden_states.data_ptr(), shuffle_weight(w1_ref, layout=(16, 16)).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) + else: + hidden_states_q, hidden_states_scale = quant_func( + hidden_states, + scale=None, + quant_dtype=weight_type, + num_rows=None, + ) + assert 0, 'add fp8 up kernel' + if 1: + moe_2stage_down([1, sorted_expert_ids.shape[0]], [256], + w1_ref.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + gemm1_out.data_ptr(), shuffle_weight(w2_ref, layout=(16, 16)).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: + gemm1_out_q, gemm1_out_scale = quant_func( + gemm1_out.view(B * TOPK, -1), + scale=None, + quant_dtype=torch.weight_type, + num_rows=None, + ) + assert 0, 'add fp8 down kernel' + else: assert 0, f'not support kernel type "{kernel_type}"' return cur_out @@ -670,20 +694,23 @@ def test_small_batch_perf(batch, HIDDEN_SIZE=4096, INTER_SIZE=2048, TP=8): show_perf(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=2048, TP=8, 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)) + if 'aiter' in test_sets: + 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)) # 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)) + if 'mxn_2s' in test_sets: + perf.append(entry_common('mxn_2s', batch=batch, prec=[torch.bfloat16, get_fp8type()], TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP)) + if 'mxn_splitk_2s' in test_sets: + # 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 __name__ == '__main__': @@ -702,7 +729,7 @@ 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: + if 0: test_acc(TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP) batch = [1, 2, 4, 8, 12, 16, 32] #batch = [1, 2, 4] @@ -710,3 +737,7 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE=2048, 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 = [8192] + 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) + test_perf(batch, TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, test_sets=['aiter', 'mxn_2s']) From 15ec3aa894f03979b2d2364d7ff59532f92da53b Mon Sep 17 00:00:00 2001 From: "Li,Tingqian" Date: Tue, 12 May 2026 09:09:56 +0000 Subject: [PATCH 02/24] [jit][moe] moe_2stage_down support fp8 ptpc --- src/contrib/moe.py | 141 +++++++++++++++++++++++++++++++--- tests/contrib/moe/test_moe.py | 47 ++++++++---- 2 files changed, 162 insertions(+), 26 deletions(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index 73aff4d4..c07d48ca 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -15,6 +15,7 @@ "moe_2stage_splitk", "moe_2stage_gateup", "moe_2stage_down", + "moe_2stage_down_ref", "moe_1stage_splitk", "moe_2stage_down_loopn", ] @@ -1429,6 +1430,69 @@ def ds_write(self, index, lds_base): J.global_store_dwordx4(voffset, vdata, p_output) + +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_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, weight_dtype, @@ -1443,10 +1507,17 @@ 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*", + pt_scale: "float*", + pc_scale:"float*", M:"int",): sizeof_w = J.sizeof(weight_dtype) dtype_A = "bf16" + is_fp8_ptpc = str(weight_dtype).startswith("torch.float8_e4m3") + assert is_fp8_ptpc or str(weight_dtype) == "torch.bfloat16", str(weight_dtype) + if is_fp8_ptpc: + fp8_ptpc = {} + else: + fp8_ptpc = None e_idx = J.blockIdx.y s_e_id = J.gpr(1, 'su32') @@ -1462,6 +1533,7 @@ def moe_2stage_down(J:JIT, 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) + pc_scale[:] += s_e_id * (N * J.sizeof_DW) # per-channel scales for weights mfma_MN = 16 mfma_K = (64//mfma_MN) * (J.sizeof_DW4//J.sizeof(dtype_A)) @@ -1483,6 +1555,15 @@ def moe_2stage_down(J:JIT, 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 + if is_fp8_ptpc: + 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 + fp8_ptpc["pc_scale"] = pc_scale + 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): @@ -1508,10 +1589,10 @@ def moe_2stage_down(J:JIT, J.s_barrier() 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) + p_weight, p_output, M, fp8_ptpc) 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 +1600,49 @@ 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): + + 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, "vbf16x2") # 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: + vaddr_pc_scale = J.gpr("vu32", (J.lane_id // 16) * J.sizeof_DW4) + scales_pc = J.gpr(2, num_mfma_n, 4, "vf32") + 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: + # 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 + J.global_load_dwordx4(scales_pc[index, n], vaddr_pc_scale, fp8_ptpc["pc_scale"], mod=f"offset:{n*16*J.sizeof_DW}") + vaddr_pc_scale[0] += num_mfma_n * 16 * J.sizeof_DW def mfma_generator(index): for k in range(num_mfma_k): @@ -1553,12 +1650,32 @@ 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] + 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] # prelog0, load Bn0 J.emit(loadB_generator(0)) diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index 20abf9a2..836362d4 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -486,8 +486,8 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt 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) + w1.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + gemm1_out, w2, cur_out, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, None, w2_scale, B) elif fp8_ptpc: BLOCK_TILE_SIZE_M = TILE_M BLOCK_TILE_SIZE_N = TILE_N @@ -495,7 +495,12 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt if 1: moe_2stage_gateup([N1 // BLOCK_TILE_SIZE_N, sorted_expert_ids.shape[0]], [256], w1_ref.dtype, TOPK, K1, N1, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - hidden_states.data_ptr(), shuffle_weight(w1_ref, layout=(16, 16)).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) + hidden_states, shuffle_weight(w1_ref, layout=(16, 16)), + gemm1_out, + sorted_ids, + sorted_expert_ids, + num_valid_ids, + w1_scale, B) else: hidden_states_q, hidden_states_scale = quant_func( hidden_states, @@ -504,18 +509,31 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt num_rows=None, ) assert 0, 'add fp8 up kernel' - if 1: - moe_2stage_down([1, sorted_expert_ids.shape[0]], [256], - w1_ref.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - gemm1_out.data_ptr(), shuffle_weight(w2_ref, layout=(16, 16)).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) + # down + gemm1_out_q, gemm1_out_scale = quant_func( + gemm1_out.view(B * TOPK, -1), + scale=None, + quant_dtype=w2.dtype, + num_rows=None, + ) + if 0: + moe_2stage_down_ref(gemm1_out_q.view(B, TOPK, -1), gemm1_out_scale, + w2, w2_scale, + cur_out, + TILE_M, + sorted_ids, sorted_expert_ids, sorted_weights, num_valid_ids) else: - gemm1_out_q, gemm1_out_scale = quant_func( - gemm1_out.view(B * TOPK, -1), - scale=None, - quant_dtype=torch.weight_type, - num_rows=None, - ) - assert 0, 'add fp8 down kernel' + moe_2stage_down([1, sorted_expert_ids.shape[0]], [256], + w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + gemm1_out_q, w2, + cur_out, + sorted_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + gemm1_out_scale, + w2_scale, + B) else: assert 0, f'not support kernel type "{kernel_type}"' @@ -647,6 +665,7 @@ 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): init_env() From 21521c1ba9e5b87131a845b9897640d1ff8a04f0 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Wed, 13 May 2026 01:36:29 +0000 Subject: [PATCH 03/24] [jit][moe] moe_2stage_gateup support fp8 ptpc --- src/contrib/common/gemm.py | 38 ++++++++++++++++-------- src/contrib/moe.py | 54 ++++++++++++++++++++++++++++------- tests/contrib/moe/test_moe.py | 14 +++++++-- 3 files changed, 80 insertions(+), 26 deletions(-) diff --git a/src/contrib/common/gemm.py b/src/contrib/common/gemm.py index 21186a06..13811d3a 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_ptpc): 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_ptpc: + assert mfma_MN == 16, 'fp8 ptpc 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_ptpc = is_fp8_ptpc + if is_fp8_ptpc: + 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=}") @@ -204,20 +216,20 @@ def swizzle(row, col): for m in range(self.wave_nCM): 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 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") @@ -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_ptpc: + 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] @@ -331,7 +345,7 @@ def mfma_generator(k): 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) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index c07d48ca..4cad9c12 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1154,7 +1154,8 @@ def moe_2stage_gateup(J:JIT, p_sorted_ids:"void*", p_sorted_expert_ids:"void*", p_num_valid_ids:"void*", - p_w_scale:"float*", + pt_scale: "float*", + pc_scale:"float*", M:"int",): class MFMA_DW4Loader: def __init__(self, J:JIT, ptr, buff_size, row_index, @@ -1296,11 +1297,15 @@ def ds_write(self, index, lds_base): BLOCK_TILE_SIZE_N_HALF = BLOCK_TILE_SIZE_N // 2 + is_fp8_ptpc = str(weight_dtype).startswith("torch.float8_e4m3") + assert is_fp8_ptpc or str(weight_dtype) == "torch.bfloat16", str(weight_dtype) + 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 @@ -1347,10 +1352,14 @@ def ds_write(self, index, lds_base): 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) + for n in range(A_vert_write): + J.global_load_dword(v_sorted_id_write[n], v_sorted_id_write_off, p_sorted_ids, mod=f'offset:{n * write_row_lanes_per_wave * sizeof_f32}') + + if is_fp8_ptpc: + v_sorted_id_scale = J.gpr(BLOCK_TILE_SIZE_M // 2 // 16, 'vu32') + v_sorted_id_scale_off = J.gpr((lane_mod_16 + J.warp_id // 2 * (BLOCK_TILE_SIZE_M // 2)) * 4) + for n in range(BLOCK_TILE_SIZE_M // 2 // 16): + J.global_load_dword(v_sorted_id_scale[n], v_sorted_id_scale_off, p_sorted_ids, mod=f'offset:{n * 16 * sizeof_f32}') p_output[:] = p_output[:] + BLOCK_TILE_SIZE_N_HALF * sizeof_bf16 * J.blockIdx.x @@ -1362,18 +1371,36 @@ def ds_write(self, index, lds_base): p_weight[:] = p_weight[:] + s_e_id * (N * K * sizeof_w) - gemm = UGEMM(J, 16, [BLOCK_TILE_SIZE_M // 2, BLOCK_TILE_SIZE_N // 2], [2, 2], K, N) + gemm = UGEMM(J, 16, [BLOCK_TILE_SIZE_M // 2, BLOCK_TILE_SIZE_N // 2], [2, 2], K, N, is_fp8_ptpc) 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, + 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_weight, N * K * sizeof_bf16, 16, - gemm.wg_N, sizeof_bf16*gemm.wg_K, stride_B, + loaderB = MFMA_DW4Loader_preshuffled(J, p_weight, N * K * sizeof_w, 16, + gemm.wg_N, sizeof_w*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) + if is_fp8_ptpc: + 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) + pc_scale[:] += s_e_id * (N * J.sizeof_DW) + BLOCK_TILE_SIZE_N_HALF * J.sizeof_DW * J.blockIdx.x + 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, 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, pc_scale, mod=f"offset:{n // 2 * 16 * J.sizeof_DW}") + # 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 @@ -1398,6 +1425,11 @@ def ds_write(self, index, lds_base): 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_ptpc: + tmp[0] *= v_pt_scales[m] + tmp[1] *= v_pt_scales[m] + tmp[0] *= scales_pc[2 * n + 0, i] + tmp[1] *= scales_pc[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") diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index 836362d4..5226097a 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -484,7 +484,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt 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) + 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) 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, w2, cur_out, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, None, w2_scale, B) @@ -492,7 +492,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt BLOCK_TILE_SIZE_M = TILE_M BLOCK_TILE_SIZE_N = TILE_N quant_func = aiter.get_hip_quant(aiter.QuantType.per_Token) - if 1: + if 0: moe_2stage_gateup([N1 // BLOCK_TILE_SIZE_N, sorted_expert_ids.shape[0]], [256], w1_ref.dtype, TOPK, K1, N1, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, hidden_states, shuffle_weight(w1_ref, layout=(16, 16)), @@ -508,7 +508,15 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt quant_dtype=weight_type, num_rows=None, ) - assert 0, 'add fp8 up kernel' + 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_q, w1, + gemm1_out, + sorted_ids, + sorted_expert_ids, + num_valid_ids, + hidden_states_scale, + w1_scale, B) # down gemm1_out_q, gemm1_out_scale = quant_func( gemm1_out.view(B * TOPK, -1), From 2ffbfdbf904b6f33b5638939fe34fe73bc069c66 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Wed, 13 May 2026 02:32:55 +0000 Subject: [PATCH 04/24] [jit][moe] add profiling info per kernel --- tests/contrib/moe/test_moe.py | 70 +++++++++++++++++------------------ 1 file changed, 33 insertions(+), 37 deletions(-) diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index 5226097a..fe044bf2 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -492,38 +492,32 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt BLOCK_TILE_SIZE_M = TILE_M BLOCK_TILE_SIZE_N = TILE_N quant_func = aiter.get_hip_quant(aiter.QuantType.per_Token) - if 0: - moe_2stage_gateup([N1 // BLOCK_TILE_SIZE_N, sorted_expert_ids.shape[0]], [256], - w1_ref.dtype, TOPK, K1, N1, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - hidden_states, shuffle_weight(w1_ref, layout=(16, 16)), - gemm1_out, - sorted_ids, - sorted_expert_ids, - num_valid_ids, - w1_scale, B) - else: + + 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, ) + 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([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_q, w1, - gemm1_out, - sorted_ids, - sorted_expert_ids, - num_valid_ids, - hidden_states_scale, - w1_scale, B) + w1.dtype, TOPK, K1, N1, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + hidden_states_q, w1, + gemm1_out, + sorted_ids, + sorted_expert_ids, + num_valid_ids, + hidden_states_scale, + w1_scale, B) # down - gemm1_out_q, gemm1_out_scale = quant_func( - gemm1_out.view(B * TOPK, -1), - scale=None, - quant_dtype=w2.dtype, - num_rows=None, - ) + 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 0: moe_2stage_down_ref(gemm1_out_q.view(B, TOPK, -1), gemm1_out_scale, w2, w2_scale, @@ -531,17 +525,18 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt TILE_M, sorted_ids, sorted_expert_ids, sorted_weights, num_valid_ids) else: - moe_2stage_down([1, sorted_expert_ids.shape[0]], [256], - w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - gemm1_out_q, w2, - cur_out, - sorted_ids, - sorted_weights, - sorted_expert_ids, - num_valid_ids, - gemm1_out_scale, - w2_scale, - B) + with cudaPerf(2 * B * TOPK * HIDDEN_SIZE * INTER_SIZE_TP, HIDDEN_SIZE * INTER_SIZE_TP * access_expert * ele_size, name=f"down") as p: + moe_2stage_down([1, sorted_expert_ids.shape[0]], [256], + w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + gemm1_out_q, w2, + cur_out, + sorted_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + gemm1_out_scale, + w2_scale, + B) else: assert 0, f'not support kernel type "{kernel_type}"' @@ -610,7 +605,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt quantype = " @ PTPC" else: quantype = " @ blockwise" - print(f"{kernel_type}[{B=} {weight_type=}{quantype}] acc OK") + 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 @@ -765,6 +760,7 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE=2048, TP # 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 = [8192] + batch = [i * 8192 for i in range(1, 11)] + #batch = [i * 6554 for i in range(1, 11)] 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) test_perf(batch, TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, test_sets=['aiter', 'mxn_2s']) From 7c18c2224665841e78043511508cff0b6f4bf568 Mon Sep 17 00:00:00 2001 From: "Li,Tingqian" Date: Wed, 13 May 2026 05:39:55 +0000 Subject: [PATCH 05/24] [jit][moe] fix accuracy bug in down_kernel --- src/contrib/moe.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index 4cad9c12..9b0619e2 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1634,7 +1634,7 @@ def down_kernel(J, mfma_MN, num_mfma_n, BM, N, K, pC:"void*", 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//sizeof_w) # each DW4 vgpr holds a mfma_K=32(bf16) or 64(fp8) @@ -1657,7 +1657,7 @@ def down_kernel(J, mfma_MN, num_mfma_n, BM, N, K, soff_b[0] = 0 if fp8_ptpc is not None: - vaddr_pc_scale = J.gpr("vu32", (J.lane_id // 16) * J.sizeof_DW4) + vaddr_pc_scale = J.gpr("vu32", (J.lane_id // 16) * J.sizeof_DW4 + J.warp_id * (num_mfma_n * 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") def loadB_generator(index): @@ -1674,7 +1674,7 @@ def loadB_generator(index): for n in range(num_mfma_n): yield 1 J.global_load_dwordx4(scales_pc[index, n], vaddr_pc_scale, fp8_ptpc["pc_scale"], mod=f"offset:{n*16*J.sizeof_DW}") - vaddr_pc_scale[0] += num_mfma_n * 16 * 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): From 1f57749700d420b913122680b6dd2d0c73d7d014 Mon Sep 17 00:00:00 2001 From: "Li,Tingqian" Date: Wed, 13 May 2026 08:13:11 +0000 Subject: [PATCH 06/24] [jit][moe] fix crash bug in moe_2stage_down --- src/contrib/moe.py | 32 +++++++++++++++++++++++++++++--- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index 9b0619e2..85375690 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1548,6 +1548,7 @@ def moe_2stage_down(J:JIT, assert is_fp8_ptpc or str(weight_dtype) == "torch.bfloat16", str(weight_dtype) if is_fp8_ptpc: fp8_ptpc = {} + dtype_A = "fp8" else: fp8_ptpc = None @@ -1600,7 +1601,7 @@ def moe_2stage_down(J:JIT, 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): @@ -1620,8 +1621,32 @@ def moe_2stage_down(J:JIT, J.s_waitcnt(mod=f"lgkmcnt(0)") J.s_barrier() - num_mfma_n = 1 if BLOCK_TILE_SIZE_M > 64 else 2 + # 0,1,2,3,4,5,6,7, num_mfma_m + if 0: + s_sorted_ids = J.gpr(num_mfma_m, 'su32') + for m in range(num_mfma_m): + J.v_readfirstlane_b32(s_sorted_ids[m], v_sorted_id[m]) + J.s_nop(8) + if 0: + with J.If((s_sorted_ids[num_mfma_m//4]>>24) == TOPK): + VALIE_TILE_SIZE_M = (num_mfma_m//4)*mfma_MN + num_mfma_n = 1 if VALIE_TILE_SIZE_M > 64 else 2 + num_mfma_n = 1 + down_kernel(J, mfma_MN, num_mfma_n, VALIE_TILE_SIZE_M, N, K, + A, lds_token_ids, lds_weights, + p_weight, p_output, M, fp8_ptpc) + J.s_endpgm() + + with J.If((s_sorted_ids[4]>>24) == TOPK): + VALIE_TILE_SIZE_M = 4*16 + num_mfma_n = 1 if VALIE_TILE_SIZE_M > 64 else 2 + num_mfma_n = 1 + down_kernel(J, mfma_MN, num_mfma_n, VALIE_TILE_SIZE_M, N, K, + A, lds_token_ids, lds_weights, + p_weight, p_output, M, 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_weight, p_output, M, fp8_ptpc) @@ -1660,6 +1685,7 @@ def down_kernel(J, mfma_MN, num_mfma_n, BM, N, K, vaddr_pc_scale = J.gpr("vu32", (J.lane_id // 16) * J.sizeof_DW4 + J.warp_id * (num_mfma_n * 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): @@ -1673,7 +1699,7 @@ def loadB_generator(index): # fp8_ptpc["pc_scale"] = pc_scale for n in range(num_mfma_n): yield 1 - J.global_load_dwordx4(scales_pc[index, n], vaddr_pc_scale, fp8_ptpc["pc_scale"], mod=f"offset:{n*16*J.sizeof_DW}") + buff_sb.load_dwordx4(scales_pc[index, n], vaddr_pc_scale, 0, offset12=n*16*J.sizeof_DW) vaddr_pc_scale[0] += num_warps * num_mfma_n * mfma_MN * J.sizeof_DW def mfma_generator(index): From 5af69089ecddd01354748a5305f6e3369a18b192 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Wed, 13 May 2026 08:15:25 +0000 Subject: [PATCH 07/24] [jit][moe] add ref for moe_2stage_gateup --- src/contrib/moe.py | 36 +++++++++++++++++++++++++++++++++++ tests/contrib/moe/test_moe.py | 27 ++++++++++++++++---------- 2 files changed, 53 insertions(+), 10 deletions(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index 85375690..c6629c82 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -14,6 +14,7 @@ "moe_gemm_batch", "moe_2stage_splitk", "moe_2stage_gateup", + "moe_2stage_gateup_ref", "moe_2stage_down", "moe_2stage_down_ref", "moe_1stage_splitk", @@ -1490,6 +1491,41 @@ def de_shuffle_weight(weight_, mfma_MN = 16): 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, diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index fe044bf2..d2937343 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -500,16 +500,23 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt quant_dtype=weight_type, num_rows=None, ) - 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([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_q, w1, - gemm1_out, - sorted_ids, - sorted_expert_ids, - num_valid_ids, - hidden_states_scale, - w1_scale, B) + if 0: + moe_2stage_gateup_ref(hidden_states_q, hidden_states_scale, + w1, w1_scale, + gemm1_out, + TILE_M, + sorted_ids, sorted_expert_ids, sorted_weights, num_valid_ids, TOPK) + else: + 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([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_q, w1, + gemm1_out, + sorted_ids, + sorted_expert_ids, + num_valid_ids, + hidden_states_scale, + w1_scale, B) # down with cudaPerf(0, 0, name=f"quant_down") as p: gemm1_out_q, gemm1_out_scale = quant_func( From ee99a84ffed462f7a717401b877b6f51e980d35f Mon Sep 17 00:00:00 2001 From: "Li,Tingqian" Date: Thu, 14 May 2026 06:51:40 +0000 Subject: [PATCH 08/24] [jit][moe] support dynamic block-size-M to moe_2stage_down --- src/contrib/moe.py | 35 +++++++++++++++-------------------- src/core/asmjit.py | 9 +++++++++ 2 files changed, 24 insertions(+), 20 deletions(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index c6629c82..d8e20e0f 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1657,32 +1657,25 @@ def moe_2stage_down(J:JIT, J.s_waitcnt(mod=f"lgkmcnt(0)") J.s_barrier() + J.get_sgpr_const(0x8000) + J.get_sgpr_const(0x3020706) # 0,1,2,3,4,5,6,7, num_mfma_m - if 0: - s_sorted_ids = J.gpr(num_mfma_m, 'su32') - for m in range(num_mfma_m): - J.v_readfirstlane_b32(s_sorted_ids[m], v_sorted_id[m]) - J.s_nop(8) - if 0: - with J.If((s_sorted_ids[num_mfma_m//4]>>24) == TOPK): - VALIE_TILE_SIZE_M = (num_mfma_m//4)*mfma_MN - num_mfma_n = 1 if VALIE_TILE_SIZE_M > 64 else 2 + if 1: + VALID_TILE_SIZES = [s for s in [32, 64] 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 - down_kernel(J, mfma_MN, num_mfma_n, VALIE_TILE_SIZE_M, N, K, + down_kernel(J, mfma_MN, num_mfma_n, TILE_SIZE, N, K, A, lds_token_ids, lds_weights, p_weight, p_output, M, fp8_ptpc) J.s_endpgm() - with J.If((s_sorted_ids[4]>>24) == TOPK): - VALIE_TILE_SIZE_M = 4*16 - num_mfma_n = 1 if VALIE_TILE_SIZE_M > 64 else 2 - num_mfma_n = 1 - down_kernel(J, mfma_MN, num_mfma_n, VALIE_TILE_SIZE_M, N, K, - A, lds_token_ids, lds_weights, - p_weight, p_output, M, fp8_ptpc) - J.s_endpgm() - num_mfma_n = 1 if BLOCK_TILE_SIZE_M > 64 else 2 + num_mfma_n = 1 # num_mfma_n has accuracy issue when using fp8, need to investigate more 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, fp8_ptpc) @@ -1720,8 +1713,8 @@ def down_kernel(J, mfma_MN, num_mfma_n, BM, N, K, if fp8_ptpc is not None: vaddr_pc_scale = J.gpr("vu32", (J.lane_id // 16) * J.sizeof_DW4 + J.warp_id * (num_mfma_n * 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) - 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): @@ -1896,6 +1889,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/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 From ff4697dc97a114c01c3057994b81c83aebfe4433 Mon Sep 17 00:00:00 2001 From: "Li,Tingqian" Date: Thu, 14 May 2026 09:21:53 +0000 Subject: [PATCH 09/24] [jit][moe] use store_dw4 + reduce instead of atomic_add --- src/contrib/moe.py | 26 ++++++++++++++------ src/contrib/moe_gemm_mxfp4.py | 10 ++++++-- tests/contrib/moe/test_moe.py | 45 +++++++++++++++++++++++++++++++---- 3 files changed, 67 insertions(+), 14 deletions(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index d8e20e0f..bbf4ecdf 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1652,16 +1652,18 @@ def moe_2stage_down(J:JIT, 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, mod=f"offset:{lds_token_ids + m*16*4}") + 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() + M[0] *= TOPK + J.get_sgpr_const(0x8000) J.get_sgpr_const(0x3020706) # 0,1,2,3,4,5,6,7, num_mfma_m if 1: - VALID_TILE_SIZES = [s for s in [32, 64] if s < BLOCK_TILE_SIZE_M] + VALID_TILE_SIZES = [s for s in [64] 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]) @@ -1777,7 +1779,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 @@ -1798,7 +1800,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}") @@ -1813,11 +1815,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})") @@ -1867,16 +1875,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") # 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) + 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 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/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index d2937343..48b68889 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -485,9 +485,17 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt 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(), None, w1_scale, B) + gemm2_out = torch.empty(B, TOPK, N2, dtype=torch.bfloat16, device=hidden_states.device) 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, w2, cur_out, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, None, w2_scale, B) + gemm1_out, w2, gemm2_out, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, None, w2_scale, B) + 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_ptpc: BLOCK_TILE_SIZE_M = TILE_M BLOCK_TILE_SIZE_N = TILE_N @@ -525,6 +533,18 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt quant_dtype=w2.dtype, num_rows=None, ) + + if 0: + print(sorted_expert_ids) + total_wasted = 0 + for k in range(0, sorted_ids.numel(), BLOCK_TILE_SIZE_M): + wasted = ((sorted_ids[k:k+BLOCK_TILE_SIZE_M] >> 24) == TOPK).sum().item() + if wasted > 0: + print(f"{wasted/BLOCK_TILE_SIZE_M*100:.2f} %") + total_wasted += wasted + print(f"Total wasted: {total_wasted/sorted_ids.numel()*100:.2f} %") + assert 0 + if 0: moe_2stage_down_ref(gemm1_out_q.view(B, TOPK, -1), gemm1_out_scale, w2, w2_scale, @@ -532,11 +552,12 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt 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) with cudaPerf(2 * B * TOPK * HIDDEN_SIZE * INTER_SIZE_TP, HIDDEN_SIZE * INTER_SIZE_TP * access_expert * ele_size, name=f"down") as p: moe_2stage_down([1, sorted_expert_ids.shape[0]], [256], w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, gemm1_out_q, w2, - cur_out, + gemm2_out, #cur_out, sorted_ids, sorted_weights, sorted_expert_ids, @@ -544,6 +565,17 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt gemm1_out_scale, w2_scale, B) + #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) else: assert 0, f'not support kernel type "{kernel_type}"' @@ -568,6 +600,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt tflops_res.append(p.tflops()) latencies.append(p.dt()) bw.append(p.bw()) + diff = 0 else: if weight_type == torch.float4_e2m1fn_x2: # fp4 no shuffle @@ -596,7 +629,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt #print(f">>>>>>>>>>>>>>> {calc_diff(cur_out, ref_out)=} ") #print(f">>>>>>>>>>>>>>> {calc_diff(aiter_out, cur_out)=} ") - diff = calc_diff(ref_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) @@ -616,7 +649,8 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt 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 + 'bw': sum(bw[1:]) / len(bw[1:]), + "diff" : diff} # GB/s def is_arch_type(arch): props = torch.cuda.get_device_properties() @@ -707,7 +741,7 @@ def show_perf(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') + 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}') @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): @@ -768,6 +802,7 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE=2048, TP # test_perf(batch, TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP) else: batch = [i * 8192 for i in range(1, 11)] + #batch = [8192] #batch = [i * 6554 for i in range(1, 11)] 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) test_perf(batch, TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, test_sets=['aiter', 'mxn_2s']) From 78ff984ebcc08ca641ece412d0122e5b7e83af0e Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Thu, 14 May 2026 09:28:43 +0000 Subject: [PATCH 10/24] [jit][moe] support xcd swizzle --- src/contrib/common/gemm.py | 14 +++---- src/contrib/moe.py | 75 ++++++++++++++++++++++++++--------- tests/contrib/moe/test_moe.py | 31 ++++++++------- 3 files changed, 80 insertions(+), 40 deletions(-) diff --git a/src/contrib/common/gemm.py b/src/contrib/common/gemm.py index 13811d3a..8c1af018 100644 --- a/src/contrib/common/gemm.py +++ b/src/contrib/common/gemm.py @@ -135,7 +135,7 @@ def __init__(self, J, wave_size:list, wave_cnt:list, K, N, - is_fp8_ptpc): + is_fp8_ptpc=False): self.K = K self.N = N @@ -206,21 +206,21 @@ 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 * 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 * 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 * self.sizeof_b) // J.sizeof_DW4 @@ -234,10 +234,10 @@ def swizzle(row, col): 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)) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index bbf4ecdf..18d4e9a3 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1140,6 +1140,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 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, @@ -1157,7 +1197,8 @@ def moe_2stage_gateup(J:JIT, p_num_valid_ids:"void*", pt_scale: "float*", pc_scale:"float*", - M:"int",): + M:"int", + num_blocks:"int",): class MFMA_DW4Loader: def __init__(self, J:JIT, ptr, buff_size, row_index, wg_M:int, row_bytes:int, stride_bytes:int, @@ -1216,7 +1257,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 @@ -1239,7 +1280,7 @@ 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 + self.prefetch_voffset = lane_id * sizeof_DWORDX4 + BLOCK_TILE_SIZE_N // 2 * stride_bytes * blockn_idx prefetch_id = 0 for wave in range(wave_cnt): with J.If(warp_id[0] == wave): @@ -1312,12 +1353,10 @@ def ds_write(self, index, lds_base): 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 + e_idx, blockn_idx = xcd_swizzle(J, J.blockIdx.x, num_blocks, N // BLOCK_TILE_SIZE_N, 8, 10) 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') @@ -1327,7 +1366,7 @@ def ds_write(self, index, lds_base): 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)) + J.debug_setup((J.warp_id[0] == 0) & (blockn_idx[0] == 0) & (e_idx[0] == 0)) # hide following initialization into s_waitcnt p_sorted_ids[:] += e_idx * (BLOCK_TILE_SIZE_M * 4) @@ -1337,15 +1376,12 @@ def ds_write(self, index, lds_base): warp_id = J.gpr("su32") J.v_readfirstlane_b32(warp_id, J.threadIdx.x[0] // 64) + assert BLOCK_TILE_SIZE_M % (4 * read_row_lanes_per_wave) == 0 + A_vert = 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) - 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 n in range(A_vert): + J.global_load_dword(v_sorted_id[n], v_sorted_id_off, p_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 @@ -1362,7 +1398,7 @@ def ds_write(self, index, lds_base): for n in range(BLOCK_TILE_SIZE_M // 2 // 16): J.global_load_dword(v_sorted_id_scale[n], v_sorted_id_scale_off, p_sorted_ids, mod=f'offset:{n * 16 * sizeof_f32}') - p_output[:] = p_output[:] + BLOCK_TILE_SIZE_N_HALF * sizeof_bf16 * J.blockIdx.x + p_output[:] = p_output[:] + BLOCK_TILE_SIZE_N_HALF * sizeof_bf16 * blockn_idx # wait for v_sorted_id J.s_waitcnt(mod=f"vmcnt(0)") @@ -1380,7 +1416,7 @@ def ds_write(self, index, lds_base): total_wave_cnt, swizzle_row_div, False) loaderB = MFMA_DW4Loader_preshuffled(J, p_weight, N * K * sizeof_w, 16, gemm.wg_N, sizeof_w*gemm.wg_K, stride_B, - total_wave_cnt, swizzle_row_div, False, N) + 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) @@ -1391,7 +1427,7 @@ def ds_write(self, index, lds_base): 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) - pc_scale[:] += s_e_id * (N * J.sizeof_DW) + BLOCK_TILE_SIZE_N_HALF * J.sizeof_DW * J.blockIdx.x + J.warp_id % 2 * BLOCK_TILE_SIZE_N_HALF // 2 * J.sizeof_DW # per-channel scales for weights + 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): @@ -1577,7 +1613,8 @@ def moe_2stage_down(J:JIT, p_num_valid_ids:"void*", # [2] value: [65536, 8192] pt_scale: "float*", pc_scale:"float*", - M:"int",): + M:"int", + num_blocks:"int",): sizeof_w = J.sizeof(weight_dtype) dtype_A = "bf16" is_fp8_ptpc = str(weight_dtype).startswith("torch.float8_e4m3") @@ -1588,7 +1625,7 @@ def moe_2stage_down(J:JIT, else: fp8_ptpc = None - e_idx = J.blockIdx.y + e_idx, _ = xcd_swizzle(J, J.blockIdx.y, num_blocks, 1, 8, 10) 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') diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index 48b68889..b5bee673 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -482,13 +482,14 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt 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], + 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(), None, w1_scale, B) + 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]) gemm2_out = torch.empty(B, TOPK, N2, dtype=torch.bfloat16, device=hidden_states.device) 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, w2, gemm2_out, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, None, w2_scale, B) + gemm1_out, w2, gemm2_out, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, None, w2_scale, B, sorted_expert_ids.shape[0]) num_WG = 80 * 4 num_tokens_wg = B // num_WG num_extra_tokens = B % num_WG @@ -516,7 +517,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt sorted_ids, sorted_expert_ids, sorted_weights, num_valid_ids, TOPK) else: 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([N1 // BLOCK_TILE_SIZE_N, sorted_expert_ids.shape[0]], [256], + 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_q, w1, gemm1_out, @@ -524,7 +525,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt sorted_expert_ids, num_valid_ids, hidden_states_scale, - w1_scale, B) + w1_scale, B, N1 // BLOCK_TILE_SIZE_N * sorted_expert_ids.shape[0]) # down with cudaPerf(0, 0, name=f"quant_down") as p: gemm1_out_q, gemm1_out_scale = quant_func( @@ -553,7 +554,8 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt 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) - with cudaPerf(2 * B * TOPK * HIDDEN_SIZE * INTER_SIZE_TP, HIDDEN_SIZE * INTER_SIZE_TP * access_expert * ele_size, name=f"down") as p: + down_mem_size = HIDDEN_SIZE * INTER_SIZE_TP * access_expert * ele_size + B * TOPK * INTER_SIZE_TP * 2 + B * TOPK * HIDDEN_SIZE * 2 + with cudaPerf(2 * B * TOPK * HIDDEN_SIZE * INTER_SIZE_TP, down_mem_size, name=f"down") as p: moe_2stage_down([1, sorted_expert_ids.shape[0]], [256], w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, gemm1_out_q, w2, @@ -564,8 +566,9 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt num_valid_ids, gemm1_out_scale, w2_scale, - B) - #with cudaPerf(rw_bytes=B*(TOPK+1)*N2*2, name="reduce") as p: + B, + sorted_expert_ids.shape[0]) + with cudaPerf(rw_bytes=B*(TOPK+1)*N2*2, name="reduce") as p: if 0: cur_out = gemm2_out.sum(dim=1) else: @@ -642,9 +645,9 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_pt quantype="" if wei_is_fp8(weight_type): if fp8_ptpc: - quantype = " @ PTPC" + quantype = "@PTPC" else: - quantype = " @ blockwise" + 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 @@ -697,9 +700,9 @@ def entry_common(kernel_type, batch, prec=[torch.bfloat16], TILE_M=32, TILE_N=64 quan_type="" if wei_is_fp8(weight_type): if fp8_ptpc: - quan_type = " @ PTPC" + quan_type = "@PTPC" else: - quan_type = " @ blockwise" + quan_type = "@blockwise" perf[kernel_type][str(weight_type)+quan_type] = perf_prec INTER_SIZE = org_INTER_SIZE @@ -768,7 +771,7 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE=2048, TP # 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)) if 'mxn_2s' in test_sets: - perf.append(entry_common('mxn_2s', batch=batch, prec=[torch.bfloat16, get_fp8type()], TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP)) + perf.append(entry_common('mxn_2s', batch=batch, prec=[torch.bfloat16, get_fp8type()], TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP)) if 'mxn_splitk_2s' in test_sets: # 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)) @@ -804,5 +807,5 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE=2048, TP batch = [i * 8192 for i in range(1, 11)] #batch = [8192] #batch = [i * 6554 for i in range(1, 11)] - 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) + entry_common('mxn_2s', batch=batch, prec=[torch.bfloat16, get_fp8type()], TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=0) test_perf(batch, TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, test_sets=['aiter', 'mxn_2s']) From ab0bba0672844e74ea9666788b25c7bf6bb8b4f0 Mon Sep 17 00:00:00 2001 From: "Li,Tingqian" Date: Fri, 15 May 2026 05:35:55 +0000 Subject: [PATCH 11/24] [jit][moe] allow TILE_M=256 in moe_2stage_down --- src/contrib/moe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index 18d4e9a3..4ab80fd8 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1700,7 +1700,7 @@ def moe_2stage_down(J:JIT, J.get_sgpr_const(0x3020706) # 0,1,2,3,4,5,6,7, num_mfma_m if 1: - VALID_TILE_SIZES = [s for s in [64] if s < BLOCK_TILE_SIZE_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]) @@ -1739,7 +1739,7 @@ def down_kernel(J, mfma_MN, num_mfma_n, BM, N, K, # 4 warps work in parallel along N dimension buff_b = J.Buffer(pB, N * K * sizeof_w) # ping-pong buffer - B = J.gpr(2, num_mfma_n, num_mfma_k, 4, "vbf16x2") # 4 x (8,bf16) or (16,fp8) + 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 From 85da1f8e2687ae8e5373954a751410c551beff24 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Fri, 15 May 2026 09:27:12 +0000 Subject: [PATCH 12/24] [jit][moe] support dynamic block-size-M to moe_2stage_gateup --- src/contrib/moe.py | 250 ++++++++++++++++++++++++--------------------- 1 file changed, 134 insertions(+), 116 deletions(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index 4ab80fd8..64b98a5b 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1361,6 +1361,12 @@ def ds_write(self, index, lds_base): 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) + + 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_sorted_ids, VALID_TILE_SIZES[n] * J.sizeof_DW) + J.s_waitcnt(mod=f"lgkmcnt(0)") # invalid padding section J.Jump("continue_following", e_idx * BLOCK_TILE_SIZE_M < max_id) @@ -1376,128 +1382,140 @@ def ds_write(self, index, lds_base): warp_id = J.gpr("su32") J.v_readfirstlane_b32(warp_id, J.threadIdx.x[0] // 64) - assert BLOCK_TILE_SIZE_M % (4 * read_row_lanes_per_wave) == 0 - A_vert = 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_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 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) - for n in range(A_vert_write): - J.global_load_dword(v_sorted_id_write[n], v_sorted_id_write_off, p_sorted_ids, mod=f'offset:{n * write_row_lanes_per_wave * sizeof_f32}') - - if is_fp8_ptpc: - v_sorted_id_scale = J.gpr(BLOCK_TILE_SIZE_M // 2 // 16, 'vu32') - v_sorted_id_scale_off = J.gpr((lane_mod_16 + J.warp_id // 2 * (BLOCK_TILE_SIZE_M // 2)) * 4) - for n in range(BLOCK_TILE_SIZE_M // 2 // 16): - J.global_load_dword(v_sorted_id_scale[n], v_sorted_id_scale_off, p_sorted_ids, mod=f'offset:{n * 16 * sizeof_f32}') - p_output[:] = p_output[:] + BLOCK_TILE_SIZE_N_HALF * sizeof_bf16 * blockn_idx - # 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_weight[:] = p_weight[:] + s_e_id * (N * K * sizeof_w) - - gemm = UGEMM(J, 16, [BLOCK_TILE_SIZE_M // 2, BLOCK_TILE_SIZE_N // 2], [2, 2], K, N, is_fp8_ptpc) - 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_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_ptpc: - 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) - 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, 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, pc_scale, mod=f"offset:{n // 2 * 16 * J.sizeof_DW}") - # vaddr_pc_scale[0] += num_mfma_n * 16 * J.sizeof_DW + 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_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_sorted_ids, mod=f'offset:{n * write_row_lanes_per_wave * sizeof_f32}') + + if is_fp8_ptpc: + 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_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_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_ptpc) + 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_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_ptpc: + 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) + 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, 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, pc_scale, mod=f"offset:{n // 2 * 16 * J.sizeof_DW}") + # 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_ptpc: + tmp[0] *= v_pt_scales[m] + tmp[1] *= v_pt_scales[m] + tmp[0] *= scales_pc[2 * n + 0, i] + tmp[1] *= scales_pc[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}") + J.s_waitcnt(mod=f"lgkmcnt(0)") + with J.ExecMask(cur_m < M[0], early_skip=False): + J.global_store_dwordx4(voffset, vdata, p_output) - s_cvt_bf16_bias = J.gpr(1, "su32") - s_cvt_bf16_bias[0] = 0x00008000 + J.free_lds(lds_out) - # 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_ptpc: - tmp[0] *= v_pt_scales[m] - tmp[1] *= v_pt_scales[m] - tmp[0] *= scales_pc[2 * n + 0, i] - tmp[1] *= scales_pc[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}") - J.s_waitcnt(mod=f"lgkmcnt(0)") - with J.ExecMask(cur_m < M[0]): - J.global_store_dwordx4(voffset, vdata, p_output) + J.get_sgpr_const(0x8000) + J.get_sgpr_const(0x3020706) + if 1: + for i, TILE_SIZE in enumerate(VALID_TILE_SIZES): + with J.If((s_sorted_id_tile[i] >> 24) == TOPK): + kernel(TILE_SIZE) + J.s_endpgm() + kernel(BLOCK_TILE_SIZE_M) def de_shuffle_weight(weight_, mfma_MN = 16): From 5cf98c154d57df898602c37da6492a63584ad2de Mon Sep 17 00:00:00 2001 From: "Li,Tingqian" Date: Fri, 15 May 2026 12:23:02 +0000 Subject: [PATCH 13/24] [jit][moe] TILE_M=64 instead of 128, finetune down --- src/contrib/moe.py | 12 ++++++------ tests/contrib/moe/test_moe.py | 12 +++++++----- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index 64b98a5b..c4fa59a9 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1725,14 +1725,13 @@ def moe_2stage_down(J:JIT, for i, TILE_SIZE in enumerate(VALID_TILE_SIZES): with J.If((s_sorted_ids[i] >> 24) == TOPK): - num_mfma_n = 1 + 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_weight, p_output, M, fp8_ptpc) J.s_endpgm() num_mfma_n = 1 if BLOCK_TILE_SIZE_M > 64 else 2 - num_mfma_n = 1 # num_mfma_n has accuracy issue when using fp8, need to investigate more 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, fp8_ptpc) @@ -1768,7 +1767,7 @@ def down_kernel(J, mfma_MN, num_mfma_n, BM, N, K, soff_b[0] = 0 if fp8_ptpc is not None: - vaddr_pc_scale = J.gpr("vu32", (J.lane_id // 16) * J.sizeof_DW4 + J.warp_id * (num_mfma_n * mfma_MN * J.sizeof_DW)) # point to the start of pc scales for this WG + 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) @@ -1785,7 +1784,7 @@ def loadB_generator(index): # 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*16*J.sizeof_DW) + 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): @@ -1816,6 +1815,7 @@ def mfma_generator(index): 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 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] @@ -1931,13 +1931,13 @@ def loop_body(ni): J.s_waitcnt(mod=f"lgkmcnt({min(15,num_loads - i - 1)})") if vmem_lane_size == J.sizeof_DW4: with J.ExecMask(voff_vmem_row[i] < M[0], early_skip=False): - J.global_store_dwordx4(voff_vmem[i], temp_c[i], "off") # this is fast: (48us) + 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]) diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index b5bee673..5b37b937 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -804,8 +804,10 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE=2048, TP # 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 = [i * 8192 for i in range(1, 11)] - #batch = [8192] - #batch = [i * 6554 for i in range(1, 11)] - entry_common('mxn_2s', batch=batch, prec=[torch.bfloat16, get_fp8type()], TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=0) - test_perf(batch, TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, test_sets=['aiter', 'mxn_2s']) + #batch = [i * 8192 for i in range(1, 11)] + batch = [i * 1024 for i in range(1, 9)] + [8192*2, 8192*4, 8192*8] + #batch = [i * 6553 for i in range(1, 11)] + #entry_common('mxn_2s', batch=batch, prec=[get_fp8type()], TILE_M=64, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=10) + test_perf(batch, TILE_M=64, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, test_sets=['aiter', 'mxn_2s']) + + From f71a5d8e34977030096eb3376c381860ff0d52ad Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Mon, 18 May 2026 01:50:46 +0000 Subject: [PATCH 14/24] [jit][moe] fix incorrect dynamic block-size-M selection for gateup --- src/contrib/moe.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index c4fa59a9..50eceac9 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1356,17 +1356,12 @@ def ds_write(self, index, lds_base): # grid: N // 32, sorted_expert_ids.shape[0] # expert index in p_sorted_expert_ids - e_idx, blockn_idx = xcd_swizzle(J, J.blockIdx.x, num_blocks, N // BLOCK_TILE_SIZE_N, 8, 10) + e_idx, blockn_idx = xcd_swizzle(J, J.blockIdx.x, num_blocks, N // BLOCK_TILE_SIZE_N, 4, 20) 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) - 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_sorted_ids, VALID_TILE_SIZES[n] * J.sizeof_DW) - J.s_waitcnt(mod=f"lgkmcnt(0)") # invalid padding section J.Jump("continue_following", e_idx * BLOCK_TILE_SIZE_M < max_id) @@ -1376,6 +1371,13 @@ def ds_write(self, index, lds_base): # hide following initialization into s_waitcnt 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_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 lane_mod_16 = get_lane_id_mod(J, 16) lane_div_16 = get_lane_id_div(J, 16) @@ -1643,7 +1645,7 @@ def moe_2stage_down(J:JIT, else: fp8_ptpc = None - e_idx, _ = xcd_swizzle(J, J.blockIdx.y, num_blocks, 1, 8, 10) + e_idx, _ = xcd_swizzle(J, J.blockIdx.y, num_blocks, 1, 4, 20) 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') From 2c03a580e403c3600484fe35337736cd0f3c18eb Mon Sep 17 00:00:00 2001 From: "Li, Tingqian" Date: Wed, 20 May 2026 06:49:27 +0000 Subject: [PATCH 15/24] [jit][moe] refactor test_moe, call ref when prec_fp8_t --- tests/contrib/moe/test_moe.py | 1344 +++++++++++++++++---------------- 1 file changed, 709 insertions(+), 635 deletions(-) diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index 5b37b937..62b84522 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -14,6 +14,7 @@ USE_FP4_SHUFFLE_WEIGHT = 1 +DDD = int(os.getenv("DDD", "0")) def _run_aiter(hidden_states, w1, # [expert(local_expert:EP), inter_dim*2, dim] N,K @@ -22,8 +23,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 +35,7 @@ def _run_aiter(hidden_states, quant_type = QuantType.per_128x128 else: quant_type = QuantType.per_Token + """ return fused_moe( hidden_states, w1, @@ -80,633 +83,697 @@ 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 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 + +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 + 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 + + # adjust INTER_SIZE_TP + 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 + 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 + + 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: + 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) + 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) + 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) - 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, 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_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') - 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 - 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(), - None, w1_scale, B, N1 // BLOCK_TILE_SIZE_N * sorted_expert_ids.shape[0]) - gemm2_out = torch.empty(B, TOPK, N2, dtype=torch.bfloat16, device=hidden_states.device) - 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, w2, gemm2_out, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, None, w2_scale, B, sorted_expert_ids.shape[0]) - 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_ptpc: - BLOCK_TILE_SIZE_M = TILE_M - BLOCK_TILE_SIZE_N = TILE_N - 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: - moe_2stage_gateup_ref(hidden_states_q, hidden_states_scale, - w1, w1_scale, - gemm1_out, - TILE_M, - sorted_ids, sorted_expert_ids, sorted_weights, num_valid_ids, TOPK) - else: - 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([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_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]) - # 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 0: - print(sorted_expert_ids) - total_wasted = 0 - for k in range(0, sorted_ids.numel(), BLOCK_TILE_SIZE_M): - wasted = ((sorted_ids[k:k+BLOCK_TILE_SIZE_M] >> 24) == TOPK).sum().item() - if wasted > 0: - print(f"{wasted/BLOCK_TILE_SIZE_M*100:.2f} %") - total_wasted += wasted - print(f"Total wasted: {total_wasted/sorted_ids.numel()*100:.2f} %") - assert 0 - - if 0: - moe_2stage_down_ref(gemm1_out_q.view(B, TOPK, -1), gemm1_out_scale, - w2, w2_scale, - 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 - with cudaPerf(2 * B * TOPK * HIDDEN_SIZE * INTER_SIZE_TP, down_mem_size, name=f"down") as p: - moe_2stage_down([1, sorted_expert_ids.shape[0]], [256], - w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - 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]) - 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 + 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 + 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(), + None, w1_scale, B, N1 // BLOCK_TILE_SIZE_N * sorted_expert_ids.shape[0]) + gemm2_out = torch.empty(B, TOPK, N2, dtype=torch.bfloat16, device=hidden_states.device) + 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, w2, gemm2_out, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, None, w2_scale, B, sorted_expert_ids.shape[0]) + 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 + quant_func = aiter.get_hip_quant(fp8_quant_type) + + 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 fp8_quant_type == aiter.QuantType.per_Tensor: + hs_s = hidden_states_scale + w1_s = w1_scale + if fp8_quant_type == aiter.QuantType.per_Tensor: + hs_s = hs_s.expand(B).view(B, 1) + 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: + 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([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_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]) + # 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 + if fp8_quant_type == aiter.QuantType.per_Tensor: + g1_s = gemm1_out_scale + w2_s = w2_scale + if fp8_quant_type == aiter.QuantType.per_Tensor: + g1_s = g1_s.expand(B*TOPK).view(B*TOPK, 1) + 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 + with cudaPerf(2 * B * TOPK * HIDDEN_SIZE * INTER_SIZE_TP, down_mem_size, name=f"down") as p: + moe_2stage_down([1, sorted_expert_ids.shape[0]], [256], + w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + 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]) + 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) + + 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: - 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()) - diff = 0 - 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)#, 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_ptpc: - quantype = "@PTPC" - 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 + 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 is_arch_type(arch): - props = torch.cuda.get_device_properties() - return arch in props.gcnArchName + perf[kernel_type][str(weight_type)+"@"+str(quant_type)] = perf_prec + self.perf.append(perf) -def get_fp8type(): - return torch.float8_e4m3fn if is_arch_type('950') else torch.float8_e4m3fnuz + def show_perf(self): + print('\nsummary:') + 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}') -def get_fp4type_if_valid(): - return torch.float4_e2m1fn_x2 if is_arch_type('950') else None -# 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 +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) def init_env(): torch.set_printoptions(linewidth=3000, sci_mode=False, edgeitems=8, ) @@ -714,77 +781,70 @@ def init_env(): 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, diff : {data["diff"]:.6f}') @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, test_sets=['aiter', 'mxn_2s', 'mxn_splitk_2s']): +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 = [] + test = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK) if 'aiter' in test_sets: - 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)) + 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)) if 'mxn_2s' in test_sets: - perf.append(entry_common('mxn_2s', batch=batch, prec=[torch.bfloat16, get_fp8type()], TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP)) + test.entry_common('mxn_2s', batch=batch, prec=[ prec_fp8_ptpc]) if 'mxn_splitk_2s' in test_sets: # 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) + 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() @@ -796,18 +856,32 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE=2048, 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 0: - test_acc(TILE_M=TILE_M, TILE_N=TILE_N, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP) + 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 = [i * 8192 for i in range(1, 11)] - batch = [i * 1024 for i in range(1, 9)] + [8192*2, 8192*4, 8192*8] - #batch = [i * 6553 for i in range(1, 11)] - #entry_common('mxn_2s', batch=batch, prec=[get_fp8type()], TILE_M=64, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=10) - test_perf(batch, TILE_M=64, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, test_sets=['aiter', 'mxn_2s']) + #batch = [i * 1024 for i in range(1, 9)] + [8192*2, 8192*4, 8192*8] + batch = [i * 128 for i in range(1, 9)] + batch = [64,128,256,512,1024,2048,4096,8192, 16384, 32768, 65536, 131072] + batch = [8192, 16384, 32768, 65536, 131072] + batch = [1024, 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] + + TILE_M, TILE_N = 128, 128 + HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK = 4096, 192, 192, 8 + test = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK) + test.entry_common('aiter', batch, prec=[prec_fp8_t]) + test.entry_common('mxn_2s', batch, prec=[prec_fp8_t]) + test.show_perf() + #entry_common('mxn_2s', batch=batch, prec=[get_fp8type()], E=E, TOPK=TOPK, TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=10) + #test_perf(batch, TILE_M=128, TILE_N=256, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE_TP=INTER_SIZE_TP, E=E, TOPK=TOPK, test_sets=['aiter', 'mxn_2s']) From c3006fda87682f2deba5d79cac5f5dabf6285d20 Mon Sep 17 00:00:00 2001 From: "Li, Tingqian" Date: Wed, 20 May 2026 08:20:50 +0000 Subject: [PATCH 16/24] [jit][moe] support prec_fp8_t in prefill --- src/contrib/common/gemm.py | 12 ++--- src/contrib/moe.py | 96 ++++++++++++++++++++++++----------- tests/contrib/moe/test_moe.py | 32 ++++++------ 3 files changed, 88 insertions(+), 52 deletions(-) diff --git a/src/contrib/common/gemm.py b/src/contrib/common/gemm.py index 8c1af018..580f2025 100644 --- a/src/contrib/common/gemm.py +++ b/src/contrib/common/gemm.py @@ -135,7 +135,7 @@ def __init__(self, J, wave_size:list, wave_cnt:list, K, N, - is_fp8_ptpc=False): + is_fp8=False): self.K = K self.N = N @@ -145,8 +145,8 @@ 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 - if is_fp8_ptpc: - assert mfma_MN == 16, 'fp8 ptpc only support 16x16x16' + 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) @@ -176,8 +176,8 @@ def __init__(self, J, self.wave_nCM = wave_nCM self.wave_nCN = wave_nCN self.wave_nCK = wave_nCK - self.is_fp8_ptpc = is_fp8_ptpc - if is_fp8_ptpc: + self.is_fp8 = is_fp8 + if is_fp8: self.sizeof_a = 1 self.sizeof_b = 1 else: @@ -309,7 +309,7 @@ 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_ptpc: + 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] diff --git a/src/contrib/moe.py b/src/contrib/moe.py index 50eceac9..305413cb 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1188,6 +1188,7 @@ 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_input:"void*", p_weight:"void*", p_output:"void*", @@ -1339,8 +1340,9 @@ def ds_write(self, index, lds_base): BLOCK_TILE_SIZE_N_HALF = BLOCK_TILE_SIZE_N // 2 - is_fp8_ptpc = str(weight_dtype).startswith("torch.float8_e4m3") - assert is_fp8_ptpc or str(weight_dtype) == "torch.bfloat16", str(weight_dtype) + 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 @@ -1403,7 +1405,7 @@ def kernel(CUR_BLOCK_TILE_SIZE_M): for n in range(A_vert_write): J.global_load_dword(v_sorted_id_write[n], v_sorted_id_write_off, p_sorted_ids, mod=f'offset:{n * write_row_lanes_per_wave * sizeof_f32}') - if is_fp8_ptpc: + 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): @@ -1417,7 +1419,7 @@ def kernel(CUR_BLOCK_TILE_SIZE_M): p_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_ptpc) + 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, @@ -1430,20 +1432,29 @@ def kernel(CUR_BLOCK_TILE_SIZE_M): # C_reg: [wave_nCM, wave_nCN] C_reg = gemm.run(loaderA, loaderB, None, M, 0, 0) - if is_fp8_ptpc: + 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) - 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, 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, pc_scale, mod=f"offset:{n // 2 * 16 * J.sizeof_DW}") + + if quant_type_w == "QuantType.per_Token": + 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, 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, 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('vu32', 0) + pc_scale[:] += 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)") @@ -1471,11 +1482,15 @@ def kernel(CUR_BLOCK_TILE_SIZE_M): 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_ptpc: + if is_fp8: tmp[0] *= v_pt_scales[m] tmp[1] *= v_pt_scales[m] - tmp[0] *= scales_pc[2 * n + 0, i] - tmp[1] *= scales_pc[2 * n + 1, i] + 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") @@ -1624,6 +1639,7 @@ def moe_2stage_down(J:JIT, with_silu, # BLOCK_TILE_SIZE_M, # 32 BLOCK_TILE_SIZE_N, # 64 + quant_type_w, p_input:"void*", # [8192, 8, 128] p_weight:"void*", # [128, 2048, 128] p_output:"void*", # [8192, 2048] @@ -1637,10 +1653,10 @@ def moe_2stage_down(J:JIT, num_blocks:"int",): sizeof_w = J.sizeof(weight_dtype) dtype_A = "bf16" - is_fp8_ptpc = str(weight_dtype).startswith("torch.float8_e4m3") - assert is_fp8_ptpc or str(weight_dtype) == "torch.bfloat16", str(weight_dtype) - if is_fp8_ptpc: - fp8_ptpc = {} + 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 @@ -1659,7 +1675,10 @@ def moe_2stage_down(J:JIT, 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) - pc_scale[:] += s_e_id * (N * J.sizeof_DW) # per-channel scales for weights + if quant_type_w == "QuantType.per_Tensor": + pc_scale[:] += s_e_id * (1 * J.sizeof_DW) # per-tensor scales for weights + else: + pc_scale[:] += s_e_id * (N * J.sizeof_DW) # per-channel scales for weights mfma_MN = 16 mfma_K = (64//mfma_MN) * (J.sizeof_DW4//J.sizeof(dtype_A)) @@ -1681,14 +1700,22 @@ def moe_2stage_down(J:JIT, 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 - if is_fp8_ptpc: + 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 - fp8_ptpc["pc_scale"] = pc_scale + + if quant_type_w == "QuantType.per_Token": + fp8_ptpc["pc_scale"] = 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, 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): @@ -1769,9 +1796,10 @@ def down_kernel(J, mfma_MN, num_mfma_n, BM, N, K, soff_b[0] = 0 if fp8_ptpc is not None: - 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) + 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): @@ -1780,7 +1808,7 @@ def loadB_generator(index): buff_b.load_dwordx4(B[index,n,k], voff_b, soff_b) 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: + 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 @@ -1818,10 +1846,16 @@ def mfma_generator(index): C[index,m,n,2] *= fp8_ptpc["v_pt_scales"][m] C[index,m,n,3] *= fp8_ptpc["v_pt_scales"][m] yield 16 - 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_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)) diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index 62b84522..f2801b67 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -141,11 +141,13 @@ def __call__(self, kernel_type, weight_type, quant_type, B=1, run_count=0): INTER_SIZE_TP = self.INTER_SIZE_TP # adjust INTER_SIZE_TP - 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 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 @@ -540,7 +542,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu 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, + w1.dtype, TOPK, K1, N1, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, str(fp8_quant_type), 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]) gemm2_out = torch.empty(B, TOPK, N2, dtype=torch.bfloat16, device=hidden_states.device) @@ -557,7 +559,9 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu 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 - quant_func = aiter.get_hip_quant(fp8_quant_type) + # 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( @@ -566,13 +570,11 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu quant_dtype=weight_type, num_rows=None, ) - if fp8_quant_type == aiter.QuantType.per_Tensor: + if 0: hs_s = hidden_states_scale w1_s = w1_scale if fp8_quant_type == aiter.QuantType.per_Tensor: - hs_s = hs_s.expand(B).view(B, 1) w1_s = w1_s[:,None,None].expand(-1,N1,1) - moe_2stage_gateup_ref(hidden_states_q, hs_s, w1, w1_s, gemm1_out, @@ -581,7 +583,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu else: 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([N1 // BLOCK_TILE_SIZE_N * sorted_expert_ids.shape[0]], [256], - w1.dtype, TOPK, K1, N1, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + w1.dtype, TOPK, K1, N1, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, str(fp8_quant_type), hidden_states_q, w1, gemm1_out, sorted_ids, @@ -612,11 +614,10 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu print(f"Total wasted: {total_wasted/sorted_ids.numel()*100:.2f} %") assert 0 - if fp8_quant_type == aiter.QuantType.per_Tensor: + if 0: g1_s = gemm1_out_scale w2_s = w2_scale if fp8_quant_type == aiter.QuantType.per_Tensor: - g1_s = g1_s.expand(B*TOPK).view(B*TOPK, 1) 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, @@ -628,7 +629,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu down_mem_size = HIDDEN_SIZE * INTER_SIZE_TP * access_expert * ele_size + B * TOPK * INTER_SIZE_TP * 2 + B * TOPK * HIDDEN_SIZE * 2 with cudaPerf(2 * B * TOPK * HIDDEN_SIZE * INTER_SIZE_TP, down_mem_size, name=f"down") as p: moe_2stage_down([1, sorted_expert_ids.shape[0]], [256], - w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, str(fp8_quant_type), gemm1_out_q, w2, gemm2_out, #cur_out, sorted_ids, @@ -868,8 +869,9 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE_TP=2048, #batch = [i * 1024 for i in range(1, 9)] + [8192*2, 8192*4, 8192*8] batch = [i * 128 for i in range(1, 9)] batch = [64,128,256,512,1024,2048,4096,8192, 16384, 32768, 65536, 131072] - batch = [8192, 16384, 32768, 65536, 131072] - batch = [1024, 8192,] + #batch = [8192, 16384, 32768, 65536, 131072] + 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) @@ -878,7 +880,7 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE_TP=2048, TILE_M, TILE_N = 128, 128 HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK = 4096, 192, 192, 8 test = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK) - test.entry_common('aiter', batch, prec=[prec_fp8_t]) + test.entry_common('aiter', batch, prec=[prec_fp8_ptpc]) test.entry_common('mxn_2s', batch, prec=[prec_fp8_t]) test.show_perf() #entry_common('mxn_2s', batch=batch, prec=[get_fp8type()], E=E, TOPK=TOPK, TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=10) From 8dab15b8509a45226815191592909cfada9dbb47 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Wed, 20 May 2026 08:29:53 +0000 Subject: [PATCH 17/24] [jit][moe] support pref_fp_t in decoding --- src/contrib/common/gemm_splitk.py | 32 +++++++--- src/contrib/moe.py | 98 ++++++++++++++++++------------- tests/contrib/moe/test_moe.py | 72 ++++++++++++++--------- 3 files changed, 122 insertions(+), 80 deletions(-) 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 305413cb..a29ec21b 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -55,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 @@ -105,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 @@ -121,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 @@ -208,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 @@ -286,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 @@ -380,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: @@ -551,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] @@ -596,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 diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index f2801b67..70ed32d4 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -16,6 +16,29 @@ 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 w2, # [expert(local_expert:EP), dim, inter_dim] @@ -91,15 +114,6 @@ def wei_is_fp8(weight_type): from aiter.ops.quant import pertoken_quant ext_topk_ids = None -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 def quant_expert_weights(w1, quant_type, dtype): if quant_type == aiter.QuantType.per_Token: @@ -289,6 +303,7 @@ def weight_per_128x128_quant(weight, quant_dtype): 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] @@ -298,8 +313,8 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu 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) + 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( @@ -318,7 +333,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu 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) + 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) @@ -345,7 +360,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu 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, fp8_ptpc) + 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: @@ -770,12 +785,6 @@ def show_perf(self): 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}') -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) - def init_env(): torch.set_printoptions(linewidth=3000, sci_mode=False, edgeitems=8, ) torch.set_default_device('cuda') @@ -865,11 +874,6 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE_TP=2048, # 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 = [i * 8192 for i in range(1, 11)] - #batch = [i * 1024 for i in range(1, 9)] + [8192*2, 8192*4, 8192*8] - batch = [i * 128 for i in range(1, 9)] - batch = [64,128,256,512,1024,2048,4096,8192, 16384, 32768, 65536, 131072] - #batch = [8192, 16384, 32768, 65536, 131072] batch = [512, 1024, 2048] #batch = [8192,] if 0: @@ -877,12 +881,22 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE_TP=2048, print(ext_topk_ids.shape) batch[0] = ext_topk_ids.shape[0] - TILE_M, TILE_N = 128, 128 HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK = 4096, 192, 192, 8 - test = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK) - test.entry_common('aiter', batch, prec=[prec_fp8_ptpc]) - test.entry_common('mxn_2s', batch, prec=[prec_fp8_t]) - test.show_perf() + #decoding + 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_fp8_t]) + test_dec.entry_common('16x32_2s_b1', [1], prec=[prec_fp8_t]) + test_dec.entry_common('16x32_2s_b', batch, prec=[prec_fp8_t]) + + TILE_M, TILE_N = 128, 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_fp8_t]) + test_prefill.entry_common('mxn_2s', batch, prec=[prec_fp8_t]) + test_dec.show_perf() + test_prefill.show_perf() #entry_common('mxn_2s', batch=batch, prec=[get_fp8type()], E=E, TOPK=TOPK, TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=10) #test_perf(batch, TILE_M=128, TILE_N=256, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE_TP=INTER_SIZE_TP, E=E, TOPK=TOPK, test_sets=['aiter', 'mxn_2s']) From e4c8c1314d62ca014e7756e7cddf1b7f51743ed5 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Wed, 20 May 2026 09:26:25 +0000 Subject: [PATCH 18/24] [jit][moe] default use preshuffle on for aiter --- tests/contrib/moe/test_moe.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index 70ed32d4..d0142208 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -214,10 +214,12 @@ def __call__(self, kernel_type, weight_type, quant_type, B=1, run_count=0): 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)] + 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)] + 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): From 7b3c6cbc2d4033b6fe642a2db96930812d67a910 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Wed, 20 May 2026 12:15:39 +0000 Subject: [PATCH 19/24] [jit][moe] dynamic schedule gateup --- src/contrib/common/gemm.py | 9 +- src/contrib/moe.py | 357 ++++++++++++++++++---------------- tests/contrib/moe/test_moe.py | 25 ++- 3 files changed, 216 insertions(+), 175 deletions(-) diff --git a/src/contrib/common/gemm.py b/src/contrib/common/gemm.py index 580f2025..c6169431 100644 --- a/src/contrib/common/gemm.py +++ b/src/contrib/common/gemm.py @@ -364,6 +364,7 @@ def mfma_generator(k): J.emit(mfma0, 16*2) # 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() @@ -372,17 +373,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*3) loaderA.prefetch(r) - J.emit([mfma0, mfma1], 16*8) + J.emit([mfma0, mfma1], 16*3) 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*3) loaderB.prefetch(r) - J.emit([mfma0, mfma1], 16*8) + J.emit([mfma0, mfma1], 16*3) # 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) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index a29ec21b..b38886c1 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1175,7 +1175,7 @@ def xcd_swizzle(J, blk1d, num_blocks, num_oc_blocks, NUM_XCD, NUM_CU_PER_XCD): If.Else() blk_m[0] = blk1d // num_oc_blocks blk_n[0] = blk1d - blk_m * num_oc_blocks - elif num_oc_blocks <= 4: + 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 @@ -1203,6 +1203,7 @@ def moe_2stage_gateup(J:JIT, 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*", @@ -1370,183 +1371,207 @@ def ds_write(self, index, lds_base): read_col_lanes_per_wave = 4 * 16 * 2 // 16 read_row_lanes_per_wave = 64 // read_col_lanes_per_wave - # grid: N // 32, sorted_expert_ids.shape[0] - # expert index in p_sorted_expert_ids - e_idx, blockn_idx = xcd_swizzle(J, J.blockIdx.x, num_blocks, N // BLOCK_TILE_SIZE_N, 4, 20) - 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) & (blockn_idx[0] == 0) & (e_idx[0] == 0)) - - # hide following initialization into s_waitcnt - 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_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 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) - p_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_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_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_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_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_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": - 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, 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, 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('vu32', 0) - pc_scale[:] += 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 + 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)") + J.v_readfirstlane_b32(idx, v_idx) + J.s_barrier() + + # 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) + 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)") + # 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) & (blockn_idx[0] == 0) & (e_idx[0] == 0)) + + # 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 - 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_output) + vdata = J.gpr(4, "vu32", align=4) - J.free_lds(lds_out) + 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) - J.get_sgpr_const(0x8000) - J.get_sgpr_const(0x3020706) - if 1: - for i, TILE_SIZE in enumerate(VALID_TILE_SIZES): - with J.If((s_sorted_id_tile[i] >> 24) == TOPK): - kernel(TILE_SIZE) - J.s_endpgm() + J.free_lds(lds_out) - kernel(BLOCK_TILE_SIZE_M) + J.get_sgpr_const(0x8000) + J.get_sgpr_const(0x3020706) + + 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) def de_shuffle_weight(weight_, mfma_MN = 16): diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index d0142208..82acf5a6 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -558,10 +558,12 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu 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, str(fp8_quant_type), - 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]) + 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([80], [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]) gemm2_out = torch.empty(B, TOPK, N2, dtype=torch.bfloat16, device=hidden_states.device) moe_2stage_down([1, sorted_expert_ids.shape[0]], [256], w1.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, @@ -598,9 +600,11 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu 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([N1 // BLOCK_TILE_SIZE_N * sorted_expert_ids.shape[0]], [256], + moe_2stage_gateup([80], [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, @@ -892,13 +896,24 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE_TP=2048, test_dec.entry_common('16x32_2s_b1', [1], prec=[prec_fp8_t]) test_dec.entry_common('16x32_2s_b', batch, prec=[prec_fp8_t]) + # prefill per tensor TILE_M, TILE_N = 128, 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_fp8_t]) test_prefill.entry_common('mxn_2s', batch, prec=[prec_fp8_t]) + + # prefill ptpc + TILE_M, TILE_N = 128, 256 + HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK = 4096, 128, 512, 10 + batch = [i * 1024 for i in range(1, 9)] + [16 * 1024, 32 * 1024, 64 * 1024] + test_ptpc = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK) + test_ptpc.entry_common('aiter', batch, prec=[prec_fp8_ptpc]) + test_ptpc.entry_common('mxn_2s', batch, prec=[prec_fp8_ptpc]) + test_dec.show_perf() test_prefill.show_perf() + test_ptpc.show_perf() #entry_common('mxn_2s', batch=batch, prec=[get_fp8type()], E=E, TOPK=TOPK, TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=10) #test_perf(batch, TILE_M=128, TILE_N=256, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE_TP=INTER_SIZE_TP, E=E, TOPK=TOPK, test_sets=['aiter', 'mxn_2s']) From 45ed5b212eb1b56e4056fd7709919d93b940e033 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Thu, 21 May 2026 00:28:06 +0000 Subject: [PATCH 20/24] [jit][moe] dynamic schedule down --- src/contrib/moe.py | 182 ++++++++++++++++++++-------------- tests/contrib/moe/test_moe.py | 49 +++++---- 2 files changed, 134 insertions(+), 97 deletions(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index b38886c1..0ca92ffd 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1677,8 +1677,10 @@ def moe_2stage_down(J:JIT, TOPK, K, N, # 8, 128, 2048 with_silu, # BLOCK_TILE_SIZE_M, # 32 + SUB_M, 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] @@ -1700,28 +1702,12 @@ def moe_2stage_down(J:JIT, else: fp8_ptpc = None - e_idx, _ = xcd_swizzle(J, J.blockIdx.y, num_blocks, 1, 4, 20) - 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) - if quant_type_w == "QuantType.per_Tensor": - pc_scale[:] += s_e_id * (1 * J.sizeof_DW) # per-tensor scales for weights - else: - pc_scale[:] += s_e_id * (N * J.sizeof_DW) # per-channel scales for weights 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") @@ -1729,80 +1715,124 @@ 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) - - J.s_waitcnt(mod=f"vmcnt(0)") + M[0] *= TOPK - 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 + J.get_sgpr_const(0x8000) + J.get_sgpr_const(0x3020706) - 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 + 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() - if quant_type_w == "QuantType.per_Token": - fp8_ptpc["pc_scale"] = pc_scale + e_idx, sub_m = xcd_swizzle(J, idx, num_blocks, BLOCK_TILE_SIZE_M // SUB_M, 4, 20) + # 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)") + # invalid padding section + J.Jump("continue_following", e_idx * BLOCK_TILE_SIZE_M < max_id) + J.s_endpgm() + J.Label("continue_following") + 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": - v_pc_scales = J.gpr('vf32') - v_offset_zero = J.gpr('vu32', 0) - J.global_load_dword(v_pc_scales[0], v_offset_zero, pc_scale) - fp8_ptpc["v_pc_scales"] = v_pc_scales + 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 - 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_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) - 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) + J.s_waitcnt(mod=f"vmcnt(0)") - J.s_waitcnt(mod=f"vmcnt(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 - lds_weights = J.alloc_lds(256*4) - lds_token_ids = J.alloc_lds(256*4) + 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 - 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): + 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): - 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}") + 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)) - J.s_waitcnt(mod=f"lgkmcnt(0)") - J.s_barrier() + 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) - M[0] *= TOPK + J.s_waitcnt(mod=f"vmcnt(0)") - J.get_sgpr_const(0x8000) - J.get_sgpr_const(0x3020706) - # 0,1,2,3,4,5,6,7, num_mfma_m - if 1: - 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_weight, p_output, M, fp8_ptpc) - J.s_endpgm() + 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}") - 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, fp8_ptpc) + J.s_waitcnt(mod=f"lgkmcnt(0)") + J.s_barrier() + + # 0,1,2,3,4,5,6,7, num_mfma_m + if 0: + 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_weight, p_output, M, fp8_ptpc) + J.s_endpgm() + + if 1: + 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, fp8_ptpc) def down_kernel(J, mfma_MN, num_mfma_n, BM, N, K, A, # A = J.gpr(num_mfma_m, num_mfma_k, 4, "abf16x2") diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index 82acf5a6..9d8ab57e 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -565,9 +565,12 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu 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]) gemm2_out = torch.empty(B, TOPK, N2, dtype=torch.bfloat16, device=hidden_states.device) - 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, w2, gemm2_out, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, None, w2_scale, B, sorted_expert_ids.shape[0]) + BLOCK_TILE_SIZE_N = 128 + BLOCK_SUB_TILE_SIZE_M = 64 + id_buf2 = torch.zeros(64, dtype=torch.int32) + moe_2stage_down([160], [256], + w1.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_SUB_TILE_SIZE_M, BLOCK_TILE_SIZE_N, + 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]) num_WG = 80 * 4 num_tokens_wg = B // num_WG num_extra_tokens = B % num_WG @@ -646,12 +649,15 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu TILE_M, sorted_ids, sorted_expert_ids, sorted_weights, num_valid_ids) else: + BLOCK_TILE_SIZE_N = 128 + BLOCK_SUB_TILE_SIZE_M = 64 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([1, sorted_expert_ids.shape[0]], [256], - w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_TILE_SIZE_N, str(fp8_quant_type), - gemm1_out_q, w2, + moe_2stage_down([160], [256], + w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_SUB_TILE_SIZE_M, BLOCK_TILE_SIZE_N, str(fp8_quant_type), + id_buf2, gemm1_out_q, w2, gemm2_out, #cur_out, sorted_ids, sorted_weights, @@ -889,30 +895,31 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE_TP=2048, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK = 4096, 192, 192, 8 #decoding - 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_fp8_t]) - test_dec.entry_common('16x32_2s_b1', [1], prec=[prec_fp8_t]) - test_dec.entry_common('16x32_2s_b', batch, prec=[prec_fp8_t]) - - # prefill per tensor - TILE_M, TILE_N = 128, 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_fp8_t]) - test_prefill.entry_common('mxn_2s', batch, prec=[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_fp8_t]) + # test_dec.entry_common('16x32_2s_b1', [1], prec=[prec_fp8_t]) + # test_dec.entry_common('16x32_2s_b', batch, prec=[prec_fp8_t]) + + # # prefill per tensor + # TILE_M, TILE_N = 128, 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_fp8_t]) + # test_prefill.entry_common('mxn_2s', batch, prec=[prec_fp8_t]) # prefill ptpc TILE_M, TILE_N = 128, 256 HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK = 4096, 128, 512, 10 batch = [i * 1024 for i in range(1, 9)] + [16 * 1024, 32 * 1024, 64 * 1024] + #batch = [ 8192] test_ptpc = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK) test_ptpc.entry_common('aiter', batch, prec=[prec_fp8_ptpc]) test_ptpc.entry_common('mxn_2s', batch, prec=[prec_fp8_ptpc]) - test_dec.show_perf() - test_prefill.show_perf() + # test_dec.show_perf() + # test_prefill.show_perf() test_ptpc.show_perf() #entry_common('mxn_2s', batch=batch, prec=[get_fp8type()], E=E, TOPK=TOPK, TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=10) #test_perf(batch, TILE_M=128, TILE_N=256, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE_TP=INTER_SIZE_TP, E=E, TOPK=TOPK, test_sets=['aiter', 'mxn_2s']) From a23151c10386ab29747d264bbc8ac5663f211793 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Thu, 21 May 2026 01:16:10 +0000 Subject: [PATCH 21/24] [jit][moe] refactor to support dyn and static schedule --- src/contrib/common/gemm.py | 40 ++++++++---- src/contrib/moe.py | 120 ++++++++++++++++++++-------------- tests/contrib/moe/test_moe.py | 96 +++++++++++++++++---------- 3 files changed, 160 insertions(+), 96 deletions(-) diff --git a/src/contrib/common/gemm.py b/src/contrib/common/gemm.py index c6169431..637fce60 100644 --- a/src/contrib/common/gemm.py +++ b/src/contrib/common/gemm.py @@ -341,6 +341,24 @@ 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): @@ -358,10 +376,10 @@ 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) @@ -373,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*3) + J.emit([mfma0, mfma1], 16*mfma_cnt['ds_write']) loaderA.prefetch(r) - J.emit([mfma0, mfma1], 16*3) + 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*3) + J.emit([mfma0, mfma1], 16*mfma_cnt['ds_write']) loaderB.prefetch(r) - J.emit([mfma0, mfma1], 16*3) + 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) @@ -394,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 @@ -471,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: @@ -518,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) @@ -680,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/moe.py b/src/contrib/moe.py index 0ca92ffd..2388dc58 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1214,7 +1214,8 @@ def moe_2stage_gateup(J:JIT, pt_scale: "float*", pc_scale:"float*", M:"int", - num_blocks:"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, @@ -1378,23 +1379,7 @@ def ds_write(self, index, lds_base): warp_id = J.gpr("su32") J.v_readfirstlane_b32(warp_id, J.threadIdx.x[0] // 64) - 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)") - J.v_readfirstlane_b32(idx, v_idx) - J.s_barrier() - + 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) @@ -1573,6 +1558,27 @@ def kernel(CUR_BLOCK_TILE_SIZE_M): 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)") + 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 @@ -1677,7 +1683,6 @@ def moe_2stage_down(J:JIT, TOPK, K, N, # 8, 128, 2048 with_silu, # BLOCK_TILE_SIZE_M, # 32 - SUB_M, BLOCK_TILE_SIZE_N, # 64 quant_type_w, p_id:"void*", @@ -1691,7 +1696,13 @@ def moe_2stage_down(J:JIT, pt_scale: "float*", pc_scale:"float*", M:"int", - num_blocks:"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") @@ -1715,29 +1726,15 @@ def moe_2stage_down(J:JIT, # collect token id row = J.lane_id % mfma_MN col = J.lane_id // mfma_MN - M[0] *= TOPK + M_TOPK = J.gpr(M[0] * TOPK) J.get_sgpr_const(0x8000) J.get_sgpr_const(0x3020706) - 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() - + 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 # 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) @@ -1810,8 +1807,16 @@ def moe_2stage_down(J:JIT, J.s_waitcnt(mod=f"lgkmcnt(0)") J.s_barrier() - # 0,1,2,3,4,5,6,7, num_mfma_m - if 0: + 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): @@ -1822,17 +1827,34 @@ def moe_2stage_down(J:JIT, 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_weight, p_output, M, fp8_ptpc) + p_cur_weight, p_output, M_TOPK, fp8_ptpc) J.s_endpgm() - if 1: - 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, fp8_ptpc) + 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") diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index 9d8ab57e..80f92b31 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -141,6 +141,8 @@ class TestCase: 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 @@ -153,6 +155,9 @@ def __call__(self, kernel_type, weight_type, quant_type, B=1, run_count=0): 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: @@ -558,19 +563,26 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu 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([80], [256], + 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]) + 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) - BLOCK_TILE_SIZE_N = 128 - BLOCK_SUB_TILE_SIZE_M = 64 id_buf2 = torch.zeros(64, dtype=torch.int32) - moe_2stage_down([160], [256], - w1.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_SUB_TILE_SIZE_M, BLOCK_TILE_SIZE_N, - 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]) + 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 @@ -581,6 +593,13 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu 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) @@ -605,7 +624,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu 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([80], [256], + 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, @@ -614,7 +633,8 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu sorted_expert_ids, num_valid_ids, hidden_states_scale, - w1_scale, B, N1 // BLOCK_TILE_SIZE_N * sorted_expert_ids.shape[0]) + 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( @@ -649,14 +669,12 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu TILE_M, sorted_ids, sorted_expert_ids, sorted_weights, num_valid_ids) else: - BLOCK_TILE_SIZE_N = 128 - BLOCK_SUB_TILE_SIZE_M = 64 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([160], [256], - w2.dtype, TOPK, K2, N2, False, BLOCK_TILE_SIZE_M, BLOCK_SUB_TILE_SIZE_M, BLOCK_TILE_SIZE_N, str(fp8_quant_type), + 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, @@ -666,7 +684,8 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu gemm1_out_scale, w2_scale, B, - sorted_expert_ids.shape[0]) + 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) @@ -895,31 +914,36 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE_TP=2048, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK = 4096, 192, 192, 8 #decoding - # 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_fp8_t]) - # test_dec.entry_common('16x32_2s_b1', [1], prec=[prec_fp8_t]) - # test_dec.entry_common('16x32_2s_b', batch, prec=[prec_fp8_t]) - - # # prefill per tensor - # TILE_M, TILE_N = 128, 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_fp8_t]) - # test_prefill.entry_common('mxn_2s', batch, prec=[prec_fp8_t]) - - # prefill ptpc - TILE_M, TILE_N = 128, 256 + 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_fp8_t]) + test_dec.entry_common('16x32_2s_b1', [1], prec=[prec_fp8_t]) + test_dec.entry_common('16x32_2s_b', batch, prec=[prec_fp8_t]) + + # prefill per tensor + TILE_M, TILE_N = 128, 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_fp8_t]) + test_prefill.entry_common('mxn_2s', batch, prec=[prec_fp8_t]) + + # prefill ptpc(with dynamic scheduler) + TILE_M, TILE_N, STAGE2_TILE_N = 128, 256, 128 HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK = 4096, 128, 512, 10 batch = [i * 1024 for i in range(1, 9)] + [16 * 1024, 32 * 1024, 64 * 1024] - #batch = [ 8192] - test_ptpc = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK) - test_ptpc.entry_common('aiter', batch, prec=[prec_fp8_ptpc]) - test_ptpc.entry_common('mxn_2s', batch, prec=[prec_fp8_ptpc]) + test_ptpc_dyn = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK, DYN_SCHEDULE=True, STAGE2_TILE_N=STAGE2_TILE_N) + test_ptpc_dyn.entry_common('aiter', batch, prec=[prec_bf16, prec_fp8_ptpc]) + test_ptpc_dyn.entry_common('mxn_2s', batch, prec=[prec_bf16, prec_fp8_ptpc]) + + # prefill ptpc(with default scheduler) + TILE_M, TILE_N = 64, 128 + test_ptpc = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK, DYN_SCHEDULE=False) + test_ptpc.entry_common('mxn_2s', batch, prec=[prec_bf16, prec_fp8_ptpc]) - # test_dec.show_perf() - # test_prefill.show_perf() + test_dec.show_perf() + test_prefill.show_perf() + test_ptpc_dyn.show_perf() test_ptpc.show_perf() #entry_common('mxn_2s', batch=batch, prec=[get_fp8type()], E=E, TOPK=TOPK, TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=10) #test_perf(batch, TILE_M=128, TILE_N=256, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE_TP=INTER_SIZE_TP, E=E, TOPK=TOPK, test_sets=['aiter', 'mxn_2s']) From 05939a80b0edc19642689203872990a104ec6322 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Mon, 25 May 2026 07:51:55 +0000 Subject: [PATCH 22/24] [jit][moe] test case for qwen3.5 & hunyuan --- tests/contrib/moe/test_moe.py | 80 ++++++++++++++++++----------------- 1 file changed, 41 insertions(+), 39 deletions(-) diff --git a/tests/contrib/moe/test_moe.py b/tests/contrib/moe/test_moe.py index 80f92b31..2b9a2916 100644 --- a/tests/contrib/moe/test_moe.py +++ b/tests/contrib/moe/test_moe.py @@ -346,7 +346,7 @@ def run(hidden_states, w1, w2, topk_weight, topk_ids, w1_scale, w2_scale, fp8_qu # 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: + 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 @@ -807,8 +807,8 @@ def entry_common(self, kernel_type, batch, prec): perf[kernel_type][str(weight_type)+"@"+str(quant_type)] = perf_prec self.perf.append(perf) - def show_perf(self): - print('\nsummary:') + 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(): @@ -912,40 +912,42 @@ def test_perf(batch, TILE_M=32, TILE_N=64, HIDDEN_SIZE=4096, INTER_SIZE_TP=2048, print(ext_topk_ids.shape) batch[0] = ext_topk_ids.shape[0] - HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK = 4096, 192, 192, 8 - #decoding - 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_fp8_t]) - test_dec.entry_common('16x32_2s_b1', [1], prec=[prec_fp8_t]) - test_dec.entry_common('16x32_2s_b', batch, prec=[prec_fp8_t]) - - # prefill per tensor - TILE_M, TILE_N = 128, 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_fp8_t]) - test_prefill.entry_common('mxn_2s', batch, prec=[prec_fp8_t]) - - # prefill ptpc(with dynamic scheduler) - TILE_M, TILE_N, STAGE2_TILE_N = 128, 256, 128 + # 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 - batch = [i * 1024 for i in range(1, 9)] + [16 * 1024, 32 * 1024, 64 * 1024] - test_ptpc_dyn = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK, DYN_SCHEDULE=True, STAGE2_TILE_N=STAGE2_TILE_N) - test_ptpc_dyn.entry_common('aiter', batch, prec=[prec_bf16, prec_fp8_ptpc]) - test_ptpc_dyn.entry_common('mxn_2s', batch, prec=[prec_bf16, prec_fp8_ptpc]) - - # prefill ptpc(with default scheduler) - TILE_M, TILE_N = 64, 128 - test_ptpc = TestCase(TILE_M, TILE_N, HIDDEN_SIZE, INTER_SIZE_TP, E, TOPK, DYN_SCHEDULE=False) - test_ptpc.entry_common('mxn_2s', batch, prec=[prec_bf16, prec_fp8_ptpc]) - - test_dec.show_perf() - test_prefill.show_perf() - test_ptpc_dyn.show_perf() - test_ptpc.show_perf() - #entry_common('mxn_2s', batch=batch, prec=[get_fp8type()], E=E, TOPK=TOPK, TILE_M=128, TILE_N=128, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE=INTER_SIZE, TP=TP, run_count=10) - #test_perf(batch, TILE_M=128, TILE_N=256, HIDDEN_SIZE=HIDDEN_SIZE, INTER_SIZE_TP=INTER_SIZE_TP, E=E, TOPK=TOPK, test_sets=['aiter', 'mxn_2s']) - - + 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=}') From 65a0262b033c888b141d3540ddb42da9b65cffb7 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Thu, 28 May 2026 07:05:52 +0000 Subject: [PATCH 23/24] [jit][moe] fix using sgpr in v_mul_u32_u24 to compute prefetch offset --- src/contrib/moe.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index 2388dc58..7fe9c8e2 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1297,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 * blockn_idx + 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): From 8a58deaf004e41f61091e147207a39dbf71b5123 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Tue, 2 Jun 2026 02:00:31 +0000 Subject: [PATCH 24/24] [jit][moe] fix out of bound access expert_id if dyn scheduler is true --- src/contrib/moe.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/contrib/moe.py b/src/contrib/moe.py index 7fe9c8e2..cd4ba509 100644 --- a/src/contrib/moe.py +++ b/src/contrib/moe.py @@ -1380,19 +1380,20 @@ def ds_write(self, index, lds_base): 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)") 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) - 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)") # 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) + + J.s_waitcnt(mod=f"lgkmcnt(0)") J.debug_setup((J.warp_id[0] == 0) & (blockn_idx[0] == 0) & (e_idx[0] == 0)) # hide following initialization into s_waitcnt @@ -1732,19 +1733,20 @@ def moe_2stage_down(J:JIT, J.get_sgpr_const(0x8000) J.get_sgpr_const(0x3020706) + J.s_waitcnt(mod=f"lgkmcnt(0)") 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 - # 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)") # 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)") 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))