diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index bda806cc15b..169299b580a 100755 --- a/.github/scripts/aiter_test.sh +++ b/.github/scripts/aiter_test.sh @@ -36,12 +36,12 @@ fi skip_tests=( "op_tests/multigpu_tests/bench_mega_moe_v2.py" + "op_tests/multigpu_tests/test_mega_moe_v2.py" "op_tests/multigpu_tests/test_wide_ep_moe.py" "op_tests/multigpu_tests/test_dispatch_combine.py" "op_tests/multigpu_tests/test_communication.py" "op_tests/multigpu_tests/test_mori_all2all.py" "op_tests/multigpu_tests/test_fused_ar_rms.py" - "op_tests/multigpu_tests/test_mega_moe_v2.py" "op_tests/multigpu_tests/triton_test/test_reduce_scatter_all_gather.py" "op_tests/multigpu_tests/triton_test/test_fused_rs_rmsnorm_quant_ag.py" ) diff --git a/aiter/aot/flydsl/common.py b/aiter/aot/flydsl/common.py index 8018f28302f..d23c6c117a4 100755 --- a/aiter/aot/flydsl/common.py +++ b/aiter/aot/flydsl/common.py @@ -140,9 +140,10 @@ def _collect_aot_jobs_for(kind: OpKind) -> list[dict[str, Any]]: ``flydsl.expr``). Job collection is therefore not free in the parent process, just shifted once out of every child.""" if kind is OpKind.MEGA_MOE: - from .mega_moe import default_jobs + from .mega_moe import production_jobs - return default_jobs() + # Ships both V4-Pro and Kimi-K3 so neither JITs online at serving time. + return production_jobs() if kind is OpKind.FMHA_FP8: from .fmha_fp8 import default_jobs diff --git a/aiter/aot/flydsl/mega_moe.py b/aiter/aot/flydsl/mega_moe.py index 3d7dad6833f..b5e6a21a694 100755 --- a/aiter/aot/flydsl/mega_moe.py +++ b/aiter/aot/flydsl/mega_moe.py @@ -21,10 +21,24 @@ ) DEFAULT_MTPRS = (8192, 16384, 32768) +DEFAULT_QUANTS = ("a8w4",) # DeepSeek-V4-Pro deployment profiles: r0, r32, and r64 redundant experts. # Keep all three in the default AOT job set so the service never falls back to # an online compile merely because EPLB changes the physical expert count. DEFAULT_EXPERTS_PER_RANKS = (48, 52, 56) +# Kimi-K3 deployment profile (EP8 / epr112, topk16, d3584 / i3072). Added to +# the default AOT job set so K3 serving never falls back to an online compile. +K3_EXPERTS_PER_RANKS = (112,) +K3_TOPK = 16 +K3_MODEL_DIM = 3584 +K3_INTER_DIM = 3072 +# K3 uses SiTUv2 (hidden_act="situ"); tanh self-saturates so no hard clamp. +# beta=4.0 / linear_beta=25.0 are the K3 checkpoint constants the MegaMoE SiTU +# kernel bakes in (SGLang kimi_k3.py asserts exactly these); they MUST match +# op_tests kimi_k3_route or the AOT bundle name will not match at serving. +K3_ACT = "situv2" +K3_SITU_BETA = 4.0 +K3_SITU_LINEAR_BETA = 25.0 WORLD_SIZE = 8 TOPK = 6 MODEL_DIM = 7168 @@ -59,11 +73,15 @@ def default_jobs( mtprs=DEFAULT_MTPRS, experts_per_ranks=DEFAULT_EXPERTS_PER_RANKS, *, + quants=DEFAULT_QUANTS, world_size=WORLD_SIZE, topk=TOPK, model_dim=MODEL_DIM, inter_dim=INTER_DIM, swiglu_limit=SWIGLU_LIMIT, + act="silu", + situ_beta=1.0, + situ_linear_beta=1.0, ): shape_suffix = "" if (world_size, topk, model_dim, inter_dim) != ( @@ -73,12 +91,16 @@ def default_jobs( INTER_DIM, ): shape_suffix = f"_w{world_size}_k{topk}_d{model_dim}_i{inter_dim}" + # Activation is baked into the Stage1 epilogue, so a non-default act needs a + # distinct AOT artifact name (mirrors the runtime kernel-name act suffix). + act_suffix = "" if act == "silu" else f"_{act}" return [ { "kernel_name": ( - f"mega_moe_stage{stage}_bundle_mtpr{mtpr}_epr{experts_per_rank}" - f"_rank{rank}{shape_suffix}" + f"mega_moe_{quant}_stage{stage}_bundle_mtpr{mtpr}_epr{experts_per_rank}" + f"_rank{rank}{shape_suffix}{act_suffix}" ), + "quant": quant, "stage": stage, "mtpr": mtpr, "experts_per_rank": experts_per_rank, @@ -88,14 +110,51 @@ def default_jobs( "model_dim": model_dim, "inter_dim": inter_dim, "swiglu_limit": swiglu_limit, + "act": act, + "situ_beta": situ_beta, + "situ_linear_beta": situ_linear_beta, } for mtpr in mtprs for experts_per_rank in experts_per_ranks + for quant in quants for rank in range(world_size) for stage in (1, 2) ] +def production_jobs( + mtprs=DEFAULT_MTPRS, + *, + quants=DEFAULT_QUANTS, + world_size=WORLD_SIZE, + swiglu_limit=SWIGLU_LIMIT, +): + """Every geometry the service ships: V4-Pro (default shape) + Kimi-K3. + + Used by the default build so both profiles are AOT-covered and neither + triggers an online JIT compile at serving time. + """ + jobs = default_jobs( + mtprs, quants=quants, world_size=world_size, swiglu_limit=swiglu_limit + ) + jobs += default_jobs( + mtprs, + K3_EXPERTS_PER_RANKS, + quants=quants, + world_size=world_size, + topk=K3_TOPK, + model_dim=K3_MODEL_DIM, + inter_dim=K3_INTER_DIM, + # SiTUv2 self-saturates -> no hard clamp; keep swiglu_limit at 0 so the + # kernel name matches the runtime kimi_k3_route (swiglu_limit=0.0). + swiglu_limit=0.0, + act=K3_ACT, + situ_beta=K3_SITU_BETA, + situ_linear_beta=K3_SITU_LINEAR_BETA, + ) + return jobs + + def _tensor(shape, dtype): return torch.empty(shape, dtype=dtype, device="cpu") @@ -106,11 +165,15 @@ def _compile_stage1( rank, plan, *, + a_dtype, world_size, topk, model_dim, inter_dim, swiglu_limit, + act="silu", + situ_beta=1.0, + situ_linear_beta=1.0, ): from aiter.ops.flydsl.kernels.mega_moe.mega_moe_prepare import ( preload_mega_moe_prepare, @@ -182,7 +245,7 @@ def _compile_stage1( # the validation/per-stage benchmark invokes it to isolate Stage1 timing. quant_rows = 2 quant_groups = quant_rows * scale_dim - _get_launcher(model_dim, "fp8")( + _get_launcher(model_dim, a_dtype)( _tensor((quant_rows, model_dim), torch.bfloat16), _tensor((quant_rows, model_dim), torch.float8_e4m3fn), _tensor((quant_rows, scale_dim), torch.uint8), @@ -206,6 +269,9 @@ def _compile_stage1( tile_state_stride=tile_state_stride, variants=plan.stage1_variants, swiglu_limit=swiglu_limit, + act=act, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, ) launch( _tensor((1, inter_dim), torch.float8_e4m3fn), @@ -236,6 +302,7 @@ def _compile_stage2( rank, plan, *, + a_dtype, world_size, topk, model_dim, @@ -278,6 +345,7 @@ def _compile_stage2( "cu_num": NUM_CU, "p2p_quant_type": key.p2p_quant, "fixed_slot_dispatch": plan.fixed_slot_dispatch, + "a_dtype": a_dtype, } residual = ( replace(stage2, skew_cu=stage2.persist_cu) @@ -296,8 +364,10 @@ def _compile_stage2( BK=residual.block_k, use_nt=residual.use_nt, g2_bhoist=residual.b_hoist, + g2_b2stage=residual.b2stage, g2_ascale_pf=residual.ascale_prefetch, g2_spart=residual.spatial_partition, + g2_deep_a_pipeline=residual.deep_a_pipeline, persist=residual.persist, persist_cu=residual.persist_cu, persist_strided=residual.persist_strided, @@ -332,7 +402,10 @@ def _compile_stage2( use_nt=stage2.use_nt, cu_num=stage2.pair_cu, g2_bhoist=stage2.b_hoist, + g2_b2stage=stage2.b2stage, g2_ascale_pf=stage2.ascale_prefetch, + g2_deep_a_pipeline=stage2.deep_a_pipeline, + a_dtype=a_dtype, ) # Stage2's production bundle includes the terminal fused combine kernels. @@ -408,6 +481,10 @@ def compile_one_config(**job): model_dim = job.get("model_dim", MODEL_DIM) inter_dim = job.get("inter_dim", INTER_DIM) swiglu_limit = job.get("swiglu_limit", SWIGLU_LIMIT) + act = job.get("act", "silu") + situ_beta = job.get("situ_beta", 1.0) + situ_linear_beta = job.get("situ_linear_beta", 1.0) + a_dtype = "fp8" plan = build_mega_moe_bundle_plan( job["mtpr"], experts_per_rank=job["experts_per_rank"], @@ -422,11 +499,15 @@ def compile_one_config(**job): job["experts_per_rank"], job["rank"], plan, + a_dtype=a_dtype, world_size=world_size, topk=topk, model_dim=model_dim, inter_dim=inter_dim, swiglu_limit=swiglu_limit, + act=act, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, ) else: _compile_stage2( @@ -434,6 +515,7 @@ def compile_one_config(**job): job["experts_per_rank"], job["rank"], plan, + a_dtype=a_dtype, world_size=world_size, topk=topk, model_dim=model_dim, @@ -460,16 +542,44 @@ def main(): parser.add_argument("--model-dim", type=int, default=MODEL_DIM) parser.add_argument("--inter-dim", type=int, default=INTER_DIM) parser.add_argument("--swiglu-limit", type=float, default=SWIGLU_LIMIT) + parser.add_argument("--act", default="silu", choices=("silu", "situv2")) + parser.add_argument("--situ-beta", type=float, default=1.0) + parser.add_argument("--situ-linear-beta", type=float, default=1.0) + parser.add_argument( + "--quant", + nargs="+", + choices=DEFAULT_QUANTS, + default=list(DEFAULT_QUANTS), + ) args = parser.parse_args() - jobs = default_jobs( - tuple(args.mtpr), - tuple(args.experts_per_rank), - world_size=args.world_size, - topk=args.topk, - model_dim=args.model_dim, - inter_dim=args.inter_dim, - swiglu_limit=args.swiglu_limit, + default_shape = ( + list(args.experts_per_rank) == list(DEFAULT_EXPERTS_PER_RANKS) + and args.topk == TOPK + and args.model_dim == MODEL_DIM + and args.inter_dim == INTER_DIM ) + if default_shape: + # Default build ships both V4-Pro and Kimi-K3 so neither JITs online. + jobs = production_jobs( + tuple(args.mtpr), + quants=tuple(args.quant), + world_size=args.world_size, + swiglu_limit=args.swiglu_limit, + ) + else: + jobs = default_jobs( + tuple(args.mtpr), + tuple(args.experts_per_rank), + quants=tuple(args.quant), + world_size=args.world_size, + topk=args.topk, + model_dim=args.model_dim, + inter_dim=args.inter_dim, + swiglu_limit=args.swiglu_limit, + act=args.act, + situ_beta=args.situ_beta, + situ_linear_beta=args.situ_linear_beta, + ) results = run_jobs_parallel(compile_one_config, jobs) failed = sum(result["compile_time"] is None for result in results) print(f"Compiled: {len(results) - failed} ok, {failed} failed") diff --git a/aiter/ops/flydsl/kernels/mega_moe/dispatch.py b/aiter/ops/flydsl/kernels/mega_moe/dispatch.py index 79ccff13962..cab5de6dd3a 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/dispatch.py +++ b/aiter/ops/flydsl/kernels/mega_moe/dispatch.py @@ -96,7 +96,7 @@ def _load_fanout_pair( # ``packed`` controls route classification, destination offsets, and the # producer task mapping. Keep the VMEM dependency explicit before those # values cross lane/control-flow boundaries. - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) enabled = (packed & fx.Int32(1 << 16)) != fx.Int32(0) pair_a = packed & fx.Int32(0xFF) pair_b = (packed >> fx.Int32(8)) & fx.Int32(0xFF) @@ -440,7 +440,7 @@ def _publish_tile_range( comm_ops.atomic_add_system(remote_queue_tail, fx.Int32(1)) ) ptr_buf_tensor(remote_queue, fx.Int32)[ready_slot] = tile - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) comm_ops.fence_system_release() comm_ops.store_i32_system(remote_queue_epoch, ready_slot, payload_epoch) @@ -552,7 +552,7 @@ def dp(i): ptr_buf_tensor(remote_weights, fx.Int32)[payload_row] = weight_bits ptr_buf_tensor(remote_srcmap, fx.Int32)[payload_row] = source_encoding - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() if tid == fx.Int32(0): comm_ops.fence_system_release() @@ -659,7 +659,7 @@ def dp(i): work_tail[fx.Int32(0)] = ready_work max_expert_tiles_buffer[fx.Int32(0)] = max_expert_tiles - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) comm_ops.fence_system_release() for source in range(lane, fz_npes, 64): remote_ready = plan_ready_table[source] @@ -703,7 +703,7 @@ def dp(slot): destination_tile_m_lane = ptr_buf_tensor( destination_ready_rows, fx.Int32 )[fx.Int32(0)] - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) destination_tile_m = fx.Int32( fx.rocdl.readfirstlane(T.i32, destination_tile_m_lane) ) @@ -726,7 +726,7 @@ def dp(slot): cache_modifier=2, ) ) - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) group_count_lane = fx.Int32(0) group_source_prefix_lane = fx.Int32(0) for source in range_constexpr(npes): @@ -767,7 +767,7 @@ def dp(slot): cache_modifier=2, ) ) - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) normal_count = fx.Int32(0) normal_source_prefix = fx.Int32(0) group_count = group_member.select( @@ -959,7 +959,7 @@ def _derive_next_fanout_pairs( next_parity * fx.Int32(npes) + fx.Int32(destination), packed, ) - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) comm_ops.fence_agent_release() @@ -1030,7 +1030,7 @@ def dp(i): group_phase_base + fx.Int32(group_blocks), ) comm_ops.fence_agent_acquire() - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() comm_ops.fence_agent_acquire() @@ -1044,7 +1044,7 @@ def dp(i): fx.Int32(group_block * total_segments) + segment ] local_hist[segment] = segment_count - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() if warp == fx.Int32(0): @@ -1059,7 +1059,7 @@ def dp(i): payload_chunk_rows=payload_chunk_rows, dispatch_blocks=dispatch_blocks, ) - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() comm_ops.fence_agent_release() @@ -1075,7 +1075,7 @@ def dp(i): remote_count_matrix[fx.Int32(fz_rank * total_segments) + segment] = ( count ) - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() # Warp 0 plans local experts after all source matrices arrive. @@ -1323,13 +1323,13 @@ def dp(i): if lane == fx.Int32(0): num_valid_buffer[fx.Int32(0)] = row_carry max_expert_tiles_buffer[fx.Int32(0)] = max_expert_tiles - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) comm_ops.fence_system_release() for source in range(lane, fz_npes, 64): remote_ready = plan_ready_table[source] ready_index = parity * fx.Int32(fz_npes) + fx.Int32(fz_rank) comm_ops.store_i32_system(remote_ready, ready_index, expected) - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) elif warp == fx.Int32(1): # Build the global-expert exclusive prefix cooperatively. pairs_per_lane = (total_segments + 63) // 64 @@ -1360,13 +1360,13 @@ def dp(i): block_hist[block_index] = block_prefix block_prefix = block_prefix + block_count source_prefix = source_prefix + lane_counts[item] - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) comm_ops.fence_agent_release() # All compact offsets are derived locally from the exchanged histogram. # The prepare kernel owns this single synchronization edge; MegaStage1 # never recounts or regroups routes. - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() # Warp 0 observes COUNT_DONE and performs the system acquire above, but # every warp participates in the deterministic offset derivation below. @@ -1383,13 +1383,13 @@ def dp(i): total_experts=fz_total_experts, total_segments=total_segments, ) - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() if tid == fx.Int32(0): comm_ops.fence_system_release() comm_ops.store_i32_system(a_pair_ready, parity, expected) - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() if tid == fx.Int32(0): comm_ops.wait_i32_until_equals( @@ -1493,13 +1493,13 @@ def record_count(segment): if active_route: route_segment[route] = cached_segment - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() for segment in range(tid, count_segments, block_threads): block_count = fx.ptr_load(count_scratch + fx.Int64(segment)) block_index = producer_slot * fx.Int32(count_segments) + segment block_hist[block_index] = block_count - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() if tid == fx.Int32(0): comm_ops.fence_agent_release() @@ -1514,7 +1514,7 @@ def record_count(segment): block_index = producer_slot * fx.Int32(count_segments) + segment block_base = block_hist[block_index] fx.ptr_store(block_base, count_scratch + fx.Int64(segment)) - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() for token_batch in range( token_batch0, i32_cur_tok, token_stride @@ -1543,7 +1543,7 @@ def record_count(segment): pair_entry = shared_group.select(group_entry, route) pair_order[position] = pair_entry - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() if tid == fx.Int32(0): comm_ops.fence_agent_release() @@ -1916,7 +1916,7 @@ def _store_token(destination_buffer, vals): _row_body(row, wk, None) if chunk_active: - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() if tid == fx.Int32(0): if group_task: diff --git a/aiter/ops/flydsl/kernels/mega_moe/gemm1.py b/aiter/ops/flydsl/kernels/mega_moe/gemm1.py index 3e05c3e2ccf..82e027aebee 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/gemm1.py +++ b/aiter/ops/flydsl/kernels/mega_moe/gemm1.py @@ -290,7 +290,8 @@ def build_fused_gemm1(*, x_tensor, w_rsrc, sw_rsrc, sx_rsrc, m_repeat, num_acc_n, a_k_step_bytes, total_threads, k_iters, a_lds_i32, n_tiles, expert_offset, b_cache_modifier, swizzle_a, pipe_weights, mfma_amajor, async_a_copy, use_tile_resource, indirect_input, indexed_input=False, row_map_rsrc=None, - source_rows=0, swiglu_limit=0.0, pair_k=False, evec): + source_rows=0, swiglu_limit=0.0, pair_k=False, evec, + act="silu", situ_beta=1.0, situ_linear_beta=1.0): # fmt: on """Build the GEMM1 atoms and return its expert resolver and tile runner.""" sched = TileScheduler( @@ -329,7 +330,8 @@ def build_fused_gemm1(*, x_tensor, w_rsrc, sw_rsrc, sx_rsrc, epi = SiluQuantEpilogue(out_rsrc=out_rsrc, out_scale_rsrc=os_rsrc, sorted_rsrc=trb_rsrc, tokens=0, inter_dim=inter_dim, m_repeat=m_repeat, num_acc_n=num_acc_n, sort_block_m=sort_block_m, tile_n=tile_n, num_waves=num_waves, lds_out=c_tile, swiglu_limit=swiglu_limit, always_valid=True, - out_tensor=out_tensor if use_tile_resource else None, evec=evec) + out_tensor=out_tensor if use_tile_resource else None, evec=evec, + act=act, situ_beta=situ_beta, situ_linear_beta=situ_linear_beta) # fmt: on def _decode(flat): @@ -364,6 +366,7 @@ def compile_gemm1( mfma_amajor: bool = False, swizzle_a: bool = True, async_a_copy: bool = False, use_tile_resource: bool = True, waves_per_eu_hint: int = 2, b_cache_modifier: int = 0, swiglu_limit: float = 0.0, + act: str = "silu", situ_beta: float = 1.0, situ_linear_beta: float = 1.0, ): # fmt: on """Compile standalone group GEMM1 from the fused Stage1 compute body.""" @@ -381,6 +384,7 @@ def compile_gemm1( assert num_acc_n % 2 == 0 and m_repeat % 2 == 0 a_k_step_bytes = tile_k + out_row_bytes = inter_dim k_iters = model_dim // tile_k total_threads = num_waves * 64 a_lds_size = sort_block_m * a_k_step_bytes @@ -413,10 +417,11 @@ def kernel( if const_expr(use_tile_resource): out_rsrc = None else: + out_elem = fx.Int16 out_rsrc = ptr_buf_tensor( fx.get_iter(out), - fx.Int16, - num_records_bytes=num_valid * fx.Int32(inter_dim), + out_elem, + num_records_bytes=num_valid * fx.Int32(out_row_bytes), ) scale_cols = (inter_dim // 32 + 7) // 8 * 8 os_rsrc = ptr_buf_tensor( @@ -439,6 +444,8 @@ def kernel( pipe_weights=pipe_weights, mfma_amajor=mfma_amajor, async_a_copy=async_a_copy, use_tile_resource=use_tile_resource, indirect_input=False, swiglu_limit=swiglu_limit, evec=8 if use_tile_resource else 2, + act=act, situ_beta=situ_beta, situ_linear_beta=situ_linear_beta, + ) total_work = (num_valid // fx.Int32(sort_block_m)) * fx.Int32(n_tiles) for flat in range(fx.block_idx.x, total_work, grid_x): @@ -468,7 +475,8 @@ def gemm1_kernel( tile_n: int = 256, tile_k: int = 256, num_waves: int = 4, grid_mult: int = 4, pipe_weights: bool = True, mfma_amajor: bool = False, swizzle_a: bool = True, async_a_copy: bool = False, use_tile_resource: bool = True, waves_per_eu_hint: int = 2, - num_cu: int = 256, b_cache_modifier: int = 0, swiglu_limit: float = 0.0, + num_cu: int = 256, b_cache_modifier: int = 0, + swiglu_limit: float = 0.0, act: str = "silu", situ_beta: float = 1.0, situ_linear_beta: float = 1.0, ): # fmt: on """Run standalone MegaMoEV2 group GEMM1 and return ``(out, out_scale)``.""" @@ -486,10 +494,11 @@ def gemm1_kernel( pipe_weights=pipe_weights, mfma_amajor=mfma_amajor, swizzle_a=swizzle_a, async_a_copy=async_a_copy, use_tile_resource=use_tile_resource, waves_per_eu_hint=waves_per_eu_hint, b_cache_modifier=b_cache_modifier, - swiglu_limit=swiglu_limit, + swiglu_limit=swiglu_limit, act=act, situ_beta=situ_beta, situ_linear_beta=situ_linear_beta, ) _run_compiled( - launch, out, x, w.view(torch.uint8), scale_x, scale_w.view(torch.uint8), tile_row_base, expert_ids, out_scale, + launch, out, x, w.view(torch.uint8), scale_x, scale_w.view(torch.uint8), + tile_row_base, expert_ids, out_scale, fx.Int32(num_valid), fx.Int32(grid_x), stream, ) return out, out_scale diff --git a/aiter/ops/flydsl/kernels/mega_moe/gemm2.py b/aiter/ops/flydsl/kernels/mega_moe/gemm2.py index 6ca3a5eeaa1..7a1d925e14d 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/gemm2.py +++ b/aiter/ops/flydsl/kernels/mega_moe/gemm2.py @@ -182,12 +182,16 @@ def gemm2_compute_v2( SBM=None, g2_bhoist=True, g2_ascale_pf=True, + g2_b2stage=False, + g2_deep_a_pipeline=False, expert_offset=0, explicit_m_row=None, explicit_n_block=None, explicit_expert=None, ): """Run GEMM2, optionally using an explicitly selected expert row/tile.""" + assert aStages == (4 if g2_deep_a_pipeline else kStages + 1) + assert not g2_deep_a_pipeline or g2_b2stage # SBM is the sort padding unit; BM is the compute tile and must divide SBM. if SBM is None: SBM = BM @@ -219,6 +223,8 @@ def gemm2_compute_v2( num_n_blocks = N_OUT_rt // fx.Int32(BN) KH4 = K_rt // fx.Int32(8) # i32 col stride (= K_HALF//4) K_SCALE_CHUNKS_MAX = INTER_MAX // 256 + kFenceEvery = (aStages - kStages) if g2_deep_a_pipeline else 1 + assert kFenceEvery >= 1 # Padded shapes mask weight tiles beyond the real K/N extents. N_real = None @@ -516,7 +522,7 @@ def store_c_carry(state): n += 1 return n - if const_expr(BM == 64 and BN == 256): + if const_expr(BM == 64 and BN == 256 and not g2_b2stage): # BM64/BN256 uses the 1-stage B path unconditionally. for kt_iv, state in range( fx.Int32(0), @@ -650,9 +656,13 @@ def prefetch_next_b(kt_rt): kt_rt = fx.Int32(kt_iv) if const_expr(g2_bhoist): prefetch_next_b(kt_rt) - gpu.barrier() - issue_a_ds_read(kt_rt % fx.Int32(aStages)) + if const_expr(g2_deep_a_pipeline): + if (kt_rt % fx.Int32(kFenceEvery)) == fx.Int32(0): + gpu.barrier() + else: + gpu.barrier() nxt_a = kt_rt + fx.Int32(kStages) + issue_a_ds_read(kt_rt % fx.Int32(aStages)) if nxt_a < K_TILES_RT: issue_a_load_lds(nxt_a % fx.Int32(aStages), nxt_a) # A-scale from the prefetch carry (g2_ascale_pf) or loaded synchronously here. diff --git a/aiter/ops/flydsl/kernels/mega_moe/gemm_util.py b/aiter/ops/flydsl/kernels/mega_moe/gemm_util.py index 2a550dec39e..ee09fb37215 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/gemm_util.py +++ b/aiter/ops/flydsl/kernels/mega_moe/gemm_util.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors -"""A8W4 GEMM utilities for MegaMoE v2 stage1.""" +"""A8W4 GEMM utilities for MegaMoE v2 Stage1.""" import flydsl.compiler as flyc import flydsl.expr as fx @@ -8,15 +8,16 @@ from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec +from ..act import situ_mul, tanh_via_sigmoid_f32 from ..tensor_shim import buf_copy_load, buf_copy_store, ptr_buf_tensor _PACK = 2 # fp4 micro-scale pack (per-32 E8M0): pack_M = pack_N = pack_K = 2 @flyc.jit -def _load_row_map_subgroup(row_map_rsrc, row_index, lane): - """Load one source row per 16-lane K-slice group and broadcast it.""" - subgroup_lane = lane & fx.Int32(15) +def _load_row_map_subgroup(row_map_rsrc, row_index, lane, subgroup_width): + """Load and broadcast one source row per K-step copy subgroup.""" + subgroup_lane = lane % fx.Int32(subgroup_width) packed = fx.Int32(0) if subgroup_lane == fx.Int32(0): packed = row_map_rsrc[row_index] @@ -32,8 +33,7 @@ def _load_row_map_subgroup(row_map_rsrc, row_index, lane): def wait_lds_barrier(vmcnt=63): """Drain LDS writes and optionally older VMEM while preserving newer loads.""" - waitcnt = (vmcnt & 0xF) | ((vmcnt & 0x30) << 10) | (7 << 4) - fx.rocdl.s_waitcnt(waitcnt) + fx.rocdl.s_waitcnt(vmcnt=vmcnt, lgkmcnt=0) fx.barrier() @@ -90,7 +90,9 @@ def __init__( assert row_map_rsrc is not None and self._source_rows > 0 if const_expr(self._async_copy): assert total_threads % 64 == 0 - assert (sort_block_m * 16) % total_threads == 0 + total_chunks = sort_block_m * (k_step_bytes // 16) + assert total_chunks % 64 == 0 + assert total_chunks < total_threads or total_chunks % total_threads == 0 assert row_bytes % 16 == 0 and k_step_bytes % 16 == 0 self._dma_atom = fx.make_copy_atom( fx.rocdl.BufferCopyLDS128b(), @@ -147,14 +149,14 @@ def for_tile(self, tile_row_base_i32): self._row_map_rsrc, tile_row_base_i32 + row, lane, + chunks_per_row, ) source_row = packed & fx.Int32(0xFFFFFF) row_byte = source_row * fx.Int32(self._row_bytes) if const_expr(self._swizzle): col_i32 = chunk * fx.Int32(4) - swz = row * fx.Int32(row_stride_i32) + ( - col_i32 ^ ((row & fx.Int32(15)) << fx.Int32(2)) - ) + swizzle_mask = (row & fx.Int32(15)) << fx.Int32(2) + swz = row * fx.Int32(row_stride_i32) + (col_i32 ^ swizzle_mask) lds_byte = swz * fx.Int32(4) else: lds_byte = lin * fx.Int32(16) @@ -193,47 +195,61 @@ def prefetch_to_lds(self, k_step_byte_off, lds_dst, base_i32=0): """Issue swizzled direct global-to-LDS copies with a wave-uniform LDS base.""" koff = fx.Int32(k_step_byte_off) base_bytes = fx.Int32(base_i32) * fx.Int32(4) - lds_f8 = fx.recast_iter( - fx.Float8E4M3FN, - lds_dst.ptr, - ) - total_chunks = self._sort_block_m * 16 + lds_elem = fx.recast_iter(fx.Float8E4M3FN, lds_dst.ptr) + chunks_per_row = self._k_step_bytes // 16 + total_chunks = self._sort_block_m * chunks_per_row + copy_waves = min(total_chunks, self._total_threads) // 64 for round_base in range_constexpr( 0, total_chunks, self._total_threads, ): - physical = fx.Int32(round_base) + fx.Int32(self._tx) - row = physical // fx.Int32(16) - physical_chunk = physical % fx.Int32(16) - if const_expr(self._swizzle): - logical_chunk = physical_chunk ^ (row & fx.Int32(15)) - else: - logical_chunk = physical_chunk - source_row = row - if const_expr(self._indexed_input): - source_index = round_base // self._total_threads - source_byte = self._chunks[source_index][1] - source_row = source_byte // fx.Int32(self._row_bytes) - src_byte = ( - source_row * fx.Int32(self._row_bytes) - + koff - + logical_chunk * fx.Int32(16) - ) - src = fx.slice( - self._tile_dma, - (None, src_byte), - ) - wave_base = base_bytes + fx.Int32((round_base + self._wave * 64) * 16) - dst = fx.make_view( - fx.add_offset(lds_f8, wave_base), - fx.make_layout(1, 1), + copy_active = ( + self._wave < fx.Int32(copy_waves) + if const_expr(total_chunks < self._total_threads) + else fx.Boolean(True) ) - fx.copy(self._dma_atom, src, dst) + + # issue_copy is invoked below in this same iteration and never + # escapes it, so the late-binding B023 warns about cannot happen. + @flyc.jit + def issue_copy(active): + if active: + physical = fx.Int32(round_base) + fx.Int32(self._tx) # noqa: B023 + row = physical // fx.Int32(chunks_per_row) + physical_chunk = physical % fx.Int32(chunks_per_row) + if const_expr(self._swizzle): + logical_chunk = physical_chunk ^ (row & fx.Int32(15)) + else: + logical_chunk = physical_chunk + source_row = row + if const_expr(self._indexed_input): + source_index = round_base // self._total_threads # noqa: B023 + source_byte = self._chunks[source_index][1] + source_row = source_byte // fx.Int32(self._row_bytes) + src_byte = ( + source_row * fx.Int32(self._row_bytes) + + koff + + logical_chunk * fx.Int32(16) + ) + src = fx.slice( + self._tile_dma, + (None, src_byte), + ) + wave_base = base_bytes + fx.Int32( + (round_base + self._wave * 64) * 16 # noqa: B023 + ) + dst = fx.make_view( + fx.add_offset(lds_elem, wave_base), + fx.make_layout(1, 1), + ) + fx.copy(self._dma_atom, src, dst) + + issue_copy(copy_active) class AS2RLoader: - """Load one FP8 MFMA A operand from row-major LDS.""" + """Load one FP8 scaled-MFMA A operand from row-major LDS.""" def __init__(self, *, k_step_bytes, swizzle=False): self._k_step_bytes = k_step_bytes @@ -249,7 +265,7 @@ def _load_16b(self, lds_src, i32_off): return Vec(v).bitcast(fx.Int32) def load_operand(self, lds_src, mi, ksub, base_i32=0): - """A operand for m-tile mi, K=128 sub-block ksub. base_i32 = ping/pong. fp8 32-per-lane = 16@K + 16@K+64.""" + """Load one K=128 A operand; FP4 uses one 16B chunk and FP8 uses two.""" row = fx.Int32(mi * 16) + fx.Int32(self._lane % 16) row_i32 = row * fx.Int32(self._k_step_bytes // 4) + fx.Int32(base_i32) klane4 = fx.Int32(self._lane // 16) * fx.Int32(4) @@ -339,7 +355,7 @@ def load_step(self, base_row_i32, kstep_i32): class AScaleLoader: - """Per-1x32 E8M0 A scales STAGED to LDS once per tile (stage), read via ds_read in K-loop (kills VMEM flood).""" + """Stage per-1x32 E8M0 A scales to LDS once per tile.""" def __init__( self, @@ -365,7 +381,7 @@ def __init__( assert row_map_rsrc is not None def stage(self, lds_ascale, tile_row_base_i32): - """Coalesced gmem->LDS copy of this tile's e8m0 A-scale block [sort_block_m, n_scale]. Call before K-loop.""" + """Copy this tile's E8M0 scale block from global memory to LDS.""" total = self._sort_block_m * self._n_scale assert total % 16 == 0, "A-scale tile must contain whole 16-byte copy chunks" base = tile_row_base_i32 * fx.Int32(self._n_scale) @@ -407,6 +423,7 @@ def copy_indexed_slot(slot: fx.Int32): self._row_map_rsrc, tile_row_base_i32 + row, lane, + padded_groups_per_row, ) source_row = packed & fx.Int32(0xFFFFFF) lin = row * fx.Int32(groups_per_row) + group @@ -432,11 +449,9 @@ def copy_chunk(lin: fx.Int32): copy_chunk(lin) def load_step(self, lds_ascale, kstep_i32): - """One packed i32 per pack-group, read from the LDS-staged A-scale (ds_read).""" - lane_row = fx.Int32(self._lane % 16) - col0 = kstep_i32 * fx.Int32(8) + fx.Int32( - self._lane // 16 - ) # e8m0 col = kstep*8 + ksub*4 + KLane + """Read one packed i32 per scale group from LDS.""" + lane_row = self._lane % 16 + col0 = kstep_i32 * 8 + self._lane // 16 out = [] for g in range_constexpr(self._n_groups): r0 = fx.Int32(g * 32) + lane_row @@ -444,9 +459,7 @@ def load_step(self, lds_ascale, kstep_i32): b = [] for ksub in range_constexpr(_PACK): for rr in (r0, r1): - b.append( - self._read_scale_lds(lds_ascale, rr, col0 + fx.Int32(ksub * 4)) - ) + b.append(self._read_scale_lds(lds_ascale, rr, col0 + ksub * 4)) packed = ( b[0] | (b[1] << fx.Int32(8)) @@ -633,12 +646,12 @@ def _e8m0_scale(amax): class SiluQuantEpilogue: - """SwiGLU followed by FP8 quantization and per-32 E8M0 output scales.""" + """silu(gate)*up -> fp8 + per-32 E8M0 out-scale.""" # fmt: off def __init__(self, *, out_rsrc, out_scale_rsrc, sorted_rsrc, tokens, inter_dim, m_repeat, num_acc_n, - sort_block_m, tile_n, num_waves, lds_out, swiglu_limit=0.0, always_valid=False, out_tensor=None, - evec): + sort_block_m, tile_n, num_waves, lds_out, evec, swiglu_limit=0.0, always_valid=False, + out_tensor=None, act="silu", situ_beta=1.0, situ_linear_beta=1.0): # fmt: on self._out_rsrc = out_rsrc self._out_scale_rsrc = out_scale_rsrc @@ -652,8 +665,16 @@ def __init__(self, *, out_rsrc, out_scale_rsrc, sorted_rsrc, tokens, inter_dim, self._num_waves = num_waves self._lds_out = lds_out self._swiglu_limit = float(swiglu_limit) + self._act = act + self._situ_beta = float(situ_beta) + self._situ_linear_beta = float(situ_linear_beta) + assert act in ("silu", "situv2"), f"unsupported mega_moe act {act!r}" + if act == "situv2": + assert self._situ_beta > 0.0 and self._situ_linear_beta > 0.0, \ + "situv2 requires positive beta / linear_beta" self._always_valid = always_valid self._out_tensor = out_tensor + self._out_row_bytes = inter_dim # Columns quantized per lane in the store pass. 2: one 16-bit store and a # 16-lane max per pair; 8: one 8-byte store and a 4-lane max per 8 values. # Without out_tensor, out_rsrc is the whole output in that store's width: @@ -673,12 +694,28 @@ def _combine(self, acc): for mi in range_constexpr(self._m_repeat): for ni in range_constexpr(gui_n): g_idx = mi * self._num_acc_n + ni * _PACK - out.append(self._silu_mul(acc[g_idx], acc[g_idx + 1])) + out.append(self._act_mul(acc[g_idx], acc[g_idx + 1])) return out - def _silu_mul(self, gate_v4, up_v4): + def _act_mul(self, gate_v4, up_v4): gv = Vec(gate_v4) uv = Vec(up_v4) + if const_expr(self._act == "situv2"): + # SiTUv2 (Kimi-K3 hidden_act="situ"): + # situ(g) = beta * tanh(g / beta) * sigmoid(g) + # situ_up(u) = linear_beta * tanh(u / linear_beta) + # No hard clamp -- tanh self-saturates, matching SGLang SituAndMul and + # the split-K silu_and_mul_fq situv2 path. beta/linear_beta are baked + # per network (one K3 value), so they stay compile-time here. + b = fx.Float32(self._situ_beta) + br = fx.Float32(1.0 / self._situ_beta) + lb = fx.Float32(self._situ_linear_beta) + lbr = fx.Float32(1.0 / self._situ_linear_beta) + elems = [ + situ_mul(gv[i], uv[i], b, br, lb, lbr, tanh=tanh_via_sigmoid_f32) + for i in range_constexpr(4) + ] + return Vec.from_elements(elems, fx.Float32) if self._swiglu_limit <= 0: elems = [self._silu(gv[i]) * uv[i] for i in range_constexpr(4)] return Vec.from_elements(elems, fx.Float32) @@ -781,10 +818,10 @@ def store(self, acc, tile_i32, tile_row_base_i32, n_tile_base_i32): scaled0 = v0 * quant_scale scaled1 = v1 * quant_scale packed = rocdl.cvt_pk_fp8_f32(T.i32, scaled0, scaled1, fx.Int32(0), 0) - short_raw = fx.Int32(packed).to(fx.Int16) + payload = fx.Int32(packed).to(fx.Int16) out_byte = out_row_base + gcol out_byte = valid.select(out_byte, fx.Int32(0x40000000)) - out_rsrc[out_byte // fx.Int32(2)] = short_raw + out_rsrc[out_byte // fx.Int32(2)] = payload col_s = gcol >> fx.Int32(5) is_writer = (gcol & fx.Int32(31)) == fx.Int32(0) diff --git a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py index db77b4bac41..995ba2194b4 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py +++ b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_config.py @@ -10,6 +10,7 @@ TOKEN_BUCKETS = ( 1, + 2, 4, 8, 16, @@ -31,16 +32,21 @@ # 128; AITER_MEGA_FIXED_SLOT_MAX_MTPR extends it (see envs.py). FIXED_SLOT_MAX_MTPR = envs.AITER_MEGA_FIXED_SLOT_MAX_MTPR MAX_MTPR_CLASS = 32768 -# Source-indexed payload storage cuts the maximum-capacity activation buffer -# by roughly ``topk``. Keep every smaller capacity on the historical layout. INDEXED_PAYLOAD_MIN_MTPR = MAX_MTPR_CLASS INDEXED_PAYLOAD_MIN_SBM = 128 +# Source-indexed payload storage cuts the maximum-capacity activation buffer +# by roughly ``topk``. Keep every smaller capacity on the historical layout. REFERENCE_EXPERTS_PER_RANK = 48 +KIMI_K3_EXPERTS_PER_RANK = 112 # Compact route metadata dedicates ten bits to the global expert/group segment. # Under the EP8 protocol this admits 8 * 127 expert segments plus 8 group # segments. The next expert would require segment 1024 and cannot be encoded. MAX_FANOUT_SEGMENTS = 1024 MAX_FANOUT_EXPERTS_PER_RANK = 256 +P2P_QUANT_AUTO = "auto" +P2P_QUANT_NONE = "none" +P2P_QUANT_FP8_BLOCKWISE = "fp8_blockwise_1x32" +SUPPORTED_P2P_QUANT_MODES = (P2P_QUANT_NONE, P2P_QUANT_FP8_BLOCKWISE) def fixed_stage1_epoch_slot(grid_mult: int, num_dispatch_cu: int, num_cu: int) -> int: @@ -94,6 +100,7 @@ class Stage2Config: skew_cu: int = 0 block_k: int = 256 b_hoist: bool = True + b2stage: bool = False ascale_prefetch: bool = True spatial_partition: int = 402 bf16_lds: bool = False @@ -101,6 +108,7 @@ class Stage2Config: pair_cu: int = 0 pair_block_m: int = 32 pair_block_n: int = 256 + deep_a_pipeline: bool = False @dataclass(frozen=True, slots=True) @@ -120,6 +128,8 @@ def __post_init__(self): raise ValueError(f"unsupported p2p_quant={self.p2p_quant!r}") if self.p2p_quant != "none" and self.stage2.bf16_lds: raise ValueError("FP8 P2P requires Stage2 bf16_lds=False") + if self.stage2.deep_a_pipeline and not self.stage2.b2stage: + raise ValueError("Stage2 deep_a_pipeline requires b2stage=True") @dataclass(frozen=True, slots=True) @@ -383,15 +393,31 @@ def _select_bucket_config( return MegaMoEConfig(stage1=stage1, stage2=stage2, p2p_quant="none") +def _replace_config( + config: MegaMoEConfig, + *, + stage1: dict[str, object] | None = None, + stage2: dict[str, object] | None = None, +) -> MegaMoEConfig: + return replace( + config, + stage1=replace(config.stage1, **stage1) if stage1 else config.stage1, + stage2=replace(config.stage2, **stage2) if stage2 else config.stage2, + ) + + def select_mega_moe_config( tokens: int, mtpr: int, *, + p2p_quant: str = P2P_QUANT_AUTO, experts_per_rank: int = REFERENCE_EXPERTS_PER_RANK, model_dim: int = 7168, inter_dim: int = 3072, world_size: int = 8, ) -> MegaMoEConfig: + if p2p_quant not in (P2P_QUANT_AUTO, *SUPPORTED_P2P_QUANT_MODES): + raise ValueError(f"unsupported p2p_quant={p2p_quant!r}") if mtpr <= 0 or mtpr & (mtpr - 1): raise ValueError(f"mtpr={mtpr} must be a positive power of two") if tokens > mtpr: @@ -422,19 +448,78 @@ def select_mega_moe_config( f"MegaMoE v2 fanout needs {total_segments} segments, exceeding " f"the {MAX_FANOUT_SEGMENTS}-segment route metadata limit" ) - return _select_bucket_config( - bucket, + # Bucket 2 is an exact tuning identity, while its untuned A8 base remains + # bucket 1 so adding it does not change existing A8 geometry. + base_bucket = 1 if bucket == 2 else bucket + config = _select_bucket_config( + base_bucket, mtpr_class, model_dim, inter_dim, fixed_slot_dispatch, ) + if p2p_quant == P2P_QUANT_AUTO: + desired_p2p = ( + P2P_QUANT_FP8_BLOCKWISE if mtpr >= P2P_FP8_MIN_MTPR else config.p2p_quant + ) + else: + desired_p2p = p2p_quant + if desired_p2p != config.p2p_quant: + config = replace(config, p2p_quant=desired_p2p) + if experts_per_rank == KIMI_K3_EXPERTS_PER_RANK: + # Kimi-K3 packs 112 experts per rank, so the dispatch carries far more + # work than the generic num_dispatch_cu=32 (sized for 32-expert ranks) + # can keep fed. Use 160 CUs through the middle buckets and 96 at the + # ends (bs=1 cannot fill 160; the largest buckets spend more on + # scheduling than they recover). + config = _replace_config( + config, + stage1={"num_dispatch_cu": 96 if bucket == 1 or bucket >= 2048 else 160}, + ) + # Stage2 runs in a different phase than the dispatch, so its tuning + # composes with the CU change above. + if bucket >= 2048: + config = _replace_config( + config, + stage2={"persist_cu": 192 if bucket <= 4096 else 256}, + ) + # At this inter_dim Stage2 is a large fraction of e2e, so widen its + # N-tile to BN256 off the generic BN128. block_m=64 fills the SBM64 + # sort block for 256-2048; the largest buckets already run BM64. + if bucket >= 4096: + config = _replace_config( + config, + stage2={"block_n": 256, "b2stage": True}, + ) + elif 256 <= bucket <= 2048: + config = _replace_config( + config, + stage2={"block_m": 64, "block_n": 256, "b2stage": True}, + ) + elif 8 <= bucket <= 128: + # BN256 extends down to bs 8-128; bs 1-4 stay BN128 (too few tokens + # to fill the wider tile). SBM32 here pins block_m to 32. b2stage + # stays off here: it is neutral for these small buckets and only + # pays off once Stage2 dominates at bs>=256. + config = _replace_config( + config, + stage2={"block_n": 256}, + ) + + if config.stage2.aligned_pair and config.p2p_quant != P2P_QUANT_FP8_BLOCKWISE: + # The aligned-pair Stage2 kernel only emits the FP8 P2P row layout, so + # reject the combination here instead of at launch. + raise ValueError( + f"p2p_quant={config.p2p_quant!r} conflicts with aligned-pair Stage2" + ) + return config @cache def build_mega_moe_bundle_plan( mtpr: int, *, + p2p_quant: str = P2P_QUANT_AUTO, experts_per_rank: int = REFERENCE_EXPERTS_PER_RANK, model_dim: int = 7168, inter_dim: int = 3072, @@ -461,6 +546,7 @@ def build_mega_moe_bundle_plan( config = select_mega_moe_config( bucket, mtpr, + p2p_quant=p2p_quant, experts_per_rank=experts_per_rank, model_dim=model_dim, inter_dim=inter_dim, diff --git a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_prepare.py b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_prepare.py index 09b184e3c53..b7e9ebbffe1 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_prepare.py +++ b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_prepare.py @@ -174,7 +174,7 @@ def disp_ptr(slot): hist_rsrc = ptr_buf_tensor(local_hist, fx.Int32) for segment in range(tid, total_segments, block_threads): hist_rsrc[segment] = fx.Int32(0) - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() if tid == fx.Int32(0): comm_ops.store_i32_system( diff --git a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage1.py b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage1.py index 1277daa5109..87f4b500101 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage1.py +++ b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage1.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors -"""Fused stage1 with low-ID dispatch producers and oversubscribed FP8xFP4 grouped-GEMM1 consumers.""" +"""Fused Stage1 dispatch with FP8 activation and FP4-weight GEMM1.""" import functools @@ -13,6 +13,7 @@ from flydsl.runtime.device import get_rocm_arch from .. import communication_ops_utils as comm_ops +from ..kernels_common import ceildiv from ..tensor_shim import ( _preload_compiled, _run_compiled, @@ -49,10 +50,6 @@ def __init__(self, kernel, grid_x, block_x, waves_per_eu_hint): self.waves_per_eu_hint = int(waves_per_eu_hint) -def ceildiv(a, b): - return (a + b - 1) // b - - def _validate_fixed_slot_geometry( npes, experts_per_rank, max_tokens_per_rank, cap, tile_m ): @@ -100,6 +97,8 @@ def compile_mega_moe_stage1( work_shards: int | None = None, payload_chunk_rows: int = 0, tile_state_stride: int = 0, swiglu_limit: float = 0.0, + act: str = "silu", situ_beta: float = 1.0, situ_linear_beta: float = 1.0, + bounds_check: bool = False, _return_kernel_spec: bool = False, ): arch = str(get_rocm_arch() or "") @@ -150,10 +149,8 @@ def compile_mega_moe_stage1( NUM_ACC_N = n_per_wave // 16 assert NUM_ACC_N % 2 == 0 and M_REPEAT % 2 == 0 - TILE_K_BYTES = tile_k // 2 - assert TILE_K_BYTES % 128 == 0 + assert tile_k == 256, "MegaMoE v2 GEMM1 requires tile_k=256" A_K_STEP_BYTES = tile_k - assert A_K_STEP_BYTES == 256, "MegaMoE v2 GEMM1 requires tile_k=256" K_ITERS = model_dim // tile_k TOTAL_THREADS = NUM_WAVES * 64 WORK_SHARDS = 4 if work_shards is None and int(fuse_mtpr) >= 8192 else 8 @@ -195,7 +192,8 @@ def compile_mega_moe_stage1( ) # Small batches stream B; large batches cache it across M tiles. b_cache_modifier = int(b_nt) if int(b_nt) >= 0 else (3 if fz_mtpr <= 512 else 0) - fz_n_i32, fz_nbytes = model_dim // 4, model_dim + fz_nbytes = model_dim + fz_n_i32 = fz_nbytes // 4 fz_scale_bytes = int(fuse_scale_dim) fz_scale_n_i32 = (fz_scale_bytes + 3) // 4 if fz_scale_bytes > 0 else 0 if fixed_slot_dispatch and fz_scale_n_i32 > 64: @@ -226,9 +224,38 @@ class SharedStorage: EPI_EVEC = envs.AITER_MEGA_S1_EPI_EVEC dispatch_path = "fixedslot" if fixed_slot_dispatch else "compact" swiglu_suffix = "" if swiglu_limit <= 0 else f"_sl{str(float(swiglu_limit)).replace('.', 'p')}" + # Activation is baked into the epilogue, so it (and its baked beta/linear_beta) + # must key the kernel or a situv2 launch would collide with the silu cache. + act_suffix = "" if act == "silu" else ( + f"_{act}" + f"b{str(float(situ_beta)).replace('.', 'p')}" + f"lb{str(float(situ_linear_beta)).replace('.', 'p')}" + ) WORK_BATCH = 1 + + # Optional hardware bounds checking. A V# descriptor with a finite + # num_records makes the hardware drop out-of-range stores and return zero + # for out-of-range loads, so an indexing mistake surfaces as a wrong result + # instead of a page fault. That matters because amdgpu runs with + # noretry=1 here: a fault is non-retryable and can leave the queue in a + # state only a privileged reset clears. + # + # Only buffers whose exact allocation is known at this level are bounded. + # A bound that is too small would silently clamp a valid access, which is + # worse than the fault it replaces, so anything uncertain keeps the + # unbounded default. Notably tile_ready/tile_expected/ready_tile_queue are + # allocated as 2 * metadata_blocks while tile_state_stride carries only + # metadata_blocks, hence the doubling below. + WORK_HEAD_BYTES = 8 * 16 * 4 # workspace["work_head"] = zeros(8 * 16, int32) + TILE_STATE_BYTES = tile_state_stride * 2 * 4 + PER_DESTINATION_BYTES = fz_npes * 4 + + def bounded(nbytes): + """Keyword arguments pinning a view's V# extent, when enabled.""" + return {"num_records_bytes": nbytes} if bounds_check else {} kernel_name = ( f"megamoe_stage1_{dispatch_path}_t{sort_block_m}x{tile_n}x{tile_k}" + "_afp8ofp8" f"_w{NUM_WAVES}_gm{grid_mult}" f"_dcu{dispatch_blocks}_pw{int(pipe_weights)}ma{int(mfma_amajor)}sw{int(swizzle_a)}" f"_cgc{grid_x}" @@ -239,9 +266,13 @@ class SharedStorage: f"_rc31_wb{WORK_BATCH}_adaptive_bc" f"_ix{int(indexed_payload)}" f"{swiglu_suffix}" + f"{act_suffix}" f"_ev{EPI_EVEC}b" f"{'_fc' if FAST_COPY else ''}" f"{'_kp2' if PAIR_K else ''}" + # Kept in the name so a bounds-checked build never reuses a cached + # unbounded kernel, or the reverse. + f"{'_bc1' if bounds_check else ''}" ) @flyc.kernel(name=kernel_name, known_block_size=[TOTAL_THREADS, 1, 1]) @@ -258,8 +289,8 @@ def kernel( a_scale_lds = lds.A_scale c_tile = _LdsF32View(fx.recast_iter(fx.Float32, lds.pool.ptr)) disp_rsrc = ptr_buf_tensor(addr_disp, fx.Int64) - parity_rsrc = ptr_buf_tensor(addr_parity, fx.Int32) - expected_rsrc = ptr_buf_tensor(addr_expected, fx.Int32) + parity_rsrc = ptr_buf_tensor(addr_parity, fx.Int32, **bounded(4)) + expected_rsrc = ptr_buf_tensor(addr_expected, fx.Int32, **bounded(8)) def _disp_ptr(slot): return disp_rsrc[fx.Int32(int(slot))] @@ -348,7 +379,11 @@ def _disp_ptr(slot): ) comm_ops.fence_system_acquire() if tid == fx.Int32(0): - work_head_rsrc = ptr_buf_tensor(a_work_head, fx.Int32) + # work_shard indexes this at shard * 16 int32s, so an + # out-of-range WORK_SHARDS lands here first. + work_head_rsrc = ptr_buf_tensor( + a_work_head, fx.Int32, **bounded(WORK_HEAD_BYTES) + ) for shard in range_constexpr(WORK_SHARDS): work_head_rsrc[fx.Int32(shard * 16)] = fx.Int32(0) comm_ops.store_i32_system( @@ -358,13 +393,13 @@ def _disp_ptr(slot): for destination in range_constexpr(fz_npes): group_done_rsrc[fx.Int32(destination)] = fx.Int32(0) if tid == fx.Int32(0): - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) comm_ops.fence_agent_release() parity_rsrc[fx.Int32(0)] = next_parity - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) comm_ops.fence_agent_release() comm_ops.store_i32_system(gate_addr, fx.Int32(0), gate_epoch) - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) fx.barrier() else: if tid == fx.Int32(0): @@ -414,10 +449,14 @@ def _disp_ptr(slot): fx.barrier() producer_destination = producer_slot % fx.Int32(fz_npes) producers_per_destination = ptr_buf_tensor( - a_payload_blocks_per_destination, fx.Int32 + a_payload_blocks_per_destination, + fx.Int32, + **bounded(PER_DESTINATION_BYTES), )[producer_destination] chunks_per_destination = ptr_buf_tensor( - a_payload_chunks_per_destination, fx.Int32 + a_payload_chunks_per_destination, + fx.Int32, + **bounded(PER_DESTINATION_BYTES), )[producer_destination] emit_dispatch_payload( num_waves=NUM_WAVES, fz_epr=fz_epr, fz_k=fz_k, fz_mtpr=fz_mtpr, fz_rank=fz_rank, @@ -495,6 +534,7 @@ def _disp_ptr(slot): row_map_rsrc=srcmap_rsrc, source_rows=source_rows, swiglu_limit=swiglu_limit, + act=act, situ_beta=situ_beta, situ_linear_beta=situ_linear_beta, ) if tid == fx.Int32(0): @@ -517,14 +557,16 @@ def _disp_ptr(slot): total_work = num_m_tiles * fx.Int32(N_TILES) use_ready_order = fx.Int32(0) == fx.Int32(1) if const_expr(compact_dispatch): - max_expert_tiles = ptr_buf_tensor(a_max_expert_tiles, fx.Int32)[ - fx.Int32(0) - ] + max_expert_tiles = ptr_buf_tensor( + a_max_expert_tiles, fx.Int32, **bounded(4) + )[fx.Int32(0)] use_ready_order = max_expert_tiles * fx.Int32(4) >= num_m_tiles def _wait_tile_payload(flat): tile_index = flat // fx.Int32(N_TILES) - expected_tiles = ptr_buf_tensor(addr_tile_expected, fx.Int32)[tile_index] + expected_tiles = ptr_buf_tensor( + addr_tile_expected, fx.Int32, **bounded(TILE_STATE_BYTES) + )[tile_index] comm_ops.wait_i32_until_equals( addr_tile_ready + fx.Int64(tile_index) * fx.Int64(4), expected_tiles, @@ -563,7 +605,14 @@ def _run_work_batch(first_work, scheduled_first): if work < total_work: if const_expr(fixed_slot_dispatch): fx.barrier() - elif not use_ready_order: + else: + # Both the sequential and the ready-order compact paths + # must wait for the tile to be published and acquire + # before reading its per-tile metadata. The ready-order + # path previously skipped this, so a consumer could read + # tile_input_base (and the other per-tile metadata) before + # the planner's store was visible, feeding a stale/garbage + # row into the A-scale gather and faulting under skew. if tid == fx.Int32(0): _wait_tile_payload(work) fx.barrier() @@ -694,6 +743,10 @@ def compile_mega_moe_stage1_bundle( tile_state_stride: int, variants: tuple[Stage1Config, ...], swiglu_limit: float = 0.0, + act: str = "silu", + situ_beta: float = 1.0, + situ_linear_beta: float = 1.0, + bounds_check: bool = False, ): """Compile every production Stage1 variant into one profile module.""" if not variants: @@ -728,6 +781,10 @@ def compile_mega_moe_stage1_bundle( payload_chunk_rows=config.payload_chunk_rows, tile_state_stride=tile_state_stride, swiglu_limit=swiglu_limit, + act=act, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + bounds_check=bounds_check, _return_kernel_spec=True, ) for config in variants @@ -813,7 +870,7 @@ def run_mega_moe_stage1(out, x, w, scale_x, scale_w, sorted_token_ids, expert_id use_tile_resource=True, waves_per_eu_hint=2, b_nt=-1, work_shards=None, payload_chunk_rows=0, tile_state_stride=0, - swiglu_limit=0.0): + swiglu_limit=0.0, act="silu", situ_beta=1.0, situ_linear_beta=1.0, bounds_check=False): launch = compile_mega_moe_stage1( model_dim=model_dim, inter_dim=inter_dim, rank=rank, experts_per_rank=experts_per_rank, fuse_npes=fuse_npes, fuse_topk=fuse_topk, fuse_cap=fuse_cap, fuse_mtpr=fuse_mtpr, @@ -824,7 +881,8 @@ def run_mega_moe_stage1(out, x, w, scale_x, scale_w, sorted_token_ids, expert_id waves_per_eu_hint=waves_per_eu_hint, num_cu=num_cu, num_dispatch_cu=num_dispatch_cu, b_nt=b_nt, work_shards=work_shards, payload_chunk_rows=payload_chunk_rows, tile_state_stride=tile_state_stride, - swiglu_limit=swiglu_limit, + swiglu_limit=swiglu_limit, act=act, situ_beta=situ_beta, situ_linear_beta=situ_linear_beta, + bounds_check=bounds_check, ) _run_compiled( launch, out, x, w, scale_x, scale_w, sorted_token_ids, expert_ids, num_valid_ids, out_scale, diff --git a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage2.py b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage2.py index d189254f22e..4d35ebe0367 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage2.py +++ b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage2.py @@ -299,8 +299,9 @@ def compile_mega_moe_stage2(*, model_dim: int, inter_dim: int, experts: int, top max_tok: int, recv_cap: int | None = None, comb_inp_nbytes: int | None = None, BM: int = 32, BN: int = 256, BK: int = 256, use_nt: bool = True, HIDDEN_MAX: int = 8192, INTER_MAX: int = 8192, a_dtype: str = "fp8", SBM: int | None = None, - persist: bool = False, cu_num: int = 0, has_pad: bool = False, g2_bhoist=None, g2_ascale_pf=None, - g2_spart=None, persist_strided: bool = False, g2_bf16_lds: bool = False, p2p_quant_type: str = "none", + persist: bool = False, cu_num: int = 0, has_pad: bool = False, g2_bhoist=None, g2_b2stage: bool = False, + g2_ascale_pf=None, g2_spart=None, g2_deep_a_pipeline: bool = False, persist_strided: bool = False, + g2_bf16_lds: bool = False, p2p_quant_type: str = "none", fixed_slot_dispatch: bool = False, skew_cu: int = 0, runtime_pair_skip: bool = False, scatter_vec: int = 8): # fmt: on @@ -328,6 +329,8 @@ def compile_mega_moe_stage2(*, model_dim: int, inter_dim: int, experts: int, top raise AssertionError(f"a_dtype must be 'fp4' or 'fp8', got {a_dtype!r}") if persist and cu_num <= 0: raise AssertionError(f"persist=True requires cu_num>0, got {cu_num}") + if g2_deep_a_pipeline and not g2_b2stage: + raise AssertionError("g2_deep_a_pipeline requires g2_b2stage=True") if skew_cu and (not persist or not 0 < skew_cu <= cu_num): raise AssertionError(f"skew_cu={skew_cu} requires persist=True and 0 0 else cu_num @@ -660,11 +665,13 @@ def run_mega_moe_stage2(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, arg_cu model_dim=model_dim, inter_dim=inter_dim, experts=experts, topk=topk, rank=rank, npes=npes, max_tok=max_tok, recv_cap=recv_cap, comb_inp_nbytes=comb_inp_nbytes, BM=BM, BN=BN, BK=BK, use_nt=use_nt, HIDDEN_MAX=HIDDEN_MAX, INTER_MAX=INTER_MAX, SBM=SBM, persist=persist, - cu_num=launch_cu_num, g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, - g2_spart=g2_spart, persist_strided=persist_strided, g2_bf16_lds=g2_bf16_lds, + cu_num=launch_cu_num, g2_bhoist=g2_bhoist, g2_b2stage=g2_b2stage, + g2_ascale_pf=g2_ascale_pf, g2_spart=g2_spart, g2_deep_a_pipeline=g2_deep_a_pipeline, + persist_strided=persist_strided, g2_bf16_lds=g2_bf16_lds, p2p_quant_type=p2p_quant_type, fixed_slot_dispatch=fixed_slot_dispatch, skew_cu=skew_cu, runtime_pair_skip=runtime_pair_skip, scatter_vec=scatter_vec, + a_dtype=a_dtype, ) max_m_blocks = (row_capacity + BM - 1) // BM grid_blocks = launch_cu_num if persist else max_m_blocks @@ -682,10 +689,10 @@ def preload_mega_moe_stage2(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, ar arg_count_matrix, arg_pair_config, arg_parity, arg_p2p, row_capacity, i32_inter, i32_hidden, stream, *, model_dim, inter_dim, experts, topk, rank, npes, max_tok, recv_cap, comb_inp_nbytes, BM, SBM, - HIDDEN_MAX, INTER_MAX, cu_num, BN=256, BK=256, use_nt=True, g2_bhoist=True, - g2_ascale_pf=True, g2_spart=402, persist=False, persist_cu=0, persist_strided=False, + HIDDEN_MAX, INTER_MAX, cu_num, BN=256, BK=256, use_nt=True, g2_bhoist=True, g2_b2stage=False, + g2_ascale_pf=True, g2_spart=402, g2_deep_a_pipeline=False, persist=False, persist_cu=0, persist_strided=False, g2_bf16_lds=False, p2p_quant_type="none", fixed_slot_dispatch=False, skew_cu=0, - runtime_pair_skip=False, scatter_vec=8): + runtime_pair_skip=False, scatter_vec=8, a_dtype="fp8"): # fmt: on """Compile and load one fused Stage2 variant without dispatching it.""" launch_cu_num = min(cu_num, persist_cu) if persist and persist_cu > 0 else cu_num @@ -693,11 +700,13 @@ def preload_mega_moe_stage2(arg_aq, arg_ascale, arg_bq, arg_bscale, arg_eids, ar model_dim=model_dim, inter_dim=inter_dim, experts=experts, topk=topk, rank=rank, npes=npes, max_tok=max_tok, recv_cap=recv_cap, comb_inp_nbytes=comb_inp_nbytes, BM=BM, BN=BN, BK=BK, use_nt=use_nt, HIDDEN_MAX=HIDDEN_MAX, INTER_MAX=INTER_MAX, SBM=SBM, persist=persist, - cu_num=launch_cu_num, g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, - g2_spart=g2_spart, persist_strided=persist_strided, g2_bf16_lds=g2_bf16_lds, + cu_num=launch_cu_num, g2_bhoist=g2_bhoist, g2_b2stage=g2_b2stage, + g2_ascale_pf=g2_ascale_pf, g2_spart=g2_spart, g2_deep_a_pipeline=g2_deep_a_pipeline, + persist_strided=persist_strided, g2_bf16_lds=g2_bf16_lds, p2p_quant_type=p2p_quant_type, fixed_slot_dispatch=fixed_slot_dispatch, skew_cu=skew_cu, runtime_pair_skip=runtime_pair_skip, scatter_vec=scatter_vec, + a_dtype=a_dtype, ) max_m_blocks = (row_capacity + BM - 1) // BM grid_blocks = launch_cu_num if persist else max_m_blocks diff --git a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage2_aligned_pair.py b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage2_aligned_pair.py index 3f5acd473d6..af7d34e7f56 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage2_aligned_pair.py +++ b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_stage2_aligned_pair.py @@ -222,7 +222,8 @@ def compile_mega_moe_stage2_aligned_pair(*, model_dim: int, inter_dim: int, recv_cap: int, comb_inp_nbytes: int, BM: int = 32, SBM: int = 128, BN: int = 256, BK: int = 256, INTER_MAX: int = 8192, use_nt: bool = False, cu_num: int = 722, g2_bhoist=True, - g2_ascale_pf=True): + g2_b2stage: bool = False, g2_ascale_pf=True, + g2_deep_a_pipeline: bool = False, a_dtype: str = "fp8"): # fmt: on """Compile the production two-expert common-row Stage2 kernel.""" arch = str(get_rocm_arch() or "") @@ -244,8 +245,14 @@ def compile_mega_moe_stage2_aligned_pair(*, model_dim: int, inter_dim: int, raise ValueError( f"aligned-pair BN must be divisible by {ALIGNED_PAIR_SCATTER_VEC}" ) - a_stages = kStages + 1 - compute_lds_bytes = _stage2_lds_bytes(BM, BN, BK, "fp8", a_stages) + if a_dtype not in ("fp4", "fp8"): + raise ValueError(f"a_dtype must be 'fp4' or 'fp8', got {a_dtype!r}") + if g2_deep_a_pipeline and not g2_b2stage: + raise ValueError("g2_deep_a_pipeline requires g2_b2stage=True") + is_f8 = a_dtype == "fp8" + a_pack = 1 if is_f8 else 2 + a_stages = 4 if g2_deep_a_pipeline else kStages + 1 + compute_lds_bytes = _stage2_lds_bytes(BM, BN, BK, a_dtype, a_stages) second_compute_off = compute_lds_bytes lds_packed_a_off = compute_lds_bytes * 2 lds_packed_b_off = lds_packed_a_off + BM * 4 @@ -265,8 +272,8 @@ def compile_mega_moe_stage2_aligned_pair(*, model_dim: int, inter_dim: int, expert_offset = rank * experts total_experts = npes * experts total_segments = total_experts + npes - k_bytes = inter_dim - kh_tile_a = BK + k_bytes = inter_dim // a_pack + kh_tile_a = BK // a_pack g2_bhoist, g2_ascale_pf, _, _, _, _ = _resolve_g2_knobs( g2_bhoist, g2_ascale_pf, @@ -282,7 +289,8 @@ class SharedStorage: kernel_name = ( f"megamoe_stage2_aligned_pair_runtime_t{BM}x{BN}x{BK}_cu{cu_num}" f"_seqacc4_sv{ALIGNED_PAIR_SCATTER_VEC}_ms1" - f"_nt{int(use_nt)}_bh{int(g2_bhoist)}apf{int(g2_ascale_pf)}" + f"_a{a_dtype}_nt{int(use_nt)}_bh{int(g2_bhoist)}b2{int(g2_b2stage)}" + f"apf{int(g2_ascale_pf)}da{int(g2_deep_a_pipeline)}" "_rtv3_rq2_hw1" ) @@ -311,7 +319,7 @@ def kernel(arg_aq: fx.Int64, arg_ascale: fx.Int64, arg_bq: fx.Int64, packed_pair = pair_config[ active_parity * fx.Int32(npes) + fx.Int32(rank) ] - fx.rocdl.s_waitcnt(0) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0) # Every lane loads the same entry; broadcasting a lane-0 SSA value # after divergent control flow can select the wrong experts. pair_a_lane = packed_pair & fx.Int32(0xFF) @@ -390,7 +398,7 @@ def issue_all_a_loads(m_row, lds_base): m_row, wave, lane, - True, + is_f8, kh_tile_a, fx.Int32(k_bytes), BM=BM, @@ -447,8 +455,9 @@ def run_pair(m_block): lds_compute, arg_ascale, arg_bq, arg_bscale, fx.Int64(0), arg_aq, i32_max_m_blocks, fx.Int32(0), lane, wave, i32_inter, i32_hidden, fx.Int32(0), fx.Int32(0), BM=BM, BN=BN, BK=BK, use_nt=use_nt, - INTER_MAX=INTER_MAX, aStages=a_stages, a_dtype="fp8", SBM=BM, - g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, + INTER_MAX=INTER_MAX, aStages=a_stages, a_dtype=a_dtype, SBM=BM, + g2_bhoist=g2_bhoist, g2_b2stage=g2_b2stage, + g2_ascale_pf=g2_ascale_pf, g2_deep_a_pipeline=g2_deep_a_pipeline, expert_offset=expert_offset, explicit_m_row=m_row_a, explicit_n_block=n_block, explicit_expert=global_a) # fmt: on @@ -463,8 +472,9 @@ def run_pair(m_block): lds_compute_b, arg_ascale, arg_bq, arg_bscale, fx.Int64(0), arg_aq, i32_max_m_blocks, fx.Int32(0), lane, wave, i32_inter, i32_hidden, fx.Int32(0), fx.Int32(0), BM=BM, BN=BN, BK=BK, use_nt=use_nt, - INTER_MAX=INTER_MAX, aStages=a_stages, a_dtype="fp8", SBM=BM, - g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, + INTER_MAX=INTER_MAX, aStages=a_stages, a_dtype=a_dtype, SBM=BM, + g2_bhoist=g2_bhoist, g2_b2stage=g2_b2stage, + g2_ascale_pf=g2_ascale_pf, g2_deep_a_pipeline=g2_deep_a_pipeline, expert_offset=expert_offset, explicit_m_row=m_row_b, explicit_n_block=n_block, explicit_expert=global_b) _store_accumulator_tile( diff --git a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_v2.py b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_v2.py index 7bca56d9ff6..b68e9a3548d 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/mega_moe_v2.py +++ b/aiter/ops/flydsl/kernels/mega_moe/mega_moe_v2.py @@ -2,6 +2,7 @@ # Copyright (c) 2025 FlyDSL Project Contributors """MegaMoE v2 fused dispatch, GEMM1, GEMM2, and combine implementation.""" +import os from dataclasses import replace import flydsl.expr as fx @@ -37,8 +38,9 @@ class MegaMoEV2: # fmt: off def __init__(self, *, rank: int, world_size: int, model_dim: int, inter_dim: int, experts: int, topk: int, quant: str, w1: torch.Tensor, w1_scale: torch.Tensor, w2: torch.Tensor, w2_scale: torch.Tensor, - max_tok_per_rank: int, mega_scheme: str = "fixedslot", swiglu_limit: float = 0.0, - fanout_masks: tuple[int, ...] = ()): + max_tok_per_rank: int, mega_scheme: str = "fixedslot", stage2_p2p_quant: str = "auto", + swiglu_limit: float = 0.0, fanout_masks: tuple[int, ...] = (), + act: str = "silu", situ_beta: float = 1.0, situ_linear_beta: float = 1.0): # fmt: on if quant != "a8w4": raise ValueError("MegaMoEV2 currently supports quant='a8w4' only") @@ -56,9 +58,20 @@ def __init__(self, *, rank: int, world_size: int, model_dim: int, inter_dim: int if not 0 < self.topk <= 16: raise ValueError(f"MegaMoEV2 topk must be in [1, 16], got {self.topk}") self.mtpr = int(max_tok_per_rank) + self._a_dtype = "fp8" + self._a_torch_dtype = torch.float8_e4m3fn + self._a_view_dim = self.model_dim self.swiglu_limit = float(swiglu_limit) + self.act = act + self.situ_beta = float(situ_beta) + self.situ_linear_beta = float(situ_linear_beta) + if self.act not in ("silu", "situv2"): + raise ValueError(f"MegaMoEV2 act must be 'silu' or 'situv2', got {self.act!r}") + if self.act == "situv2" and (self.situ_beta <= 0.0 or self.situ_linear_beta <= 0.0): + raise ValueError("situv2 requires positive situ_beta / situ_linear_beta") self._bundle_plan = build_mega_moe_bundle_plan( self.mtpr, + p2p_quant=stage2_p2p_quant, experts_per_rank=self.epr, model_dim=self.model_dim, inter_dim=self.inter_dim, @@ -83,7 +96,7 @@ def __init__(self, *, rank: int, world_size: int, model_dim: int, inter_dim: int self.comb_cfg = FlyDSLDispatchCombineConfig(rank=self.rank, world_size=self.world_size, hidden_dim=self.model_dim, max_num_inp_token_per_rank=self.mtpr, num_experts_per_rank=self.epr, num_experts_per_token=self.topk, combine_dtype=torch.bfloat16, - dispatch_dtype=torch.float8_e4m3fn, scale_dim=self._s1_scale_dim, scale_type_size=1, + dispatch_dtype=self._a_torch_dtype, scale_dim=self._s1_scale_dim, scale_type_size=1, enable_std_moe=False, enable_group_major=True, gm_unit_size=capacity_tile_m, gm_scheme=mega_scheme, gm_compact=compact, gm_indexed_payload=compact and self.mtpr >= INDEXED_PAYLOAD_MIN_MTPR, @@ -116,6 +129,13 @@ def _build_fused_stage1(self, w1, w1_scale): ) self.sort_block_m = 32 + # Opt-in hardware bounds checking for the Stage1 control buffers. Off + # by default so production keeps today's descriptors; turn it on while + # sweeping configurations, where an indexing mistake would otherwise + # page-fault and can leave a GPU needing a privileged reset. + self._s1_bounds_check = ( + os.environ.get("AITER_MEGA_MOE_BOUNDS_CHECK", "0") == "1" + ) self._s1_w1 = w1.contiguous().view(torch.uint8) self._s1_w1_scale = w1_scale.contiguous().view(torch.uint8) op = self.comb_op._gm @@ -165,7 +185,9 @@ def _build_fused_stage1(self, w1, w1_scale): inter_dim = self.inter_dim a2rows = self._s1_nvm - self._s1_out = torch.zeros((a2rows, inter_dim), dtype=torch.float8_e4m3fn, device=self.dev) + self._s1_out = torch.zeros( + (a2rows, inter_dim), dtype=torch.float8_e4m3fn, device=self.dev + ) prows = ((a2rows + 255) // 256) * 256 pcols = (((inter_dim // 32) + 7) // 8) * 8 self._s1_osd = torch.zeros(prows * pcols + inter_dim, dtype=torch.uint8, device=self.dev) @@ -383,8 +405,6 @@ def preload_stage1_bundle(self): config.prepare_quant_cu, (quant_groups + 511) // 512, ) - # Preload fused-quant and prequantized-input variants so both public - # forward paths remain AOT-only after warmup. for preload_quant_blocks in sorted({0, quant_blocks}): self._s1_preload_prepare( fx.Int64(self._s1_disp.data_ptr()), @@ -447,6 +467,7 @@ def preload_stage1_bundle(self): tile_state_stride=self._s1_tile_state_stride, variants=self._bundle_plan.stage1_variants, swiglu_limit=self.swiglu_limit, + act=self.act, situ_beta=self.situ_beta, situ_linear_beta=self.situ_linear_beta, ) def _select_config(self, tokens: int) -> MegaMoEConfig: @@ -535,10 +556,10 @@ def _run_fused_stage1( cur_tok = int(x.shape[0]) if cur_tok > self.mtpr: raise ValueError(f"run_tokens={cur_tok} > max_tok_per_rank={self.mtpr}") - if x.dtype != torch.float8_e4m3fn or not x.is_contiguous(): - raise ValueError("x must be contiguous float8_e4m3fn") - if tuple(x.shape) != (cur_tok, self.model_dim): - raise ValueError(f"x must have shape ({cur_tok}, {self.model_dim})") + if x.dtype != self._a_torch_dtype or not x.is_contiguous(): + raise ValueError(f"x must be contiguous {self._a_torch_dtype}") + if tuple(x.shape) != (cur_tok, self._a_view_dim): + raise ValueError(f"x must have shape ({cur_tok}, {self._a_view_dim})") if wts.dtype != torch.float32 or not wts.is_contiguous(): raise ValueError("wts must be contiguous float32") if tuple(wts.shape) != (cur_tok, self.topk): @@ -590,6 +611,10 @@ def _run_fused_stage1( "fixed_slot_dispatch": self._s1_fixed_slot, "num_cu": self._s1_num_cu, "swiglu_limit": self.swiglu_limit, + "act": self.act, + "situ_beta": self.situ_beta, + "situ_linear_beta": self.situ_linear_beta, + "bounds_check": self._s1_bounds_check, } entry = self._active_bundle_entry use_bundle = ( @@ -635,7 +660,7 @@ def quantize(self, x_bf16): # A zero-token rank still participates in the downstream # collectives, but HIP rejects a zero-grid quant kernel launch. return self._s1_quant_x[:0], self._s1_quant_scale[:0] - return per_1x32_mx_quant(x_bf16, quant_mode="fp8") + return per_1x32_mx_quant(x_bf16, quant_mode=self._a_dtype) def _run_joint( self, @@ -800,6 +825,7 @@ def _build_fused_stage2(self): "comb_inp_nbytes": int(comb_cfg.max_num_inp_token_per_rank) * int(k) * p2p_row_nbytes, "HIDDEN_MAX": int(comb_cfg.hidden_dim), "INTER_MAX": int(self.inter_dim), "cu_num": int(cu_num), "p2p_quant_type": p2p_quant, "fixed_slot_dispatch": bool(self._s1_fixed_slot), + "a_dtype": self._a_dtype, } self._g2_combine_placeholder = torch.empty( 1, comb_cfg.hidden_dim, dtype=comb_cfg.combine_dtype, device=dev @@ -836,8 +862,10 @@ def _fused_stage2_call(self, launcher, config, stream, *, runtime_pair_skip, sca BK=stage2.block_k, use_nt=stage2.use_nt, g2_bhoist=stage2.b_hoist, + g2_b2stage=stage2.b2stage, g2_ascale_pf=stage2.ascale_prefetch, g2_spart=stage2.spatial_partition, + g2_deep_a_pipeline=stage2.deep_a_pipeline, persist=stage2.persist, persist_cu=stage2.persist_cu, persist_strided=stage2.persist_strided, @@ -903,7 +931,10 @@ def _aligned_pair_stage2_call(self, launcher, config: MegaMoEConfig, stream): use_nt=stage2.use_nt, cu_num=stage2.pair_cu, g2_bhoist=stage2.b_hoist, + g2_b2stage=stage2.b2stage, g2_ascale_pf=stage2.ascale_prefetch, + g2_deep_a_pipeline=stage2.deep_a_pipeline, + a_dtype=str(invariants["a_dtype"]), ) def _preload_aligned_pair_stage2(self, config: MegaMoEConfig, stream): diff --git a/aiter/ops/flydsl/kernels/mega_moe/quant.py b/aiter/ops/flydsl/kernels/mega_moe/quant.py index 507386b4b30..6e647fa4fed 100644 --- a/aiter/ops/flydsl/kernels/mega_moe/quant.py +++ b/aiter/ops/flydsl/kernels/mega_moe/quant.py @@ -72,6 +72,45 @@ def emit_per_1x32_mx_fp8_group(in_buf, out_buf, scale_buf, group_id): ) +def emit_per_1x32_mx_fp4_group(in_buf, out_buf, scale_buf, group_id): + """Quantize one contiguous 32-BF16 group to packed FP4 plus E8M0 scale.""" + in_vec = group_id * fx.Int32(GROUP * 2 // 16) + act = [] + local_max = fx.Float32(1e-10) + for chunk in range_constexpr(GROUP // 8): + raw = _load_i32x4(in_buf, in_vec + fx.Int32(chunk)) + values = fx.Vector(raw).bitcast(fx.BFloat16).to(fx.Float32) + local_max = local_max.maximumf(fmath.absf(values).reduce(ReductionOp.MAX)) + for elem in range_constexpr(8): + act.append(values[elem]) + + working = (local_max * fx.Int32(_FP4_INV_MAX_POS_BITS).bitcast(fx.Float32)).bitcast( + fx.Int32 + ) + mantissa = working & fx.Int32(0x7FFFFF) + biased_exp = (working >> fx.Int32(23)) & fx.Int32(0xFF) + e8m0 = (mantissa != fx.Int32(0)).select(biased_exp + fx.Int32(1), biased_exp) + e8m0 = (e8m0 > fx.Int32(255)).select(fx.Int32(255), e8m0) + scale_buf[group_id] = e8m0.to(fx.Uint8) + + dequant_scale = (e8m0 << fx.Int32(23)).bitcast(fx.Float32) + words = [] + for word in range_constexpr(GROUP // 8): + packed = fx.Int32(0) + for pair in range_constexpr(4): + idx = word * 8 + pair * 2 + packed = rocdl.cvt_scalef32_pk_fp4_f32( + T.i32, + packed, + act[idx], + act[idx + 1], + dequant_scale, + pair, + ) + words.append(packed) + _store_i32x4(out_buf, group_id, fx.Vector.from_elements(words, fx.Int32)) + + def build_per_1x32_mx_quant_module(n: int, quant_mode: str): """Return a @flyc.jit launcher for 1x32 MX quant of a [m, n] bf16 matrix.""" assert n % 32 == 0, f"n={n} must be divisible by 32" @@ -81,7 +120,6 @@ def build_per_1x32_mx_quant_module(n: int, quant_mode: str): ), f"quant_mode must be fp4|fp8, got {quant_mode!r}" scale_n = n // GROUP - inv_max_pos_bits = _FP4_INV_MAX_POS_BITS if need_fp4 else _FP8_E4M3_INV_MAX_POS_BITS @flyc.kernel(name=f"per_1x32_mx_quant_{quant_mode}_n{n}") def quant_kernel(x: fx.Tensor, y: fx.Tensor, scale: fx.Tensor, m: fx.Int32): @@ -92,47 +130,7 @@ def quant_kernel(x: fx.Tensor, y: fx.Tensor, scale: fx.Tensor, m: fx.Int32): group_id = fx.block_idx.x * fx.Int32(BLOCK) + fx.thread_idx.x if group_id < m * fx.Int32(scale_n): if const_expr(need_fp4): - in_vec = group_id * fx.Int32(GROUP * 2 // 16) - act = [] - local_max = fx.Float32(1e-10) - for chunk in range_constexpr(GROUP // 8): - raw = _load_i32x4(in_buf, in_vec + fx.Int32(chunk)) - values = fx.Vector(raw).bitcast(fx.BFloat16).to(fx.Float32) - local_max = local_max.maximumf( - fmath.absf(values).reduce(ReductionOp.MAX) - ) - for elem in range_constexpr(8): - act.append(values[elem]) - - working = ( - local_max * fx.Int32(inv_max_pos_bits).bitcast(fx.Float32) - ).bitcast(fx.Int32) - mantissa = working & fx.Int32(0x7FFFFF) - biased_exp = (working >> fx.Int32(23)) & fx.Int32(0xFF) - e8m0 = (mantissa != fx.Int32(0)).select( - biased_exp + fx.Int32(1), biased_exp - ) - e8m0 = (e8m0 > fx.Int32(255)).select(fx.Int32(255), e8m0) - scale_buf[group_id] = e8m0.to(fx.Uint8) - - dequant_scale = (e8m0 << fx.Int32(23)).bitcast(fx.Float32) - words = [] - for word in range_constexpr(GROUP // 8): - packed = fx.Int32(0) - for pair in range_constexpr(4): - idx = word * 8 + pair * 2 - packed = rocdl.cvt_scalef32_pk_fp4_f32( - T.i32, - packed, - act[idx], - act[idx + 1], - dequant_scale, - pair, - ) - words.append(packed) - _store_i32x4( - out_buf, group_id, fx.Vector.from_elements(words, fx.Int32) - ) + emit_per_1x32_mx_fp4_group(in_buf, out_buf, scale_buf, group_id) else: emit_per_1x32_mx_fp8_group(in_buf, out_buf, scale_buf, group_id) diff --git a/op_tests/multigpu_tests/test_mega_moe_v2.py b/op_tests/multigpu_tests/test_mega_moe_v2.py index 22321c256d1..52922fd7719 100644 --- a/op_tests/multigpu_tests/test_mega_moe_v2.py +++ b/op_tests/multigpu_tests/test_mega_moe_v2.py @@ -37,15 +37,20 @@ "topk": 6, "swiglu_limit": 10.0, }, - # Kimi-K3 routing/weight geometry. This exercises topk16 and EP8/epr112; - # the numerical reference intentionally keeps MegaMoEV2's current bounded - # SwiGLU activation while the K3 activation integration remains separate. + # Kimi-K3 routing/weight geometry. This exercises topk16 and EP8/epr112 and + # the real K3 activation: SiTUv2 (hidden_act="situ"), which self-saturates + # via tanh, so it carries no hard swiglu clamp. beta=4.0 / linear_beta=25.0 + # are the K3 checkpoint constants the MegaMoE SiTU kernel bakes in (SGLang + # kimi_k3.py asserts exactly these for the MegaMoE a2a backend). "kimi_k3_route": { "model_dim": 3584, - "inter_dim": 512, + "inter_dim": 3072, "experts": 896, "topk": 16, - "swiglu_limit": 10.0, + "swiglu_limit": 0.0, + "act": "situv2", + "situ_beta": 4.0, + "situ_linear_beta": 25.0, }, } @@ -246,6 +251,9 @@ def _reference( inter_dim, experts, swiglu_limit, + act="silu", + situ_beta=1.0, + situ_linear_beta=1.0, ): x_all, token_counts = _all_gather_variable(x) weights_all, _ = _all_gather_variable(route_weights, token_counts) @@ -270,13 +278,20 @@ def _reference( inp = x_all[rows].float() gate = inp @ w1[:inter_dim].T up = inp @ w1[inter_dim:].T - # MegaMoEV2 treats swiglu_limit <= 0 as "no clamp"; clamping here - # unconditionally turns up into zeros at the operator's own default of - # 0.0, so the reference is all-zero and relL2 divides by zero. - if swiglu_limit > 0: - gate = gate.clamp(max=swiglu_limit) - up = up.clamp(-swiglu_limit, swiglu_limit) - hidden = F.silu(gate) * up + if act == "situv2": + # SiTUv2 (Kimi-K3 hidden_act="situ"), matching SGLang SituAndMul and + # the operator's no-clamp situv2 epilogue: tanh self-saturates. + gate = situ_beta * torch.tanh(gate / situ_beta) * torch.sigmoid(gate) + up = situ_linear_beta * torch.tanh(up / situ_linear_beta) + hidden = gate * up + else: + # MegaMoEV2 treats swiglu_limit <= 0 as "no clamp"; clamping here + # unconditionally turns up into zeros at the operator's own default of + # 0.0, so the reference is all-zero and relL2 divides by zero. + if swiglu_limit > 0: + gate = gate.clamp(max=swiglu_limit) + up = up.clamp(-swiglu_limit, swiglu_limit) + hidden = F.silu(gate) * up out = (hidden @ w2.T) * weights_all[rows, slots, None] partial.index_add_(0, rows, out) del w1, w2, inp, hidden, out @@ -345,6 +360,9 @@ def _run_size(moe, x, weights, ids, ref_weights, args, rank, world, device): moe.inter_dim, moe.experts, moe.swiglu_limit, + moe.act, + moe.situ_beta, + moe.situ_linear_beta, ) rel_l2 = _rel_l2(output, reference) if rel_l2 >= args.rtol: @@ -422,6 +440,9 @@ def masked(): moe.inter_dim, moe.experts, moe.swiglu_limit, + moe.act, + moe.situ_beta, + moe.situ_linear_beta, ) # -1: every row of every rank is padding, so there is nothing to compare. rel_l2 = _rel_l2(eager_output[:real], reference) @@ -503,11 +524,17 @@ def select_config(tokens): def main(): parser = argparse.ArgumentParser() parser.add_argument("--network", choices=NETWORKS, default="v4_pro") + parser.add_argument("--quant", choices=("a8w4",), default="a8w4") + parser.add_argument( + "--stage2-p2p-quant", + choices=("auto", "none", "fp8_blockwise_1x32"), + default="auto", + ) parser.add_argument("--bs-list", default="128") parser.add_argument("--iters", type=int, default=30) parser.add_argument("--seed", type=int, default=123) parser.add_argument("--accuracy-max-bs", type=int, default=128) - parser.add_argument("--rtol", type=float, default=0.10) + parser.add_argument("--rtol", type=float) parser.add_argument("--max-tok-per-rank", type=int) parser.add_argument("--rank-tokens", default="") parser.add_argument("--config-tokens", type=int, default=0) @@ -535,6 +562,8 @@ def main(): > 1 ): raise ValueError("fanout adversarial flags are mutually exclusive") + if args.rtol is None: + args.rtol = 0.10 batch_sizes = [int(value) for value in args.bs_list.split(",")] if not batch_sizes or min(batch_sizes) <= 0: raise ValueError("--bs-list must contain positive integers") @@ -614,13 +643,14 @@ def main(): shared_moe = MegaMoEV2( rank=rank, world_size=world, - quant="a8w4", + quant=args.quant, w1=w1, w1_scale=w1_scale, w2=w2, w2_scale=w2_scale, max_tok_per_rank=max_tok_per_rank, fanout_masks=fanout_masks, + stage2_p2p_quant=args.stage2_p2p_quant, **network, ) moe = shared_moe