From 028b995f4fff7dec26c6fe2cdd039302347c9b2b Mon Sep 17 00:00:00 2001 From: yanboshao Date: Tue, 4 Aug 2026 04:29:52 +0000 Subject: [PATCH 01/23] gfx1250 MoE: fused gemm2 P2P scatter-combine (ep_scatter), rebased onto updated main Reapply the gfx1250 fused-MoE ep_scatter feature (combine_mode="scatter_fused": gemm2's TDM epilogue P2P-writes each route-weighted output row into peers' comb_inp, so combine just sums -- no gather-reduce) onto the updated origin/main, which meanwhile landed #4482 (real SiTUv2 + >512-expert chunked psum scan), SiTUv2 path). Squashed to a single commit. Per-file resolution: - dispatch_combine_v2/*: vendored cco-LSA v2 intranode dispatch/combine op-layer. - mxfp4_preshuffle_gfx1250_tdm.py: on upstream's make_lds_copy_ops LDS API + SiTUv2 epilogue, add the TDM gather-store ep epilogue (tdm_scatter with in-kernel global_view/lds_view). tdm_scatter is vendored locally (tdm_gather_shim.py, on the stock FlyDSL wheel's low-level TDM intrinsics) so this branch needs no FlyDSL-side patch; route weight hoisted per wm row (_wf_rows). - moe_contiguous_psum.py: adopt upstream's chunked scan (E>512 correct) for the non-EP remap; keep the multi-block grid-stride remap + ep_rowmap kernels for EP. - grouped_moe_gfx1250.py: keep upstream SiTUv2 (stage1_act=3, situ_beta); add the ep_scatter dispatch wiring (ep_rowmap build, _ep_gemm2_kwargs, ep_scatter return). - batched_gemm_mxfp4.py / fused_moe.py: thread both situ_* and ep_* params. - tuned_grouped_fmoe.csv: tuning points (99 rows). Dropped only the full-subtile PF prefetch pipeline (da8d794): loaders and lds_addr_keepalive stay at upstream. lalala-sh's ds-read hoist is kept. gfx1250-only; not run in this environment. Compile-verify the LDS-API migration + ep/SiTUv2 epilogue and re-run test_mega_moe --combine on hardware. Co-authored-by: lalala-sh Co-authored-by: zhimding --- aiter/configs/tuned_grouped_fmoe.csv | 28 + aiter/fused_moe.py | 48 + aiter/ops/flydsl/batched_gemm_mxfp4.py | 18 + .../ops/flydsl/dispatch_combine_v2/README.md | 92 + .../flydsl/dispatch_combine_v2/__init__.py | 17 + .../dispatch_combine_op.py | 1122 ++++++++++ .../dispatch_combine_v2/flydsl_prims.py | 169 ++ .../dispatch_combine_v2/intranode_kernels.py | 1846 +++++++++++++++++ .../dispatch_combine_v2/tuning_configs.py | 258 +++ aiter/ops/flydsl/grouped_moe_gfx1250.py | 161 +- .../ops/flydsl/kernels/moe_contiguous_psum.py | 424 +++- .../kernels/mxfp4_preshuffle_gfx1250_tdm.py | 297 ++- aiter/ops/flydsl/kernels/tdm_gather_shim.py | 573 +++++ op_tests/multigpu_tests/test_mega_moe.py | 837 ++++++++ 14 files changed, 5820 insertions(+), 70 deletions(-) create mode 100644 aiter/ops/flydsl/dispatch_combine_v2/README.md create mode 100644 aiter/ops/flydsl/dispatch_combine_v2/__init__.py create mode 100644 aiter/ops/flydsl/dispatch_combine_v2/dispatch_combine_op.py create mode 100644 aiter/ops/flydsl/dispatch_combine_v2/flydsl_prims.py create mode 100644 aiter/ops/flydsl/dispatch_combine_v2/intranode_kernels.py create mode 100644 aiter/ops/flydsl/dispatch_combine_v2/tuning_configs.py create mode 100644 aiter/ops/flydsl/kernels/tdm_gather_shim.py create mode 100644 op_tests/multigpu_tests/test_mega_moe.py diff --git a/aiter/configs/tuned_grouped_fmoe.csv b/aiter/configs/tuned_grouped_fmoe.csv index a3c6a9a95c8..8b45bc2f69c 100644 --- a/aiter/configs/tuned_grouped_fmoe.csv +++ b/aiter/configs/tuned_grouped_fmoe.csv @@ -70,3 +70,31 @@ gfx1250,256,4096,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_ gfx1250,256,8192,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,49152,128,512,256,128,256,128,1,4,2,2,1,1,0,,1,1,grouped_gemm1_fp4_gugu,grouped_gemm2_fp4,6079.7677,0,0 gfx1250,256,16384,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,98304,128,512,256,128,512,256,1,4,2,2,1,1,0,,1,0,grouped_gemm1_fp4_gugu,grouped_gemm2_fp4,11927.1248,0,0 gfx1250,256,32768,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,196608,128,512,256,128,256,128,1,4,2,2,1,1,0,,1,1,grouped_gemm1_fp4_gugu,grouped_gemm2_fp4,23452.4903,0,0 +gfx1250,256,32,7168,3072,96,-1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,256,512,32,256,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,128,7168,3072,96,-1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,512,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,256,7168,3072,96,-1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,256,512,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,512,7168,3072,96,-1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,512,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,1024,7168,3072,96,-1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,512,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,2048,7168,3072,96,-1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,512,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,4096,7168,3072,96,-1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,256,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,8192,7168,3072,96,-1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,256,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,16384,7168,3072,96,-1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,64,512,256,64,256,256,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,2033.0,0,0 +gfx1250,256,32768,7168,3072,96,-1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,64,512,256,64,256,256,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,3692.0,0,0 +gfx1250,256,32,7168,3072,96,6,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,256,512,32,256,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,128,7168,3072,96,6,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,512,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,256,7168,3072,96,6,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,256,512,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,2048,7168,3072,96,6,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,512,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,4096,7168,3072,96,6,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,256,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,8192,7168,3072,96,6,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,256,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,16384,7168,3072,96,6,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,64,512,256,64,256,256,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,2033.0,0,0 +gfx1250,256,32768,7168,3072,96,6,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,64,512,256,64,256,256,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,3692.0,0,0 +gfx1250,256,32,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,256,512,32,256,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,128,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,512,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,256,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,256,512,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,512,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,512,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,1024,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,512,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,2048,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,512,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,4096,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,256,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,8192,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,32,256,256,32,512,512,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,468.0528,0,0 +gfx1250,256,16384,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,64,512,256,64,256,256,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,2033.0,0,0 +gfx1250,256,32768,7168,3072,96,1,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,384,64,512,256,64,256,256,1,4,2,2,1,1,0,,0,0,grouped_gemm1_fp4_gguu,grouped_gemm2_fp4,3692.0,0,0 diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index c6c0a1ae7b7..fe42b2450fe 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -489,6 +489,14 @@ def fused_moe( shared_w1_scale: torch.Tensor | None = None, shared_w2_scale: torch.Tensor | None = None, shared_expert_id: int = -1, + ep_scatter: bool = False, + ep_arena_handle: int = 0, + ep_comb_inp_off: int = 0, + ep_wire_nbytes: int = 0, + ep_rank: int = 0, + ep_max_tok: int = 0, + ep_topk: int = 0, + ep_tis: torch.Tensor | None = None, ): if ( any( @@ -557,6 +565,14 @@ def fused_moe( beta=beta, linear_beta=linear_beta, gate_mode=gate_mode, + ep_scatter=ep_scatter, + ep_arena_handle=ep_arena_handle, + ep_comb_inp_off=ep_comb_inp_off, + ep_wire_nbytes=ep_wire_nbytes, + ep_rank=ep_rank, + ep_max_tok=ep_max_tok, + ep_topk=ep_topk, + ep_tis=ep_tis, ) @@ -626,6 +642,14 @@ def fused_moe_( beta: float | None = None, linear_beta: float | None = None, gate_mode: str = GateMode.SEPARATED.value, + ep_scatter: bool = False, + ep_arena_handle: int = 0, + ep_comb_inp_off: int = 0, + ep_wire_nbytes: int = 0, + ep_rank: int = 0, + ep_max_tok: int = 0, + ep_topk: int = 0, + ep_tis: torch.Tensor | None = None, ) -> torch.Tensor: return _fused_moe_impl( hidden_states=hidden_states, @@ -653,6 +677,14 @@ def fused_moe_( beta=beta, linear_beta=linear_beta, gate_mode=gate_mode, + ep_scatter=ep_scatter, + ep_arena_handle=ep_arena_handle, + ep_comb_inp_off=ep_comb_inp_off, + ep_wire_nbytes=ep_wire_nbytes, + ep_rank=ep_rank, + ep_max_tok=ep_max_tok, + ep_topk=ep_topk, + ep_tis=ep_tis, ) @@ -682,6 +714,14 @@ def _fused_moe_impl( beta: float | None = None, linear_beta: float | None = None, gate_mode: str = GateMode.SEPARATED.value, + ep_scatter: bool = False, + ep_arena_handle: int = 0, + ep_comb_inp_off: int = 0, + ep_wire_nbytes: int = 0, + ep_rank: int = 0, + ep_max_tok: int = 0, + ep_topk: int = 0, + ep_tis: torch.Tensor | None = None, *, _q_dtype_a: torch.dtype | None = None, _metadata_transform: Callable | None = None, @@ -804,6 +844,14 @@ def _fused_moe_impl( num_local_tokens=num_local_tokens, situ_beta=1.0 if beta is None else float(beta), situ_linear_beta=1.0 if linear_beta is None else float(linear_beta), + ep_scatter=ep_scatter, + ep_arena_handle=ep_arena_handle, + ep_comb_inp_off=ep_comb_inp_off, + ep_wire_nbytes=ep_wire_nbytes, + ep_rank=ep_rank, + ep_max_tok=ep_max_tok, + ep_topk=ep_topk, + ep_tis=ep_tis, ) if grouped_a8w4_out is not None: diff --git a/aiter/ops/flydsl/batched_gemm_mxfp4.py b/aiter/ops/flydsl/batched_gemm_mxfp4.py index 233d6a711b6..82cab72244f 100644 --- a/aiter/ops/flydsl/batched_gemm_mxfp4.py +++ b/aiter/ops/flydsl/batched_gemm_mxfp4.py @@ -51,6 +51,13 @@ def flydsl_grouped_gemm_a8w4_masked( stage1_quant_out=0, quant_scale=None, quant_wmma_rep=1, + ep_p2p_write=0, + ep_off_comb_inp=0, + ep_wire_nbytes=0, + ep_slot_stride=0, + ep_arena_handle=0, + ep_world=0, + ep_rowmap=None, situ_beta=1.0, situ_linear_beta=1.0, ): @@ -97,6 +104,10 @@ def flydsl_grouped_gemm_a8w4_masked( quant_scale_tensor = out # dummy, never written else: quant_scale_tensor = quant_scale.view(torch.uint8) + # EP gemm2-fused P2P scatter: the down-proj epilogue writes each weighted row + # into peers' comb_inp instead of TDM-storing arg_c. ep_rowmap carries the + # per-row (dst_packed, f32 weight) map; when off, pass a dummy tensor (unread). + ep_rowmap_tensor = ep_rowmap if ep_rowmap is not None else out launch_gemm_a8w4_tdm( out, ptr_arg(a), @@ -126,6 +137,13 @@ def flydsl_grouped_gemm_a8w4_masked( quant_scale_tensor, float(situ_beta), float(situ_linear_beta), + ep_p2p_write=int(ep_p2p_write), + ep_off_comb_inp=int(ep_off_comb_inp), + ep_wire_nbytes=int(ep_wire_nbytes), + ep_slot_stride=int(ep_slot_stride), + ep_arena_handle=int(ep_arena_handle), + ep_world=int(ep_world), + arg_ep_rowmap=ep_rowmap_tensor, ) return out diff --git a/aiter/ops/flydsl/dispatch_combine_v2/README.md b/aiter/ops/flydsl/dispatch_combine_v2/README.md new file mode 100644 index 00000000000..d707c72459f --- /dev/null +++ b/aiter/ops/flydsl/dispatch_combine_v2/README.md @@ -0,0 +1,92 @@ +# cco-LSA intranode MoE dispatch / combine (ops v2, FlyDSL) + +Intranode (single-node, EP8) MoE dispatch + combine built on **mori-cco LSA** +(intra-node P2P over the flat symmetric VA) and **FlyDSL** device kernels. A +mori-parity reimplementation that swaps the mori-shmem provider for cco-LSA: +peer addresses are computed in-kernel via `cco.Window(arena).lsa_ptr(pe, off)`, +no host P2P tables. Reference = ROCm/FlyDSL PR #522 +(`dispatch_combine_intranode_{kernel,op}.py`). + +Supported token dtypes: **bf16**, **f32**, **fp8** (gather-only; OCP e4m3 on +gfx950, e4m3**fnuz** max 240 on gfx942) and **fp4** (e2m1, gather-only, +**gfx950-only** — the `cvt_scalef32_*_fp4` intrinsics don't exist on gfx942). +Combine: gather (UseP2PRead) **and** scatter (`_nop2p`); weighted combine +(`out_weights`); StdMoE (ConvertDispatchOutput / ConvertCombineInput, standalone ++ wired into the op); fp8 combine-wire **quant** (`fp8_direct_cast` **and** +`fp8_blockwise`, scatter-only — distinct from the plain fp8 token dtype, which +keeps a bf16 external payload); per-token scales forwarding; +`max_total_recv_tokens` cap; mori-parity host op-layer + per-device, +dtype-aware tuning table. Not done: `skip_stage1` (FlyDSL-only). + +> **Test-only, not a mori API (yet).** These modules import each other by +> top-level name (`from intranode_kernels import ...`, `import flydsl_prims as P`) +> and have no `__init__.py`; they only resolve after the tests do +> `sys.path.insert(0, )`. There is no `mori.ops.dispatch_combine_v2` +> package export, so `import mori.ops.dispatch_combine_v2.dispatch_combine_op` +> will fail. Use it via the test/bench harnesses below. Wiring it in as a real +> package (relative imports + `__init__.py` + `ops/__init__.py` export) is a +> follow-up. + +## Layout + +| file | role | +|---|---| +| `flydsl_prims.py` | device primitives: system atomics / ordered stores / fences / volatile-spin waits | +| `intranode_kernels.py` | all FlyDSL intranode kernel factories: `make_dispatch` (+scales/replay), `make_combine` (gather) / `make_combine_scatter` (`_nop2p`, bf16/f32/fp8/fp4), `make_convert_dispatch_output` / `make_convert_combine_input` (StdMoE), `make_local_expert_count` | +| `dispatch_combine_op.py` | `SymmArena` + `EpDispatchCombineOp` / `EpDispatchCombineConfig` (+`.tuned()`) / `EpDispatchRoutingHandle` — mori-parity host op-layer (scales, scatter/quant combine, StdMoE, recv cap, LEC, reset, replay) | +| `tuning_configs.py` | per-(world,hidden,topk) block/warp lookup | + +Tests/bench live under `tests/python/ops/dispatch_combine_v2/`: + +| file | role | +|---|---| +| `test_dispatch_combine_v2_intranode.py` | pytest wrapper: runs `test_op.py` under torchrun for the representative modes and asserts every line PASS | +| `test_op.py` | EP8 op-layer test (gather/scatter, quant, StdMoE, recv-cap, scales, LEC, reset, replay) | +| `bench_dispatch_combine.py` | eager + CUDA-graph perf bench + e2e correctness. Envs: `DTYPE=bf16\|f32\|fp8\|fp4`, `COMBINE=gather\|scatter`, `QUANT=none\|fp8_direct_cast\|fp8_blockwise`, `STDMOE=1`, `SCALE_DIM`, `SWEEP`, `DISP_BLOCK`/`COMB_BLOCK`, `WARP_NUM`/`COMB_WARP`, `MODE`, `TUNED` | +| `run_bench.sh` | bench launcher (runs `bench_dispatch_combine.py` in the container) | + +(Each script inlines a tiny torchrun/gloo `Dist` bootstrap — gloo only carries the cco unique-id and pass/fail counts.) + +## Run (inside the container, 8 GPUs) + +`torchrun --standalone` uses a localhost rendezvous, so no socket-iface env is +needed. Intranode only (no GDA/RDMA). + +```bash +cd tests/python/ops/dispatch_combine_v2 + +pytest test_dispatch_combine_v2_intranode.py -v # EP8 correctness (all modes) +torchrun --standalone --nproc_per_node=8 test_op.py # op-layer correctness (env-driven) +torchrun --standalone --nproc_per_node=8 bench_dispatch_combine.py # perf + e2e correctness +``` + +Config via env: `HIDDEN`, `TOPK`, `EPR`, `SWEEP`, `DTYPE`, `COMBINE`, `QUANT`, +`DISP_BLOCK`/`COMB_BLOCK`, `WARP_NUM`/`COMB_WARP`, `MODE=eager|graph|both`, `TUNED`. + +## Design notes + +- **dispatch**: per (token, k) dedup same-dest-PE via ballot; lane0 remote + `atomic_add` allocates a recv slot; publish origin id + idx/wts + 16B dual-issue + token copy to the peer; grid barrier; per-peer count signal; collect `total_recv`. +- **combine** (gather, = mori `UseP2PRead`): cross-device entry barrier, then each + local token gathers its k expert outputs **remotely** from `peer.out_tok[dest_tok_id]` + and reduces in f32. Register-light i32 reads (2 bf16 / `v2f32` accumulate) + 2-way + unroll keep VGPRs low so 16 warps/block run at high occupancy to hide xGMI read + latency; remote reads are latency-bound so combine needs ~128 blocks, while + dispatch's posted writes saturate at ~64 blocks (half the CUs). +- Self-written volatile/atomic spin-waits (`flydsl_prims.spin_until_*`) — mori-shmem's + `wait_until_*` assert on a cco-only stack. Counters self-reset in-kernel → CUDAGraph-safe. + +## Perf (EP8, hidden=7168, top-k=8, 256 experts; dispatch 64blk / combine 128blk × 16warp, CUDA-graph, bf16) + +Per-rank bandwidth = `recv_tok * per_token_bytes / time` (the bench sizes the +payload per dtype, `hidden*2` for bf16). Indicative bf16 numbers on **MI308X +(gfx942)** xGMI: + +| tok/rank | dispatch | combine | +|---:|---:|---:| +| 512 | 268 GB/s | 213 GB/s | +| 2048 | 306 GB/s | 294 GB/s | +| 8192 | 314 GB/s | 323 GB/s | + +Cross-impl (v2 vs mori v1) latency tables for fp8/fp4 are in PR ROCm/mori#448. diff --git a/aiter/ops/flydsl/dispatch_combine_v2/__init__.py b/aiter/ops/flydsl/dispatch_combine_v2/__init__.py new file mode 100644 index 00000000000..055febed3bb --- /dev/null +++ b/aiter/ops/flydsl/dispatch_combine_v2/__init__.py @@ -0,0 +1,17 @@ +# Vendored from mori python/mori/ops/dispatch_combine_v2 (cco-LSA v2 EP dispatch/combine). +# The op/kernel layer lives here so gfx1250 kernel development happens in aiter; the +# underlying cco communication substrate (mori.cco.Communicator + libmori_cco*.{so,bc}) +# and mori.jit/mori.tensor_utils remain installed-mori runtime dependencies. +from .dispatch_combine_op import ( + SymmArena, + EpDispatchCombineConfig, + EpDispatchCombineOp, + EpDispatchRoutingHandle, +) + +__all__ = [ + "SymmArena", + "EpDispatchCombineConfig", + "EpDispatchCombineOp", + "EpDispatchRoutingHandle", +] diff --git a/aiter/ops/flydsl/dispatch_combine_v2/dispatch_combine_op.py b/aiter/ops/flydsl/dispatch_combine_v2/dispatch_combine_op.py new file mode 100644 index 00000000000..fb83d5089bc --- /dev/null +++ b/aiter/ops/flydsl/dispatch_combine_v2/dispatch_combine_op.py @@ -0,0 +1,1122 @@ +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Host op-layer for the cco-LSA intranode dispatch/combine kernels. One SymmArena +window holds the symmetric staging; per-rank metadata are plain device tensors.""" +from dataclasses import dataclass +import os + +import torch + +import flydsl.expr as fx +from mori.tensor_utils import from_gpu_ptr + +from .intranode_kernels import ( + make_dispatch, + make_combine, + make_combine_scatter, + make_combine_fused_reduce, + make_convert_dispatch_output, + make_convert_combine_input, + make_local_expert_count, +) + +_COMBINE_QUANT_TYPES = ("none", "fp8_direct_cast", "fp8_blockwise") + +_DT = { + torch.bfloat16: 2, + torch.float32: 4, + torch.float8_e4m3fnuz: 1, + torch.float8_e4m3fn: 1, +} +_FP8_DTYPES = (torch.float8_e4m3fnuz, torch.float8_e4m3fn) + + +def _align_up(x, a): + return (x + a - 1) // a * a + + +class SymmArena: + """One cco symmetric window carved into named, aligned sub-regions. A kernel + reaches peer pe's copy of region R via cco.Window(handle).lsa_ptr(pe, off_R).""" + + _ALIGN = 256 + + def __init__(self, comm, regions): + self._comm = comm + self._offsets = {} + self._sizes = {} + off = 0 + for name, nbytes in regions: + off = _align_up(off, self._ALIGN) + self._offsets[name] = off + self._sizes[name] = nbytes + off += nbytes + self._total = max(_align_up(off, self._ALIGN), self._ALIGN) + self._mem = comm.alloc_mem(self._total) + self._win = comm.register_window(self._mem.ptr, self._total) + + @property + def handle(self): + return self._win.handle + + @property + def total_bytes(self): + return self._total + + def offset(self, name): + return self._offsets[name] + + def local_ptr(self, name): + return self._win.local_ptr + self._offsets[name] + + def zero(self, name=None): + """Zero the whole window, or just region `name` (zero-copy int8 torch view).""" + if name is None: + ptr, nbytes = self._win.local_ptr, self._total + else: + ptr, nbytes = self.local_ptr(name), self._sizes[name] + from_gpu_ptr(ptr, (nbytes,), torch.int8).zero_() + + def close(self): + """Free the symmetric window (deregister before freeing the backing mem).""" + self._win.close() + self._mem.close() + + +@dataclass +class EpDispatchCombineConfig: + rank: int + world_size: int + hidden_dim: int + max_num_inp_token_per_rank: int + num_experts_per_rank: int + num_experts_per_token: int + # Base token dtype; dispatch_data_type / combine_data_type override per-op + # (None => data_type). Asymmetric (fp8 dispatch -> bf16 combine) is gather-only. + data_type: torch.dtype = torch.bfloat16 + dispatch_data_type: torch.dtype = None + combine_data_type: torch.dtype = None + # Per-token quant scales forwarded verbatim to dest out_scales (0 disables). + scale_dim: int = 0 + scale_type_size: int = 0 + # "gather" (UseP2PRead) or "scatter" (mori _nop2p, fp8 compression home). + combine_mode: str = "gather" + # Combine-side wire quantization (scatter-only; distinct from the dispatch + # token dtype). Forcing it != "none" switches combine_mode to scatter. + combine_quant_type: str = "none" # none | fp8_direct_cast | fp8_blockwise + # Geometry: None => auto (tuned schedule from tuning_configs in __post_init__); + # pin any to opt out. Combine's warp count is kept separate from dispatch's. + dispatch_block_num: int = None + combine_block_num: int = None + dispatch_warp_num_per_block: int = None + combine_warp_num_per_block: int = None + # Optional per-token plan: (max_tok_inclusive | None, disp_block, disp_warp, + # comb_block, comb_warp) buckets; the op precompiles the distinct (block, warp) + # variants and picks one at runtime. None => auto or single-shot fallback. + schedule: tuple = None + enable_std_moe: bool = False + max_total_recv_tokens: int = 0 # mori maxTotalRecvTokens; 0 = worst-case ws*M + + def __post_init__(self): + # all-or-none: setting only one silently defaults the other to data_type. + if (self.dispatch_data_type is None) != (self.combine_data_type is None): + raise ValueError( + "dispatch_data_type / combine_data_type must be set together " + "(all-or-none): got dispatch_data_type=" + f"{self.dispatch_data_type}, combine_data_type={self.combine_data_type}. " + "Set data_type alone for a symmetric op, or set both explicitly for " + "asymmetric dispatch/combine dtypes." + ) + if self.combine_quant_type not in _COMBINE_QUANT_TYPES: + raise ValueError( + f"combine_quant_type must be one of {_COMBINE_QUANT_TYPES}, " + f"got {self.combine_quant_type!r}" + ) + if self.combine_mode not in ("gather", "scatter", "scatter_fused"): + raise ValueError( + "combine_mode must be gather|scatter|scatter_fused, " + f"got {self.combine_mode!r}" + ) + if self.combine_quant_type != "none": + self.combine_mode = "scatter" + # The dispatch grid barrier iterates peers as range(lane, npes, wave) and + # resets inside the loop, correct only when npes <= wavefront. + if self.world_size > 64: + raise ValueError( + f"intranode op supports world_size <= 64, got {self.world_size}" + ) + # Token copy moves whole 16 B (vec4) chunks; per-token size must be 16 B aligned. + if self.token_nbytes % 16 != 0: + raise ValueError( + f"per-token transport bytes must be 16 B aligned (vec4 copy); " + f"hidden_dim={self.hidden_dim}, dispatch_dtype={self.dispatch_dtype} -> " + f"token_nbytes={self.token_nbytes}" + ) + if self.is_asymmetric_dtype: + # asymmetric dtype (separate disp_out/comb_stg buffers): gather/non-quant/ + # non-StdMoE only. + if ( + self.combine_mode != "gather" + or self.combine_quant_type != "none" + or self.enable_std_moe + ): + raise ValueError( + "combine_data_type (asymmetric dtype) requires combine_mode=gather, " + "combine_quant_type=none, enable_std_moe=False" + ) + if torch.float4_e2m1fn_x2 in (self.dispatch_dtype, self.combine_dtype): + raise ValueError( + "asymmetric dispatch/combine dtype does not support fp4" + ) + if self.combine_token_nbytes % 16 != 0: + raise ValueError( + f"combine per-token bytes must be 16 B aligned; combine_data_type=" + f"{self.combine_data_type} -> {self.combine_token_nbytes}" + ) + self._resolve_geometry() + + def _resolve_geometry(self): + """Fill block/warp/schedule. Tuned-by-default: if the caller pinned nothing, + pull the tuned geometry from tuning_configs.lookup; if any field is pinned, + honor it and fill the rest with the single-shot fallback (no schedule).""" + pinned = self.schedule is not None or any( + g is not None + for g in ( + self.dispatch_block_num, + self.combine_block_num, + self.dispatch_warp_num_per_block, + self.combine_warp_num_per_block, + ) + ) + if not pinned: + from .tuning_configs import lookup + + t = lookup( + self.world_size, + self.hidden_dim, + self.num_experts_per_token, + dtype=self.dtype_str, + ) + self.dispatch_block_num = t["dispatch_block_num"] + self.combine_block_num = t["combine_block_num"] + self.dispatch_warp_num_per_block = t["dispatch_warp_num_per_block"] + self.combine_warp_num_per_block = t["combine_warp_num_per_block"] + self.schedule = t["schedule"] + else: + # explicit geometry: fill any unset field with the single-shot default + if self.dispatch_block_num is None: + self.dispatch_block_num = 64 + if self.combine_block_num is None: + self.combine_block_num = 80 + if self.dispatch_warp_num_per_block is None: + self.dispatch_warp_num_per_block = 16 + if self.combine_warp_num_per_block is None: + self.combine_warp_num_per_block = 4 + + @property + def is_scatter(self): + return self.combine_mode in ("scatter", "scatter_fused") + + @property + def is_fused(self): + """gemm2-fused scatter: gemm2 P2P-writes weighted per-(token,k) results into + comb_inp; combine only barriers + sums. No Stage-1 write, no gather_reduce.""" + return self.combine_mode == "scatter_fused" + + @property + def fp8_direct_cast(self): + return self.combine_quant_type == "fp8_direct_cast" + + @property + def fp8_blockwise(self): + return self.combine_quant_type == "fp8_blockwise" + + @property + def combine_scale_dim(self): + """Per-token block count for fp8_blockwise combine (block_elems=128).""" + return self.hidden_dim // 128 if self.fp8_blockwise else 0 + + @property + def wire_elem_size(self): + """comb_inp transport element size: 1 byte for fp8 paths, else elem_size.""" + return ( + 1 + if self.combine_quant_type in ("fp8_direct_cast", "fp8_blockwise") + else self.elem_size + ) + + @property + def tdm_gather_scatter(self): + """gemm2-fused P2P scatter via TDM gather-store. The felix TDM GEMM2's + fused epilogue always scatters via the TDM gather-store, which requires a + pow2-padded comb_inp slot so the per-slot stride divides the 4GB-aligned + per-rank symmetric window (lets gemm2 fold (pe,slot) into a single TDM + gather row index). Tied to the same switch that selects the felix TDM + GEMM2 (AITER_EP_SCATTER_TDM, default on); set it to 0 to fall back onto + the grouped/contiguous mxscale gemm2 (unpadded buffer_store scatter).""" + return self.is_fused and os.environ.get("AITER_EP_SCATTER_TDM", "1") in ( + "1", + "true", + "True", + "yes", + "on", + ) + + @property + def comb_inp_slot_nbytes(self): + """Per-(token,k) comb_inp slot stride in bytes. Padded up to a power of + two for the TDM gather-store fused path; else the natural hidden row + size (hidden_dim * wire_elem_size).""" + base = self.hidden_dim * self.wire_elem_size + if self.tdm_gather_scatter: + p = 1 + while p < base: + p <<= 1 + return p + return base + + @classmethod + def tuned(cls, **kwargs): + """Build a config with geometry from tuning_configs (unless overridden in + kwargs). Back-compat alias; the plain constructor is also tuned-by-default.""" + from .tuning_configs import lookup + + dt = kwargs.get("data_type", torch.bfloat16) + dtype = ( + "fp4" + if dt == torch.float4_e2m1fn_x2 + else ("fp8" if dt in _FP8_DTYPES else "bf16") + ) + t = lookup( + kwargs["world_size"], + kwargs["hidden_dim"], + kwargs["num_experts_per_token"], + dtype=dtype, + ) + for k, v in t.items(): + kwargs.setdefault(k, v) + return cls(**kwargs) + + @property + def dispatch_dtype(self): + """Dispatch transport dtype (== data_type unless dispatch_data_type set).""" + return ( + self.dispatch_data_type + if self.dispatch_data_type is not None + else self.data_type + ) + + @property + def is_fp4(self): + return self.dispatch_dtype == torch.float4_e2m1fn_x2 + + @property + def is_fp8(self): + return self.dispatch_dtype in _FP8_DTYPES + + @property + def elem_size(self): + # fp4 is 0.5 B/elem; return a nominal 1 (kernels use the fp4 flag + + # token_nbytes for actual sizing, never elem_size for fp4 token buffers). + return 1 if self.is_fp4 else _DT[self.dispatch_dtype] + + @property + def token_nbytes(self): + """Per-token transport bytes (fp4 packs 2 e2m1/byte -> hidden/2).""" + return self.hidden_dim // 2 if self.is_fp4 else self.hidden_dim * self.elem_size + + @property + def combine_dtype(self): + """Combine transport dtype (== data_type unless combine_data_type set).""" + return ( + self.combine_data_type + if self.combine_data_type is not None + else self.data_type + ) + + @property + def combine_elem_size(self): + cdt = self.combine_dtype + return 1 if cdt == torch.float4_e2m1fn_x2 else _DT[cdt] + + @property + def combine_token_nbytes(self): + cdt = self.combine_dtype + return ( + self.hidden_dim // 2 + if cdt == torch.float4_e2m1fn_x2 + else self.hidden_dim * self.combine_elem_size + ) + + @property + def is_asymmetric_dtype(self): + return self.dispatch_dtype != self.combine_dtype + + @property + def dtype_str(self): + """Token/dispatch dtype key for tuning_configs.lookup (fp4/fp8/default).""" + if self.is_fp4: + return "fp4" + if self.is_fp8: + return "fp8" + return "bf16" + + @property + def max_recv(self): + """Sentinel / sender-side atomic-add allocation bound (always ws*M).""" + return self.world_size * self.max_num_inp_token_per_rank + + @property + def effective_max_recv_per_rank(self): + if self.max_total_recv_tokens <= 0: + return self.max_num_inp_token_per_rank + per = (self.max_total_recv_tokens + self.world_size - 1) // self.world_size + return min(per, self.max_num_inp_token_per_rank) + + @property + def effective_max_recv(self): + """Recv-slot cap passed to the kernels as max_recv (mori MaxNumTokensToRecv).""" + return self.world_size * self.effective_max_recv_per_rank + + +class EpDispatchRoutingHandle: + """Per-call routing snapshot. + + disp_dest_tok_id_map: forward (src_tok,k)->dest flat slot (tok_map). + disp_tok_id_to_src_tok_id_local: reverse recv-slot->src token (tis). + inter_node_*: empty placeholders (intranode-only; kept for 5-tensor shape parity). + + The reverse map is materialized LAZILY on first access: recv_to_src_token is + written by peers via P2P during dispatch and is only visible after the caller's + post-dispatch barrier, so cloning it eagerly would race those writes. + """ + + def __init__( + self, + disp_dest_tok_id_map, + inter_node_disp_dest_tok_id_map, + inter_node_disp_send_map, + total_recv_token_num, + disp_tok_id_to_src_tok_id_local=None, + cur_rank_num_token=0, + *, + reverse_src_view=None, + ): + self.disp_dest_tok_id_map = disp_dest_tok_id_map + self.inter_node_disp_dest_tok_id_map = inter_node_disp_dest_tok_id_map + self.inter_node_disp_send_map = inter_node_disp_send_map + self.total_recv_token_num = total_recv_token_num + self.cur_rank_num_token = cur_rank_num_token + # Either an already-materialized reverse map (from_tensors round-trip) or + # a live arena view to clone on first access (dispatch()). + self._reverse_cache = disp_tok_id_to_src_tok_id_local + self._reverse_src_view = reverse_src_view + + @property + def disp_tok_id_to_src_tok_id_local(self): + if self._reverse_cache is None: + # First access (post-barrier): clone off the arena before the next + # dispatch overwrites the region. + self._reverse_cache = self._reverse_src_view.clone() + return self._reverse_cache + + def tensors(self): + return ( + self.disp_dest_tok_id_map, + self.inter_node_disp_dest_tok_id_map, + self.inter_node_disp_send_map, + self.total_recv_token_num, + self.disp_tok_id_to_src_tok_id_local, + ) + + @classmethod + def from_tensors(cls, tensors, cur_rank_num_token=0): + return cls(*tensors, cur_rank_num_token=cur_rank_num_token) + + +class EpDispatchCombineOp: + def __init__(self, cfg: EpDispatchCombineConfig, comm): + self.cfg = cfg + self.comm = comm + device = torch.device("cuda", torch.cuda.current_device()) + self.dev = device + elem_size = cfg.elem_size + is_fp4 = cfg.is_fp4 + is_fp8 = cfg.is_fp8 + token_nbytes = cfg.token_nbytes # per-token transport bytes (fp4 = hidden/2) + if (is_fp4 or is_fp8) and cfg.is_scatter: + raise ValueError( + "plain fp4/fp8 token dtype is gather-only " + "(fp8 quant uses combine_quant_type=fp8_direct_cast, not data_type)" + ) + topk = cfg.num_experts_per_token + hidden_dim = cfg.hidden_dim + max_tok_per_rank = cfg.max_num_inp_token_per_rank + recv_cap = cfg.effective_max_recv # recv-slot cap (== ws*M unless capped) + self._recv_cap = recv_cap + + self._scale_bytes = cfg.scale_dim * cfg.scale_type_size + self._scale_num_i32 = (self._scale_bytes + 3) // 4 + self._enable_scales = self._scale_bytes > 0 + + regions = [ + ("tok_off", 4), + ("recv_num", cfg.world_size * 4), + ("recv_to_src_token", recv_cap * 4), + ("out_idx", recv_cap * topk * 4), + ("out_wts", recv_cap * topk * 4), + # disp_out: dispatch dest / expert-GEMM input, kept separate from comb_stg + # so combine's copy-in never clobbers the dispatched tokens. + ("disp_out", recv_cap * token_nbytes), + # comb_stg: combine staging (post-expert results that peers gather). + ("comb_stg", recv_cap * cfg.combine_token_nbytes), + ("cross_device_barrier", cfg.world_size * 8), + ] + if self._enable_scales: + regions.append(("out_scales", recv_cap * self._scale_num_i32 * 4)) + # scatter combine needs its own staging regions + if cfg.is_scatter: + wire_elem_size = cfg.wire_elem_size + if cfg.is_fused: + # per-(origin_token, k): this rank's own M tokens x topk slots; each + # (token,k) is written exactly once by the owner rank's gemm2 (the + # expert k of that token lives on one rank) -> no cross-rank collision. + # Weights are pre-multiplied in gemm2, so no comb_wts region. + regions.append( + ("comb_inp", max_tok_per_rank * topk * cfg.comb_inp_slot_nbytes) + ) + else: + regions.append( + ( + "comb_inp", + cfg.world_size * max_tok_per_rank * hidden_dim * wire_elem_size, + ) + ) + regions.append( + ("comb_wts", cfg.world_size * max_tok_per_rank * topk * 4) + ) + if cfg.fp8_blockwise: + regions.append( + ( + "comb_scales", + cfg.world_size + * max_tok_per_rank + * cfg.combine_scale_dim + * 4, + ) + ) + self.arena = SymmArena(comm, regions) + self.arena.zero() + + self.token_dest_map = torch.full( + (max_tok_per_rank * topk,), -1, dtype=torch.int32, device=device + ) + self._empty_i32 = torch.empty( + 0, dtype=torch.int32, device=device + ) # inter-node placeholders + self.dest_pe_counter = torch.zeros( + cfg.world_size, dtype=torch.int32, device=device + ) + self.dispatch_barrier = torch.zeros(1, dtype=torch.int32, device=device) + self.total_recv = torch.zeros(1, dtype=torch.int32, device=device) + self.combine_barrier = torch.zeros(1, dtype=torch.int32, device=device) + self.cross_device_flag = torch.ones(1, dtype=torch.int64, device=device) + c_dt = cfg.combine_dtype # combine output dtype + c_elem = cfg.combine_elem_size + if c_dt == torch.float4_e2m1fn_x2: # fp4 combine outputs fp4 (hidden/2 B/token) + self.combine_out = torch.zeros( + max_tok_per_rank * (hidden_dim // 2), dtype=torch.int8, device=device + ) + elif c_dt in _FP8_DTYPES: # fp8 combine outputs fp8 (1 byte/elem) + self.combine_out = torch.zeros( + max_tok_per_rank * hidden_dim, dtype=torch.int8, device=device + ) + else: + self.combine_out = torch.zeros( + max_tok_per_rank * hidden_dim, + dtype=torch.int16 if c_elem == 2 else torch.int32, + device=device, + ) + self.combine_out_weights = torch.zeros( + max_tok_per_rank * topk, dtype=torch.float32, device=device + ) + + arena = self.arena + # Distinct (block, warp) variants to precompile; a per-token schedule picks + # the best at runtime, else single-shot. Scatter combine is not schedule-tuned. + schedule = cfg.schedule + if schedule: + dispatch_specs = sorted({(db, dw) for (_, db, dw, _, _) in schedule}) + combine_specs = sorted({(cb, cw) for (_, _, _, cb, cw) in schedule}) + else: + dispatch_specs = [(cfg.dispatch_block_num, cfg.dispatch_warp_num_per_block)] + combine_specs = [(cfg.combine_block_num, cfg.combine_warp_num_per_block)] + if cfg.is_scatter: + # Scatter combine hits an all-block grid barrier mid-kernel (Stage 1 + # write -> barrier -> Stage 3 read). With the gather-tuned (block=80, + # warp=4) that barrier serializes badly (occupancy-bound: too many + # blocks to co-reside, so the grid barrier spins). Far fewer blocks with + # more warps per block (8x16) keeps all blocks resident and roughly + # halves the scatter combine kernel time. Override via SCATTER_COMB_BW + # (e.g. "16,16") to retune for a different token count / hidden dim. + _b, _w = 8, 16 + _bw = os.environ.get("SCATTER_COMB_BW") + if _bw: + _b, _w = (int(x) for x in _bw.split(",")) + combine_specs = [(_b, _w)] + elif cfg.is_fused: + # The fused combine is an xdb cross-device barrier + a light + # per-(token,k) topk sum. That sum is *tiny* (tens of MB) so the + # kernel is dominated by the cross-BLOCK grid barrier (block 0 signals + # every block via comb_bar, then waits for all of them to report). That + # handshake serializes with block_num, so once block_num exceeds what + # co-resides on the 256-CU GPU the kernel falls off a cliff: at bs512 a + # 512-block grid ran ~1255 us/call vs ~94 us/call for a 16-block grid + # (~13x). The reduce is not big enough to be HBM bound here, so the + # right geometry is the FEWEST blocks that still saturate HBM -- + # empirically 8-32 blocks across bs128..bs2048 (best per_layer, and + # scatter_fused then beats gather). Token-adaptive tiers + # (count_upper_bound, (block, warp)); block count grows very gently + # with tokens and stays small. Override via SCATTER_COMB_BW to retune. + self._fused_comb_tiers = [ + (256, (8, 16)), + (1024, (16, 16)), + (None, (32, 16)), + ] + combine_specs = [s for _, s in self._fused_comb_tiers] + else: + combine_specs = [(_b, _w)] + self._dispatch_specs = dispatch_specs + self._combine_specs = combine_specs + + self._dispatch_kwargs = dict( + rank=cfg.rank, + npes=cfg.world_size, + experts_per_rank=cfg.num_experts_per_rank, + experts_per_token=topk, + hidden_dim=hidden_dim, + hidden_elem_size=elem_size, + max_tok_per_rank=max_tok_per_rank, + max_recv=recv_cap, + off_tok_off=arena.offset("tok_off"), + off_recv_num=arena.offset("recv_num"), + off_tis=arena.offset("recv_to_src_token"), + off_out_idx=arena.offset("out_idx"), + off_out_wts=arena.offset("out_wts"), + off_out_tok=arena.offset("disp_out"), + off_out_scales=arena.offset("out_scales") if self._enable_scales else 0, + scale_dim=cfg.scale_dim, + scale_type_size=cfg.scale_type_size, + fp4=is_fp4, + ) + # (block, warp) -> compiled dispatch / combine kernel. + self._dispatch_variants = { + (b, w): make_dispatch( + block_num=b, warp_num_per_block=w, **self._dispatch_kwargs + ) + for (b, w) in dispatch_specs + } + self._dispatch_replay_variants = {} # lazily compiled per (block, warp) + if cfg.is_fused: + # gemm2 already P2P-wrote weighted per-(token,k) results into comb_inp; + # combine only barriers (wait for peers' writes) + sums the topk slots. + self._combine_variants = { + (b, w): make_combine_fused_reduce( + rank=cfg.rank, + npes=cfg.world_size, + experts_per_token=topk, + hidden_dim=hidden_dim, + hidden_elem_size=elem_size, + max_tok_per_rank=max_tok_per_rank, + block_num=b, + warp_num_per_block=w, + off_comb_inp=arena.offset("comb_inp"), + off_xdb_mem=arena.offset("cross_device_barrier"), + slot_stride_nbytes=cfg.comb_inp_slot_nbytes, + ) + for (b, w) in combine_specs + } + elif cfg.is_scatter: + self._combine_variants = { + (b, w): make_combine_scatter( + rank=cfg.rank, + npes=cfg.world_size, + experts_per_token=topk, + hidden_dim=hidden_dim, + hidden_elem_size=elem_size, + max_tok_per_rank=max_tok_per_rank, + max_recv=recv_cap, + block_num=b, + warp_num_per_block=w, + off_out_tok=arena.offset("comb_stg"), + off_comb_inp=arena.offset("comb_inp"), + off_tis=arena.offset("recv_to_src_token"), + off_xdb_mem=arena.offset("cross_device_barrier"), + off_out_wts=arena.offset("out_wts"), + off_comb_wts=arena.offset("comb_wts"), + off_comb_scales=( + arena.offset("comb_scales") if cfg.fp8_blockwise else 0 + ), + fp8_direct_cast=cfg.fp8_direct_cast, + fp8_blockwise=cfg.fp8_blockwise, + scale_dim=cfg.combine_scale_dim, + reset_total_recv=False, + ) + for (b, w) in combine_specs + } + else: + self._combine_variants = { + (b, w): make_combine( + rank=cfg.rank, + npes=cfg.world_size, + experts_per_token=topk, + hidden_dim=hidden_dim, + hidden_elem_size=cfg.combine_elem_size, + max_tok_per_rank=max_tok_per_rank, + max_recv=recv_cap, + block_num=b, + warp_num_per_block=w, + off_out_tok=arena.offset("comb_stg"), + off_xdb_mem=arena.offset("cross_device_barrier"), + off_out_wts=arena.offset("out_wts"), + reset_total_recv=True, + fp4=(cfg.combine_dtype == torch.float4_e2m1fn_x2), + ) + for (b, w) in combine_specs + } + + self._local_expert_count_buf = torch.zeros( + cfg.num_experts_per_rank, dtype=torch.int32, device=device + ) + self._local_expert_count = make_local_expert_count( + rank=cfg.rank, + experts_per_rank=cfg.num_experts_per_rank, + experts_per_token=topk, + block_num=cfg.dispatch_block_num, + warp_num_per_block=cfg.dispatch_warp_num_per_block, + ) + + if cfg.enable_std_moe: + assert elem_size == 2, "StdMoE convert path is bf16-only" + experts_per_rank = cfg.num_experts_per_rank + max_tok_per_expert = cfg.world_size * max_tok_per_rank + self._max_tok_per_expert = max_tok_per_expert + self.packed_x = torch.zeros( + experts_per_rank * max_tok_per_expert * hidden_dim, + dtype=torch.int16, + device=device, + ) + self.packed_count = torch.zeros( + experts_per_rank, dtype=torch.int32, device=device + ) + self.packed_src = torch.zeros( + experts_per_rank * max_tok_per_expert, dtype=torch.int32, device=device + ) + self.slot_map = torch.full( + (recv_cap * topk,), -1, dtype=torch.int64, device=device + ) + self._convert_dispatch = make_convert_dispatch_output( + rank=cfg.rank, + experts_per_rank=experts_per_rank, + experts_per_token=topk, + hidden_dim=hidden_dim, + hidden_elem_size=elem_size, + max_tok_per_expert=max_tok_per_expert, + block_num=cfg.dispatch_block_num, + warp_num_per_block=cfg.dispatch_warp_num_per_block, + ) + self._convert_combine = make_convert_combine_input( + rank=cfg.rank, + experts_per_rank=experts_per_rank, + experts_per_token=topk, + hidden_dim=hidden_dim, + hidden_elem_size=elem_size, + max_tok_per_expert=max_tok_per_expert, + block_num=cfg.combine_block_num, + warp_num_per_block=cfg.combine_warp_num_per_block, + ) + self._closed = False + + def close(self): + """Free this op's symmetric arena window. Call (or use as a context + manager) when the op is discarded but its Communicator lives on — + otherwise the arena stays in comm._resources until the comm is destroyed.""" + if self._closed: + return + self._closed = True + self.arena.close() + + def __enter__(self): + return self + + def __exit__(self, *exc): + self.close() + + def recv_tokens(self): + """Arena disp_out [max_recv, hidden] (dispatch dest / expert-GEMM input). + Separate from comb_stg, so combine's copy-in never overwrites it — the + expert can read this in place without a defensive .clone(). + fp4 packs 2 e2m1 per float4_e2m1fn_x2 element -> last dim is hidden/2.""" + cols = self.cfg.hidden_dim // 2 if self.cfg.is_fp4 else self.cfg.hidden_dim + return from_gpu_ptr( + self.arena.local_ptr("disp_out"), + (self._recv_cap, cols), + self.cfg.dispatch_dtype, + ) + + def combine_in_view(self): + """comb_stg as combine dtype [max_recv, hidden] — combine()'s copy target.""" + cdt = self.cfg.combine_dtype + cols = ( + self.cfg.hidden_dim // 2 + if cdt == torch.float4_e2m1fn_x2 + else self.cfg.hidden_dim + ) + return from_gpu_ptr( + self.arena.local_ptr("comb_stg"), (self._recv_cap, cols), cdt + ) + + def convert_dispatch_output(self): + """mori ConvertDispatchOutput: repack recv tokens into per-local-expert + buckets. Returns (packed_x, packed_count, packed_src); GEMM overwrites + packed_x in place.""" + assert self.cfg.enable_std_moe, "op built without enable_std_moe" + self.packed_count.zero_() + self.slot_map.fill_(-1) + arena = self.arena + stream = fx.Stream(torch.cuda.current_stream()) + self._convert_dispatch( + arena.local_ptr("disp_out"), + arena.local_ptr("out_idx"), + arena.local_ptr("recv_to_src_token"), + self.total_recv.data_ptr(), + self.packed_x.data_ptr(), + self.packed_count.data_ptr(), + self.packed_src.data_ptr(), + self.slot_map.data_ptr(), + stream, + ) + experts_per_rank = self.cfg.num_experts_per_rank + max_tok_per_expert = self._max_tok_per_expert + hidden_dim = self.cfg.hidden_dim + packed_x_view = from_gpu_ptr( + self.packed_x.data_ptr(), + (experts_per_rank, max_tok_per_expert, hidden_dim), + self.cfg.dispatch_dtype, + ) + return packed_x_view, self.packed_count, self.packed_src + + def convert_combine_input(self, routing): + """mori ConvertCombineInput: weighted-reduce each recv token's local-expert + outputs from packed_x back into comb_stg. Run after GEMM, before combine.""" + assert self.cfg.enable_std_moe, "op built without enable_std_moe" + arena = self.arena + stream = fx.Stream(torch.cuda.current_stream()) + self._convert_combine( + arena.local_ptr("comb_stg"), + arena.local_ptr("out_wts"), + routing.total_recv_token_num.data_ptr(), + self.packed_x.data_ptr(), + self.slot_map.data_ptr(), + stream, + ) + + def recv_weights(self): + """Arena out_wts as [max_recv, topk] f32 (forwarded per-token weights).""" + return from_gpu_ptr( + self.arena.local_ptr("out_wts"), + (self._recv_cap, self.cfg.num_experts_per_token), + torch.float32, + ) + + def recv_indices(self): + """Arena out_idx as [max_recv, topk] i32 (forwarded expert indices).""" + return from_gpu_ptr( + self.arena.local_ptr("out_idx"), + (self._recv_cap, self.cfg.num_experts_per_token), + torch.int32, + ) + + def recv_scales(self): + """Forwarded per-token scales as opaque i32 dwords [max_recv, scale_num_i32], + or None if built without scales.""" + if not self._enable_scales: + return None + return from_gpu_ptr( + self.arena.local_ptr("out_scales"), + (self._recv_cap, self._scale_num_i32), + torch.int32, + ) + + def _pick(self, num_tokens): + """((disp_block, disp_warp), (comb_block, comb_warp)) for a runtime token + count via the per-token schedule; falls back to the single-shot specs + otherwise. Returned specs are clamped to the precompiled variants.""" + schedule = self.cfg.schedule + disp_spec = comb_spec = None + if schedule: + for bucket in schedule: + max_tok = bucket[0] + if max_tok is None or num_tokens <= max_tok: + disp_spec, comb_spec = (bucket[1], bucket[2]), ( + bucket[3], + bucket[4], + ) + break + if disp_spec is None: + last = schedule[-1] + disp_spec, comb_spec = (last[1], last[2]), (last[3], last[4]) + else: + disp_spec, comb_spec = self._dispatch_specs[0], self._combine_specs[0] + if disp_spec not in self._dispatch_variants: + disp_spec = self._dispatch_specs[-1] + if comb_spec not in self._combine_variants: + comb_spec = self._combine_specs[-1] + return disp_spec, comb_spec + + def dispatch( + self, input, weights, scales, indices, *, routing=None, return_routing=False + ): + """mori-parity dispatch. input [n_tok,hidden], weights [n_tok,topk] f32, + scales [n_tok,scale_dim] (or None), indices [n_tok,topk] i32. + + routing=: replay a prior handle (reuse cached dest-slot layout, skip + atomic routing). return_routing=: also return the handle. Mutually + exclusive. Returns (out, out_weights, out_scales, out_indices, + total_recv[, routing]); out == arena disp_out (safe to read without + .clone() — combine stages into a separate comb_stg buffer). + """ + if routing is not None and return_routing: + raise ValueError( + "pass either routing= (replay) or return_routing=True, not both" + ) + num_input_tokens = input.shape[0] + disp_spec, _ = self._pick(num_input_tokens) + # total_recv is self-reset inside the dispatch kernel (warp 0, Phase 2). + scale_ptr = ( + scales.data_ptr() if (scales is not None and self._enable_scales) else 0 + ) + stream = fx.Stream(torch.cuda.current_stream()) + weight_ptr = weights.data_ptr() if weights is not None else 0 + if routing is not None: + kern = self._dispatch_replay_variants.get(disp_spec) + if kern is None: + kern = self._dispatch_replay_variants[disp_spec] = make_dispatch( + replay=True, + block_num=disp_spec[0], + warp_num_per_block=disp_spec[1], + **self._dispatch_kwargs, + ) + dest_map_ptr = routing.disp_dest_tok_id_map.data_ptr() + kern( + self.arena.handle, + input.data_ptr(), + indices.data_ptr(), + weight_ptr, + dest_map_ptr, + self.dest_pe_counter.data_ptr(), + self.dispatch_barrier.data_ptr(), + self.total_recv.data_ptr(), + scale_ptr, + self.cfg.rank, + num_input_tokens, + stream, + ) + else: + self._dispatch_variants[disp_spec]( + self.arena.handle, + input.data_ptr(), + indices.data_ptr(), + weight_ptr, + self.token_dest_map.data_ptr(), + self.dest_pe_counter.data_ptr(), + self.dispatch_barrier.data_ptr(), + self.total_recv.data_ptr(), + scale_ptr, + self.cfg.rank, + num_input_tokens, + stream, + ) + + out = self.recv_tokens() + out_weights = self.recv_weights() + out_scales = self.recv_scales() + out_indices = self.recv_indices() + base = (out, out_weights, out_scales, out_indices, self.total_recv) + if not return_routing: + return base + + # Pass a live arena view; the reverse map is cloned lazily on first + # access (post-barrier), see EpDispatchRoutingHandle. + recv_to_src_view = from_gpu_ptr( + self.arena.local_ptr("recv_to_src_token"), (self._recv_cap,), torch.int32 + ) + routing = EpDispatchRoutingHandle( + disp_dest_tok_id_map=self.token_dest_map.clone(), + inter_node_disp_dest_tok_id_map=self._empty_i32, + inter_node_disp_send_map=self._empty_i32, + total_recv_token_num=self.total_recv, + cur_rank_num_token=num_input_tokens, + reverse_src_view=recv_to_src_view, + ) + return base + (routing,) + + def combine(self, input, weights=None, indices=None, *, routing): + """mori-parity combine. input [<=max_recv,hidden] post-expert tokens + (copied into arena comb_stg if not already there). weights/indices are + accepted for API parity but unused (weights come from forwarded out_wts, + routing carries the mapping). Returns (out [ct,hidden], out_weights [ct,topk]). + """ + if self.cfg.is_fused: + return self._combine_fused(routing) + comb_stg_ptr = self.arena.local_ptr("comb_stg") + if self.cfg.is_scatter: + # Scatter Stage 1 only LOCAL-reads its source, so read the post-expert + # tokens (GEMM output) directly and skip the copy into symmetric + # comb_stg. StdMoE already staged comb_stg, so use that. + src_tok_ptr = comb_stg_ptr if self.cfg.enable_std_moe else input.data_ptr() + else: + # Gather peers REMOTE-read comb_stg, so the post-expert tokens must be + # copied into symmetric memory first. StdMoE already staged comb_stg. + if not self.cfg.enable_std_moe and input.data_ptr() != comb_stg_ptr: + # copy in the combine-dtype layout (not recv_tokens()'s dispatch view) + dst = self.combine_in_view().view(-1)[: input.numel()] + dst.copy_(input.reshape(-1)) + self.combine_out.zero_() + stream = fx.Stream(torch.cuda.current_stream()) + _, comb_spec = self._pick(routing.cur_rank_num_token) + if self.cfg.is_scatter: + self._combine_variants[comb_spec]( + self.arena.handle, + routing.disp_dest_tok_id_map.data_ptr(), + self.combine_barrier.data_ptr(), + self.cross_device_flag.data_ptr(), + routing.total_recv_token_num.data_ptr(), + self.combine_out.data_ptr(), + self.combine_out_weights.data_ptr(), + src_tok_ptr, + self.cfg.rank, + routing.cur_rank_num_token, + stream, + ) + else: + self._combine_variants[comb_spec]( + self.arena.handle, + routing.disp_dest_tok_id_map.data_ptr(), + self.combine_barrier.data_ptr(), + self.cross_device_flag.data_ptr(), + routing.total_recv_token_num.data_ptr(), + self.combine_out.data_ptr(), + self.combine_out_weights.data_ptr(), + self.cfg.rank, + routing.cur_rank_num_token, + stream, + ) + count = routing.cur_rank_num_token + hidden_dim = self.cfg.hidden_dim + topk = self.cfg.num_experts_per_token + cdt = self.cfg.combine_dtype + cols = ( + hidden_dim // 2 if cdt == torch.float4_e2m1fn_x2 else hidden_dim + ) # fp4 out is hidden/2 float4 elems + out = self.combine_out[: count * cols].view(cdt).view(count, cols) + return out, self.combine_out_weights[: count * topk].view(count, topk) + + # ---- gemm2-fused scatter (combine_mode="scatter_fused") ---- # + def zero_fused_staging(self): + """Zero the per-(token,k) comb_inp before the gemm2 P2P writes of a layer. + Needed because gemm2 writes only the (token,k) slots whose expert survived + dispatch/capacity; dropped slots must read 0 in the combine topk sum.""" + assert self.cfg.is_fused, "zero_fused_staging is scatter_fused-only" + self.arena.zero("comb_inp") + + def ep_scatter_params(self): + """Handles the moe gemm2 epilogue needs to P2P-write its weighted per-(token,k) + results straight into peers' comb_inp (combine_mode='scatter_fused').""" + assert self.cfg.is_fused, "ep_scatter_params is scatter_fused-only" + return dict( + ep_arena_handle=self.arena.handle, + ep_comb_inp_off=self.arena.offset("comb_inp"), + ep_tis_off=self.arena.offset("recv_to_src_token"), + ep_wire_nbytes=self.cfg.comb_inp_slot_nbytes, + ep_rank=self.cfg.rank, + ep_max_tok=self.cfg.max_num_inp_token_per_rank, + ep_topk=self.cfg.num_experts_per_token, + ) + + def _combine_fused(self, routing): + """gemm2 already P2P-wrote weighted per-(token,k) results into comb_inp; just + barrier (wait for peers' writes) + sum the topk slots per origin token.""" + stream = fx.Stream(torch.cuda.current_stream()) + count = routing.cur_rank_num_token + # Token-adaptive geometry: few blocks (barrier-bound) at low token counts, + # more blocks (sum/BW-bound) at high token counts. Falls back to _pick when + # a single spec is compiled (e.g. SCATTER_COMB_BW pinned). + _tiers = getattr(self, "_fused_comb_tiers", None) + if _tiers is not None: + comb_spec = _tiers[-1][1] + for _ub, _spec in _tiers: + if _ub is not None and count <= _ub: + comb_spec = _spec + break + else: + _, comb_spec = self._pick(count) + self.combine_out.zero_() + self._combine_variants[comb_spec]( + self.arena.handle, + self.combine_barrier.data_ptr(), + self.cross_device_flag.data_ptr(), + self.combine_out.data_ptr(), + self.cfg.rank, + count, + stream, + ) + hidden_dim = self.cfg.hidden_dim + topk = self.cfg.num_experts_per_token + cdt = self.cfg.combine_dtype + cols = hidden_dim // 2 if cdt == torch.float4_e2m1fn_x2 else hidden_dim + out = self.combine_out[: count * cols].view(cdt).view(count, cols) + return out, self.combine_out_weights[: count * topk].view(count, topk) + + def local_expert_count(self): + """[num_experts_per_rank] i32: recv tokens per local expert. Call after + dispatch, before combine (gather resets total_recv).""" + self._local_expert_count_buf.zero_() + stream = fx.Stream(torch.cuda.current_stream()) + self._local_expert_count( + self.arena.local_ptr("out_idx"), + self.total_recv.data_ptr(), + self._local_expert_count_buf.data_ptr(), + stream, + ) + return self._local_expert_count_buf + + def reset(self): + """Zero arena staging + per-rank counters/barriers (mori LaunchReset). + Kernels self-reset counters already; this forces a clean slate.""" + self.arena.zero() + self.token_dest_map.fill_(-1) + self.dest_pe_counter.zero_() + self.dispatch_barrier.zero_() + self.total_recv.zero_() + self.combine_barrier.zero_() + self.cross_device_flag.fill_(1) diff --git a/aiter/ops/flydsl/dispatch_combine_v2/flydsl_prims.py b/aiter/ops/flydsl/dispatch_combine_v2/flydsl_prims.py new file mode 100644 index 00000000000..1d7c0aa963d --- /dev/null +++ b/aiter/ops/flydsl/dispatch_combine_v2/flydsl_prims.py @@ -0,0 +1,169 @@ +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""FlyDSL device primitives for the cco-LSA dispatch/combine kernels: system-scope +atomics, ordered stores, fences, uncached/non-temporal loads and volatile spin-waits, +on top of cco peer pointers (cco.Window(h).lsa_ptr(pe, off)). + +The atomic / ordered-store / fence / volatile-load ops stay on +flydsl._mlir.dialects.llvm: the high-level FlyDSL API exposes no memory ordering, +sync-scope, volatile, or non-temporal control, which these primitives require. +""" +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm as _llvm_d +from flydsl._mlir.dialects import scf +from flydsl.expr import arith +from flydsl.expr.typing import T +import flydsl.expr as fx + + +def _gptr(addr_i64): + return _llvm_d.IntToPtrOp( + _llvm_d.PointerType.get(address_space=1), arith.unwrap(addr_i64) + ).result + + +def _ptr_plus(base_i64, offset, elem_bytes): + """Global pointer for base + offset*elem_bytes (offset may be i32 or i64).""" + addr = fx.Int64(arith.unwrap(base_i64)) + fx.Int64(arith.unwrap(offset)) * elem_bytes + return _gptr(addr) + + +def atomic_add_global(addr_i64, val): + """Monotonic remote global fetch-and-add at addr_i64; returns old value.""" + return _llvm_d.AtomicRMWOp( + _llvm_d.AtomicBinOp.add, + _gptr(addr_i64), + arith.unwrap(val), + _llvm_d.AtomicOrdering.monotonic, + ).res + + +def store_i32_system(addr_i64, offset, val): + """System-release i32 store at addr + offset*4.""" + _llvm_d.StoreOp( + arith.unwrap(val), + _ptr_plus(addr_i64, offset, 4), + alignment=4, + ordering=_llvm_d.AtomicOrdering.release, + syncscope="one-as", + ) + + +def store_i64_system(addr_i64, offset, val): + """System-release i64 store at addr + offset*8.""" + _llvm_d.StoreOp( + arith.unwrap(val), + _ptr_plus(addr_i64, offset, 8), + alignment=8, + ordering=_llvm_d.AtomicOrdering.release, + syncscope="one-as", + ) + + +def fence_system_acquire(): + _llvm_d.FenceOp(_llvm_d.AtomicOrdering.acquire, syncscope="one-as") + + +def fence_system_release(): + _llvm_d.FenceOp(_llvm_d.AtomicOrdering.release, syncscope="one-as") + + +def _unwrap(v): + return v.ir_value() if hasattr(v, "ir_value") else v + + +def load_i32_acquire(addr_i64): + """Volatile monotonic i32 load: volatile+ordering keeps the spin re-read from + being hoisted/CSE'd out of the wait loop (LICM would otherwise spin on a stale + value).""" + return _llvm_d.LoadOp( + T.i32, + _gptr(addr_i64), + alignment=4, + volatile_=True, + ordering=_llvm_d.AtomicOrdering.monotonic, + syncscope="one-as", + ).res + + +def load_i64_acquire(addr_i64): + """Volatile monotonic i64 load.""" + return _llvm_d.LoadOp( + T.i64, + _gptr(addr_i64), + alignment=8, + volatile_=True, + ordering=_llvm_d.AtomicOrdering.monotonic, + syncscope="one-as", + ).res + + +def load_i32_nt(base_i64, offset): + """Non-temporal global i32 load at base + offset*4. A raw global load (VGPR + address) avoids the per-expert buffer-descriptor waterfall; caller ensures the + address is in-bounds.""" + return _llvm_d.LoadOp( + T.i32, _ptr_plus(base_i64, offset, 4), alignment=4, nontemporal=True + ).res + + +def load_v4i32_nt(base_i64, offset): + """Non-temporal global vector<4xi32> load at base + offset*4 (global_load_dwordx4). + offset is in i32 units; alignment=4 since the per-warp offset is only 4B-aligned.""" + return _llvm_d.LoadOp( + T.i32x4, _ptr_plus(base_i64, offset, 4), alignment=4, nontemporal=True + ).res + + +def _spin(addr_i64, keep_waiting, *, width=32): + """Spin on a volatile/atomic load at addr_i64 until keep_waiting(cur) is false; + returns the awaited value. Self-contained (mori's wait_until_* need ShmemStates + and cannot run on a cco-only stack).""" + if width == 64: + ty, load, wrap = T.i64, load_i64_acquire, fx.Int64 + else: + ty, load, wrap = T.i32, load_i32_acquire, fx.Int32 + loop = scf.WhileOp([ty], [_unwrap(load(addr_i64))]) + cond = ir.Block.create_at_start(loop.before, [ty]) + body = ir.Block.create_at_start(loop.after, [ty]) + with ir.InsertionPoint(cond): + scf.ConditionOp( + _unwrap(keep_waiting(wrap(cond.arguments[0]))), [cond.arguments[0]] + ) + with ir.InsertionPoint(body): + scf.YieldOp([_unwrap(load(addr_i64))]) + return wrap(loop.results[0]) + + +def spin_until_eq_i64(addr_i64, val): + """Spin until *addr (i64) == val.""" + return _spin(addr_i64, lambda cur: cur != fx.Int64(val), width=64) + + +def spin_until_eq_i32(addr_i64, val): + """Spin until *addr == val.""" + return _spin(addr_i64, lambda cur: cur != fx.Int32(val)) + + +def spin_until_gt_i32(addr_i64, val): + """Spin until *addr > val (signed); returns the value seen.""" + return _spin(addr_i64, lambda cur: cur <= fx.Int32(val)) diff --git a/aiter/ops/flydsl/dispatch_combine_v2/intranode_kernels.py b/aiter/ops/flydsl/dispatch_combine_v2/intranode_kernels.py new file mode 100644 index 00000000000..5dd3492b798 --- /dev/null +++ b/aiter/ops/flydsl/dispatch_combine_v2/intranode_kernels.py @@ -0,0 +1,1846 @@ +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""FlyDSL intranode (single-node cco-LSA P2P) device kernels for the dispatch/combine +op: dispatch (+scales/replay), combine (gather + scatter/quant), StdMoE convert and +local expert count. Each is a compile-time-parameterised @flyc.jit factory; peer +addressing goes through cco.Window(handle).lsa_ptr(pe, off). + +Conventions: `rsrc_*` = buffer resource descriptor; `safe_*` = real value on live +lanes / in-bounds fallback (0 or self-rank) on dropped lanes; "sentinel" = tok_map +dropped-slot marker (dest PE == npes); "tis" = recv-slot -> global source-token map; +"xdb" = cross-device barrier (bumped i64 flag handshaked across all ranks). +""" +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import arith, const_expr, range_constexpr, T, vector +from flydsl.expr.buffer_ops import ( + buffer_load, + buffer_store, + create_buffer_resource_from_addr, +) +from flydsl.expr.rocdl import ( + ballot, + readlane, + ds_bpermute, + fmed3, + cvt_pk_f32_fp8, + cvt_pk_fp8_f32, + cvt_scalef32_pk_f32_fp4, + cvt_scalef32_pk_fp4_f32, + cvt_scale_pk8_f32_fp4, + cvt_scalef32_pk8_fp4_f32, +) +from flydsl.expr.typing import Int32, Int64 + +import mori.cco.device.flydsl as cco +from . import flydsl_prims as P + +# Wavefront size: gfx9 = 64, gfx12 (gfx1250) = 32. Override with MORI_WAVE_SIZE; +# key off the arch string since get_warp_size(gfx1250) wrongly reports 64. +import os as _os + + +def _detect_wave_size(): + v = _os.environ.get("MORI_WAVE_SIZE") + if v: + return int(v) + try: + from mori.jit.config import detect_gpu_arch + + return 32 if detect_gpu_arch().startswith("gfx12") else 64 + except Exception: + return 64 + + +WAVE = _detect_wave_size() +LANE_MASK = WAVE - 1 +LOG2_WAVE = WAVE.bit_length() - 1 +_BALLOT_INT = T.i64 if WAVE == 64 else T.i32 +_LANE_STRIDE_I32 = WAVE * 4 # one wave of lanes, vec4 (16B) each +_MAIN_STRIDE_I32 = 2 * _LANE_STRIDE_I32 +_BUTTERFLY_OFFSETS = tuple(WAVE >> i for i in range(1, LOG2_WAVE + 1)) + +# NOTE: the cross-device xdb barrier is kept inlined per kernel (dispatch Phase 2, +# combine gather/scatter/fused) rather than a shared helper: FlyDSL only AST-rewrites +# dynamic `if` in the @flyc.kernel body, not in called helpers, so a helper's +# `if :` would raise "cannot evaluate dynamic Boolean" at trace time. + + +# ── dispatch ────────────────────────────────────────────────────────────── + + +def make_dispatch( + *, + rank, + npes, + experts_per_rank, + experts_per_token, + hidden_dim, + hidden_elem_size, + max_tok_per_rank, + max_recv, + block_num, + warp_num_per_block, + off_tok_off, + off_recv_num, + off_tis, + off_out_idx, + off_out_wts, + off_out_tok, + off_out_scales=0, + scale_dim=0, + scale_type_size=0, + enable_signal=True, + replay=False, + fp4=False, +): + # fp4 (e2m1) packs 2 values/byte -> token is hidden_dim/2 bytes; dispatch is a + # pure byte mover (no fp4 decode). + nbytes = hidden_dim // 2 if fp4 else hidden_dim * hidden_elem_size + n_i32 = nbytes // 4 + # sentinel: tok_map dropped-slot marker whose dest_pe (value // max_recv) == npes. + sentinel_val = npes * max_recv + # Optional per-token scales forwarded verbatim to the dest peer's out_scales. + scale_bytes = scale_dim * scale_type_size + scale_num_i32 = (scale_bytes + 3) // 4 + enable_scales = scale_bytes > 0 + + @flyc.kernel(known_block_size=[warp_num_per_block * WAVE, 1, 1]) + def ep_dispatch( + arena: Int64, + addr_inp_tok: Int64, + addr_inp_idx: Int64, + addr_inp_wts: Int64, + addr_tok_map: Int64, + addr_dest_pe_ctr: Int64, + addr_disp_bar: Int64, + addr_total_recv: Int64, + addr_inp_scales: Int64, + my_lsa_rank: Int32, + inp_cur_tok: Int32, + ): + tid = fx.thread_idx.x + bid = fx.block_idx.x + lane = tid & LANE_MASK + warp = tid >> LOG2_WAVE + global_warp_id = bid * warp_num_per_block + warp + global_warp_num = block_num * warp_num_per_block + work_limit = inp_cur_tok * experts_per_token + + window = cco.Window(arena) + rsrc_inp_idx = create_buffer_resource_from_addr(addr_inp_idx) + rsrc_inp_wts = create_buffer_resource_from_addr(addr_inp_wts) + rsrc_tok_map = create_buffer_resource_from_addr(addr_tok_map) + rsrc_dest_ctr = create_buffer_resource_from_addr(addr_dest_pe_ctr) + rsrc_disp_bar = create_buffer_resource_from_addr(addr_disp_bar) + + # ── Phase 1: P2P-scatter each (src_tok, k_slot) to its dest PE ── + for work_idx in range(global_warp_id, work_limit, global_warp_num): + src_tok = work_idx // experts_per_token + k_slot = work_idx % experts_per_token + dest_expert = buffer_load( + rsrc_inp_idx, work_idx, vec_width=1, dtype=T.i32() + ) + # Dedup: a token routed to several experts on the SAME dest PE is sent + # once. If a lower lane (< k_slot) already targets our dest_pe, drop this + # (src_tok, k_slot); safe_lane keeps the probe in-bounds for lanes >= k_slot. + safe_lane = arith.select(lane < k_slot, lane, 0) + lane_expert = buffer_load( + rsrc_inp_idx, + src_tok * experts_per_token + safe_lane, + vec_width=1, + dtype=T.i32(), + ) + dest_pe = dest_expert // experts_per_rank + lane_dest_pe = lane_expert // experts_per_rank + dup_per_lane = arith.select( + lane_dest_pe == dest_pe, arith.select(lane < k_slot, lane, WAVE), WAVE + ) + dup_ballot = ballot(_BALLOT_INT(), dup_per_lane < WAVE) + is_dup = dup_ballot != 0 + + if const_expr(replay): + # decode dest_tok_id from cached tok_map (skip atomic alloc) + cached = buffer_load(rsrc_tok_map, work_idx, vec_width=1, dtype=T.i32()) + is_dup_or_overflow = cached >= sentinel_val + do_publish = cached < sentinel_val + dest_tok_id = cached - dest_pe * max_recv + else: + dest_tok_lane0 = arith.constant(0) + if lane == 0: + if dup_ballot == 0: + peer_tok_off = fx.Int64(window.lsa_ptr(dest_pe, off_tok_off)) + dest_tok_lane0 = P.atomic_add_global(peer_tok_off, fx.Int32(1)) + dest_tok_id = readlane(T.i32(), dest_tok_lane0, 0) + overflow = dest_tok_id >= max_recv + is_dup_or_overflow = arith.select(is_dup, is_dup, overflow) + no_dup = dup_ballot == 0 + in_cap = dest_tok_id < max_recv + do_publish = arith.select(no_dup, in_cap, no_dup) + tok_map_entry = arith.select( + is_dup_or_overflow, sentinel_val, dest_pe * max_recv + dest_tok_id + ) + if lane == 0: + buffer_store(tok_map_entry, rsrc_tok_map, work_idx) + + if lane == 0: + if do_publish: + # publish this recv slot's origin (global source token id) into + # the dest peer's tis, for combine routing. + src_tok_encoded = rank * max_tok_per_rank + src_tok + peer_tis = fx.Int64(window.lsa_ptr(dest_pe, off_tis)) + buffer_store( + src_tok_encoded, + create_buffer_resource_from_addr(peer_tis), + dest_tok_id, + ) + dest_ctr_addr = fx.Int64(addr_dest_pe_ctr) + fx.Int64( + dest_pe + ) * fx.Int64(4) + P.atomic_add_global(dest_ctr_addr, fx.Int32(1)) + + # Per-lane (weight, expert-idx) scatter (lanes < k). + if lane < experts_per_token: + if do_publish: + weight_src_off = src_tok * experts_per_token + lane + weight_val = buffer_load( + rsrc_inp_wts, weight_src_off, vec_width=1, dtype=T.f32() + ) + idx_val = buffer_load( + rsrc_inp_idx, weight_src_off, vec_width=1, dtype=T.i32() + ) + dest_slot = dest_tok_id * experts_per_token + lane + peer_wts = fx.Int64(window.lsa_ptr(dest_pe, off_out_wts)) + buffer_store( + arith.bitcast(T.i32(), weight_val), + create_buffer_resource_from_addr(peer_wts), + dest_slot, + ) + peer_idx = fx.Int64(window.lsa_ptr(dest_pe, off_out_idx)) + buffer_store( + idx_val, create_buffer_resource_from_addr(peer_idx), dest_slot + ) + + # Per-token scales scatter: forward the src token's scale_num_i32 dwords + # (lane-strided) to the dest peer's out_scales[dest_tok_id], verbatim. + if const_expr(enable_scales): + if do_publish: + rsrc_inp_scales = create_buffer_resource_from_addr(addr_inp_scales) + peer_scales = fx.Int64(window.lsa_ptr(dest_pe, off_out_scales)) + rsrc_peer_scales = create_buffer_resource_from_addr(peer_scales) + for k_off in range(lane, scale_num_i32, WAVE): + scale_val = buffer_load( + rsrc_inp_scales, + src_tok * scale_num_i32 + k_off, + vec_width=1, + dtype=T.i32(), + ) + buffer_store( + scale_val, + rsrc_peer_scales, + dest_tok_id * scale_num_i32 + k_off, + ) + + # Token-embedding scatter: each lane owns 4 i32 (16B). Two vec4 streams + # for memory-level parallelism, one-stream tail for the remainder; + # dropped slots set copy_end == lane_i32_off (no-op). + peer_tok_base = fx.Int64(window.lsa_ptr(dest_pe, off_out_tok)) + remote_tok_addr = peer_tok_base + fx.Int64(dest_tok_id) * fx.Int64(nbytes) + local_tok_addr = fx.Int64(addr_inp_tok) + fx.Int64(src_tok) * fx.Int64( + nbytes + ) + rsrc_src = create_buffer_resource_from_addr(local_tok_addr) + rsrc_dst = create_buffer_resource_from_addr(remote_tok_addr) + lane_i32_off = lane * 4 + safe_end_i32 = (n_i32 // _MAIN_STRIDE_I32) * _MAIN_STRIDE_I32 + if const_expr(n_i32 >= _MAIN_STRIDE_I32 and safe_end_i32 > 0): + copy_end_main = arith.select( + is_dup_or_overflow, lane_i32_off, safe_end_i32 + ) + for chunk in range(lane_i32_off, copy_end_main, _MAIN_STRIDE_I32): + vec_a = buffer_load(rsrc_src, chunk, vec_width=4, dtype=T.i32()) + vec_b = buffer_load( + rsrc_src, chunk + _LANE_STRIDE_I32, vec_width=4, dtype=T.i32() + ) + buffer_store(vec_a, rsrc_dst, chunk) + buffer_store(vec_b, rsrc_dst, chunk + _LANE_STRIDE_I32) + if const_expr(safe_end_i32 < n_i32): + copy_end_tail = arith.select(is_dup_or_overflow, lane_i32_off, n_i32) + for chunk in range( + lane_i32_off + safe_end_i32, copy_end_tail, _LANE_STRIDE_I32 + ): + vec_a = buffer_load(rsrc_src, chunk, vec_width=4, dtype=T.i32()) + buffer_store(vec_a, rsrc_dst, chunk) + elif const_expr(n_i32 < _MAIN_STRIDE_I32): + copy_end_small = arith.select(is_dup_or_overflow, lane_i32_off, n_i32) + for chunk in range(lane_i32_off, copy_end_small, _LANE_STRIDE_I32): + vec_a = buffer_load(rsrc_src, chunk, vec_width=4, dtype=T.i32()) + buffer_store(vec_a, rsrc_dst, chunk) + + if const_expr(enable_signal): + # Self-reset total_recv (warp 0 zeros it here + release-fence; only warp 0 + # accumulates into it in Phase 3, so no cross-block race). CUDAGraph-safe. + if global_warp_id == 0: + if lane == 0: + buffer_store( + arith.constant(0), + create_buffer_resource_from_addr(addr_total_recv), + 0, + ) + P.fence_system_release() + + # ── Phase 2: grid barrier + per-peer count signal ── + fx.barrier() + if tid == 0: + P.atomic_add_global(fx.Int64(addr_disp_bar), arith.constant(1)) + + local_recv_num = fx.Int64(window.lsa_ptr(my_lsa_rank, off_recv_num)) + for dest_pe in range(lane, npes, WAVE): + if global_warp_id == 0: + P.spin_until_eq_i32(fx.Int64(addr_disp_bar), block_num) + P.fence_system_acquire() + buffer_store(arith.constant(0), rsrc_disp_bar, 0) + signal_value = ( + buffer_load(rsrc_dest_ctr, dest_pe, vec_width=1, dtype=T.i32()) + + 1 + ) + peer_recv_num = fx.Int64(window.lsa_ptr(dest_pe, off_recv_num)) + recv_num_remote_addr = peer_recv_num + fx.Int64(rank) * fx.Int64(4) + P.spin_until_eq_i32(recv_num_remote_addr, 0) + P.store_i32_system( + recv_num_remote_addr, arith.constant(0), signal_value + ) + + # ── Phase 3: collect per-source counts into total_recv ── + for src_pe in range(lane, npes, WAVE): + if global_warp_id == 0: + recv_num_src_addr = local_recv_num + fx.Int64(src_pe) * fx.Int64(4) + signal_value = P.spin_until_gt_i32(recv_num_src_addr, 0) + peer_recv_count = signal_value - 1 + P.store_i32_system( + recv_num_src_addr, arith.constant(0), arith.constant(0) + ) + P.atomic_add_global(fx.Int64(addr_total_recv), peer_recv_count) + buffer_store(arith.constant(0), rsrc_dest_ctr, src_pe) + + if global_warp_id == 0: + if lane == 0: + local_tok_off = fx.Int64(window.lsa_ptr(my_lsa_rank, off_tok_off)) + P.store_i32_system( + local_tok_off, arith.constant(0), arith.constant(0) + ) + + @flyc.jit + def run( + arena: Int64, + addr_inp_tok: Int64, + addr_inp_idx: Int64, + addr_inp_wts: Int64, + addr_tok_map: Int64, + addr_dest_pe_ctr: Int64, + addr_disp_bar: Int64, + addr_total_recv: Int64, + addr_inp_scales: Int64, + my_lsa_rank: Int32, + inp_cur_tok: Int32, + stream=fx.Stream(None), + ): + ep_dispatch( + arena, + addr_inp_tok, + addr_inp_idx, + addr_inp_wts, + addr_tok_map, + addr_dest_pe_ctr, + addr_disp_bar, + addr_total_recv, + addr_inp_scales, + my_lsa_rank, + inp_cur_tok, + ).launch( + grid=(block_num, 1, 1), + block=[warp_num_per_block * WAVE, 1, 1], + stream=stream, + ) + + return run + + +# ── combine (gather + scatter/quant) ────────────────────────────────────── + + +def _V2BF16(): + return T.VectorType.get([2], T.bf16()) + + +def _V2F32(): + return T.VectorType.get([2], T.f32()) + + +def _V4F32(): + return T.VectorType.get([4], T.f32()) + + +def _V8F32(): + return T.VectorType.get([8], T.f32()) + + +def _V1I32(): + return T.VectorType.get([1], T.i32()) + + +def _accum_funcs(hidden_elem_size, fp8_direct_cast=False, fp4=False): + if fp4: # fp4 e2m1: i32 = 8 packed fp4 -> v8f32 + if WAVE == 32: + # gfx1250: one pack-8 instr converts all 8 fp4/i32 (the gfx950 pk-2 + # variants don't exist here). Plain fp4 => scale 1.0, scale_sel 0. + def to_accum(i32_scalar): + # dequant uses the E8M0 block-scale; 1.0 == 2^(127-127) => 127. + return cvt_scale_pk8_f32_fp4( + res=_V8F32(), + src=i32_scalar, + scale=arith.constant(127, type=T.i32()), + scale_sel=0, + ) + + def from_accum(acc): + return cvt_scalef32_pk8_fp4_f32( + res=T.i32(), src=acc, scale=arith.constant(1.0, type=T.f32()) + ) + + def zero_accum(): + return arith.constant_vector(0.0, _V8F32()) + + return to_accum, from_accum, zero_accum + + # cvt_scalef32_pk_*_fp4 are gfx950-only (fail on gfx942); opt-in (fp4=True). + def to_accum(i32_scalar): + one = arith.constant(1.0, type=T.f32()) + pairs = [ + cvt_scalef32_pk_f32_fp4( + res=_V2F32(), src=i32_scalar, scale=one, src_sel_index=s + ) + for s in range(4) + ] + lo4 = vector.shuffle(pairs[0], pairs[1], [0, 1, 2, 3]) + hi4 = vector.shuffle(pairs[2], pairs[3], [0, 1, 2, 3]) + return vector.shuffle(lo4, hi4, [0, 1, 2, 3, 4, 5, 6, 7]) + + def from_accum(acc): + one = arith.constant(1.0, type=T.f32()) + old = arith.constant(0, type=T.i32()) + for s in range(4): + fa = vector.extract(acc, static_position=[s * 2]) + fb = vector.extract(acc, static_position=[s * 2 + 1]) + old = cvt_scalef32_pk_fp4_f32( + res=T.i32(), + old_vdst=old, + src0=fa, + src1=fb, + scale=one, + dst_sel_index=s, + ) + return old + + def zero_accum(): + return arith.constant_vector(0.0, _V8F32()) + + return to_accum, from_accum, zero_accum + return _accum_funcs_int(hidden_elem_size, fp8_direct_cast) + + +def _accum_funcs_int(hidden_elem_size, fp8_direct_cast=False): + """Return (to_accum, from_accum, zero_accum) for one i32 'unit' of the + transport dtype, mirroring the FlyDSL reference's per-dtype branches. + + Each i32 packs: 2 bf16 (v2f32), 1 f32 (scalar), or 4 fp8 (v4f32). + """ + if hidden_elem_size == 2: # bf16: i32 = 2 bf16 + + def to_accum(i32_scalar): + return vector.bitcast( + _V2BF16(), vector.from_elements(_V1I32(), [i32_scalar]) + ).extf(_V2F32()) + + def from_accum(acc): + return vector.extract( + vector.bitcast(_V1I32(), acc.truncf(_V2BF16())), static_position=[0] + ) + + def zero_accum(): + return to_accum(arith.constant(0)) + + elif hidden_elem_size == 4: # f32: i32 = 1 f32 + + def to_accum(i32_scalar): + return fx.Float32(arith.bitcast(T.f32(), arith.unwrap(i32_scalar))) + + def from_accum(acc): + return fx.Int32(arith.bitcast(T.i32(), arith.unwrap(acc))) + + def zero_accum(): + return fx.Float32(arith.constant(0.0, type=T.f32())) + + elif hidden_elem_size == 1: # fp8 (OCP e4m3): i32 = 4 fp8 + + def to_accum(i32_scalar): + lo = cvt_pk_f32_fp8(res=_V2F32(), src=i32_scalar, word_sel=False) + hi = cvt_pk_f32_fp8(res=_V2F32(), src=i32_scalar, word_sel=True) + return vector.shuffle(lo, hi, [0, 1, 2, 3]) + + def from_accum(acc): + f0 = vector.extract(acc, static_position=[0]) + f1 = vector.extract(acc, static_position=[1]) + f2 = vector.extract(acc, static_position=[2]) + f3 = vector.extract(acc, static_position=[3]) + if fp8_direct_cast: # wire fp8 -> external bf16: v4f32 -> v4bf16 -> 2 i32 + v4bf16 = acc.truncf(T.VectorType.get([4], T.bf16())) + return vector.bitcast(T.VectorType.get([2], T.i32()), v4bf16) + zero = arith.constant(0, type=T.i32()) + lo = cvt_pk_fp8_f32( + res=T.i32(), src_a=f0, src_b=f1, old=zero, word_sel=False + ) + return cvt_pk_fp8_f32( + res=T.i32(), src_a=f2, src_b=f3, old=lo, word_sel=True + ) + + def zero_accum(): + return arith.constant_vector(0.0, _V4F32()) + + else: + raise ValueError(f"unsupported hidden_elem_size {hidden_elem_size}") + return to_accum, from_accum, zero_accum + + +def make_combine( + *, + rank, + npes, + experts_per_token, + hidden_dim, + hidden_elem_size, + max_tok_per_rank, + max_recv, + block_num, + warp_num_per_block, + off_out_tok, + off_xdb_mem, + off_out_wts=0, + enable_weights=True, + fp8_direct_cast=False, + fp4=False, + reset_total_recv=True, + _unroll=4, +): + # Transport dtype = external dtype; fp8_direct_cast wires fp8 while the output + # stays bf16 (2 i32 per fp8 unit). fp4: i32 = 8 fp4. + _to_accum2, _from_accum2, _zero_accum = _accum_funcs( + hidden_elem_size, fp8_direct_cast, fp4 + ) + nbytes = hidden_dim // 2 if fp4 else hidden_dim * hidden_elem_size + n_i32 = hidden_dim // 8 if fp4 else nbytes // 4 + out_n_i32 = (hidden_dim * 2) // 4 if fp8_direct_cast else n_i32 + # vec4 gather: 4 i32 per load (global_load_dwordx4) when output is 1:1. + _use_vec4 = (n_i32 % 4 == 0) and not fp8_direct_cast + out_step_mult = 2 if fp8_direct_cast else 1 + + @flyc.kernel(known_block_size=[warp_num_per_block * WAVE, 1, 1]) + def ep_combine( + arena: Int64, + addr_tok_map: Int64, + addr_comb_bar: Int64, + addr_xdb_flag: Int64, + addr_total_recv: Int64, + addr_out: Int64, + addr_out_wts: Int64, + my_lsa_rank: Int32, + cur_rank_num_token: Int32, + ): + tid = fx.thread_idx.x + bid = fx.block_idx.x + lane = tid & LANE_MASK + warp = tid >> LOG2_WAVE + global_warp_id = bid * warp_num_per_block + warp + global_warp_num = block_num * warp_num_per_block + grid_thread_id = bid * (warp_num_per_block * WAVE) + tid + + window = cco.Window(arena) + rsrc_tok_map = create_buffer_resource_from_addr(addr_tok_map) + rsrc_comb_bar = create_buffer_resource_from_addr(addr_comb_bar) + rsrc_total_recv = create_buffer_resource_from_addr(addr_total_recv) + rsrc_out = create_buffer_resource_from_addr(addr_out) + xdb_cur_flag = P.load_i64_acquire(fx.Int64(addr_xdb_flag)) + + # ── Stage 1: cross-device entry barrier (block 0 only) ── + # Block 0 handles the full xdb handshake, then signals other blocks + # via comb_bar. Eliminates the grid barrier and per-block xdb spin. + if grid_thread_id < npes: + xdb_remote = fx.Int64( + window.lsa_ptr(grid_thread_id, off_xdb_mem) + ) + fx.Int64(rank) * fx.Int64(8) + P.store_i64_system(xdb_remote, arith.constant(0), xdb_cur_flag) + if grid_thread_id == 0: + P.atomic_add_global( + fx.Int64(addr_xdb_flag), arith.constant(1, type=T.i64()) + ) + if grid_thread_id < npes: + xdb_peer_slot = fx.Int64( + window.lsa_ptr(my_lsa_rank, off_xdb_mem) + ) + fx.Int64(grid_thread_id) * fx.Int64(8) + P.spin_until_eq_i64(xdb_peer_slot, xdb_cur_flag) + if bid == 0: + fx.barrier() + if tid == 0: + P.store_i32_system( + fx.Int64(addr_comb_bar), arith.constant(0), arith.constant(1) + ) + if bid != 0: + if tid == 0: + P.spin_until_gt_i32(fx.Int64(addr_comb_bar), 0) + fx.barrier() + if tid == 0: + P.atomic_add_global(fx.Int64(addr_comb_bar), arith.constant(1)) + if const_expr(reset_total_recv): + if tid == 0: + buffer_store(arith.constant(0), rsrc_total_recv, 0) + + rsrc_out_wts = create_buffer_resource_from_addr(addr_out_wts) + + # ── Stage 2: warp-partitioned remote gather + f32 accumulate ── + # Register-light i32 reads + `_unroll`-way unroll (each lane keeps _unroll + # loads/accumulators in flight per k) hide xGMI latency with fewer warps; + # each token's hidden is split across warps_per_tok warps to fill the grid. + STEP = _unroll * WAVE + safe_tok = arith.select( + cur_rank_num_token == arith.constant(0), + arith.constant(1), + cur_rank_num_token, + ) + warps_per_tok = ( + arith.constant(global_warp_num) + safe_tok - arith.constant(1) + ) // safe_tok + units_per_warp = ( + arith.constant(n_i32) + warps_per_tok - arith.constant(1) + ) // warps_per_tok + stage3_total = cur_rank_num_token * warps_per_tok + for stage3_idx in range(global_warp_id, stage3_total, global_warp_num): + tok_id = stage3_idx // warps_per_tok + part_id = stage3_idx % warps_per_tok + unit_base = part_id * units_per_warp + tok_map_base = tok_id * experts_per_token + expert_bases = [] + expert_valids = [] + expert_pes = [] + expert_toks = [] + # One coalesced tok_map load per lane, then broadcast each k_slot + # via readlane (replaces K per-k_slot buffer_loads). + encoded_my = buffer_load( + rsrc_tok_map, tok_map_base + lane, vec_width=1, dtype=T.i32() + ) + for k_slot in range_constexpr(experts_per_token): + encoded_k = readlane(T.i32(), encoded_my, k_slot) + dest_pe_k = encoded_k // max_recv # sentinel: dest_pe == npes + dest_tok_k = encoded_k % max_recv # peer-local recv slot + valid_k = dest_pe_k < npes + safe_pe = arith.select(valid_k, dest_pe_k, arith.constant(rank)) + safe_tok_k = arith.select(valid_k, dest_tok_k, arith.constant(0)) + # Remote base peer[dest_pe].out_tok[dest_tok]; global-load gather + # (no per-expert descriptor) keeps all K loads in flight. + slot_addr = fx.Int64(window.lsa_ptr(safe_pe, off_out_tok)) + fx.Int64( + safe_tok_k + ) * fx.Int64(nbytes) + expert_bases.append(slot_addr) + expert_valids.append(valid_k) + expert_pes.append(safe_pe) + expert_toks.append(safe_tok_k) + + # Weights: once per token (part 0), reduce the K forwarded weights -> + # out_weights[tok][e]; overlaps with this warp's hidden gather. + if const_expr(enable_weights): + if part_id == arith.constant(0): + if lane < experts_per_token: + weight_acc = arith.constant(0.0, type=T.f32()) + for k_slot in range_constexpr(experts_per_token): + weight_addr = fx.Int64( + window.lsa_ptr(expert_pes[k_slot], off_out_wts) + ) + ( + fx.Int64(expert_toks[k_slot]) + * fx.Int64(experts_per_token) + + fx.Int64(lane) + ) * fx.Int64( + 4 + ) + weight_val = buffer_load( + create_buffer_resource_from_addr(weight_addr), + 0, + vec_width=1, + dtype=T.f32(), + ) + weight_acc = weight_acc + arith.select( + expert_valids[k_slot], + weight_val, + arith.constant(0.0, type=T.f32()), + ) + buffer_store(weight_acc, rsrc_out_wts, tok_map_base + lane) + rem = arith.constant(n_i32) - unit_base + eff = arith.select( + rem < units_per_warp, rem, units_per_warp + ) # i32 units this warp + out_base = tok_id * out_n_i32 + + # Nested fn: closure over expert_bases/valids (lists can't be loop-carried). + def _one(off): # reduce k contributions for one i32 unit + vals = [] + for k_slot in range_constexpr(experts_per_token): + v = P.load_i32_nt(expert_bases[k_slot], off) + vals.append( + arith.select(expert_valids[k_slot], v, arith.constant(0)) + ) + acc = _zero_accum() + for k_slot in range_constexpr(experts_per_token): + acc = acc + _to_accum2(vals[k_slot]) + buffer_store( + _from_accum2(acc), rsrc_out, out_base + off * out_step_mult + ) + + if const_expr(_use_vec4): + # Vec4 path: load 4 i32 (16B, global_load_dwordx4) per expert + # per iteration, reducing xGMI load count 4×. + VEC = 4 + STEP_CHUNK = WAVE * VEC + STEP_V4 = _unroll * STEP_CHUNK + + def _load_expert_vecs(off): + vecs = [] + valids = [] + for k_slot in range_constexpr(experts_per_token): + vecs.append(P.load_v4i32_nt(expert_bases[k_slot], off)) + valids.append(expert_valids[k_slot]) + return vecs, valids + + def _accum_loop(): + main_end = (eff // STEP_V4) * STEP_V4 + for u in range(lane * VEC, main_end, STEP_V4): + base = unit_base + u + # Issue all remote loads first (both unroll rounds) so stores + # in the j/r nest cannot block the next load batch. + pre_vecs = [] + pre_valids = [] + for r in range_constexpr(_unroll): + off_r = base + r * STEP_CHUNK + vecs_r, valids_r = _load_expert_vecs(off_r) + pre_vecs.append(vecs_r) + pre_valids.append(valids_r) + # Reduce+store: j outer, r inner (unroll innermost). + for j in range_constexpr(VEC): + for r in range_constexpr(_unroll): + off = base + r * STEP_CHUNK + acc = _zero_accum() + for k_slot in range_constexpr(experts_per_token): + elem = vector.extract( + pre_vecs[r][k_slot], static_position=[j] + ) + v = arith.select( + pre_valids[r][k_slot], elem, arith.constant(0) + ) + acc = acc + _to_accum2(v) + buffer_store( + _from_accum2(acc), rsrc_out, out_base + off + j + ) + for u in range(main_end + lane, eff, WAVE): + _one(unit_base + u) + + else: + + def _accum_loop(): + main_end = (eff // STEP) * STEP + for u in range(lane, main_end, STEP): + base = unit_base + u + pre_vals = [] + for r in range_constexpr(_unroll): + off_r = base + r * WAVE + vals_r = [] + for k_slot in range_constexpr(experts_per_token): + v = P.load_i32_nt(expert_bases[k_slot], off_r) + vals_r.append( + arith.select( + expert_valids[k_slot], v, arith.constant(0) + ) + ) + pre_vals.append(vals_r) + for r in range_constexpr(_unroll): + off = base + r * WAVE + acc = _zero_accum() + for k_slot in range_constexpr(experts_per_token): + acc = acc + _to_accum2(pre_vals[r][k_slot]) + buffer_store( + _from_accum2(acc), + rsrc_out, + out_base + off * out_step_mult, + ) + for u in range(main_end + lane, eff, WAVE): + _one(unit_base + u) + + _accum_loop() + + # Epilogue: block 0 waits for every other block to finish (comb_bar + # increments from Stage-1), then resets comb_bar for the next call. + if global_warp_id == 0: + if lane == 0: + P.spin_until_gt_i32(fx.Int64(addr_comb_bar), block_num - 1) + buffer_store(arith.constant(0), rsrc_comb_bar, 0) + + @flyc.jit + def run( + arena: Int64, + addr_tok_map: Int64, + addr_comb_bar: Int64, + addr_xdb_flag: Int64, + addr_total_recv: Int64, + addr_out: Int64, + addr_out_wts: Int64, + my_lsa_rank: Int32, + cur_rank_num_token: Int32, + stream=fx.Stream(None), + ): + ep_combine( + arena, + addr_tok_map, + addr_comb_bar, + addr_xdb_flag, + addr_total_recv, + addr_out, + addr_out_wts, + my_lsa_rank, + cur_rank_num_token, + ).launch( + grid=(block_num, 1, 1), + block=[warp_num_per_block * WAVE, 1, 1], + stream=stream, + ) + + return run + + +_FP8_MAX = 240.0 # gfx942 native fp8 is e4m3fnuz: max finite 240 (NOT OCP 448); +# clamping above 240 yields NaN from cvt_pk_fp8_f32 on this arch + + +def _fabs(f): + return arith.maximumf(f, arith.negf(f)) + + +def _bf16x2(i32_scalar): # i32 (2 bf16) -> v2f32 + return vector.bitcast(_V2BF16(), vector.from_elements(_V1I32(), [i32_scalar])).extf( + _V2F32() + ) + + +def _warp_amax(lane, v): + """Max-reduce an f32 across the wavefront (butterfly via ds_bpermute, per-lane + gather index). Every lane returns the wavefront max; raw values in/out.""" + for off in _BUTTERFLY_OFFSETS: + idx = arith.unwrap((lane ^ off) * 4) # byte addr = lane*4 + o = ds_bpermute(T.i32(), idx, arith.bitcast(T.i32(), arith.unwrap(v))) + v = arith.maximumf(v, arith.bitcast(T.f32(), o)) + return v + + +def make_combine_scatter( + *, + rank, + npes, + experts_per_token, + hidden_dim, + hidden_elem_size, + max_tok_per_rank, + max_recv, + block_num, + warp_num_per_block, + off_out_tok, + off_comb_inp, + off_tis, + off_xdb_mem, + off_out_wts=0, + off_comb_wts=0, + off_comb_scales=0, + enable_weights=True, + fp8_direct_cast=False, + fp8_blockwise=False, + scale_dim=0, + reset_total_recv=True, + _s3_cache=2, + _unroll=2, +): + """Scatter combine (mori _nop2p path): 2 passes vs gather, but compresses the + transport to fp8 (the natural home for fp8_direct_cast). + + Stage 1 each computing rank P2P-WRITES post-expert tokens to the ORIGIN's + comb_inp[computing_rank*M + origin_lid] (bf16->fp8 if cast). + Stage 2 cross-device barrier. + Stage 3 origin LOCAL-reads comb_inp for its token's k expert PEs and reduces.""" + if fp8_blockwise and WAVE != 64: + raise NotImplementedError( + "fp8_blockwise combine's coalesced path assumes wave64 (one wave == one " + "128-elem block); wave32 (gfx12) port is TODO" + ) + _fp8_out = fp8_direct_cast or fp8_blockwise + wire_elem_size = 1 if _fp8_out else hidden_elem_size + to_acc, from_acc, zero_acc = _accum_funcs(wire_elem_size, _fp8_out) + inp_nbytes = hidden_dim * hidden_elem_size # source out_tok (bf16/f32) + wire_nbytes = hidden_dim * wire_elem_size # comb_inp transport + wire_n_i32 = wire_nbytes // 4 + out_n_i32 = (hidden_dim * 2) // 4 if _fp8_out else wire_n_i32 + out_step_mult = 2 if _fp8_out else 1 + if fp8_blockwise: + block_elems = hidden_dim // scale_dim + assert ( + hidden_dim % scale_dim == 0 and block_elems == 128 + ), "blockwise (coalesced path): block_elems must be 128 (scale_dim = hidden/128)" + block_i32_fp8 = block_elems // 4 # fp8 i32 units per block (=32) + + @flyc.kernel(known_block_size=[warp_num_per_block * WAVE, 1, 1]) + def ep_combine_s( + arena: Int64, + addr_tok_map: Int64, + addr_comb_bar: Int64, + addr_xdb_flag: Int64, + addr_total_recv: Int64, + addr_out: Int64, + addr_out_wts: Int64, + addr_src_tok: Int64, + my_lsa_rank: Int32, + cur_rank_num_token: Int32, + ): + tid = fx.thread_idx.x + bid = fx.block_idx.x + lane = tid & LANE_MASK + warp = tid >> LOG2_WAVE + global_warp_id = bid * warp_num_per_block + warp + global_warp_num = block_num * warp_num_per_block + grid_thread_id = bid * (warp_num_per_block * WAVE) + tid + + window = cco.Window(arena) + rsrc_tok_map = create_buffer_resource_from_addr(addr_tok_map) + rsrc_comb_bar = create_buffer_resource_from_addr(addr_comb_bar) + rsrc_total_recv = create_buffer_resource_from_addr(addr_total_recv) + rsrc_tis = create_buffer_resource_from_addr( + fx.Int64(window.lsa_ptr(my_lsa_rank, off_tis)) + ) + rsrc_out = create_buffer_resource_from_addr(addr_out) + rsrc_out_wts = create_buffer_resource_from_addr(addr_out_wts) + total_recv = buffer_load(rsrc_total_recv, 0, vec_width=1, dtype=T.i32()) + + # ── Stage 1: scatter post-expert tokens back to origin's comb_inp ── + # addr_src_tok is a raw LOCAL pointer to the post-expert tokens (the GEMM + # output). Scatter only ever LOCAL-reads its source, so the caller can pass + # the GEMM output directly and skip the copy into symmetric comb_stg that + # gather requires (gather peers remote-read it, so it must be symmetric). + src_tok_base = fx.Int64(addr_src_tok) + for recv_slot in range(global_warp_id, total_recv, global_warp_num): + # tis encodes origin = src_pe*max_tok_per_rank + local_id + encoded_origin = buffer_load( + rsrc_tis, recv_slot, vec_width=1, dtype=T.i32() + ) + origin_pe = encoded_origin // max_tok_per_rank + origin_lid = encoded_origin % max_tok_per_rank + dst = fx.Int64(window.lsa_ptr(origin_pe, off_comb_inp)) + ( + fx.Int64(rank * max_tok_per_rank + origin_lid) + ) * fx.Int64(wire_nbytes) + src = src_tok_base + fx.Int64(recv_slot) * fx.Int64(inp_nbytes) + rsrc_src = create_buffer_resource_from_addr(src) + rsrc_dst = create_buffer_resource_from_addr(dst) + if const_expr(fp8_blockwise): + # Blockwise fp8 quant, coalesced + warp-reduce (block_elems==128, one + # i32/lane spans a block): ds_bpermute butterfly gives every lane the + # block amax; scale = (amax>MAX)? amax/MAX : 1, quant = clamp(v*MAX/amax); + # sign sentinel: if ANY block scaled, negate block-0 scale. + scale_dst = fx.Int64( + window.lsa_ptr(origin_pe, off_comb_scales) + ) + fx.Int64(rank * max_tok_per_rank + origin_lid) * fx.Int64( + scale_dim + ) * fx.Int64( + 4 + ) + rsrc_scales = create_buffer_resource_from_addr(scale_dst) + fp8max = arith.constant(_FP8_MAX, type=T.f32()) + nlim = arith.constant(-_FP8_MAX, type=T.f32()) + any_scaled = arith.constant(0) != arith.constant(0) # False (uniform) + for scale_block in range_constexpr(scale_dim): + v2 = _bf16x2( + buffer_load( + rsrc_src, + scale_block * 64 + lane, + vec_width=1, + dtype=T.i32(), + ) + ) + e0 = vector.extract(v2, static_position=[0]) + e1 = vector.extract(v2, static_position=[1]) + amax = _warp_amax(lane, arith.maximumf(_fabs(e0), _fabs(e1))) + scaled = amax > fp8max + any_scaled = arith.select(scaled, scaled, any_scaled) + scale = arith.select( + scaled, + arith.divf(amax, fp8max), + arith.constant(1.0, type=T.f32()), + ) + inv = arith.select( + scaled, + arith.divf(fp8max, amax), + arith.constant(1.0, type=T.f32()), + ) + if lane == 0: + buffer_store(scale, rsrc_scales, scale_block) + f0 = fmed3(T.f32(), arith.mulf(e0, inv), fp8max, nlim) + f1 = fmed3(T.f32(), arith.mulf(e1, inv), fp8max, nlim) + my_packed = cvt_pk_fp8_f32( + res=T.i32(), + src_a=f0, + src_b=f1, + old=arith.constant(0, type=T.i32()), + word_sel=False, + ) + # neighbour (lane^1)'s 2 fp8 (its low 16) via ds_bpermute. + nbr_packed = ds_bpermute( + T.i32(), + arith.unwrap((lane ^ arith.constant(1)) * arith.constant(4)), + arith.unwrap(my_packed), + ) + my_lo16 = my_packed & arith.constant(0xFFFF) + packed_pair = my_lo16 | ( + (nbr_packed & arith.constant(0xFFFF)) << arith.constant(16) + ) + if (lane & arith.constant(1)) == arith.constant(0): + buffer_store( + packed_pair, + rsrc_dst, + scale_block * block_i32_fp8 + (lane >> arith.constant(1)), + ) + if any_scaled: + if lane == 0: + s0 = buffer_load(rsrc_scales, 0, vec_width=1, dtype=T.f32()) + buffer_store(arith.negf(s0), rsrc_scales, 0) + elif const_expr(fp8_direct_cast): + # 2 bf16 i32 -> v4f32 -> cvt_pk_fp8 x2 -> 1 fp8 i32 + for elem in range(lane, wire_n_i32, WAVE): + bf = buffer_load(rsrc_src, elem * 2, vec_width=2, dtype=T.i32()) + v4 = vector.bitcast(T.VectorType.get([4], T.bf16()), bf).extf( + _V4F32() + ) + f0 = vector.extract(v4, static_position=[0]) + f1 = vector.extract(v4, static_position=[1]) + f2 = vector.extract(v4, static_position=[2]) + f3 = vector.extract(v4, static_position=[3]) + z = arith.constant(0, type=T.i32()) + lo = cvt_pk_fp8_f32( + res=T.i32(), src_a=f0, src_b=f1, old=z, word_sel=False + ) + fp8 = cvt_pk_fp8_f32( + res=T.i32(), src_a=f2, src_b=f3, old=lo, word_sel=True + ) + buffer_store(fp8, rsrc_dst, elem) + else: + # bf16/f32 (no compression): plain local-read -> remote-write copy. + # vec4 (global_load/store_dwordx4, 16B) + _unroll-way prefetch to keep + # multiple remote stores in flight and hide xGMI write latency. + _v4n = wire_n_i32 // 4 + _step = _unroll * WAVE + _main_end = (_v4n // _step) * _step + for _b in range(lane, _main_end, _step): + _pre = [] + for _r in range_constexpr(_unroll): + _pre.append(P.load_v4i32_nt(src, (_b + _r * WAVE) * 4)) + for _r in range_constexpr(_unroll): + buffer_store(_pre[_r], rsrc_dst, (_b + _r * WAVE) * 4) + for _u in range(_main_end + lane, _v4n, WAVE): + buffer_store(P.load_v4i32_nt(src, _u * 4), rsrc_dst, _u * 4) + for elem in range(_v4n * 4 + lane, wire_n_i32, WAVE): + v = buffer_load(rsrc_src, elem, vec_width=1, dtype=T.i32()) + buffer_store(v, rsrc_dst, elem) + if const_expr(enable_weights): + # forward this recv slot's weights to the ORIGIN's + # comb_wts[computing_rank*M + lid] (dedicated region). + weight_src = fx.Int64( + window.lsa_ptr(my_lsa_rank, off_out_wts) + ) + fx.Int64(recv_slot) * fx.Int64(experts_per_token) * fx.Int64(4) + weight_dst = fx.Int64( + window.lsa_ptr(origin_pe, off_comb_wts) + ) + fx.Int64(rank * max_tok_per_rank + origin_lid) * fx.Int64( + experts_per_token + ) * fx.Int64( + 4 + ) + if lane < experts_per_token: + weight_val = buffer_load( + create_buffer_resource_from_addr(weight_src), + lane, + vec_width=1, + dtype=T.i32(), + ) + buffer_store( + weight_val, create_buffer_resource_from_addr(weight_dst), lane + ) + + xdb_cur_flag = P.load_i64_acquire(fx.Int64(addr_xdb_flag)) + + # ── Stage 2: cross-device barrier ── + # Grid sync ensures all blocks finished Stage-1 scatter writes. + # Block 0 then handles xdb and signals completion via comb_bar. + P.fence_system_release() + fx.barrier() + if tid == 0: + P.atomic_add_global(fx.Int64(addr_comb_bar), arith.constant(1)) + if grid_thread_id < npes: + P.spin_until_eq_i32(fx.Int64(addr_comb_bar), block_num) + P.fence_system_acquire() + xdb_remote = fx.Int64( + window.lsa_ptr(grid_thread_id, off_xdb_mem) + ) + fx.Int64(rank) * fx.Int64(8) + P.store_i64_system(xdb_remote, arith.constant(0), xdb_cur_flag) + if grid_thread_id == 0: + P.atomic_add_global( + fx.Int64(addr_xdb_flag), arith.constant(1, type=T.i64()) + ) + if grid_thread_id < npes: + xdb_slot = fx.Int64(window.lsa_ptr(my_lsa_rank, off_xdb_mem)) + fx.Int64( + grid_thread_id + ) * fx.Int64(8) + P.spin_until_eq_i64(xdb_slot, xdb_cur_flag) + if bid == 0: + fx.barrier() + if tid == 0: + P.store_i32_system( + fx.Int64(addr_comb_bar), + arith.constant(0), + arith.constant(block_num + 1), + ) + if bid != 0: + if tid == 0: + P.spin_until_gt_i32(fx.Int64(addr_comb_bar), block_num) + fx.barrier() + if tid == 0: + P.atomic_add_global(fx.Int64(addr_comb_bar), arith.constant(1)) + if const_expr(reset_total_recv): + if tid == 0: + buffer_store(arith.constant(0), rsrc_total_recv, 0) + + # ── Stage 3: local read of comb_inp + reduce ── + comb_inp_base = fx.Int64(window.lsa_ptr(my_lsa_rank, off_comb_inp)) + safe_tok = arith.select( + cur_rank_num_token == arith.constant(0), + arith.constant(1), + cur_rank_num_token, + ) + warps_per_tok = ( + arith.constant(global_warp_num) + safe_tok - arith.constant(1) + ) // safe_tok + units_per_warp = ( + arith.constant(wire_n_i32) + warps_per_tok - arith.constant(1) + ) // warps_per_tok + stage3_total = cur_rank_num_token * warps_per_tok + for stage3_idx in range(global_warp_id, stage3_total, global_warp_num): + tok_id = stage3_idx // warps_per_tok + part_id = stage3_idx % warps_per_tok + unit_base = part_id * units_per_warp + tok_map_base = tok_id * experts_per_token + expert_rsrcs = [] + expert_addrs = [] + expert_valids = [] + expert_pes = [] + expert_scales = [] + for k_slot in range_constexpr(experts_per_token): + encoded_k = buffer_load( + rsrc_tok_map, tok_map_base + k_slot, vec_width=1, dtype=T.i32() + ) + dest_pe = encoded_k // max_recv + valid = dest_pe < npes + safe_pe = arith.select(valid, dest_pe, arith.constant(rank)) + # LOCAL comb_inp[computing_pe*M + tok_id] + src_addr = comb_inp_base + ( + fx.Int64(safe_pe) * fx.Int64(max_tok_per_rank) + fx.Int64(tok_id) + ) * fx.Int64(wire_nbytes) + expert_rsrcs.append(create_buffer_resource_from_addr(src_addr)) + expert_addrs.append(src_addr) + expert_valids.append(valid) + expert_pes.append(safe_pe) + if const_expr(fp8_blockwise): + scale_addr = fx.Int64( + window.lsa_ptr(my_lsa_rank, off_comb_scales) + ) + ( + fx.Int64(safe_pe) * fx.Int64(max_tok_per_rank) + + fx.Int64(tok_id) + ) * fx.Int64( + scale_dim + ) * fx.Int64( + 4 + ) + expert_scales.append(create_buffer_resource_from_addr(scale_addr)) + if const_expr(enable_weights): + if part_id == arith.constant(0): + if lane < experts_per_token: + weight_acc = arith.constant(0.0, type=T.f32()) + for k_slot in range_constexpr(experts_per_token): + # LOCAL comb_wts[computing_pe*M + tok_id] (scattered in S1) + weight_addr = ( + fx.Int64(window.lsa_ptr(my_lsa_rank, off_comb_wts)) + + ( + fx.Int64(expert_pes[k_slot]) + * fx.Int64(max_tok_per_rank) + + fx.Int64(tok_id) + ) + * fx.Int64(experts_per_token) + * fx.Int64(4) + + fx.Int64(lane) * fx.Int64(4) + ) + weight_val = buffer_load( + create_buffer_resource_from_addr(weight_addr), + 0, + vec_width=1, + dtype=T.f32(), + ) + weight_acc = weight_acc + arith.select( + expert_valids[k_slot], + weight_val, + arith.constant(0.0, type=T.f32()), + ) + buffer_store(weight_acc, rsrc_out_wts, tok_map_base + lane) + rem = arith.constant(wire_n_i32) - unit_base + eff = arith.select(rem < units_per_warp, rem, units_per_warp) + out_base = tok_id * out_n_i32 + + def _one(off): + acc = zero_acc() + # blockwise: 4 fp8 per i32 unit are 4 consecutive elements in the + # same block -> one scale per unit. scale_block = (off*4)//block_elems. + if const_expr(fp8_blockwise): + scale_block = (off * arith.constant(4)) // arith.constant( + block_elems + ) + is_b0 = scale_block == arith.constant(0) + for k_slot in range_constexpr(experts_per_token): + v = buffer_load( + expert_rsrcs[k_slot], + off, + vec_width=1, + dtype=T.i32(), + cache_modifier=_s3_cache, + ) + v = arith.select(expert_valids[k_slot], v, arith.constant(0)) + if const_expr(fp8_blockwise): + scale = buffer_load( + expert_scales[k_slot], + scale_block, + vec_width=1, + dtype=T.f32(), + ) + scale = arith.select( + is_b0, _fabs(scale), scale + ) # undo block-0 sign sentinel + # invalid expert -> v already 0; force finite scale so a + # stale/garbage comb_scales slot can't make 0*NaN = NaN. + scale = arith.select( + expert_valids[k_slot], + scale, + arith.constant(1.0, type=T.f32()), + ) + acc = acc + to_acc(v) * scale + else: + acc = acc + to_acc(v) + buffer_store(from_acc(acc), rsrc_out, out_base + off * out_step_mult) + + def _loop(): + if const_expr(not _fp8_out): + # Non-quant (bf16/f32) fast path: vec4 local reads (4 i32/lane) + # per expert + per-i32 reduce + vec4 store, WAVE-strided so the + # wave stays coalesced. fp8 paths keep the scalar _one loop. + _v4eff = eff // arith.constant(4) + for _u in range(lane, _v4eff, WAVE): + _off = unit_base + _u * arith.constant(4) + _accs = [zero_acc(), zero_acc(), zero_acc(), zero_acc()] + for k_slot in range_constexpr(experts_per_token): + _vv = P.load_v4i32_nt(expert_addrs[k_slot], _off) + for _j in range_constexpr(4): + _e = vector.extract(_vv, static_position=[_j]) + _e = arith.select( + expert_valids[k_slot], _e, arith.constant(0) + ) + _accs[_j] = _accs[_j] + to_acc(_e) + _res = vector.from_elements( + T.VectorType.get([4], T.i32()), + [ + from_acc(_accs[0]), + from_acc(_accs[1]), + from_acc(_accs[2]), + from_acc(_accs[3]), + ], + ) + buffer_store(_res, rsrc_out, out_base + _off) + for _u in range(_v4eff * arith.constant(4) + lane, eff, WAVE): + _one(unit_base + _u) + else: + for u in range(lane, eff, WAVE): + _one(unit_base + u) + + _loop() + + # Epilogue: block 0 waits for every other block's Stage-3 completion + # (comb_bar reached 2*block_num), then resets comb_bar for the next call. + if global_warp_id == 0: + if lane == 0: + P.spin_until_gt_i32(fx.Int64(addr_comb_bar), 2 * block_num - 1) + buffer_store(arith.constant(0), rsrc_comb_bar, 0) + + @flyc.jit + def run( + arena: Int64, + addr_tok_map: Int64, + addr_comb_bar: Int64, + addr_xdb_flag: Int64, + addr_total_recv: Int64, + addr_out: Int64, + addr_out_wts: Int64, + addr_src_tok: Int64, + my_lsa_rank: Int32, + cur_rank_num_token: Int32, + stream=fx.Stream(None), + ): + ep_combine_s( + arena, + addr_tok_map, + addr_comb_bar, + addr_xdb_flag, + addr_total_recv, + addr_out, + addr_out_wts, + addr_src_tok, + my_lsa_rank, + cur_rank_num_token, + ).launch( + grid=(block_num, 1, 1), + block=[warp_num_per_block * WAVE, 1, 1], + stream=stream, + ) + + return run + + +def make_combine_fused_reduce( + *, + rank, + npes, + experts_per_token, + hidden_dim, + hidden_elem_size, + max_tok_per_rank, + block_num, + warp_num_per_block, + off_comb_inp, + off_xdb_mem, + slot_stride_nbytes=None, + _s3_cache=2, +): + """Combine for the gemm2-fused scatter path (combine_mode='scatter_fused'). + + gemm2 has already P2P-written each token's WEIGHTED per-expert result into this + rank's comb_inp[origin_lid*topk + k] (one contiguous topk-block per token). So + combine only: + Stage A cross-device barrier -- wait until every peer's gemm2 P2P writes are + globally visible before reading comb_inp. + Stage B each origin token: out[t] = sum_{k