diff --git a/CLAUDE.md b/CLAUDE.md index 94b451817ef..e5a686414b0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -160,6 +160,10 @@ def test_hopper_attention(): ## Benchmarking +For SM90 NVFP4 fused/split and MXFP4 fused MegaMoE, use +[`benchmarks/bench_megamoe_fp4_sm90.py`](benchmarks/bench_megamoe_fp4_sm90.py) for timing +and [`tests/moe/test_megamoe_fp4_sm90.py`](tests/moe/test_megamoe_fp4_sm90.py) for correctness checks. + FlashInfer provides a unified benchmarking framework in `benchmarks/flashinfer_benchmark.py`. **Key features:** diff --git a/benchmarks/bench_megamoe_fp4_sm90.py b/benchmarks/bench_megamoe_fp4_sm90.py new file mode 100644 index 00000000000..b6913c56d4d --- /dev/null +++ b/benchmarks/bench_megamoe_fp4_sm90.py @@ -0,0 +1,368 @@ +"""Four/eight-rank SM90 MegaMoE CUDA-graph timing with optional source baselines. + +The reference runner exports _module(format, family, plan, buffer) and +_BaselineBuffer(buffer, offsets), as in the source-regression runner supplied +with the two pinned source checkouts. Without it, this benchmarks FlashInfer. +Each graph queues an unmeasured NCCL GPU gate before the timing events. Raw +rank vectors are saved; per-trial rank MAX is reduced before medians/ratios. +""" + +import argparse +import datetime +import functools +import gc +import hashlib +import importlib.util +import json +import os +from pathlib import Path +import statistics + +import torch +import torch.distributed as dist + + +def _capture(run, repeats): + gate = torch.zeros(1, dtype=torch.int32, device="cuda") + start = torch.cuda.Event(enable_timing=True, external=True) + end = torch.cuda.Event(enable_timing=True, external=True) + graph = torch.cuda.CUDAGraph() + dist.barrier() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + dist.all_reduce(gate) + start.record() + for _ in range(repeats): + run() + end.record() + return graph, start, end, gate + + +def _worker(rank, args): + from flashinfer import megamoe as api + from flashinfer.jit import env + from flashinfer.jit.megamoe import ( + gen_megamoe_sm90_nvfp4_module, + gen_megamoe_sm90_nvfp4_split_module, + gen_megamoe_sm90_mxfp4_module, + ) + + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", + rank=rank, + world_size=args.ep, + device_id=torch.device("cuda", rank), + timeout=datetime.timedelta(minutes=15), + ) + reference = None + if args.reference_runner: + spec = importlib.util.spec_from_file_location( + "source_baseline", args.reference_runner + ) + reference = importlib.util.module_from_spec(spec) + spec.loader.exec_module(reference) + shapes = { + "flash": (4096, 2048, 256, 6), + "mimo": (6144, 2048, 384, 8), + "pro": (7168, 3072, 384, 6), + } + cache_flush = ( + torch.empty(8 * 1024**3, dtype=torch.uint8, device="cuda") + if args.cold_l2 + else None + ) + results = [] + for shape_name in args.shapes: + h, ih, experts, topk = shapes[shape_name] + buffer = api.MegaMoESymmBuffer( + dist.group.WORLD, experts, args.capacity, topk, h, ih + ) + epr = experts // args.ep + generator = torch.Generator(device="cuda").manual_seed(7919 + rank) + for fmt in args.formats: + gran = 16 if fmt == "nvfp4" else 32 + # Identical packed weights on both sides; packing is outside timing. + raw = [] + for n, k in ((ih * 2, h), (h, ih)): + p = torch.randint( + 0, + 256, + (epr, n, k // 2), + device="cuda", + dtype=torch.uint8, + generator=generator, + ) + sf = torch.full( + (epr, n, k // gran), + 16 if gran == 16 else 122, + dtype=torch.uint8, + device="cuda", + ) + raw.append((p, sf)) + transform = ( + api.transform_nvfp4_weights_for_mega_moe + if fmt == "nvfp4" + else api.transform_mxfp4_weights_for_mega_moe + ) + w1, w2 = transform(*raw) + del raw + for m in args.tokens: + x = (torch.randn(m, h, device="cuda", generator=generator) * 32).to( + torch.float8_e4m3fn + ) + sf = torch.full((m, h // 128), 1 / 128, device="cuda") + idx = ( + torch.randn(m, experts, device="cuda", generator=generator) + .topk(topk, dim=-1) + .indices + ) + weights = torch.softmax( + torch.randn(m, topk, device="cuda", generator=generator), dim=-1 + ) + buffer.stage(x, sf, idx, weights) + y = torch.empty(m, h, device="cuda", dtype=torch.bfloat16) + family = "split" if fmt == "nvfp4" and m > 256 else "fused" + select = ( + api.select_split_plan + if family == "split" + else ( + api.select_fused_plan + if fmt == "nvfp4" + else api.select_mxfp4_fused_plan + ) + ) + plan = select( + buffer.num_sms, args.ep, experts, topk, h, ih, args.capacity, m + ) + if fmt == "mxfp4": + spec = gen_megamoe_sm90_mxfp4_module( + plan, + buffer.num_sms, + args.ep, + args.capacity, + experts, + topk, + h, + ih, + 10.0, + True, + ) + elif family == "split": + spec = gen_megamoe_sm90_nvfp4_split_module( + plan, + buffer.num_sms, + args.ep, + args.capacity, + experts, + topk, + h, + ih, + 10.0, + True, + ) + else: + spec = gen_megamoe_sm90_nvfp4_module( + plan, + buffer.num_sms, + args.capacity, + experts, + topk, + h, + ih, + 10.0, + True, + num_ranks=args.ep, + ) + module = spec.build_and_load() + func = getattr(module, f"megamoe_sm90_{fmt}_{family}") + + run_new = functools.partial( + func, + y, + buffer.buffer, + buffer.handle.buffer_ptrs, + rank, + w1, + w2, + buffer._offsets, + m, + None, + None, + None, + ) + + run_new() + torch.cuda.synchronize() + assert torch.isfinite(y).all() + runs = {"new": run_new} + bb = None + if reference is not None: + old = reference._module(fmt, family, plan, buffer) + offsets = ( + list(old.baseline_layout()) + if fmt == "mxfp4" + else list(old.megamoe_sm90_nvfp4_symm_buffer_layout()) + if family == "fused" + else buffer._offsets + ) + bb = reference._BaselineBuffer(buffer, offsets) + bb.stage(buffer, m) + old_y = torch.empty_like(y) + old_func = getattr(old, f"megamoe_sm90_{fmt}_{family}") + + run_old = functools.partial( + old_func, + old_y, + bb.storage, + bb.handle.buffer_ptrs, + rank, + w1, + w2, + bb.offsets, + m, + None, + None, + None, + ) + + run_old() + torch.cuda.synchronize() + assert torch.isfinite(old_y).all() + runs["old"] = run_old + # Freeze binary identity before any timed graph replay. + if rank == 0: + binary = spec.jit_library_path + proof = { + "name": spec.name, + "library": str(binary), + "sha256": hashlib.sha256(binary.read_bytes()).hexdigest(), + "sources": [str(p) for p in spec.sources], + "include_root": str(env.FLASHINFER_INCLUDE_DIR.resolve()), + } + if reference: + old_name = spec.name + "_original_github" + old_library = ( + env.FLASHINFER_JIT_DIR / old_name / (old_name + ".so") + ) + proof["baseline"] = { + "name": old_name, + "library": str(old_library), + "sha256": hashlib.sha256( + old_library.read_bytes() + ).hexdigest(), + } + print("JIT_READY_PROOF " + json.dumps(proof), flush=True) + captures = { + name: _capture(run, args.repeats) for name, run in runs.items() + } + for _ in range(3): + for graph, _, end, _ in captures.values(): + graph.replay() + end.synchronize() + samples = [] + order = ("old", "new", "new", "old") if reference else ("new",) + for round_idx in range(args.rounds): + for position, name in enumerate(order): + graph, start, end, _ = captures[name] + if cache_flush is not None: + cache_flush.zero_() + graph.replay() + end.synchronize() + us = start.elapsed_time(end) * 1000 / args.repeats + value = torch.tensor([us], dtype=torch.float64, device="cuda") + gathered = [torch.empty_like(value) for _ in range(args.ep)] + dist.all_gather(gathered, value) + rank_us = [v.item() for v in gathered] + if rank == 0: + samples.append( + { + "round": round_idx, + "position": position, + "variant": name, + "rank_us": rank_us, + "max_us": max(rank_us), + } + ) + if rank == 0: + medians = { + name: statistics.median( + s["max_us"] for s in samples if s["variant"] == name + ) + for name in runs + } + ratio = medians["new"] / medians["old"] if reference else None + record = { + "shape": shape_name, + "hidden": h, + "intermediate": ih, + "experts": experts, + "topk": topk, + "tokens": m, + "format": fmt, + "family": family, + "plan": vars(plan), + "median_us": medians, + "new_over_old": ratio, + "samples": samples, + "jit": proof, + } + results.append(record) + print( + f"BENCH {shape_name} {fmt} {family} M={m} {medians} ratio={ratio}", + flush=True, + ) + Path(args.output).write_text( + json.dumps( + { + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "repeats": args.repeats, + "rounds": args.rounds, + "ep": args.ep, + "cold_l2": args.cold_l2, + "gate": "NCCL GPU all-reduce before timing events inside graph", + "results": results, + }, + indent=2, + ) + + "\n" + ) + del captures, runs, y, graph, start, end, run_new + if bb is not None: + del bb, old_y, run_old + gc.collect() + torch.cuda.synchronize() + buffer.destroy() + gc.collect() + torch.cuda.synchronize() + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--shapes", nargs="+", choices=["flash", "mimo", "pro"], default=["flash"] + ) + parser.add_argument( + "--tokens", nargs="+", type=int, default=[1, 24, 32, 64, 96, 384] + ) + parser.add_argument("--ep", type=int, choices=[4, 8], default=8) + parser.add_argument( + "--formats", nargs="+", choices=["nvfp4", "mxfp4"], default=["nvfp4", "mxfp4"] + ) + parser.add_argument("--cold-l2", action="store_true") + parser.add_argument("--capacity", type=int, default=384) + parser.add_argument("--repeats", type=int, default=40) + parser.add_argument("--rounds", type=int, default=5) + parser.add_argument("--reference-runner", type=str) + parser.add_argument("--output", type=str, default="megamoe_fp4_benchmark.json") + args = parser.parse_args() + if args.repeats < 1 or args.rounds < 1: + parser.error("repeats and rounds must be positive") + if args.cold_l2 and args.repeats != 1: + parser.error("cold L2 timing requires --repeats 1") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29650") + torch.multiprocessing.spawn(_worker, args=(args,), nprocs=args.ep) diff --git a/csrc/megamoe/megamoe_sm90_common.cuh b/csrc/megamoe/megamoe_sm90_common.cuh new file mode 100644 index 00000000000..fcfa47929ec --- /dev/null +++ b/csrc/megamoe/megamoe_sm90_common.cuh @@ -0,0 +1,84 @@ +// Host utilities shared by all SM90 FP4 MegaMoE launchers and bindings. +#pragma once + +#include +#include + +#include + +#include "tvm_ffi_utils.h" + +namespace flashinfer::megamoe::detail { + +inline CUtensorMapDataType tma_dtype(DLDataType dt) { + if (dt.code == kDLFloat8_e4m3fn || (dt.code == kDLUInt && dt.bits == 8)) + return CU_TENSOR_MAP_DATA_TYPE_UINT8; + if (dt.code == kDLFloat && dt.bits == 32) return CU_TENSOR_MAP_DATA_TYPE_FLOAT32; + TVM_FFI_ICHECK(false) << "unsupported TMA dtype"; + return CU_TENSOR_MAP_DATA_TYPE_UINT8; +} + +inline CUtensorMapSwizzle tma_swizzle(int mode) { + switch (mode) { + case 0: + return CU_TENSOR_MAP_SWIZZLE_NONE; + case 32: + return CU_TENSOR_MAP_SWIZZLE_32B; + case 64: + return CU_TENSOR_MAP_SWIZZLE_64B; + case 128: + return CU_TENSOR_MAP_SWIZZLE_128B; + } + TVM_FFI_ICHECK(false) << "unsupported swizzle mode " << mode; + return CU_TENSOR_MAP_SWIZZLE_NONE; +} + +inline CUtensorMap make_tma_2d_desc(void* ptr, DLDataType dtype, int inner, int outer, + int box_inner, int box_outer, int outer_stride, int swizzle) { + const int elem_size = dtype.bits * dtype.lanes / 8; + if (swizzle) box_inner = swizzle / elem_size; + CUtensorMap tm{}; + const cuuint64_t dims[2] = {static_cast(inner), static_cast(outer)}; + const cuuint32_t box[2] = {static_cast(box_inner), + static_cast(box_outer)}; + const cuuint64_t strides[1] = { + static_cast(static_cast(outer_stride) * elem_size)}; + const cuuint32_t elem_strides[2] = {1, 1}; + const auto result = + cuTensorMapEncodeTiled(&tm, tma_dtype(dtype), 2, ptr, dims, strides, box, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, tma_swizzle(swizzle), + CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_ICHECK_EQ(result, CUDA_SUCCESS) << "cuTensorMapEncodeTiled failed: " << result; + return tm; +} + +inline CUtensorMap make_tma_sf_desc(void* ptr, DLDataType dtype, int mn, int k, int block_mn, + int gran_k) { + const int alignment = 16 / (dtype.bits * dtype.lanes / 8); + mn = (mn + alignment - 1) / alignment * alignment; + return make_tma_2d_desc(ptr, dtype, mn, (k + gran_k - 1) / gran_k, block_mn, 1, mn, 0); +} + +inline tvm::ffi::Array symm_buffer_layout(int ranks, int experts, int capacity, int topk, + int hidden, int intermediate, int pool, + int sf_pool) { + const layout::MegaMoEBuffers b(nullptr, ranks, experts, capacity, topk, hidden, intermediate, + pool, sf_pool); + auto off = [](const void* p) { return reinterpret_cast(p); }; + return tvm::ffi::Array( + {off(b.combine_token_buffer.get_end_ptr()), off(b.input_token_buffer.base), + off(b.input_sf_buffer.base), off(b.input_topk_idx_buffer.base), + off(b.input_topk_weights_buffer.base), off(b.l1_token_buffer.base), off(b.l1_sf_buffer.base), + off(b.l2_token_buffer.base), off(b.l2_sf_buffer.base), pool, sf_pool}); +} + +template +inline T* optional_tensor_ptr(const tvm::ffi::Optional& tensor) { + return tensor.has_value() ? static_cast(tensor.value().data_ptr()) : nullptr; +} + +inline void* offset_ptr(tvm::ffi::TensorView t, int64_t off) { + return static_cast(t.data_ptr()) + off; +} + +} // namespace flashinfer::megamoe::detail diff --git a/csrc/megamoe/megamoe_sm90_layout.cu b/csrc/megamoe/megamoe_sm90_layout.cu new file mode 100644 index 00000000000..b6f17aa2834 --- /dev/null +++ b/csrc/megamoe/megamoe_sm90_layout.cu @@ -0,0 +1,12 @@ +// Query the common buffer layout without compiling a GEMM specialization. +#include "megamoe_sm90_common.cuh" + +tvm::ffi::Array megamoe_sm90_symm_buffer_layout(int64_t ranks, int64_t experts, + int64_t capacity, int64_t topk, + int64_t hidden, int64_t intermediate, + int64_t pool, int64_t sf_pool) { + return flashinfer::megamoe::detail::symm_buffer_layout(ranks, experts, capacity, topk, hidden, + intermediate, pool, sf_pool); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(megamoe_sm90_symm_buffer_layout, megamoe_sm90_symm_buffer_layout); diff --git a/csrc/megamoe/megamoe_sm90_mxfp4_config.jinja b/csrc/megamoe/megamoe_sm90_mxfp4_config.jinja new file mode 100644 index 00000000000..d0c640f63f1 --- /dev/null +++ b/csrc/megamoe/megamoe_sm90_mxfp4_config.jinja @@ -0,0 +1,30 @@ +// Generated by flashinfer/jit/megamoe.py. +#pragma once +#define FI_MEGAMOE_NVLINK_BARRIER_TRAP_ONLY_TIMEOUT 1 +#include + +namespace flashinfer::megamoe::config { +constexpr uint32_t kNumSMs = {{ num_sms }}; +constexpr uint32_t kNumRanks = {{ num_ranks }}; +constexpr uint32_t kNumMaxTokensPerRank = {{ num_max_tokens_per_rank }}; +constexpr uint32_t kNumExperts = {{ num_experts }}; +constexpr uint32_t kNumTopk = {{ num_topk }}; +constexpr uint32_t kHidden = {{ hidden }}; +constexpr uint32_t kIntermediateHidden = {{ intermediate_hidden }}; +constexpr uint32_t kBlockM = {{ plan.block_m }}; +constexpr uint32_t kBlockN = {{ plan.block_n }}; +constexpr uint32_t kNumMaxPoolTokens = {{ plan.num_max_pool_tokens }}; +constexpr uint32_t kNumPaddedSFPoolTokens = {{ plan.num_padded_sf_pool_tokens }}; +constexpr uint32_t kSmemSize = {{ plan.smem_size }}; +constexpr bool kSwapAB = {{ 'true' if plan.swap_ab else 'false' }}; +constexpr bool kRS = {{ 'true' if plan.rs_swap_ab else 'false' }}; +constexpr uint32_t kNumThreads = sm90_mxfp4_num_threads(kSwapAB, kRS, kBlockM); +constexpr auto kKernel = &sm90_mxfp4_mega_moe_fused_impl< + kNumSMs, kNumRanks, kHidden, kIntermediateHidden, kNumExperts, kNumTopk, + kNumMaxTokensPerRank, {{ plan.num_experts_per_wave }}, kBlockM, kBlockN, + kNumMaxPoolTokens, kNumPaddedSFPoolTokens, {{ plan.num_stages }}, + {{ activation_clamp }}f, {{ 'true' if fast_math else 'false' }}, kSwapAB, + {{ 'true' if plan.single_active_dispatch_warp else 'false' }}, + {{ 'true' if plan.use_mode2_row_decoder else 'false' }}, + {{ 'true' if plan.use_interleaved_scheduler else 'false' }}, kRS>; +} // namespace flashinfer::megamoe::config diff --git a/csrc/megamoe/megamoe_sm90_mxfp4_fused.cu b/csrc/megamoe/megamoe_sm90_mxfp4_fused.cu new file mode 100644 index 00000000000..119675c110f --- /dev/null +++ b/csrc/megamoe/megamoe_sm90_mxfp4_fused.cu @@ -0,0 +1,62 @@ +// MXFP4 fused launch/binding. Weights use contiguous chunk-major decode tiles; +// activation descriptors, buffer ABI and communication primitives are shared. +#include "megamoe_sm90_common.cuh" +#include "megamoe_sm90_mxfp4_config.inc" + +namespace flashinfer::megamoe { +namespace cfg = config; + +void megamoe_sm90_mxfp4_fused( + tvm::ffi::TensorView y, tvm::ffi::TensorView symm_buffer, + tvm::ffi::Array symm_buffer_ptrs, int64_t rank_idx, tvm::ffi::TensorView l1_weights, + tvm::ffi::TensorView l2_weights, tvm::ffi::Array offsets, int64_t num_tokens, + tvm::ffi::Optional cumulative_local_expert_recv_stats, + tvm::ffi::Optional l1_global_scales, + tvm::ffi::Optional l2_global_scales) { + TVM_FFI_ICHECK_EQ(symm_buffer_ptrs.size(), cfg::kNumRanks); + TVM_FFI_ICHECK_GE(rank_idx, 0); + TVM_FFI_ICHECK_LT(rank_idx, cfg::kNumRanks); + TVM_FFI_ICHECK_GE(num_tokens, 1); + TVM_FFI_ICHECK_LE(num_tokens, cfg::kNumMaxTokensPerRank); + const DLDataType fp8{kDLFloat8_e4m3fn, 8, 1}, f32{kDLFloat, 32, 1}; + using detail::make_tma_2d_desc; + using detail::make_tma_sf_desc; + using detail::offset_ptr; + + auto tm_l1 = make_tma_2d_desc(offset_ptr(symm_buffer, offsets[5]), fp8, cfg::kHidden, + cfg::kNumMaxPoolTokens, 128, cfg::kBlockM, cfg::kHidden, 128); + auto tm_l1_sf = make_tma_sf_desc(offset_ptr(symm_buffer, offsets[6]), f32, + cfg::kNumPaddedSFPoolTokens, cfg::kHidden, cfg::kBlockM, 128); + auto tm_out = make_tma_2d_desc(offset_ptr(symm_buffer, offsets[7]), fp8, cfg::kIntermediateHidden, + cfg::kNumMaxPoolTokens, cfg::kBlockN / 2, cfg::kBlockM, + cfg::kIntermediateHidden, 0); + auto tm_l2 = + make_tma_2d_desc(offset_ptr(symm_buffer, offsets[7]), fp8, cfg::kIntermediateHidden, + cfg::kNumMaxPoolTokens, 128, cfg::kBlockM, cfg::kIntermediateHidden, 128); + auto tm_l2_sf = + make_tma_sf_desc(offset_ptr(symm_buffer, offsets[8]), f32, cfg::kNumPaddedSFPoolTokens, + cfg::kIntermediateHidden, cfg::kBlockM, cfg::kBlockN / 2); + const std::vector ptrs(symm_buffer_ptrs.begin(), symm_buffer_ptrs.end()); + layout::SymBuffer sym_buffer(ptrs, static_cast(rank_idx)); + void* output = y.data_ptr(); + auto* stats = detail::optional_tensor_ptr(cumulative_local_expert_recv_stats); + const auto* w1 = static_cast(l1_weights.data_ptr()); + const auto* w2 = static_cast(l2_weights.data_ptr()); + const float* gs1 = detail::optional_tensor_ptr(l1_global_scales); + const float* gs2 = detail::optional_tensor_ptr(l2_global_scales); + uint32_t tokens = static_cast(num_tokens); + unsigned long long* stamps = nullptr; + void* args[] = {&output, &stats, &tokens, &sym_buffer, &tm_l1, &tm_l1_sf, &tm_out, + &tm_l2, &tm_l2_sf, &w1, &w2, &gs1, &gs2, &stamps}; + auto* kernel = reinterpret_cast(cfg::kKernel); + TVM_FFI_ICHECK_EQ( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, cfg::kSmemSize), + cudaSuccess); + TVM_FFI_ICHECK_EQ(cudaLaunchKernel(kernel, dim3(cfg::kNumSMs), dim3(cfg::kNumThreads), args, + cfg::kSmemSize, get_stream(y.device())), + cudaSuccess); +} +} // namespace flashinfer::megamoe + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(megamoe_sm90_mxfp4_fused, + flashinfer::megamoe::megamoe_sm90_mxfp4_fused); diff --git a/csrc/megamoe/megamoe_sm90_nvfp4_config.jinja b/csrc/megamoe/megamoe_sm90_nvfp4_config.jinja new file mode 100644 index 00000000000..bd031f86834 --- /dev/null +++ b/csrc/megamoe/megamoe_sm90_nvfp4_config.jinja @@ -0,0 +1,49 @@ +// Generated by flashinfer/jit/megamoe.py -- do not edit. +// +// Every value here is a compile-time template parameter of the MegaMoE kernel: +// the shared-memory layout, the task schedule and the dispatch protocol are all +// resolved at compile time, so one module is built per (shape, plan) pair. +#pragma once + +#define MEGAMOE_NUM_SMS {{ num_sms }} +#define FI_MEGAMOE_NVLINK_BARRIER_TRAP_ONLY_TIMEOUT 1 +{% if spin_debug %}#define DG_SPIN_DEBUG_TAGS 1{% endif %} + +#include + +namespace flashinfer::megamoe::config { + +constexpr uint32_t kNumRanks = {{ num_ranks }}; +constexpr uint32_t kNumMaxTokensPerRank = {{ num_max_tokens_per_rank }}; +constexpr uint32_t kNumExperts = {{ num_experts }}; +constexpr uint32_t kNumTopk = {{ num_topk }}; +constexpr uint32_t kHidden = {{ hidden }}; +constexpr uint32_t kIntermediateHidden = {{ intermediate_hidden }}; +constexpr uint32_t kNumExpertsPerWave = {{ num_experts_per_wave }}; +constexpr uint32_t kBlockM = {{ block_m }}; +constexpr uint32_t kBlockN = {{ block_n }}; +constexpr uint32_t kNumMaxPoolTokens = {{ num_max_pool_tokens }}; +constexpr uint32_t kNumPaddedSFPoolTokens = {{ num_padded_sf_pool_tokens }}; +constexpr uint32_t kNumStages = {{ num_stages }}; +constexpr uint32_t kSmemSize = {{ smem_size }}; +constexpr float kActivationClamp = {{ activation_clamp }}f; +constexpr bool kFastMath = {{ 'true' if fast_math else 'false' }}; +constexpr bool kSwapAB = {{ 'true' if swap_ab else 'false' }}; +constexpr bool kRSSwapAB = {{ 'true' if rs_swap_ab else 'false' }}; +constexpr bool kSingleActiveDispatchWarp = {{ 'true' if single_active_dispatch_warp else 'false' }}; +constexpr bool kUseMode2RowDecoder = {{ 'true' if use_mode2_row_decoder else 'false' }}; +constexpr bool kUseInterleavedScheduler = {{ 'true' if use_interleaved_scheduler else 'false' }}; +constexpr bool kPushDispatch = {{ 'true' if push_dispatch else 'false' }}; +constexpr bool kNoCleanBarrier = {{ 'true' if no_clean_barrier else 'false' }}; + +// The kernel this module launches, fully specialised. +constexpr auto kKernel = &sm90_nvfp4_mega_moe_fused_impl< + kNumMaxTokensPerRank, kNumExperts, kNumTopk, kHidden, kIntermediateHidden, + kNumExpertsPerWave, kBlockM, kBlockN, kNumMaxPoolTokens, kNumPaddedSFPoolTokens, + kNumStages, kActivationClamp, kFastMath, + kSwapAB, kRSSwapAB, kSingleActiveDispatchWarp, kUseMode2RowDecoder, + kUseInterleavedScheduler, kPushDispatch, kNoCleanBarrier, + /* kStridedPoolDebug */ false, /* kPushProxyFence */ false, + /* kPushGenericRows */ false, /* kPhaseStamps */ false{% if num_ranks != 8 %}, kNumRanks{% endif %}>; + +} // namespace flashinfer::megamoe::config diff --git a/csrc/megamoe/megamoe_sm90_nvfp4_fused.cu b/csrc/megamoe/megamoe_sm90_nvfp4_fused.cu new file mode 100644 index 00000000000..55941eadb78 --- /dev/null +++ b/csrc/megamoe/megamoe_sm90_nvfp4_fused.cu @@ -0,0 +1,181 @@ +/* + * SM90 NVFP4 MegaMoE -- host launcher. + * + * One persistent kernel fuses the cross-rank token dispatch, both FFN GEMMs + * (NVFP4 weights x FP8 activations), SwiGLU and the combine reduction over an + * rank-parameterized NVLink symmetric buffer. Everything that shapes shared memory or the + * task schedule is a template parameter, so this file is compiled once per + * (shape, plan) pair against the generated config header. + */ +#include +#include + +#include "megamoe_sm90_common.cuh" +#include "megamoe_sm90_nvfp4_config.inc" + +namespace flashinfer::megamoe { +namespace { + +namespace cfg = ::flashinfer::megamoe::config; + +using detail::make_tma_2d_desc; +using detail::make_tma_sf_desc; + +} // namespace + +// Launch one MegaMoE call. Pointers are raw device addresses; `sym_buffer_ptrs` +// holds every rank's base address of the symmetric buffer (NVLink peer mapped). +void sm90_nvfp4_mega_moe_fused_launch(void* y, int* cumulative_local_expert_recv_stats, + int num_tokens, const std::vector& sym_buffer_ptrs, + int rank_idx, void* l1_acts, void* l1_acts_sf, void* l2_acts, + void* l2_acts_sf, void* l1_weights, void* l2_weights, + int64_t l1_acts_stride, int64_t l2_acts_stride, + int64_t l1_weights_inner, int64_t l1_weights_stride, + int64_t l2_weights_inner, int64_t l2_weights_stride, + const float* l1_global_scales, const float* l2_global_scales, + cudaStream_t stream) { + constexpr int kBlockK = 128; + constexpr int kSwizzleActs = 128; + constexpr int kBStoragePerKBlock = 80; + constexpr int kL1ScaleGranK = 128; + constexpr int kNumExpertsPerRank = cfg::kNumExperts / cfg::kNumRanks; + + const DLDataType fp8{kDLFloat8_e4m3fn, 8, 1}; + const DLDataType f32{kDLFloat, 32, 1}; + const DLDataType u8{kDLUInt, 8, 1}; + + const auto tm_l1_acts = + make_tma_2d_desc(l1_acts, fp8, cfg::kHidden, cfg::kNumMaxPoolTokens, kBlockK, cfg::kBlockM, + static_cast(l1_acts_stride), kSwizzleActs); + const auto tm_l1_acts_sf = make_tma_sf_desc(l1_acts_sf, f32, cfg::kNumPaddedSFPoolTokens, + cfg::kHidden, cfg::kBlockM, kL1ScaleGranK); + const auto tm_l1_weights = + make_tma_2d_desc(l1_weights, u8, static_cast(l1_weights_inner), + kNumExpertsPerRank * cfg::kIntermediateHidden * 2, kBStoragePerKBlock, + cfg::kBlockN, static_cast(l1_weights_stride), 0); + // L1 writes its SwiGLU output as half-width blocks into the L2 activation pool. + const auto tm_l1_output = + make_tma_2d_desc(l2_acts, fp8, cfg::kIntermediateHidden, cfg::kNumMaxPoolTokens, + cfg::kBlockN / 2, cfg::kBlockM, static_cast(l2_acts_stride), 0); + const auto tm_l2_acts = + make_tma_2d_desc(l2_acts, fp8, cfg::kIntermediateHidden, cfg::kNumMaxPoolTokens, kBlockK, + cfg::kBlockM, static_cast(l2_acts_stride), kSwizzleActs); + const auto tm_l2_acts_sf = + make_tma_sf_desc(l2_acts_sf, f32, cfg::kNumPaddedSFPoolTokens, cfg::kIntermediateHidden, + cfg::kBlockM, cfg::kBlockN / 2); + const auto tm_l2_weights = make_tma_2d_desc(l2_weights, u8, static_cast(l2_weights_inner), + kNumExpertsPerRank * cfg::kHidden, kBStoragePerKBlock, + cfg::kBlockN, static_cast(l2_weights_stride), 0); + + const layout::SymBuffer sym_buffer(sym_buffer_ptrs, + static_cast(rank_idx)); + + auto* kernel = reinterpret_cast(cfg::kKernel); + TVM_FFI_ICHECK_EQ( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, cfg::kSmemSize), + cudaSuccess) + << "failed to raise the dynamic shared-memory limit to " << cfg::kSmemSize; + + const uint32_t tokens = static_cast(num_tokens); + unsigned long long* phase_stamps = nullptr; + void* args[] = {&y, + &cumulative_local_expert_recv_stats, + &phase_stamps, + const_cast(&tokens), + const_cast*>(&sym_buffer), + const_cast(&tm_l1_acts), + const_cast(&tm_l1_acts_sf), + const_cast(&tm_l1_weights), + const_cast(&tm_l1_output), + const_cast(&tm_l2_acts), + const_cast(&tm_l2_acts_sf), + const_cast(&tm_l2_weights), + const_cast(&l1_global_scales), + const_cast(&l2_global_scales)}; + + // Persistent kernel: exactly one CTA per SM. + TVM_FFI_ICHECK_EQ( + cudaLaunchKernel(kernel, dim3(MEGAMOE_NUM_SMS), dim3(384), args, cfg::kSmemSize, stream), + cudaSuccess) + << "MegaMoE kernel launch failed"; +} + +} // namespace flashinfer::megamoe + +using tvm::ffi::Array; +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +namespace flashinfer::megamoe { + +using detail::offset_ptr; + +/*! + * \brief Byte offsets of every region inside the symmetric buffer. + * + * Returns, in order: + * [0] total bytes required + * [1] x (fp8, [num_max_tokens_per_rank, hidden]) + * [2] x_sf (f32, [num_max_tokens_per_rank, hidden/128]) + * [3] topk_idx (i64, [num_max_tokens_per_rank, num_topk]) + * [4] topk_weights (f32, [num_max_tokens_per_rank, num_topk]) + * [5] l1_acts (fp8, [num_max_pool_tokens, hidden]) + * [6] l1_acts_sf (f32, [num_padded_sf_pool_tokens, hidden/128], MN-major) + * [7] l2_acts (fp8, [num_max_pool_tokens, intermediate_hidden]) + * [8] l2_acts_sf (f32, [num_padded_sf_pool_tokens, intermediate_hidden/64], MN-major) + * [9] num_max_pool_tokens + * [10] num_padded_sf_pool_tokens + */ +Array megamoe_sm90_nvfp4_symm_buffer_layout() { + return detail::symm_buffer_layout(cfg::kNumRanks, cfg::kNumExperts, cfg::kNumMaxTokensPerRank, + cfg::kNumTopk, cfg::kHidden, cfg::kIntermediateHidden, + cfg::kNumMaxPoolTokens, cfg::kNumPaddedSFPoolTokens); +} + +/*! \brief Plan constants this module was compiled for (for assertions in Python). */ +Array megamoe_sm90_nvfp4_plan() { + return Array( + {static_cast(cfg::kBlockM), static_cast(cfg::kBlockN), + static_cast(cfg::kNumStages), static_cast(cfg::kNumExpertsPerWave), + static_cast(cfg::kPushDispatch), static_cast(cfg::kNoCleanBarrier), + static_cast(cfg::kUseInterleavedScheduler), static_cast(cfg::kRSSwapAB)}); +} + +void megamoe_sm90_nvfp4_fused(TensorView y, TensorView symm_buffer, Array symm_buffer_ptrs, + int64_t rank_idx, TensorView l1_weights, TensorView l2_weights, + Array offsets, int64_t num_tokens, + Optional cumulative_local_expert_recv_stats, + Optional l1_global_scales, + Optional l2_global_scales) { + TVM_FFI_ICHECK_EQ(symm_buffer_ptrs.size(), cfg::kNumRanks); + TVM_FFI_ICHECK_GE(rank_idx, 0); + TVM_FFI_ICHECK_LT(rank_idx, cfg::kNumRanks); + TVM_FFI_ICHECK_GE(num_tokens, 1); + TVM_FFI_ICHECK_LE(num_tokens, cfg::kNumMaxTokensPerRank); + + const std::vector ptrs(symm_buffer_ptrs.begin(), symm_buffer_ptrs.end()); + + int* stats = detail::optional_tensor_ptr(cumulative_local_expert_recv_stats); + const float* gs1 = detail::optional_tensor_ptr(l1_global_scales); + const float* gs2 = detail::optional_tensor_ptr(l2_global_scales); + + sm90_nvfp4_mega_moe_fused_launch( + y.data_ptr(), stats, static_cast(num_tokens), ptrs, static_cast(rank_idx), + offset_ptr(symm_buffer, offsets[5]), offset_ptr(symm_buffer, offsets[6]), + offset_ptr(symm_buffer, offsets[7]), offset_ptr(symm_buffer, offsets[8]), + l1_weights.data_ptr(), l2_weights.data_ptr(), + /* l1_acts_stride */ cfg::kHidden, + /* l2_acts_stride */ cfg::kIntermediateHidden, + /* l1_weights_inner */ l1_weights.size(2), l1_weights.stride(1), + /* l2_weights_inner */ l2_weights.size(2), l2_weights.stride(1), gs1, gs2, + get_stream(y.device())); +} + +} // namespace flashinfer::megamoe + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(megamoe_sm90_nvfp4_fused, + flashinfer::megamoe::megamoe_sm90_nvfp4_fused); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(megamoe_sm90_nvfp4_symm_buffer_layout, + flashinfer::megamoe::megamoe_sm90_nvfp4_symm_buffer_layout); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(megamoe_sm90_nvfp4_plan, + flashinfer::megamoe::megamoe_sm90_nvfp4_plan); diff --git a/csrc/megamoe/megamoe_sm90_nvfp4_split.cu b/csrc/megamoe/megamoe_sm90_nvfp4_split.cu new file mode 100644 index 00000000000..354b45f7428 --- /dev/null +++ b/csrc/megamoe/megamoe_sm90_nvfp4_split.cu @@ -0,0 +1,146 @@ +/* + * SM90 NVFP4 MegaMoE -- split (BN128) launcher. + * + * Two kernels per call: L1 dequantises the dispatched rows and runs the first + * GEMM into the intermediate pool; L2 runs the second GEMM and scatters the + * result back. L1 runs with a 2-CTA cluster, so it needs the extended launch. + */ +#include +#include + +#include "megamoe_sm90_common.cuh" +#include "megamoe_sm90_nvfp4_split_config.inc" + +namespace flashinfer::megamoe { +namespace { + +namespace cfg = ::flashinfer::megamoe::config; + +using detail::make_tma_2d_desc; +using detail::make_tma_sf_desc; + +void launch_clustered(void* kernel, int threads, int smem, int cluster, void** args, + cudaStream_t stream) { + TVM_FFI_ICHECK_EQ(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem), + cudaSuccess) + << "failed to raise dynamic shared memory to " << smem; + cudaLaunchConfig_t lc{}; + lc.gridDim = dim3(cfg::kNumSMs); + lc.blockDim = dim3(threads); + lc.dynamicSmemBytes = smem; + lc.stream = stream; + cudaLaunchAttribute attrs[1]; + int n = 0; + if (cluster > 1) { + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim = {static_cast(cluster), 1, 1}; + n = 1; + } + lc.attrs = attrs; + lc.numAttrs = n; + TVM_FFI_ICHECK_EQ(cudaLaunchKernelExC(&lc, kernel, args), cudaSuccess) + << "MegaMoE split kernel launch failed"; +} + +} // namespace + +void sm90_nvfp4_mega_moe_split_launch(void* y, int* stats, int num_tokens, + const std::vector& sym_buffer_ptrs, int rank_idx, + void* l1_acts, void* l1_acts_sf, void* l2_acts, + void* l2_acts_sf, void* l1_weights, void* l2_weights, + int64_t l1_weights_inner, int64_t l1_weights_stride, + int64_t l2_weights_inner, int64_t l2_weights_stride, + const float* gs1, const float* gs2, cudaStream_t stream) { + constexpr int kEPR = cfg::kNumExperts / cfg::kNumRanks; + const DLDataType fp8{kDLFloat8_e4m3fn, 8, 1}; + const DLDataType f32{kDLFloat, 32, 1}; + const DLDataType u8{kDLUInt, 8, 1}; + + const auto tm_l1_acts = + make_tma_2d_desc(l1_acts, fp8, cfg::kHidden, cfg::kNumMaxPoolTokens, cfg::kBlockK, + cfg::kBlockM, cfg::kHidden, cfg::kSwizzleActsMode); + const auto tm_l1_acts_sf = make_tma_sf_desc(l1_acts_sf, f32, cfg::kNumPaddedSFPoolTokens, + cfg::kHidden, cfg::kBlockM, 128); + const auto tm_l1_weights = make_tma_2d_desc( + l1_weights, u8, static_cast(l1_weights_inner), kEPR * cfg::kIntermediateHidden * 2, + cfg::kWeightStoragePerKBlock, cfg::kBlockN, static_cast(l1_weights_stride), 0); + // The two M warpgroups each store a 64x64 post-SwiGLU tile, so the L1 output + // box is 64x64 -- not (BLOCK_N/2, BLOCK_M). + constexpr int kL1OutputStoreBlockM = 64, kL1OutputStoreBlockN = 64; + const auto tm_l1_output = + make_tma_2d_desc(l2_acts, fp8, cfg::kIntermediateHidden, cfg::kNumMaxPoolTokens, + kL1OutputStoreBlockN, kL1OutputStoreBlockM, cfg::kIntermediateHidden, 0); + const auto tm_l2_acts = + make_tma_2d_desc(l2_acts, fp8, cfg::kIntermediateHidden, cfg::kNumMaxPoolTokens, cfg::kBlockK, + cfg::kBlockM, cfg::kIntermediateHidden, cfg::kSwizzleActsMode); + const auto tm_l2_acts_sf = make_tma_sf_desc(l2_acts_sf, f32, cfg::kNumPaddedSFPoolTokens, + cfg::kIntermediateHidden, cfg::kBlockM, 128); + const auto tm_l2_weights = make_tma_2d_desc(l2_weights, u8, static_cast(l2_weights_inner), + kEPR * cfg::kHidden, cfg::kWeightStoragePerKBlock, + cfg::kBlockN, static_cast(l2_weights_stride), 0); + + const layout::SymBuffer sym_buffer(sym_buffer_ptrs, + static_cast(rank_idx)); + const uint32_t tokens = static_cast(num_tokens); + + void* l1_args[] = {&stats, + const_cast(&tokens), + const_cast*>(&sym_buffer), + const_cast(&tm_l1_acts), + const_cast(&tm_l1_acts_sf), + const_cast(&tm_l1_weights), + const_cast(&tm_l1_output), + const_cast(&gs1)}; + launch_clustered(reinterpret_cast(cfg::kL1Kernel), cfg::kL1NumThreads, cfg::kL1SmemSize, + cfg::kL1ClusterSize, l1_args, stream); + + void* l2_args[] = {&y, + &stats, + const_cast(&tokens), + const_cast*>(&sym_buffer), + const_cast(&tm_l2_acts), + const_cast(&tm_l2_acts_sf), + const_cast(&tm_l2_weights), + const_cast(&gs2)}; + launch_clustered(reinterpret_cast(cfg::kL2Kernel), cfg::kL2NumThreads, cfg::kL2SmemSize, + cfg::kL2ClusterSize, l2_args, stream); +} + +} // namespace flashinfer::megamoe + +using tvm::ffi::Array; +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +namespace flashinfer::megamoe { + +using detail::offset_ptr; + +void megamoe_sm90_nvfp4_split(TensorView y, TensorView symm_buffer, Array symm_buffer_ptrs, + int64_t rank_idx, TensorView l1_weights, TensorView l2_weights, + Array offsets, int64_t num_tokens, + Optional cumulative_local_expert_recv_stats, + Optional l1_global_scales, + Optional l2_global_scales) { + namespace cfg = ::flashinfer::megamoe::config; + TVM_FFI_ICHECK_EQ(symm_buffer_ptrs.size(), cfg::kNumRanks); + TVM_FFI_ICHECK_GE(num_tokens, 1); + TVM_FFI_ICHECK_LE(num_tokens, cfg::kNumMaxTokensPerRank); + + const std::vector ptrs(symm_buffer_ptrs.begin(), symm_buffer_ptrs.end()); + int* stats = detail::optional_tensor_ptr(cumulative_local_expert_recv_stats); + const float* gs1 = detail::optional_tensor_ptr(l1_global_scales); + const float* gs2 = detail::optional_tensor_ptr(l2_global_scales); + + sm90_nvfp4_mega_moe_split_launch( + y.data_ptr(), stats, static_cast(num_tokens), ptrs, static_cast(rank_idx), + offset_ptr(symm_buffer, offsets[5]), offset_ptr(symm_buffer, offsets[6]), + offset_ptr(symm_buffer, offsets[7]), offset_ptr(symm_buffer, offsets[8]), + l1_weights.data_ptr(), l2_weights.data_ptr(), l1_weights.size(2), l1_weights.stride(1), + l2_weights.size(2), l2_weights.stride(1), gs1, gs2, get_stream(y.device())); +} + +} // namespace flashinfer::megamoe + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(megamoe_sm90_nvfp4_split, + flashinfer::megamoe::megamoe_sm90_nvfp4_split); diff --git a/csrc/megamoe/megamoe_sm90_nvfp4_split_config.jinja b/csrc/megamoe/megamoe_sm90_nvfp4_split_config.jinja new file mode 100644 index 00000000000..d0cd0051610 --- /dev/null +++ b/csrc/megamoe/megamoe_sm90_nvfp4_split_config.jinja @@ -0,0 +1,59 @@ +// Generated by flashinfer/jit/megamoe.py -- do not edit. +// +// Split (BN128) family: two kernels per call. L1 does dispatch-dequant and the +// first GEMM; L2 does the second GEMM and scatters. Their pipeline depth and +// shared-memory budget come from the plan module. +#pragma once + +#define FI_MEGAMOE_NVLINK_BARRIER_TRAP_ONLY_TIMEOUT 1 +// The production split has one L1 implementation: half-stream mode4 with the +// calibrated physical-warp remap. Emitting these as defines keeps the JIT key +// truthful instead of letting an environment variable pick an unvalidated arm. +#define FI_MEGAMOE_NVFP4_L1_DEQUANT_HALF_STREAM 4 +#define FI_MEGAMOE_NVFP4_L1_DEQUANT_WARP_REMAP 1 + +#include + +namespace flashinfer::megamoe::config { + +constexpr uint32_t kNumMaxTokensPerRank = {{ num_max_tokens_per_rank }}; +constexpr uint32_t kNumExperts = {{ num_experts }}; +constexpr uint32_t kNumTopk = {{ num_topk }}; +constexpr uint32_t kHidden = {{ hidden }}; +constexpr uint32_t kIntermediateHidden = {{ intermediate_hidden }}; +constexpr uint32_t kNumExpertsPerWave = {{ num_experts_per_wave }}; +constexpr uint32_t kNumMaxPoolTokens = {{ num_max_pool_tokens }}; +constexpr uint32_t kNumPaddedSFPoolTokens = {{ num_padded_sf_pool_tokens }}; +constexpr uint32_t kNumSMs = {{ num_sms }}; +constexpr uint32_t kNumRanks = {{ num_ranks }}; +constexpr float kActivationClamp = {{ activation_clamp }}f; +constexpr bool kFastMath = {{ 'true' if fast_math else 'false' }}; + +constexpr uint32_t kL1NumStages = {{ l1_num_stages }}; +constexpr uint32_t kL1SmemSize = {{ l1_smem_size }}; +constexpr uint32_t kL2NumStages = {{ l2_num_stages }}; +constexpr uint32_t kL2SmemSize = {{ l2_smem_size }}; + +// Block / thread geometry is fixed for the split family. +constexpr uint32_t kBlockM = 128; +constexpr uint32_t kBlockN = 128; +constexpr uint32_t kBlockK = 128; +constexpr uint32_t kWeightStoragePerKBlock = 80; +constexpr uint32_t kSwizzleActsMode = 128; +constexpr uint32_t kL1ClusterSize = 2, kL2ClusterSize = 1; +constexpr uint32_t kL1NumThreads = 128 + 128 + 256; +constexpr uint32_t kL2NumThreads = 0 + 128 + 256; + +constexpr auto kL1Kernel = &sm90_nvfp4_mega_moe_split_l1_impl< + kNumMaxTokensPerRank, kHidden, kIntermediateHidden, kNumExperts, kNumTopk, + kNumExpertsPerWave, kNumMaxPoolTokens, kNumPaddedSFPoolTokens, kL1NumStages, + kNumSMs, kNumRanks, kActivationClamp, kFastMath, + /* kL2ArrivalCounterRequested */ {{ 'true' if l2_arrival_counter else 'false' }}, + /* kDispatchDequantRequested */ {{ 'true' if dispatch_dequant else 'false' }}>; + +constexpr auto kL2Kernel = &sm90_nvfp4_mega_moe_split_l2_impl< + kNumMaxTokensPerRank, kHidden, kIntermediateHidden, kNumExperts, kNumTopk, + kNumExpertsPerWave, kNumMaxPoolTokens, kNumPaddedSFPoolTokens, kL2NumStages, + kNumSMs, kNumRanks, kActivationClamp, kFastMath>; + +} // namespace flashinfer::megamoe::config diff --git a/csrc/megamoe/megamoe_sm90_nvfp4_split_plan.cu b/csrc/megamoe/megamoe_sm90_nvfp4_split_plan.cu new file mode 100644 index 00000000000..c929b98ad3e --- /dev/null +++ b/csrc/megamoe/megamoe_sm90_nvfp4_split_plan.cu @@ -0,0 +1,50 @@ +/* + * Plan query for the split (BN128) MegaMoE family. + * + * Built once per shape and asked for the L1/L2 pipeline depth and shared-memory + * budget; Python then renders those into the kernel module's config header. + * Kept as a separate, template-free module so the (slow) kernel build only + * happens once the plan is known. + */ +#include "megamoe_sm90_nvfp4_split_plan.cuh" + +using tvm::ffi::Array; + +namespace flashinfer::megamoe { + +/*! \brief [l1_stages, l1_smem, l2_stages, l2_smem, experts_per_wave, + * num_max_pool_tokens, dispatch_dequant, l2_arrival_counter] */ +Array megamoe_sm90_nvfp4_split_plan(int64_t num_sms, int64_t num_ranks, + int64_t num_experts, int64_t num_max_tokens_per_rank, + int64_t num_tokens, int64_t num_topk, int64_t hidden, + int64_t intermediate_hidden, + int64_t num_padded_sf_pool_tokens) { + const SM90NVFP4MegaMoEInput input{ + static_cast(num_sms), + static_cast(num_ranks), + static_cast(num_experts), + static_cast(num_experts / num_ranks), + static_cast(num_max_tokens_per_rank), + static_cast(num_tokens), + static_cast(num_topk), + static_cast(hidden), + static_cast(intermediate_hidden), + static_cast(num_padded_sf_pool_tokens), + }; + const auto plan = select_sm90_nvfp4_split_mega_moe(input); + return Array({ + plan.l1_config.num_stages, + plan.l1_config.smem_size, + plan.l2_config.num_stages, + plan.l2_config.smem_size, + plan.l1_config.num_experts_per_wave, + plan.l1_config.num_max_pool_tokens, + static_cast(plan.dispatch_dequant), + static_cast(plan.l2_arrival_counter), + }); +} + +} // namespace flashinfer::megamoe + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(megamoe_sm90_nvfp4_split_plan, + flashinfer::megamoe::megamoe_sm90_nvfp4_split_plan); diff --git a/csrc/megamoe/megamoe_sm90_nvfp4_split_plan.cuh b/csrc/megamoe/megamoe_sm90_nvfp4_split_plan.cuh new file mode 100644 index 00000000000..82ffc9f534b --- /dev/null +++ b/csrc/megamoe/megamoe_sm90_nvfp4_split_plan.cuh @@ -0,0 +1,301 @@ +/* + * Plan selection for the SM90 NVFP4 MegaMoE split (BN128) family. + * + * The split family runs two kernels per call (L1 then L2), each with its own + * pipeline depth and shared-memory budget. Those two numbers are template + * parameters, so they must be known before the module is built -- this header + * is compiled into a small plan module that Python queries first. + * + * The sizing arithmetic is carried over unchanged from the reference + * implementation: it mirrors the kernel's shared-memory layout exactly, and any + * paraphrase of it would silently mis-size the pipeline. + */ +#pragma once + +#include +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +namespace flashinfer::megamoe { + +template +constexpr T ceil_div(T a, T b) { + return (a + b - 1) / b; +} + +struct SM90ArchSpec { + static constexpr int smem_capacity = 232448; +}; + +static int get_num_experts_per_wave_for_mega_moe(const int& num_experts_per_rank, + const int& num_tokens, const int& num_topk, + const int& intermediate_hidden, const int& block_m, + const int& block_n, const int& num_sms) { + float expected_tokens_per_expert = + static_cast(num_tokens) * num_topk / num_experts_per_rank; + if (expected_tokens_per_expert < 1) { + // Most experts don't have tokens, calculate all experts at once + return num_experts_per_rank; + } + + // Reduce per-expert block count by this factor since uneven routing leaves some experts with + // fewer tokens + constexpr int kImbalanceFactor = 2; + + // Count L1 blocks per expert assuming tokens are evenly spread across experts + const int num_m_blocks = flashinfer::megamoe::ceil_div( + static_cast(std::ceil(expected_tokens_per_expert)), block_m); + const int num_n_blocks = (2 * intermediate_hidden) / block_n; + const int num_l1_blocks_per_expert = num_m_blocks * num_n_blocks; + + // Pick the smallest value whose total blocks (after imbalance reduction) can keep all SMs busy + int num_experts_per_wave = + num_l1_blocks_per_expert > 0 + ? flashinfer::megamoe::ceil_div(kImbalanceFactor * num_sms, num_l1_blocks_per_expert) + : 1; + num_experts_per_wave = std::min(num_experts_per_wave, num_experts_per_rank); + + // Round up to the nearest divisor of num_experts_per_rank so every wave processes the same count + while (num_experts_per_wave < num_experts_per_rank and + num_experts_per_rank % num_experts_per_wave != 0) + ++num_experts_per_wave; + + return num_experts_per_wave; +} + +struct SM90NVFP4MegaMoEConfig { + static constexpr int kBlockM = 128; + static constexpr int kBlockN = 128; + static constexpr int kBlockK = 128; + static constexpr int kWeightStoragePerKBlock = 80; + static constexpr int kSwizzleActsMode = 128; + static constexpr int kL1ClusterSize = 2; + static constexpr int kL2ClusterSize = 1; + static constexpr int kL1NumDispatchThreads = 128; + static constexpr int kL2NumDispatchThreads = 0; + static constexpr int kL1NumActiveDispatchWarps = 2; + static constexpr int kL2NumActiveDispatchWarps = 0; + static constexpr int kL1ActsSFReservationGranK = 64; + static constexpr int kL2ActsSFReservationGranK = 128; + static constexpr int kNumNonEpilogueThreads = 128; + static constexpr int kNumEpilogueThreads = 256; + + int cluster_size; + + int num_max_pool_tokens; + int num_padded_sf_pool_tokens; + + int num_experts_per_wave; + + int num_stages, smem_size; + + int num_dispatch_threads; + + friend std::ostream& operator<<(std::ostream& os, const SM90NVFP4MegaMoEConfig& config) { + os << "SM90NVFP4MegaMoEConfig(" + << "block_m=" << kBlockM << ", block_n=" << kBlockN << ", block_k=" << kBlockK + << ", cluster_size=" << config.cluster_size + << ", num_max_pool_tokens=" << config.num_max_pool_tokens + << ", num_padded_sf_pool_tokens=" << config.num_padded_sf_pool_tokens + << ", swizzle_acts_mode=" << kSwizzleActsMode + << ", num_experts_per_wave=" << config.num_experts_per_wave + << ", num_stages=" << config.num_stages << ", smem_size=" << config.smem_size + << ", num_dispatch_threads=" << config.num_dispatch_threads + << ", num_non_epilogue_threads=" << kNumNonEpilogueThreads + << ", num_epilogue_threads=" << kNumEpilogueThreads << ")"; + return os; + } +}; + +struct SM90NVFP4MegaMoEInput { + int launch_num_sms; + + int num_ranks; + int num_experts; + int num_experts_per_rank; + int num_max_tokens_per_rank; + int num_tokens; + int num_topk; + int hidden; + int intermediate_hidden; + int num_padded_sf_pool_tokens; +}; + +struct SM90NVFP4MegaMoELoad { + int64_t routed_tokens; + int64_t local_experts; + float expected_tokens_per_local_expert; + + bool valid() const { return routed_tokens > 0 && local_experts > 0; } + + bool greater_equal(const int64_t value) const { return routed_tokens >= value * local_experts; } +}; + +struct SM90NVFP4MegaMoEPlan { + SM90NVFP4MegaMoEConfig l1_config; + SM90NVFP4MegaMoEConfig l2_config; + bool dispatch_dequant; + bool l2_arrival_counter; +}; + +enum class SM90NVFP4MegaMoEPhase { + L1, + L2, +}; + +static SM90NVFP4MegaMoELoad get_sm90_nvfp4_mega_moe_load(const SM90NVFP4MegaMoEInput& input) { + const int64_t routed_tokens = static_cast(input.num_tokens) * input.num_topk; + const SM90NVFP4MegaMoELoad load{ + routed_tokens, + input.num_experts_per_rank, + static_cast(routed_tokens) / input.num_experts_per_rank, + }; + TVM_FFI_ICHECK(load.valid()); + return load; +} + +static std::pair get_sm90_nvfp4_mega_moe_pipeline_config( + const SM90NVFP4MegaMoEInput& input, const SM90NVFP4MegaMoELoad& load, + const SM90NVFP4MegaMoEPhase phase) { + const auto align = [](int value, int alignment) { + return ((value + alignment - 1) / alignment) * alignment; + }; + constexpr int kSmemAlignment = 1024; + const bool is_l2 = phase == SM90NVFP4MegaMoEPhase::L2; + const int num_dispatch_threads = is_l2 ? SM90NVFP4MegaMoEConfig::kL2NumDispatchThreads + : SM90NVFP4MegaMoEConfig::kL1NumDispatchThreads; + const int num_active_dispatch_warps = is_l2 ? SM90NVFP4MegaMoEConfig::kL2NumActiveDispatchWarps + : SM90NVFP4MegaMoEConfig::kL1NumActiveDispatchWarps; + // L1 retains only the proven per-64 physical reservation. Kernel + // activation-scale semantics remain per-128. + const int acts_sf_gran_k = is_l2 ? SM90NVFP4MegaMoEConfig::kL2ActsSFReservationGranK + : SM90NVFP4MegaMoEConfig::kL1ActsSFReservationGranK; + + TVM_FFI_ICHECK(num_active_dispatch_warps <= num_dispatch_threads / 32); + const int num_dispatch_warps = num_active_dispatch_warps; + constexpr int kNumEpilogueWarps = 8; + constexpr int kNumEpilogueWarpgroups = 2; + constexpr int kWGBlockM = 64; + constexpr int kWGL1OutBlockN = 64; + + const int smem_expert_count_size = + align(input.num_experts * static_cast(sizeof(uint32_t)), kSmemAlignment); + const int smem_send_buffers_size = + align(static_cast( + layout::Buffer(layout::Data(input.hidden), num_dispatch_warps, 1).get_num_bytes()), + kSmemAlignment); + const int smem_dispatch_size = smem_expert_count_size + smem_send_buffers_size; + const int smem_nvfp4_lut = align(128 * 8, kSmemAlignment); + + const int smem_cd_l1 = kNumEpilogueWarpgroups * kWGBlockM * kWGL1OutBlockN; + // L2 stages its remote scatter through one warp-private 8-row half tile so + // that each destination row leaves as a single contiguous 16-byte-wide + // burst. Must match `kL2StageRows` / `kL2StageRowStride` in the kernel. + constexpr int kWGBlockN = 128; + constexpr int kL2StageRows = 8; + constexpr int kL2StageRowPad = 8; + const int smem_cd_l2 = kNumEpilogueWarps * kL2StageRows * (kWGBlockN + kL2StageRowPad) * + static_cast(sizeof(uint16_t)); + const int smem_cd = align(is_l2 ? smem_cd_l2 : smem_cd_l1, kSmemAlignment); + + const int num_sfa_groups_per_bk = SM90NVFP4MegaMoEConfig::kBlockK / acts_sf_gran_k; + const int smem_sfa_per_stage = align( + num_sfa_groups_per_bk * SM90NVFP4MegaMoEConfig::kBlockM * static_cast(sizeof(float)), + 128); + const int smem_per_stage = SM90NVFP4MegaMoEConfig::kBlockM * SM90NVFP4MegaMoEConfig::kBlockK + + SM90NVFP4MegaMoEConfig::kBlockN * SM90NVFP4MegaMoEConfig::kBlockK + + smem_sfa_per_stage; + const int smem_barriers_fixed = (num_dispatch_warps + 2 * kNumEpilogueWarps) * 8; + // Both phases own full, empty, and dequant barriers. Production L1 is + // always half-stream mode4 and therefore owns one additional K[64:128] + // publication mbarrier per stage. Account for it while selecting stages, + // rather than appending bytes after a max-stage decision. + const int smem_barriers_per_stage = (is_l2 ? 3 : 4) * 8; + const int smem_fixed = smem_dispatch_size + smem_nvfp4_lut + smem_cd + smem_barriers_fixed; + const int max_num_stages = + (SM90ArchSpec::smem_capacity - smem_fixed) / (smem_per_stage + smem_barriers_per_stage); + const int num_stages = + !is_l2 && load.expected_tokens_per_local_expert > 8.0f && max_num_stages > 6 ? 6 + : max_num_stages; + TVM_FFI_ICHECK(max_num_stages >= 2); + TVM_FFI_ICHECK(num_stages >= 2 && num_stages <= max_num_stages); + return { + num_stages, + smem_fixed + num_stages * (smem_per_stage + smem_barriers_per_stage), + }; +} + +static int get_num_experts_per_wave_for_sm90_nvfp4_mega_moe(const SM90NVFP4MegaMoEInput& input) { + return get_num_experts_per_wave_for_mega_moe( + input.num_experts_per_rank, input.num_tokens, input.num_topk, input.intermediate_hidden, + SM90NVFP4MegaMoEConfig::kBlockM, SM90NVFP4MegaMoEConfig::kBlockN, input.launch_num_sms); +} + +static SM90NVFP4MegaMoEConfig materialize_sm90_nvfp4_mega_moe_phase( + const SM90NVFP4MegaMoEInput& input, const SM90NVFP4MegaMoELoad& load, + const SM90NVFP4MegaMoEPhase phase) { + const bool is_l2 = phase == SM90NVFP4MegaMoEPhase::L2; + SM90NVFP4MegaMoEConfig config{ + is_l2 ? SM90NVFP4MegaMoEConfig::kL2ClusterSize : SM90NVFP4MegaMoEConfig::kL1ClusterSize, + layout::get_num_max_pool_tokens(input.num_ranks, input.num_max_tokens_per_rank, + input.num_topk, input.num_experts_per_rank), + input.num_padded_sf_pool_tokens, + get_num_experts_per_wave_for_sm90_nvfp4_mega_moe(input), + 0, + 0, + is_l2 ? SM90NVFP4MegaMoEConfig::kL2NumDispatchThreads + : SM90NVFP4MegaMoEConfig::kL1NumDispatchThreads, + }; + std::tie(config.num_stages, config.smem_size) = + get_sm90_nvfp4_mega_moe_pipeline_config(input, load, phase); + return config; +} + +static bool is_sm90_nvfp4_mega_moe_plan_legal(const SM90NVFP4MegaMoEInput& input, + const SM90NVFP4MegaMoEPlan& plan) { + const auto valid_phase = [&](const SM90NVFP4MegaMoEConfig& config) { + return config.num_experts_per_wave > 0 && + config.num_experts_per_wave <= input.num_experts_per_rank && + input.num_experts_per_rank % config.num_experts_per_wave == 0 && + config.num_stages >= 2 && config.smem_size > 0 && + config.smem_size <= SM90ArchSpec::smem_capacity; + }; + return valid_phase(plan.l1_config) && valid_phase(plan.l2_config) && + plan.l1_config.cluster_size == SM90NVFP4MegaMoEConfig::kL1ClusterSize && + plan.l1_config.num_dispatch_threads == SM90NVFP4MegaMoEConfig::kL1NumDispatchThreads && + plan.l2_config.cluster_size == SM90NVFP4MegaMoEConfig::kL2ClusterSize && + plan.l2_config.num_dispatch_threads == SM90NVFP4MegaMoEConfig::kL2NumDispatchThreads; +} + +static SM90NVFP4MegaMoEPlan select_sm90_nvfp4_split_mega_moe(const SM90NVFP4MegaMoEInput& input) { + TVM_FFI_ICHECK(input.launch_num_sms > 0); + TVM_FFI_ICHECK(input.num_ranks > 0); + TVM_FFI_ICHECK(input.num_experts_per_rank > 0); + TVM_FFI_ICHECK(input.num_experts == input.num_experts_per_rank * input.num_ranks); + TVM_FFI_ICHECK(input.num_max_tokens_per_rank > 0); + TVM_FFI_ICHECK(input.num_tokens > 0 && input.num_tokens <= input.num_max_tokens_per_rank); + TVM_FFI_ICHECK(input.num_topk > 0); + TVM_FFI_ICHECK(input.hidden > 0 && input.hidden % 128 == 0); + TVM_FFI_ICHECK(input.intermediate_hidden > 0 && input.intermediate_hidden % 128 == 0); + TVM_FFI_ICHECK(input.num_padded_sf_pool_tokens > 0); + + const auto load = get_sm90_nvfp4_mega_moe_load(input); + const SM90NVFP4MegaMoEPlan plan{ + materialize_sm90_nvfp4_mega_moe_phase(input, load, SM90NVFP4MegaMoEPhase::L1), + materialize_sm90_nvfp4_mega_moe_phase(input, load, SM90NVFP4MegaMoEPhase::L2), + // Every production BN128 split uses the accepted L1 mode4+remap + // implementation. That implementation requires the independent odd-K + // dispatch dequant team, including in the former rho<256 interval. + true, + load.expected_tokens_per_local_expert <= 32.0f || + load.expected_tokens_per_local_expert >= 128.0f, + }; + TVM_FFI_ICHECK(is_sm90_nvfp4_mega_moe_plan_legal(input, plan)); + return plan; +} + +} // namespace flashinfer::megamoe diff --git a/docs/api/megamoe.rst b/docs/api/megamoe.rst new file mode 100644 index 00000000000..45248b1646c --- /dev/null +++ b/docs/api/megamoe.rst @@ -0,0 +1,31 @@ +.. _apimegamoe: + +flashinfer.megamoe +================== + +.. currentmodule:: flashinfer.megamoe + +SM90 W4A8 expert-parallel MegaMoE with NVFP4 fused/split and MXFP4 fused kernels. +All paths use FP8 E4M3 activations and return BF16 outputs. Execution requires +NVLink-connected Hopper GPUs and a symmetric buffer allocated collectively +before CUDA graph capture. NVFP4 supports four or eight ranks; MXFP4 supports +power-of-two rank counts from one to 32. + +Buffer and execution +-------------------- + +.. autosummary:: + :toctree: ../generated + + MegaMoESymmBuffer + nvfp4_mega_moe + mxfp4_mega_moe + +Weight preparation +------------------ + +.. autosummary:: + :toctree: ../generated + + transform_nvfp4_weights_for_mega_moe + transform_mxfp4_weights_for_mega_moe diff --git a/docs/index.rst b/docs/index.rst index 3f656fca396..3364ec9b198 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -42,6 +42,7 @@ FlashInfer is a library and kernel generator for Large Language Models that prov api/gemm api/grouped_mm api/fused_moe + api/megamoe api/cascade api/comm api/sparse diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 38ffd4956c5..fb47f90f45d 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -431,3 +431,4 @@ def _log_import_version() -> None: from .kda_prefill import ( kda_prefill_supports_fp32_checkpoints as kda_prefill_supports_fp32_checkpoints, ) +from . import megamoe as megamoe diff --git a/flashinfer/jit/megamoe.py b/flashinfer/jit/megamoe.py new file mode 100644 index 00000000000..ce032362d40 --- /dev/null +++ b/flashinfer/jit/megamoe.py @@ -0,0 +1,347 @@ +"""JIT module generation for the three SM90 FP4 MegaMoE paths. + +The kernel resolves its shared-memory layout, task schedule and dispatch +protocol at compile time, so one module is built per (shape, plan) pair. The +plan is chosen in ``flashinfer.megamoe.heuristics`` and rendered into +``megamoe_sm90_nvfp4_config.inc`` here. +""" + +import os +from pathlib import Path + +import jinja2 + +from . import env as jit_env +from .core import JitSpec, gen_jit_spec, write_if_different + + +def _csrc_dir() -> Path: + installed = jit_env.FLASHINFER_CSRC_DIR / "megamoe" + if installed.exists(): + return installed + checkout = Path(__file__).resolve().parents[2] / "csrc" / "megamoe" + if checkout.exists(): + return checkout + raise FileNotFoundError( + f"MegaMoE CUDA sources not found (checked {installed}, {checkout})" + ) + + +def _gen_fp4_module(uri, template_name, filenames, render_args, use_fast_math): + """Shared config rendering, source staging and CUDA build policy.""" + csrc = _csrc_dir() + directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri + directory.mkdir(parents=True, exist_ok=True) + template = jinja2.Template((csrc / f"{template_name}.jinja").read_text()) + write_if_different( + directory / f"{template_name}.inc", template.render(**render_args) + ) + sources = [] + for filename in filenames: + source = directory / filename + write_if_different(source, (csrc / filename).read_text()) + sources.append(source) + return gen_jit_spec( + uri, + sources, + extra_cuda_cflags=[ + "-std=c++20", + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + "--ptxas-options=--register-usage-level=5", + "-DENABLE_BF16", + "-DENABLE_FP8", + ], + extra_cflags=["-std=c++20"], + extra_include_paths=[directory, csrc, jit_env.FLASHINFER_CSRC_DIR], + use_fast_math=use_fast_math, + ) + + +def get_megamoe_sm90_nvfp4_uri( + plan, + num_sms: int, + num_max_tokens_per_rank: int, + num_experts: int, + num_topk: int, + hidden: int, + intermediate_hidden: int, + activation_clamp: float, + fast_math: bool, + num_ranks: int = 8, +) -> str: + return ( + f"megamoe_sm90_nvfp4_r{num_ranks}_sm{num_sms}_e{num_experts}_k{num_topk}" + f"_h{hidden}_i{intermediate_hidden}_cap{num_max_tokens_per_rank}" + f"_bm{plan.block_m}_bn{plan.block_n}_epw{plan.num_experts_per_wave}" + f"_st{plan.num_stages}" + f"_{'swap' if plan.swap_ab else 'noswap'}" + f"{'_rs' if plan.rs_swap_ab else ''}" + f"{'_il' if plan.use_interleaved_scheduler else '_static'}" + f"{'_push' if plan.push_dispatch else ''}" + f"{'_noclean' if plan.no_clean_barrier else ''}" + f"{'_mode2' if plan.use_mode2_row_decoder else ''}" + f"{'_single' if plan.single_active_dispatch_warp else ''}" + f"_clamp{activation_clamp:g}" + f"{'_fm' if fast_math else ''}" + ) + + +def gen_megamoe_sm90_nvfp4_module( + plan, + num_sms: int, + num_max_tokens_per_rank: int, + num_experts: int, + num_topk: int, + hidden: int, + intermediate_hidden: int, + activation_clamp: float, + fast_math: bool, + num_ranks: int = 8, +) -> JitSpec: + uri = get_megamoe_sm90_nvfp4_uri( + plan, + num_sms, + num_max_tokens_per_rank, + num_experts, + num_topk, + hidden, + intermediate_hidden, + activation_clamp, + fast_math, + num_ranks=num_ranks, + ) + # Fused compiler fast-math follows the source policy. + return _gen_fp4_module( + uri, + "megamoe_sm90_nvfp4_config", + ("megamoe_sm90_nvfp4_fused.cu",), + dict( + num_sms=num_sms, + num_ranks=num_ranks, + num_max_tokens_per_rank=num_max_tokens_per_rank, + num_experts=num_experts, + num_topk=num_topk, + hidden=hidden, + intermediate_hidden=intermediate_hidden, + num_experts_per_wave=plan.num_experts_per_wave, + block_m=plan.block_m, + block_n=plan.block_n, + num_max_pool_tokens=plan.num_max_pool_tokens, + num_padded_sf_pool_tokens=plan.num_padded_sf_pool_tokens, + num_stages=plan.num_stages, + smem_size=plan.smem_size, + activation_clamp=activation_clamp, + fast_math=fast_math, + swap_ab=plan.swap_ab, + rs_swap_ab=plan.rs_swap_ab, + single_active_dispatch_warp=plan.single_active_dispatch_warp, + use_mode2_row_decoder=plan.use_mode2_row_decoder, + use_interleaved_scheduler=plan.use_interleaved_scheduler, + push_dispatch=plan.push_dispatch, + no_clean_barrier=plan.no_clean_barrier, + spin_debug=os.environ.get("FLASHINFER_MEGAMOE_SPIN_DEBUG", "0") == "1", + ), + fast_math, + ) + + +# --------------------------------------------------------------------------- +# Split (BN128) family +# --------------------------------------------------------------------------- + + +def _split_shape_uri( + num_sms, + num_ranks, + num_max_tokens_per_rank, + num_experts, + num_topk, + hidden, + intermediate_hidden, +) -> str: + return ( + f"sm{num_sms}_r{num_ranks}_e{num_experts}_k{num_topk}" + f"_h{hidden}_i{intermediate_hidden}_cap{num_max_tokens_per_rank}" + ) + + +def gen_megamoe_sm90_nvfp4_split_plan_module( + num_sms, + num_ranks, + num_max_tokens_per_rank, + num_experts, + num_topk, + hidden, + intermediate_hidden, +) -> JitSpec: + """Template-free module that reports the L1/L2 pipeline depth and shared + memory budget. Those are kernel template parameters, so they must be known + before the (slow) kernel module is built.""" + uri = "megamoe_sm90_nvfp4_split_plan_" + _split_shape_uri( + num_sms, + num_ranks, + num_max_tokens_per_rank, + num_experts, + num_topk, + hidden, + intermediate_hidden, + ) + csrc = _csrc_dir() + gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri + os.makedirs(gen_directory, exist_ok=True) + dest = gen_directory / "megamoe_sm90_nvfp4_split_plan.cu" + for filename in ( + "megamoe_sm90_nvfp4_split_plan.cu", + "megamoe_sm90_nvfp4_split_plan.cuh", + ): + with open(csrc / filename) as f: + write_if_different(gen_directory / filename, f.read()) + return gen_jit_spec( + uri, + [dest], + extra_cuda_cflags=["-std=c++20", "--expt-relaxed-constexpr"], + extra_cflags=["-std=c++20"], + extra_include_paths=[gen_directory, csrc, jit_env.FLASHINFER_CSRC_DIR], + ) + + +def get_megamoe_sm90_nvfp4_split_uri( + plan, + num_sms, + num_ranks, + num_max_tokens_per_rank, + num_experts, + num_topk, + hidden, + intermediate_hidden, + activation_clamp, + fast_math, +) -> str: + return ( + "megamoe_sm90_nvfp4_split_" + + _split_shape_uri( + num_sms, + num_ranks, + num_max_tokens_per_rank, + num_experts, + num_topk, + hidden, + intermediate_hidden, + ) + + f"_epw{plan.num_experts_per_wave}" + + f"_l1st{plan.l1_num_stages}_l2st{plan.l2_num_stages}" + + ("_dq" if plan.dispatch_dequant else "") + + ("_l2ac" if plan.l2_arrival_counter else "") + + f"_clamp{activation_clamp:g}" + + ("_fm" if fast_math else "") + ) + + +def gen_megamoe_sm90_nvfp4_split_module( + plan, + num_sms, + num_ranks, + num_max_tokens_per_rank, + num_experts, + num_topk, + hidden, + intermediate_hidden, + activation_clamp, + fast_math, +) -> JitSpec: + uri = get_megamoe_sm90_nvfp4_split_uri( + plan, + num_sms, + num_ranks, + num_max_tokens_per_rank, + num_experts, + num_topk, + hidden, + intermediate_hidden, + activation_clamp, + fast_math, + ) + # Split kFastMath selects source intrinsics; retain denormals in this family. + return _gen_fp4_module( + uri, + "megamoe_sm90_nvfp4_split_config", + ("megamoe_sm90_nvfp4_split.cu",), + dict( + num_sms=num_sms, + num_ranks=num_ranks, + num_max_tokens_per_rank=num_max_tokens_per_rank, + num_experts=num_experts, + num_topk=num_topk, + hidden=hidden, + intermediate_hidden=intermediate_hidden, + num_experts_per_wave=plan.num_experts_per_wave, + num_max_pool_tokens=plan.num_max_pool_tokens, + num_padded_sf_pool_tokens=plan.num_padded_sf_pool_tokens, + l1_num_stages=plan.l1_num_stages, + l1_smem_size=plan.l1_smem_size, + l2_num_stages=plan.l2_num_stages, + l2_smem_size=plan.l2_smem_size, + dispatch_dequant=plan.dispatch_dequant, + l2_arrival_counter=plan.l2_arrival_counter, + activation_clamp=activation_clamp, + fast_math=fast_math, + ), + False, + ) + + +def gen_megamoe_sm90_layout_module() -> JitSpec: + csrc = _csrc_dir() + return gen_jit_spec( + "megamoe_sm90_layout", + [csrc / "megamoe_sm90_layout.cu"], + extra_include_paths=[csrc, jit_env.FLASHINFER_CSRC_DIR], + extra_cuda_cflags=["-std=c++20", "--expt-relaxed-constexpr"], + extra_cflags=["-std=c++20"], + ) + + +def gen_megamoe_sm90_mxfp4_module( + plan, + num_sms, + num_ranks, + num_max_tokens_per_rank, + num_experts, + num_topk, + hidden, + intermediate_hidden, + activation_clamp, + fast_math, +) -> JitSpec: + # Include every compile-time policy in the key, including rank width. + uri = get_megamoe_sm90_nvfp4_uri( + plan, + num_sms, + num_max_tokens_per_rank, + num_experts, + num_topk, + hidden, + intermediate_hidden, + activation_clamp, + fast_math, + num_ranks=num_ranks, + ).replace("nvfp4", "mxfp4") + return _gen_fp4_module( + uri, + "megamoe_sm90_mxfp4_config", + ("megamoe_sm90_mxfp4_fused.cu",), + dict( + plan=plan, + num_sms=num_sms, + num_ranks=num_ranks, + num_max_tokens_per_rank=num_max_tokens_per_rank, + num_experts=num_experts, + num_topk=num_topk, + hidden=hidden, + intermediate_hidden=intermediate_hidden, + activation_clamp=activation_clamp, + fast_math=fast_math, + ), + fast_math, + ) diff --git a/flashinfer/megamoe/__init__.py b/flashinfer/megamoe/__init__.py new file mode 100644 index 00000000000..4e897eca47e --- /dev/null +++ b/flashinfer/megamoe/__init__.py @@ -0,0 +1,617 @@ +"""SM90 (Hopper) FP4 MegaMoE: NVFP4 fused/split and MXFP4 fused. + +The fused paths perform cross-rank token dispatch, two W4A8 GEMMs, clamped +SwiGLU and top-k reduction in one persistent kernel. NVFP4 split uses separate +L1 and L2 launches. All paths use FP8 activations and NVLink symmetric memory. +NVFP4 supports four or eight ranks; MXFP4 supports power-of-two rank counts up to 32. + +Usage:: + + buf = MegaMoESymmBuffer(group, num_experts=256, num_max_tokens_per_rank=8448, + num_topk=6, hidden=4096, intermediate_hidden=2048) + buf.stage(x_fp8, x_sf, topk_idx, topk_weights) + nvfp4_mega_moe(y, l1_weights, l2_weights, buf) +""" + +import dataclasses +import functools +from typing import Optional, Tuple + +import torch + +from ..jit.megamoe import ( + gen_megamoe_sm90_layout_module, + gen_megamoe_sm90_mxfp4_module, + gen_megamoe_sm90_nvfp4_module, + gen_megamoe_sm90_nvfp4_split_module, + gen_megamoe_sm90_nvfp4_split_plan_module, +) +from .heuristics import ( + CANDIDATE_BLOCK_M, + FusedPlan, + get_num_max_pool_tokens, + get_num_padded_sf_pool_tokens, + select_fused_plan, + select_mxfp4_fused_plan, +) +from .weights import ( + MXFP4_TILE_BYTES, + transform_mxfp4_weights_for_mega_moe, + transform_nvfp4_weights_for_mega_moe, +) + +__all__ = [ + "MegaMoESymmBuffer", + "nvfp4_mega_moe", + "mxfp4_mega_moe", + "transform_nvfp4_weights_for_mega_moe", + "transform_mxfp4_weights_for_mega_moe", + "select_fused_plan", + "select_mxfp4_fused_plan", + "FusedPlan", + "SplitPlan", + "select_split_plan", + "FAMILY_THRESHOLD", +] + + +@functools.cache +def _get_jit_module(factory, *args): + return factory(*args).build_and_load() + + +def _run_kernel(factory, entry, plan, y, w1, w2, b, m, stats, gs1, gs2): + args = (plan, b.num_sms) + # Keep the trailing fused rank argument compatible with existing JIT callers. + if factory is not gen_megamoe_sm90_nvfp4_module: + args += (b.group.size(),) + args += ( + b.num_max_tokens_per_rank, + b.num_experts, + b.num_topk, + b.hidden, + b.intermediate_hidden, + b.activation_clamp, + b.fast_math, + ) + if factory is gen_megamoe_sm90_nvfp4_module: + args += (b.group.size(),) + module = _get_jit_module(factory, *args) + getattr(module, entry)( + y, + b.buffer, + b.handle.buffer_ptrs, + b.group.rank(), + w1, + w2, + b._offsets, + m, + stats, + gs1, + gs2, + ) + return y + + +class MegaMoESymmBuffer: + """Allocate symmetric input, expert-token and output-combine storage. + + All ranks must use identical geometry and invoke kernels in matching order. + Allocation and rendezvous are collective and must precede CUDA graph capture. + + Parameters + ---------- + group : torch.distributed.ProcessGroup + Expert-parallel process group over NVLink-connected Hopper GPUs. + num_experts : int + Global expert count, divisible by the number of ranks. + num_max_tokens_per_rank : int + Maximum input token count on each rank. + num_topk : int + Number of routed expert slots per token. + hidden : int + Input and output width, a multiple of 256. + intermediate_hidden : int + Expert intermediate width, a multiple of 128. + activation_clamp : float + SwiGLU clamp value. Defaults to 10.0. + fast_math : bool + Use the family's fast-math policy. Defaults to True. + """ + + def __init__( + self, + group: torch.distributed.ProcessGroup, + num_experts: int, + num_max_tokens_per_rank: int, + num_topk: int, + hidden: int, + intermediate_hidden: int, + activation_clamp: float = 10.0, + fast_math: bool = True, + ) -> None: + import torch.distributed._symmetric_memory as symm_mem + + # Buffer geometry is format-independent; NVFP4 checks its narrower + # deployment constraints when selecting a fused/split plan. + select_mxfp4_fused_plan( + 132, + group.size(), + num_experts, + num_topk, + hidden, + intermediate_hidden, + num_max_tokens_per_rank, + 1, + ) + if torch.cuda.get_device_capability()[0] != 9: + raise ValueError("SM90 MegaMoE requires a Hopper GPU") + + self.group = group + self.num_experts = num_experts + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.num_topk = num_topk + self.hidden = hidden + self.intermediate_hidden = intermediate_hidden + self.activation_clamp = activation_clamp + self.fast_math = fast_math + self.num_sms = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + + pool = get_num_max_pool_tokens( + group.size(), num_max_tokens_per_rank, num_topk, num_experts // group.size() + ) + sf_pool = max( + get_num_padded_sf_pool_tokens(pool, bm) for bm in CANDIDATE_BLOCK_M + ) + layout = list( + _get_jit_module( + gen_megamoe_sm90_layout_module + ).megamoe_sm90_symm_buffer_layout( + group.size(), + num_experts, + num_max_tokens_per_rank, + num_topk, + hidden, + intermediate_hidden, + pool, + sf_pool, + ) + ) + self._offsets = layout + num_bytes = layout[0] + self.num_max_pool_tokens = layout[9] + self.num_padded_sf_pool_tokens = layout[10] + + self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device="cuda") + self.handle = symm_mem.rendezvous(self.buffer, group=group) + self.buffer.zero_() + group.barrier() + torch.cuda.synchronize() + + self.x = self._view( + layout[1], (num_max_tokens_per_rank, hidden), torch.float8_e4m3fn + ) + self.x_sf = self._view( + layout[2], (num_max_tokens_per_rank, hidden // 128), torch.float32 + ) + self.topk_idx = self._view( + layout[3], (num_max_tokens_per_rank, num_topk), torch.int64 + ) + self.topk_weights = self._view( + layout[4], (num_max_tokens_per_rank, num_topk), torch.float32 + ) + + def _view(self, byte_offset: int, shape: Tuple[int, ...], dtype: torch.dtype): + itemsize = torch.empty((), dtype=dtype).element_size() + numel = 1 + for s in shape: + numel *= s + flat = self.buffer[byte_offset : byte_offset + numel * itemsize] + return flat.view(dtype).view(*shape) + + def stage( + self, + x_fp8: torch.Tensor, + x_sf: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ) -> None: + """Copy this rank's tokens and routing into the symmetric buffer. + + Parameters + ---------- + x_fp8 : torch.Tensor + E4M3 activations with shape ``[tokens, hidden]``. + x_sf : torch.Tensor + FP32 activation scales with shape ``[tokens, hidden // 128]``. + topk_idx : torch.Tensor + INT64 global expert indices with shape ``[tokens, num_topk]``. + A value of -1 masks a route. + topk_weights : torch.Tensor + FP32 routing weights with shape ``[tokens, num_topk]``. + + All tensors must be on the buffer's CUDA device. The token count must + be positive and fit the allocated capacity. + """ + m = x_fp8.shape[0] + for tensor, shape, dtype in ( + (x_fp8, (m, self.hidden), torch.float8_e4m3fn), + (x_sf, (m, self.hidden // 128), torch.float32), + (topk_idx, (m, self.num_topk), torch.int64), + (topk_weights, (m, self.num_topk), torch.float32), + ): + if ( + tuple(tensor.shape) != shape + or tensor.dtype != dtype + or tensor.device != self.buffer.device + ): + raise ValueError( + f"staged input must have shape {shape}, dtype {dtype}, device {self.buffer.device}" + ) + if not 0 < m <= self.num_max_tokens_per_rank: + raise ValueError("staged token count exceeds buffer capacity") + self.x[:m].copy_(x_fp8) + self.x_sf[:m].copy_(x_sf) + self.topk_idx[:m].copy_(topk_idx) + self.topk_weights[:m].copy_(topk_weights) + self.num_staged_tokens = m + + def destroy(self) -> None: + """Release this rank's views and symmetric-memory handle.""" + self.handle = None + self.buffer = None + self.x = self.x_sf = self.topk_idx = self.topk_weights = None + + +def nvfp4_mega_moe( + y: torch.Tensor, + l1_weights: torch.Tensor, + l2_weights: torch.Tensor, + symm_buffer: MegaMoESymmBuffer, + num_tokens: Optional[int] = None, + cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, + l1_global_scales: Optional[torch.Tensor] = None, + l2_global_scales: Optional[torch.Tensor] = None, + force_push_dispatch: Optional[bool] = None, + force_no_clean_barrier: Optional[bool] = None, + kernel_family: str = "auto", + family_threshold: int = 256, +) -> torch.Tensor: + """Run NVFP4 MegaMoE using the fused or split family. + + Parameters + ---------- + y : torch.Tensor + Contiguous BF16 output with shape ``[at least num_tokens, hidden]``. + The kernel writes the first ``num_tokens`` rows in place. + l1_weights, l2_weights : torch.Tensor + Packed NVFP4 weights from :func:`transform_nvfp4_weights_for_mega_moe`. + symm_buffer : MegaMoESymmBuffer + Buffer populated by :meth:`MegaMoESymmBuffer.stage` on every rank. + num_tokens : int, optional + Tokens on this rank. Defaults to ``y.shape[0]``. + cumulative_local_expert_recv_stats : torch.Tensor, optional + INT32 vector with one entry per local expert. Receive counts are + accumulated in place; initialize the vector before first use. + l1_global_scales, l2_global_scales : torch.Tensor, optional + FP32 vectors with one multiplier per local expert. Omitted scales use 1. + force_push_dispatch, force_no_clean_barrier : bool, optional + Override the fused counter-dispatch policy. None uses the default + policy. These overrides do not apply to the split family. + kernel_family : str + ``"auto"``, ``"fused"`` or ``"split"``. Auto selects fused through + ``family_threshold`` tokens per rank, then split. + family_threshold : int + Token threshold used by auto selection. Defaults to 256. + + Returns + ------- + torch.Tensor + The output tensor ``y``. + + Notes + ----- + All tensors must be contiguous and on the buffer's CUDA device. NVFP4 + supports four or eight ranks, which must call matching families in matching order. + """ + b = symm_buffer + m = int(num_tokens if num_tokens is not None else y.shape[0]) + _validate_call( + y, + l1_weights, + l2_weights, + b, + m, + "nvfp4", + cumulative_local_expert_recv_stats, + l1_global_scales, + l2_global_scales, + ) + if b.group.size() not in (4, 8): + raise ValueError("NVFP4 MegaMoE requires 4 or 8 ranks") + + if kernel_family == "auto": + kernel_family = "fused" if m <= family_threshold else "split" + if kernel_family not in ("fused", "split"): + raise ValueError( + f"kernel_family must be 'auto', 'fused' or 'split', got {kernel_family!r}" + ) + if kernel_family == "split": + return _run_split( + y, + l1_weights, + l2_weights, + b, + m, + cumulative_local_expert_recv_stats, + l1_global_scales, + l2_global_scales, + ) + + plan = select_fused_plan( + b.num_sms, + b.group.size(), + b.num_experts, + b.num_topk, + b.hidden, + b.intermediate_hidden, + b.num_max_tokens_per_rank, + m, + force_push_dispatch=force_push_dispatch, + force_no_clean_barrier=force_no_clean_barrier, + ) + return _run_kernel( + gen_megamoe_sm90_nvfp4_module, + "megamoe_sm90_nvfp4_fused", + plan, + y, + l1_weights, + l2_weights, + b, + m, + cumulative_local_expert_recv_stats, + l1_global_scales, + l2_global_scales, + ) + + +def _validate_call(y, w1, w2, b, m, weight_format, stats, gs1, gs2): + if b.buffer is None: + raise ValueError("symmetric buffer has been destroyed") + if not 0 < m <= b.num_max_tokens_per_rank or m > getattr(b, "num_staged_tokens", 0): + raise ValueError("num_tokens must fit the staged input and buffer capacity") + epr = b.num_experts // b.group.size() + if weight_format == "nvfp4": + shapes = ( + (epr, b.intermediate_hidden * 2, b.hidden // 128 * 80), + (epr, b.hidden, b.intermediate_hidden // 128 * 80), + ) + else: + shapes = ( + ( + epr, + b.intermediate_hidden * 2 // 128 * (b.hidden // 128), + MXFP4_TILE_BYTES, + ), + (epr, b.hidden // 128 * (b.intermediate_hidden // 128), MXFP4_TILE_BYTES), + ) + if ( + y.ndim != 2 + or y.shape[0] < m + or y.shape[1] != b.hidden + or y.dtype != torch.bfloat16 + ): + raise ValueError("y must be BF16 [at least num_tokens, hidden]") + for tensor, shape, dtype in ( + (w1, shapes[0], torch.uint8), + (w2, shapes[1], torch.uint8), + ): + if tuple(tensor.shape) != shape or tensor.dtype != dtype: + raise ValueError( + f"{weight_format} weights must have shape {shape} and dtype {dtype}" + ) + for tensor, dtype in ( + (stats, torch.int32), + (gs1, torch.float32), + (gs2, torch.float32), + ): + if tensor is not None and (tensor.shape != (epr,) or tensor.dtype != dtype): + raise ValueError( + f"optional expert tensor must have shape {(epr,)} and dtype {dtype}" + ) + for tensor in (y, w1, w2, stats, gs1, gs2): + if tensor is not None and ( + tensor.device != b.buffer.device or not tensor.is_contiguous() + ): + raise ValueError( + "all tensors must be contiguous and on the buffer's CUDA device" + ) + + +def mxfp4_mega_moe( + y: torch.Tensor, + l1_weights: torch.Tensor, + l2_weights: torch.Tensor, + symm_buffer: MegaMoESymmBuffer, + num_tokens: Optional[int] = None, + cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, + l1_global_scales: Optional[torch.Tensor] = None, + l2_global_scales: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Run fused MXFP4 MegaMoE with chunk-major weights and FP8 activations. + + Parameters + ---------- + y : torch.Tensor + Contiguous BF16 output with shape ``[at least num_tokens, hidden]``. + The kernel writes the first ``num_tokens`` rows in place. + l1_weights, l2_weights : torch.Tensor + Packed MXFP4 weights from :func:`transform_mxfp4_weights_for_mega_moe`. + symm_buffer : MegaMoESymmBuffer + Buffer populated by :meth:`MegaMoESymmBuffer.stage` on every rank. + num_tokens : int, optional + Tokens on this rank. Defaults to ``y.shape[0]``. + cumulative_local_expert_recv_stats : torch.Tensor, optional + INT32 vector with one entry per local expert. Receive counts are + accumulated in place; initialize the vector before first use. + l1_global_scales, l2_global_scales : torch.Tensor, optional + FP32 vectors with one multiplier per local expert. Omitted scales use 1. + + Returns + ------- + torch.Tensor + The output tensor ``y``. + + Notes + ----- + All tensors must be contiguous and on the buffer's CUDA device. The rank + count must be a power of two from 1 to 32. Every rank must invoke the kernel + in matching order; MXFP4 always uses the fused family. + """ + b = symm_buffer + m = int(num_tokens if num_tokens is not None else y.shape[0]) + _validate_call( + y, + l1_weights, + l2_weights, + b, + m, + "mxfp4", + cumulative_local_expert_recv_stats, + l1_global_scales, + l2_global_scales, + ) + plan = select_mxfp4_fused_plan( + b.num_sms, + b.group.size(), + b.num_experts, + b.num_topk, + b.hidden, + b.intermediate_hidden, + b.num_max_tokens_per_rank, + m, + ) + return _run_kernel( + gen_megamoe_sm90_mxfp4_module, + "megamoe_sm90_mxfp4_fused", + plan, + y, + l1_weights, + l2_weights, + b, + m, + cumulative_local_expert_recv_stats, + l1_global_scales, + l2_global_scales, + ) + + +# --------------------------------------------------------------------------- +# Split (BN128) family -- selected above the fused threshold +# --------------------------------------------------------------------------- + +#: Default token threshold for selecting NVFP4 split instead of fused. +FAMILY_THRESHOLD = 256 + + +@dataclasses.dataclass(frozen=True) +class SplitPlan: + l1_num_stages: int + l1_smem_size: int + l2_num_stages: int + l2_smem_size: int + num_experts_per_wave: int + num_max_pool_tokens: int + dispatch_dequant: bool + l2_arrival_counter: bool + num_padded_sf_pool_tokens: int + + +@functools.cache +def select_split_plan( + num_sms, + num_ranks, + num_experts, + num_topk, + hidden, + intermediate_hidden, + num_max_tokens_per_rank, + num_tokens, +) -> SplitPlan: + """Ask the plan module for the L1/L2 pipeline configuration. + + The sizing arithmetic mirrors the kernel's shared-memory layout exactly, so + it is evaluated by the same C++ code the kernel is built from rather than + being re-derived here. + """ + sf_pool = max( + get_num_padded_sf_pool_tokens( + get_num_max_pool_tokens( + num_ranks, num_max_tokens_per_rank, num_topk, num_experts // num_ranks + ), + bm, + ) + for bm in CANDIDATE_BLOCK_M + ) + mod = _get_jit_module( + gen_megamoe_sm90_nvfp4_split_plan_module, + num_sms, + num_ranks, + num_max_tokens_per_rank, + num_experts, + num_topk, + hidden, + intermediate_hidden, + ) + v = list( + mod.megamoe_sm90_nvfp4_split_plan( + num_sms, + num_ranks, + num_experts, + num_max_tokens_per_rank, + num_tokens, + num_topk, + hidden, + intermediate_hidden, + sf_pool, + ) + ) + return SplitPlan( + l1_num_stages=v[0], + l1_smem_size=v[1], + l2_num_stages=v[2], + l2_smem_size=v[3], + num_experts_per_wave=v[4], + num_max_pool_tokens=v[5], + dispatch_dequant=bool(v[6]), + l2_arrival_counter=bool(v[7]), + num_padded_sf_pool_tokens=sf_pool, + ) + + +def _run_split(y, l1_weights, l2_weights, b, m, stats, gs1, gs2): + plan = select_split_plan( + b.num_sms, + b.group.size(), + b.num_experts, + b.num_topk, + b.hidden, + b.intermediate_hidden, + b.num_max_tokens_per_rank, + m, + ) + return _run_kernel( + gen_megamoe_sm90_nvfp4_split_module, + "megamoe_sm90_nvfp4_split", + plan, + y, + l1_weights, + l2_weights, + b, + m, + stats, + gs1, + gs2, + ) diff --git a/flashinfer/megamoe/heuristics.py b/flashinfer/megamoe/heuristics.py new file mode 100644 index 00000000000..a3567058cd4 --- /dev/null +++ b/flashinfer/megamoe/heuristics.py @@ -0,0 +1,367 @@ +"""Compile-time plans for the SM90 NVFP4 and MXFP4 MegaMoE kernels. + +Every field returned here becomes a template parameter of the CUDA kernel, so +the plan must be resolved on the Python side before the JIT module is built +(FlashInfer renders them into ``megamoe_sm90_nvfp4_config.inc`` via Jinja). + +Tuning buckets select tile shapes and pipeline depths for H20 and H200. +Shapes outside those buckets use the generic scheduling policy. +""" + +from dataclasses import dataclass +from typing import Optional, Tuple + +# SM90 shared-memory capacity, and the mailbox the interleaved scheduler +# appends after the dev-m barriers. +SMEM_CAPACITY = 232448 +INTERLEAVED_SCHEDULER_SMEM_BYTES = 96 + +H20_NUM_SMS = 78 +H200_NUM_SMS = 132 + +# Token-pool block-size candidates (the pool is padded so that any of these +# BLOCK_M values can address it). +CANDIDATE_BLOCK_M = (8, 16, 32, 64, 96, 128, 192) +MIN_CANDIDATE_BLOCK_M = 8 +MAX_CANDIDATE_BLOCK_M = 192 +LCM_CANDIDATE_BLOCK_M = 384 + +# (hidden, min_tokens, max_tokens, block_m, experts_per_wave, stages, +# swap_ab, mode2_row_decoder, single_active_dispatch_warp) +_H20_FUSED_BUCKETS = ( + (4096, 1, 8, 8, 16, 6, True, True, True), + (4096, 9, 16, 8, 32, 6, True, True, True), + (4096, 17, 64, 24, 32, 6, True, True, True), + (4096, 65, 128, 24, 16, 6, True, True, True), + (7168, 1, 8, 8, 48, 6, True, True, True), + (7168, 9, 16, 8, 24, 6, True, True, True), + (7168, 17, 64, 24, 48, 6, True, True, True), + (7168, 65, 191, 24, 16, 6, True, True, True), + (6144, 1, 8, 8, 48, 6, True, True, True), + (6144, 9, 16, 8, 24, 6, True, True, True), + (6144, 17, 64, 24, 48, 6, True, True, True), + (6144, 65, 144, 24, 16, 6, True, True, True), +) + +# Counter-based dispatch policy: push dispatch replaces the first NVLink +# barrier, and the parity-rotated pool additionally drops the workspace-clean +# barrier. A step is enabled only in the M range where it was measured to win +# on the random router without losing on the all-experts router. +# (hidden, num_experts, num_topk, min_tokens, max_tokens, push, rotated_pool) +_H20_COUNTER_BUCKETS = ( + (4096, 256, 6, 1, 16, True, False), + (4096, 256, 6, 17, 64, True, True), + (7168, 384, 6, 1, 16, True, True), + (6144, 384, 8, 1, 16, True, True), + (6144, 384, 8, 17, 32, True, False), +) + +# All-M arm table: which fused schedule to use. Every arm here runs the same +# fused kernel; they differ in scheduler mode and the RS epilogue. +# "dynamic-rs" -> interleaved scheduler + RS swap-AB (mode5) +# "devm-dynamic"-> interleaved scheduler, no RS +# "static-ss" -> static scheduler, no RS +# (num_sms, hidden, num_experts, num_topk, min_tokens, max_tokens, arm) +_ALLM_ARMS = ( + (78, 4096, 256, 6, 1, 128, "dynamic-rs"), + (78, 4096, 256, 6, 129, 1 << 30, "devm-dynamic"), + (78, 7168, 384, 6, 1, 191, "dynamic-rs"), + (78, 7168, 384, 6, 192, 1 << 30, "devm-dynamic"), + (78, 6144, 384, 8, 1, 144, "dynamic-rs"), + (78, 6144, 384, 8, 145, 1 << 30, "devm-dynamic"), + (132, 4096, 256, 6, 1, 35, "dynamic-rs"), + (132, 4096, 256, 6, 36, 64, "static-ss"), + (132, 4096, 256, 6, 65, 1 << 30, "devm-dynamic"), + (132, 7168, 384, 6, 1, 128, "dynamic-rs"), + (132, 7168, 384, 6, 129, 1 << 30, "devm-dynamic"), + (132, 6144, 384, 8, 1, 16, "dynamic-rs"), + (132, 6144, 384, 8, 17, 32, "devm-dynamic"), + (132, 6144, 384, 8, 33, 96, "dynamic-rs"), + (132, 6144, 384, 8, 97, 1 << 30, "devm-dynamic"), +) + + +def align(x: int, a: int) -> int: + return (x + a - 1) // a * a + + +def get_num_max_pool_tokens( + num_ranks: int, + num_max_tokens_per_rank: int, + num_topk: int, + num_experts_per_rank: int, +) -> int: + """Worst-case token-pool capacity, padded so every BLOCK_M candidate fits.""" + num_max_recv = num_ranks * num_max_tokens_per_rank + max_experts_per_token = min(num_topk, num_experts_per_rank) + return align( + num_max_recv * max_experts_per_token + + num_experts_per_rank * (MAX_CANDIDATE_BLOCK_M - 1), + LCM_CANDIDATE_BLOCK_M, + ) + + +def get_num_padded_sf_pool_tokens(num_max_pool_tokens: int, block_m: int) -> int: + return (num_max_pool_tokens // block_m) * align(block_m, 128) + + +@dataclass(frozen=True) +class FusedPlan: + block_m: int + block_n: int + num_experts_per_wave: int + num_stages: int + smem_size: int + swap_ab: bool + use_mode2_row_decoder: bool + single_active_dispatch_warp: bool + use_interleaved_scheduler: bool + rs_swap_ab: bool + push_dispatch: bool + no_clean_barrier: bool + num_max_pool_tokens: int + num_padded_sf_pool_tokens: int + + +def select_mxfp4_fused_plan( + num_sms: int, + num_ranks: int, + num_experts: int, + num_topk: int, + hidden: int, + intermediate_hidden: int, + num_max_tokens_per_rank: int, + num_tokens: int, +) -> FusedPlan: + """Select MXFP4 tile geometry, pipeline depth and decode-warp partition. + + Chunk-major weight tiles are independent of the selected BLOCK_N. + """ + if num_sms <= 1: + raise ValueError("num_sms must be greater than one") + if num_ranks <= 0 or num_ranks > 32 or num_ranks & (num_ranks - 1): + raise ValueError("num_ranks must be a power of two in [1, 32]") + if num_experts <= 0 or num_experts % num_ranks: + raise ValueError("num_experts must divide evenly across ranks") + if not 0 < num_topk <= min(32, num_experts): + raise ValueError("num_topk must be in [1, min(32, num_experts)]") + if ( + hidden <= 0 + or hidden % 256 + or intermediate_hidden <= 0 + or intermediate_hidden % 128 + ): + raise ValueError( + "hidden must be a multiple of 256 and intermediate_hidden of 128" + ) + if not 0 < num_tokens <= num_max_tokens_per_rank: + raise ValueError("num_tokens must be in (0, num_max_tokens_per_rank]") + epr = num_experts // num_ranks + + def max_tokens(bm): + return bm * num_experts // (num_ranks * num_topk) + + limit = max_tokens(24) * H20_NUM_SMS // num_sms + bm = next((b for b in (8, 24, 32) if num_tokens <= max_tokens(b) * 7 // 10), 32) + if num_tokens <= 1: + epw, stages = 24, 8 + elif num_tokens <= 8: + epw, stages = 16, 8 + elif num_tokens <= 16: + epw, stages = 24, 8 + else: + epw, stages = 48, 8 + swap = num_tokens <= max(16, limit) + rs = swap + if num_sms > H20_NUM_SMS and bm < 16 and num_tokens <= 16: + rs = False + if swap and bm >= 16: + rs = False + if not swap: + bm, stages = 64, 3 + elif not rs: + stages = 4 if bm == 8 else 3 + epw = min(epw, epr) + while epr % epw: + epw -= 1 + pool = get_num_max_pool_tokens(num_ranks, num_max_tokens_per_rank, num_topk, epr) + return FusedPlan( + block_m=bm, + block_n=256, + num_experts_per_wave=epw, + num_stages=stages, + smem_size=SMEM_CAPACITY, + swap_ab=swap, + use_mode2_row_decoder=True, + single_active_dispatch_warp=swap, + use_interleaved_scheduler=True, + rs_swap_ab=rs, + push_dispatch=False, + no_clean_barrier=False, + num_max_pool_tokens=pool, + num_padded_sf_pool_tokens=max( + get_num_padded_sf_pool_tokens(pool, b) for b in CANDIDATE_BLOCK_M + ), + ) + + +def _select_arm( + num_sms: int, hidden: int, num_experts: int, num_topk: int, m: int +) -> str: + for sms, h, e, k, lo, hi, arm in _ALLM_ARMS: + if ( + sms == num_sms + and h == hidden + and e == num_experts + and k == num_topk + and lo <= m <= hi + ): + return arm + return "devm-dynamic" + + +def _select_counter_policy( + num_sms: int, hidden: int, num_experts: int, num_topk: int, m: int +) -> Tuple[bool, bool]: + """(push_dispatch, no_clean_barrier). H200 is unmeasured and keeps barriers. + + Matching on (hidden, num_experts, num_topk) rather than hidden alone: the + payoff depends on the pushed row count (m * num_topk) and the per-expert + pool pressure, so a model that merely shares a hidden size must not inherit + a measured range. Unmatched shapes keep the always-correct barrier path. + """ + if num_sms != H20_NUM_SMS: + return False, False + for h, e, k, lo, hi, push, rotated in _H20_COUNTER_BUCKETS: + if h == hidden and e == num_experts and k == num_topk and lo <= m <= hi: + return push, push and rotated + return False, False + + +def select_fused_plan( + num_sms: int, + num_ranks: int, + num_experts: int, + num_topk: int, + hidden: int, + intermediate_hidden: int, + num_max_tokens_per_rank: int, + num_tokens: int, + force_push_dispatch: Optional[bool] = None, + force_no_clean_barrier: Optional[bool] = None, +) -> FusedPlan: + if num_sms not in (H20_NUM_SMS, H200_NUM_SMS): + raise ValueError(f"MegaMoE supports SM90 H20/H200 only, got num_sms={num_sms}") + if num_ranks not in (4, 8): + raise ValueError("NVFP4 MegaMoE requires 4 or 8 ranks") + num_experts_per_rank = num_experts // num_ranks + if num_experts_per_rank * num_ranks != num_experts: + raise ValueError("num_experts must be divisible by num_ranks") + supported_experts_per_rank = (32, 48) if num_ranks == 8 else (64, 96) + if num_experts_per_rank not in supported_experts_per_rank: + raise ValueError(f"unsupported experts per rank: {num_experts_per_rank}") + if not 0 < num_tokens <= num_max_tokens_per_rank: + raise ValueError("num_tokens must be in (0, num_max_tokens_per_rank]") + + # (block_m, block_n, experts per wave, stages, swap_ab, mode2, single warp) + tuning: Optional[Tuple[int, int, int, int, bool, bool, bool]] = None + if num_ranks == 8 and num_sms == H20_NUM_SMS: + for h, lo, hi, bm, epw, stages, swap, mode2, single in _H20_FUSED_BUCKETS: + if h == hidden and lo <= num_tokens <= hi: + tuning = (bm, 256, epw, stages, swap, mode2, single) + break + if tuning is None: + # Generic dev-m schedule: correct for any geometry, not individually tuned. + if num_tokens <= 1: + tuning = (8, 256, 24, 4, True, True, True) + elif num_tokens <= 8: + tuning = (8, 256, 16, 4, True, True, True) + elif num_tokens <= 16: + tuning = (8, 256, 24, 4, True, True, True) + elif num_tokens <= 32: + tuning = (16, 256, 48, 3, True, True, False) + elif num_tokens <= 64: + tuning = (24, 256, 48, 3, True, False, True) + elif num_tokens <= 256: + tuning = (64, 256, 48, 3, False, True, False) + else: + tuning = (128, 128, 48, 6, False, True, False) + + block_m, block_n, epw, stages, swap_ab, mode2, single_warp = tuning + + # An expert wave must divide the per-rank expert count. + epw = min(epw, num_experts_per_rank) + while epw < num_experts_per_rank and num_experts_per_rank % epw != 0: + epw += 1 + + # Pro at BLOCK_M=8 does not fit 4 stages. + is_pro = ( + num_experts == 384 + and num_topk == 6 + and hidden == 7168 + and intermediate_hidden == 3072 + ) + if is_pro and block_m == 8 and stages == 4: + stages = 3 + + # Deployment arm and counter ranges were measured at EP8. + arm = ( + _select_arm(num_sms, hidden, num_experts, num_topk, num_tokens) + if num_ranks == 8 + else "devm-dynamic" + ) + use_interleaved = arm != "static-ss" + rs_swap_ab = arm == "dynamic-rs" and swap_ab + + push, no_clean = ( + _select_counter_policy(num_sms, hidden, num_experts, num_topk, num_tokens) + if num_ranks == 8 + else (False, False) + ) + if force_push_dispatch is not None: + push = force_push_dispatch + if force_no_clean_barrier is not None: + no_clean = force_no_clean_barrier + push = push and use_interleaved + no_clean = no_clean and push + + num_max_pool_tokens = get_num_max_pool_tokens( + num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank + ) + # Push dispatch addresses the pool with a fixed stride per local expert, so + # the worst case (every routed row of every rank landing on one expert) + # must fit; otherwise fall back to the pull path. + if push: + pool_blocks = num_max_pool_tokens // block_m + slots = 2 if no_clean else 1 + blocks_per_expert = pool_blocks // (num_experts_per_rank * slots) + if num_tokens * num_ranks > blocks_per_expert * block_m: + push = False + no_clean = False + + smem_size = SMEM_CAPACITY + if use_interleaved: + smem_size = min(SMEM_CAPACITY + INTERLEAVED_SCHEDULER_SMEM_BYTES, SMEM_CAPACITY) + + return FusedPlan( + block_m=block_m, + block_n=block_n, + num_experts_per_wave=epw, + num_stages=stages, + smem_size=smem_size, + swap_ab=swap_ab, + use_mode2_row_decoder=rs_swap_ab or mode2, + single_active_dispatch_warp=single_warp, + use_interleaved_scheduler=use_interleaved, + rs_swap_ab=rs_swap_ab, + push_dispatch=push, + no_clean_barrier=no_clean, + num_max_pool_tokens=num_max_pool_tokens, + # The SF region is sized once, when the symmetric buffer is allocated, so + # it must hold the largest layout any BLOCK_M candidate can produce -- + # not the one this plan happens to use. + num_padded_sf_pool_tokens=max( + get_num_padded_sf_pool_tokens(num_max_pool_tokens, bm) + for bm in CANDIDATE_BLOCK_M + ), + ) diff --git a/flashinfer/megamoe/weights.py b/flashinfer/megamoe/weights.py new file mode 100644 index 00000000000..0e1ebb2184a --- /dev/null +++ b/flashinfer/megamoe/weights.py @@ -0,0 +1,134 @@ +"""Offline FP4 weight preparation for SM90 MegaMoE. + +Input weights are ``(packed, scales)`` in row-major ``[experts, N, K/2]`` +and ``[experts, N, K/group_size]``. Packed bytes follow Marlin order: byte b in +each four-byte word stores K[b] high and K[b+4] low. +NVFP4 scales are raw unsigned E4M3 bytes per 16 elements; MXFP4 scales are +unsigned E8M0 bytes per 32 elements. L1 inputs have gate rows followed by up. +""" + +from typing import Tuple + +import torch + +MXFP4_TILE_ROWS = 128 +MXFP4_TILE_BYTES = MXFP4_TILE_ROWS * 68 + + +def _interleave_gate_up(t: torch.Tensor) -> torch.Tensor: + e, n, *rest = t.shape + gate, up = t[:, : n // 2], t[:, n // 2 :] + return torch.stack( + (gate.reshape(e, n // 16, 8, *rest), up.reshape(e, n // 16, 8, *rest)), + dim=2, + ).reshape(e, n, *rest) + + +def _braid_signs(rows: torch.Tensor) -> torch.Tensor: + # Magnitude and sign positions are independent in the mode2 decoder. + # Braid signs before arranging the format-specific row or tile layout. + packed = rows[..., :64].reshape(*rows.shape[:-1], 16, 4) + codes = torch.cat(((packed >> 4) & 15, packed & 15), dim=-1) + signs = (codes >> 3)[..., [4, 0, 5, 1, 6, 2, 7, 3]] + braided = (codes & 7) | (signs << 3) + rows[..., :64] = (braided[..., 0::2] | (braided[..., 1::2] << 4)).reshape( + *rows.shape[:-1], 64 + ) + return rows + + +def _transform_one(weights, group_size, interleave, decode_tiles): + packed, scales = weights + if packed.dtype != torch.uint8 or scales.dtype != torch.uint8: + raise ValueError("packed weights and raw scale bytes must have dtype uint8") + if packed.ndim != 3 or scales.ndim != 3: + raise ValueError("packed weights and scales must be 3-D") + e, n, half_k = packed.shape + k = half_k * 2 + if n % 128 or k % 128 or tuple(scales.shape) != (e, n, k // group_size): + raise ValueError( + "weights must contain complete N128/K128 tiles and matching scales" + ) + if packed.device != scales.device: + raise ValueError("weights and scales must be on the same device") + if interleave: + packed, scales = _interleave_gate_up(packed), _interleave_gate_up(scales) + rows = torch.zeros((e, n, k // 128, 80), dtype=torch.uint8, device=packed.device) + rows[..., :64] = packed.reshape(e, n, k // 128, 64) + rows[..., 64 : 64 + 128 // group_size] = scales.reshape( + e, n, k // 128, 128 // group_size + ) + rows = _braid_signs(rows) + if not decode_tiles: + return rows.reshape(e, n, k // 128 * 80).contiguous() + # Chunk-major N128 tiles omit 12 padding bytes per row and keep each + # decoder's 16-byte loads aligned. Tile index = n_tile * k_blocks + k. + nt, kb = n // MXFP4_TILE_ROWS, k // 128 + tiled = rows.reshape(e, nt, MXFP4_TILE_ROWS, kb, 80).permute(0, 1, 3, 2, 4) + fp4 = tiled[..., :64].reshape(e, nt, kb, MXFP4_TILE_ROWS, 4, 16) + fp4 = fp4.permute(0, 1, 2, 4, 3, 5).reshape(e, nt * kb, MXFP4_TILE_ROWS * 64) + sf = tiled[..., 64:68].reshape(e, nt * kb, MXFP4_TILE_ROWS * 4) + return torch.cat((fp4, sf), dim=-1).contiguous() + + +def _transform(l1_weights, l2_weights, group_size, decode_tiles): + p1, _ = l1_weights + p2, _ = l2_weights + if p1.ndim != 3 or p2.ndim != 3 or p1.shape[0] != p2.shape[0]: + raise ValueError("L1 and L2 must contain the same expert count") + if p1.shape[1] != p2.shape[2] * 4 or p1.shape[2] * 2 != p2.shape[1]: + raise ValueError( + "expected L1 [E, 2*intermediate, hidden/2] and L2 [E, hidden, intermediate/2]" + ) + return ( + _transform_one(l1_weights, group_size, True, decode_tiles), + _transform_one(l2_weights, group_size, False, decode_tiles), + ) + + +def transform_nvfp4_weights_for_mega_moe( + l1_weights: Tuple[torch.Tensor, torch.Tensor], + l2_weights: Tuple[torch.Tensor, torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + """Return braided NVFP4 weights for both fused and split SM90 kernels. + + Parameters + ---------- + l1_weights : tuple[torch.Tensor, torch.Tensor] + Packed UINT8 weights of shape ``[E, 2 * intermediate, hidden // 2]`` + and raw UINT8 E4M3 scale bytes per 16 values. Gate rows precede up rows. + l2_weights : tuple[torch.Tensor, torch.Tensor] + Packed UINT8 weights of shape ``[E, hidden, intermediate // 2]`` + and raw UINT8 E4M3 scale bytes per 16 values. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + Contiguous L1 and L2 weight tensors in the kernel's packed layout, + on the input device. Weight preparation runs before graph capture. + """ + return _transform(l1_weights, l2_weights, 16, False) + + +def transform_mxfp4_weights_for_mega_moe( + l1_weights: Tuple[torch.Tensor, torch.Tensor], + l2_weights: Tuple[torch.Tensor, torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + """Return chunk-major MXFP4 weights for the SM90 fused kernel. + + Parameters + ---------- + l1_weights : tuple[torch.Tensor, torch.Tensor] + Packed UINT8 weights of shape ``[E, 2 * intermediate, hidden // 2]`` + and raw UINT8 E8M0 scale bytes per 32 values. Gate rows precede up rows. + l2_weights : tuple[torch.Tensor, torch.Tensor] + Packed UINT8 weights of shape ``[E, hidden, intermediate // 2]`` + and raw UINT8 E8M0 scale bytes per 32 values. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + Contiguous L1 and L2 weight tensors in the kernel's packed layout, + on the input device. Weight preparation runs before graph capture. + """ + return _transform(l1_weights, l2_weights, 32, True) diff --git a/include/flashinfer/megamoe/LICENSE.deep_gemm b/include/flashinfer/megamoe/LICENSE.deep_gemm new file mode 100644 index 00000000000..5c48bdc9f6e --- /dev/null +++ b/include/flashinfer/megamoe/LICENSE.deep_gemm @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 DeepSeek + +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. diff --git a/include/flashinfer/megamoe/barrier.cuh b/include/flashinfer/megamoe/barrier.cuh new file mode 100644 index 00000000000..b2ce46e75ed --- /dev/null +++ b/include/flashinfer/megamoe/barrier.cuh @@ -0,0 +1,87 @@ +#pragma once + +#include + +#include +#include +#include + +namespace flashinfer::megamoe::comm { + +CUTLASS_DEVICE void cluster_sync_with_relaxed_arrive() { + // Perform cluster_sync with `barrier.cluster.arrive.relaxed` + // This is slightly faster than `cute::cluster_sync` but has weaker memory ordering guarantee + cute::cluster_arrive_relaxed(); + cute::cluster_wait(); +} + +template +CUTLASS_DEVICE void grid_sync(const layout::Workspace& workspace, const uint32_t& sm_idx, + const uint32_t& thread_idx, const sync_scope_t& sync_scope) { + // NOTES: the implementation idea is from `cooperative_groups::this_grid().sync()` + static constexpr uint32_t kFinishSumTag = 0x80000000u; + sync_scope(); + if (thread_idx == 0) { + const auto count_ptr = workspace.get_grid_sync_count_ptr(); + const auto old_value = + ptx::atomic_add_rel(count_ptr, sm_idx == 0 ? (kFinishSumTag - (kNumSMs - 1)) : 1); + uint32_t new_value; + do { + new_value = ptx::ld_acq(count_ptr); + } while (((new_value ^ old_value) & kFinishSumTag) == 0); + } + sync_scope(); +} + +template +CUTLASS_DEVICE void nvlink_barrier(const layout::Workspace& workspace, + const layout::SymBuffer& sym_buffer, + const uint32_t& sm_idx, const uint32_t& thread_idx, + const sync_scope_t& sync_scope, const bool& sync_prologue = true, + const bool& sync_epilogue = true) { + FI_MEGAMOE_STATIC_ASSERT(kNumRanks <= kNumThreads, "Insufficient threads"); + + // Grid sync before NVLink signaling + if (sync_prologue) grid_sync(workspace, sm_idx, thread_idx, sync_scope); + + // NVLink cross-rank barrier, only SM 0 participates + if (sm_idx == 0) { + auto* counter_ptr = workspace.get_nvl_barrier_counter_ptr(); + const auto status = (*counter_ptr) & 3; + const auto signal_phase = status & 1, signal_sign = status >> 1; + auto* signal_ptr = workspace.get_nvl_barrier_signal_ptr(signal_phase); + + // Send signals to remote ranks + if (thread_idx < kNumRanks) + ptx::red_add_rel_sys(sym_buffer.map(signal_ptr, thread_idx), signal_sign ? -1 : 1); + sync_scope(); + + // Update status and wait arrival (with 30s timeout, at 2 GHz) + constexpr int64_t kNumTimeoutCycles = 30ll * 2000000000ll; + if (thread_idx == 0) { + ptx::red_add(counter_ptr, 1); + const int target = signal_sign ? 0 : static_cast(kNumRanks); + const auto start_clock = clock64(); + while (ptx::ld_acq_sys(signal_ptr) != target) { + if (clock64() - start_clock >= kNumTimeoutCycles) { +#if defined(FI_MEGAMOE_NVLINK_BARRIER_TRAP_ONLY_TIMEOUT) + FI_MEGAMOE_TRAP_ONLY_DEVICE_ASSERT(false and "NVLink barrier timeout"); +#else + printf( + "DeepGEMM NVLink barrier timeout (30s): rank=%d, counter=%d, signal=%d, target=%d, " + "phase=%d, sign=%d, tag=%d\n", + sym_buffer.rank_idx, *counter_ptr, ptx::ld_acq_sys(signal_ptr), target, signal_phase, + signal_sign, kTag); + FI_MEGAMOE_DEVICE_ASSERT(false and "NVLink barrier timeout"); +#endif + } + } + } + } + + // Grid sync after NVLink completion + if (sync_epilogue) grid_sync(workspace, sm_idx, thread_idx, sync_scope); +} + +} // namespace flashinfer::megamoe::comm diff --git a/include/flashinfer/megamoe/detail/buffers.inl b/include/flashinfer/megamoe/detail/buffers.inl new file mode 100644 index 00000000000..052df7f2ff4 --- /dev/null +++ b/include/flashinfer/megamoe/detail/buffers.inl @@ -0,0 +1,16 @@ +// Shared host/device buffer ABI; shape constants remain visible to SMEM sizing. +constexpr auto fp8_token_layout = layout::Data(kHidden); +const auto buffers = layout::MegaMoEBuffers( + sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk, kHidden, + kIntermediateHidden, kNumMaxPoolTokens, kNumPaddedSFPoolTokens); +const auto workspace = buffers.workspace; +const auto input_token_buffer = buffers.input_token_buffer; +const auto input_sf_buffer = buffers.input_sf_buffer; +const auto input_topk_idx_buffer = buffers.input_topk_idx_buffer; +const auto input_topk_weights_buffer = buffers.input_topk_weights_buffer; +const auto l1_token_buffer = buffers.l1_token_buffer; +const auto l1_sf_buffer = buffers.l1_sf_buffer; +const auto l1_topk_weights_buffer = buffers.l1_topk_weights_buffer; +const auto l2_token_buffer = buffers.l2_token_buffer; +const auto l2_sf_buffer = buffers.l2_sf_buffer; +const auto combine_token_buffer = buffers.combine_token_buffer; diff --git a/include/flashinfer/megamoe/detail/combine.inl b/include/flashinfer/megamoe/detail/combine.inl new file mode 100644 index 00000000000..b5d7f33e754 --- /dev/null +++ b/include/flashinfer/megamoe/detail/combine.inl @@ -0,0 +1,90 @@ +// Shared top-k reduction: all three SM90 FP4 MegaMoE paths. +const auto combine_load_buffer = utils::PatternVisitor([&](const uint32_t& i) { + return math::advance_ptr(smem_buffer, + (epilogue_warp_idx + i * kNumEpilogueWarps) * kNumChunkBytes); +}); +const auto combine_store_buffer = math::advance_ptr( + smem_buffer, (epilogue_warp_idx + kNumEpilogueWarps * 2) * kNumChunkBytes); + +auto combine_load_barriers = utils::PatternVisitor( + [&](const uint32_t& i) { return combine_barriers[i + epilogue_warp_idx * 2]; }); + +uint32_t combine_phase = 0; +uint32_t load_stage_idx = 0; +for (uint32_t token_idx = sm_idx * kNumEpilogueWarps + epilogue_warp_idx; token_idx < num_tokens; + token_idx += kNumSMs * kNumEpilogueWarps) { + const int stored_topk_slot_idx = + lane_idx < kNumTopk ? static_cast(__ldg(input_topk_idx_buffer.get_base_ptr() + + token_idx * kNumTopk + lane_idx)) + : -1; + const uint32_t total_mask = __ballot_sync(0xffffffff, stored_topk_slot_idx >= 0); + + for (uint32_t chunk = 0; chunk < kNumChunks; ++chunk) { + const uint32_t chunk_byte_offset = chunk * kNumChunkBytes; + + uint32_t mask = total_mask; + const auto move_mask_and_load = [&](const uint32_t& i) { + if (mask) { + const uint32_t slot_idx = __ffs(mask) - 1; + mask ^= 1 << slot_idx; + if (cute::elect_one_sync()) { + const auto src_ptr = + math::advance_ptr(combine_token_buffer.get_rank_buffer(slot_idx) + .get_data_buffer(token_idx) + .get_base_ptr(), + chunk_byte_offset); + ptx::tma_load_1d(combine_load_buffer[i], src_ptr, combine_load_barriers[i], + kNumChunkBytes); + ptx::mbarrier_arrive_and_set_tx(combine_load_barriers[i], kNumChunkBytes); + } + __syncwarp(); + return true; + } + return false; + }; + + bool do_reduce = move_mask_and_load(load_stage_idx); + + float2 reduced[kNumUint4PerLane * kNumElemsPerUint4] = {}; + while (do_reduce) { + do_reduce = move_mask_and_load(load_stage_idx ^ 1); + combine_load_barriers[load_stage_idx]->wait(combine_phase); +#pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++j) { + const auto uint4_values = combine_load_buffer[load_stage_idx][j * 32 + lane_idx]; + const auto bf16_values = reinterpret_cast(&uint4_values); +#pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++l) + ptx::accumulate(reduced[j * kNumElemsPerUint4 + l], bf16_values[l]); + } + combine_phase ^= load_stage_idx; + load_stage_idx ^= 1; + } + +#pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++j) { + uint4 casted; + auto casted_bf16 = reinterpret_cast(&casted); +#pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++l) + casted_bf16[l] = __float22bfloat162_rn(reduced[j * kNumElemsPerUint4 + l]); + + if (j == 0) { + ptx::tma_store_wait<0>(); + __syncwarp(); + } + ptx::st_shared(combine_store_buffer + j * 32 + lane_idx, casted.x, casted.y, casted.z, + casted.w); + } + __syncwarp(); + + if (cute::elect_one_sync()) { + cute::tma_store_fence(); + ptx::tma_store_1d(math::advance_ptr(y, static_cast(token_idx) * kNumHiddenBytes + + chunk_byte_offset), + combine_store_buffer, kNumChunkBytes); + cute::tma_store_arrive(); + } + __syncwarp(); + } +} diff --git a/include/flashinfer/megamoe/detail/combine_config.inl b/include/flashinfer/megamoe/detail/combine_config.inl new file mode 100644 index 00000000000..f13eea5cb9b --- /dev/null +++ b/include/flashinfer/megamoe/detail/combine_config.inl @@ -0,0 +1,24 @@ +// Shared combine sizing for the fused FP4 paths. No arithmetic policy changes. +constexpr uint32_t kNumHiddenBytes = kHidden * sizeof(nv_bfloat16); +constexpr uint32_t kNumElemsPerUint4 = sizeof(uint4) / sizeof(nv_bfloat162); + +constexpr uint32_t kNumChunkSlots = 3; +constexpr uint32_t kNumMaxRegistersForBuffer = 128; +constexpr uint32_t kNumChunks = + (kNumChunkSlots * kNumEpilogueWarps * kNumHiddenBytes <= SMEM_BEFORE_BARRIER_SIZE and + kHidden <= 32 * kNumMaxRegistersForBuffer) + ? 1 + : 2; +constexpr uint32_t kNumChunkBytes = kNumHiddenBytes / kNumChunks; +constexpr uint32_t kNumChunkUint4 = kNumChunkBytes / sizeof(uint4); +constexpr uint32_t kNumUint4PerLane = kNumChunkUint4 / 32; +FI_MEGAMOE_STATIC_ASSERT(kHidden % kNumChunks == 0, "Hidden must be divisible by number of chunks"); +FI_MEGAMOE_STATIC_ASSERT(kNumChunkSlots* kNumEpilogueWarps* kNumHiddenBytes / kNumChunks <= + SMEM_BEFORE_BARRIER_SIZE, + "Hidden is too large"); +FI_MEGAMOE_STATIC_ASSERT(kNumChunkBytes % 16 == 0, "Combine chunk must be TMA-aligned (16 bytes)"); +FI_MEGAMOE_STATIC_ASSERT(kNumChunkBytes % sizeof(uint4) == 0, + "Combine chunk must be divisible by 16 bytes"); +FI_MEGAMOE_STATIC_ASSERT(kNumChunkUint4 % 32 == 0, + "Combine chunk must be a multiple of 32 16-byte elements"); +FI_MEGAMOE_STATIC_ASSERT(kNumTopk <= 32, "Top-k must fit in a single warp"); diff --git a/include/flashinfer/megamoe/detail/common.cuh b/include/flashinfer/megamoe/detail/common.cuh new file mode 100644 index 00000000000..eaddc551b67 --- /dev/null +++ b/include/flashinfer/megamoe/detail/common.cuh @@ -0,0 +1,332 @@ +#pragma once + +#include +#include +#include +#include + +// ============ compile.cuh ============ +#if defined(__NVCC__) or (defined(__clang__) and defined(__CUDA__)) or defined(__CUDACC_RTC__) or \ + defined(__CLION_IDE__) +#define FI_MEGAMOE_IN_CUDA_COMPILATION +#endif + +#if defined(__NVCC__) || (defined(__clang__) and defined(__CUDA__)) +#define CUTLASS_HOST_DEVICE_NOINLINE __device__ __host__ +#define CUTLASS_DEVICE_NOINLINE __device__ +#elif defined(__CUDACC_RTC__) +#define CUTLASS_HOST_DEVICE_NOINLINE __device__ +#define CUTLASS_DEVICE_NOINLINE __device__ +#else +#define CUTLASS_HOST_DEVICE_NOINLINE +#define CUTLASS_DEVICE_NOINLINE +#endif + +// ============ exception.cuh ============ +// Assertions live in CUTLASS_HOST_DEVICE helpers (workspace sizing is shared by +// the launcher and the kernel), so the failure path must also be legal in a host +// pass -- a "trap;" inline asm is not. +#ifndef FI_MEGAMOE_TRAP +#if defined(__CUDA_ARCH__) +#define FI_MEGAMOE_TRAP() asm("trap;") +#else +#define FI_MEGAMOE_TRAP() __builtin_trap() +#endif +#endif + +#ifdef __CLION_IDE__ + +CUTLASS_HOST_DEVICE void host_device_printf(const char* format, ...) { FI_MEGAMOE_TRAP(); } + +#define printf host_device_printf +#endif + +#ifndef FI_MEGAMOE_DEVICE_ASSERT +#define FI_MEGAMOE_DEVICE_ASSERT(cond) \ + do { \ + if (not(cond)) { \ + printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, __LINE__, #cond); \ + FI_MEGAMOE_TRAP(); \ + } \ + } while (0) +#endif + +#ifndef FI_MEGAMOE_TRAP_ONLY_DEVICE_ASSERT +#define FI_MEGAMOE_TRAP_ONLY_DEVICE_ASSERT(cond) \ + do { \ + if (not(cond)) FI_MEGAMOE_TRAP(); \ + } while (0) +#endif + +#ifndef FI_MEGAMOE_STATIC_ASSERT +#define FI_MEGAMOE_STATIC_ASSERT(cond, ...) static_assert(cond, __VA_ARGS__) +#endif + +#ifndef FI_MEGAMOE_UNIFIED_ASSERT +#ifdef FI_MEGAMOE_IN_CUDA_COMPILATION +#define FI_MEGAMOE_UNIFIED_ASSERT(cond) FI_MEGAMOE_DEVICE_ASSERT(cond) +#else +#define FI_MEGAMOE_UNIFIED_ASSERT(cond) FI_MEGAMOE_HOST_ASSERT(cond) +#endif +#endif + +// ============ cute_tie.cuh ============ +namespace cute { + +struct ignore_t { + template + constexpr const ignore_t& operator=(T&&) const noexcept { + return *this; + } +}; + +inline constexpr ignore_t ignore{}; + +} // namespace cute + +#define CUTE_TIE_CONCAT_IMPL(A, B) A##B +#define CUTE_TIE_CONCAT(A, B) CUTE_TIE_CONCAT_IMPL(A, B) + +#define CUTE_TIE_GET_NTH_ARG(_1, _2, _3, _4, _5, _6, _7, _8, _9, _10, N, ...) N +#define CUTE_TIE_COUNT_ARGS(...) CUTE_TIE_GET_NTH_ARG(__VA_ARGS__, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0) + +#define CUTE_TIE_OP_DECL(I, TUPLE, VAR) auto VAR = ::cute::get(TUPLE) +#define CUTE_TIE_OP_ASSIGN(I, TUPLE, VAR) VAR = ::cute::get(TUPLE) + +#define CUTE_TIE_APPLY_OP_1(OP, T, V1) OP(0, T, V1); +#define CUTE_TIE_APPLY_OP_2(OP, T, V1, V2) \ + OP(0, T, V1); \ + OP(1, T, V2); +#define CUTE_TIE_APPLY_OP_3(OP, T, V1, V2, V3) \ + OP(0, T, V1); \ + OP(1, T, V2); \ + OP(2, T, V3); +#define CUTE_TIE_APPLY_OP_4(OP, T, V1, V2, V3, V4) \ + OP(0, T, V1); \ + OP(1, T, V2); \ + OP(2, T, V3); \ + OP(3, T, V4); +#define CUTE_TIE_APPLY_OP_5(OP, T, V1, V2, V3, V4, V5) \ + OP(0, T, V1); \ + OP(1, T, V2); \ + OP(2, T, V3); \ + OP(3, T, V4); \ + OP(4, T, V5); + +#define CUTE_TIE_DECL(TUPLE_EXPR, ...) \ + auto&& CUTE_TIE_CONCAT(cute_tie__temp_tuple_, __LINE__) = (TUPLE_EXPR); \ + CUTE_TIE_CONCAT(CUTE_TIE_APPLY_OP_, CUTE_TIE_COUNT_ARGS(__VA_ARGS__))( \ + CUTE_TIE_OP_DECL, CUTE_TIE_CONCAT(cute_tie__temp_tuple_, __LINE__), __VA_ARGS__) + +#define CUTE_TIE(TUPLE_EXPR, ...) \ + do { \ + auto&& CUTE_TIE_CONCAT(cute_tie__temp_tuple_, __LINE__) = (TUPLE_EXPR); \ + CUTE_TIE_CONCAT(CUTE_TIE_APPLY_OP_, CUTE_TIE_COUNT_ARGS(__VA_ARGS__))( \ + CUTE_TIE_OP_ASSIGN, CUTE_TIE_CONCAT(cute_tie__temp_tuple_, __LINE__), __VA_ARGS__); \ + } while (0) + +// ============ types.cuh ============ +namespace flashinfer::megamoe { + +enum class MmaKind { + BF16 = 0, + MXFP8FP4 = 1, +}; + +enum class GemmType { + Normal = 0, + MGroupedContiguous = 1, + MGroupedMasked = 2, + KGroupedContiguous = 3, + Batched = 4, + MGroupedContiguousWithPsumLayout = 5, +}; + +enum class KernelType { Kernel1D1D = 0, Kernel1D2D = 1, KernelNoSF = 2 }; + +} // namespace flashinfer::megamoe + +// ============ utils.cuh ============ +namespace flashinfer::megamoe::utils { + +template +struct PatternVisitor { + FuncT func; + + CUTLASS_HOST_DEVICE + explicit PatternVisitor(FuncT&& func) : func(std::forward(func)) {} + + CUTLASS_HOST_DEVICE + auto operator[](const uint32_t& i) const { return func(i); } +}; + +template +struct Vectorized { + static auto zeros() { + // TODO: add `ulonglong4` for SM100 once `__ldg` support this + if constexpr (kNumBytes > 0 and kNumBytes % 16 == 0) { + return make_uint4(0, 0, 0, 0); + } else if constexpr (kNumBytes > 0 and kNumBytes % 8 == 0) { + return make_uint2(0, 0); + } else if constexpr (kNumBytes > 0 and kNumBytes % 4 == 0) { + return 0; + } else { + FI_MEGAMOE_STATIC_ASSERT(kNumBytes > 0 and kNumBytes % 4 == 0, "Invalid vectorization"); + } + } + + using vec_t = decltype(zeros()); +}; + +} // namespace flashinfer::megamoe::utils + +// ============ math.cuh ============ +namespace flashinfer::megamoe::math { + +/// Pointer operations +template +CUTLASS_HOST_DEVICE dtype_t* advance_ptr(void* ptr, const uint64_t num_bytes) { + return reinterpret_cast(static_cast(ptr) + num_bytes); +} + +/// Math functions +template +CUTLASS_HOST_DEVICE T ceil_div(T a, T b) { + return (a + b - 1) / b; +} + +template +CUTLASS_HOST_DEVICE constexpr T constexpr_ceil_div(T a, T b) { + return (a + b - 1) / b; +} + +template +CUTLASS_HOST_DEVICE T align(T a, T b) { + return (kDoCeilAlignment ? ceil_div(a, b) : (a / b)) * b; +} + +template +CUTLASS_HOST_DEVICE constexpr T constexpr_align(T a, T b) { + return constexpr_ceil_div(a, b) * b; +} + +template +CUTLASS_HOST_DEVICE constexpr T constexpr_min(T a, T b) { + return a < b ? a : b; +} + +template +CUTLASS_DEVICE void swap(T& a, T& b) { + T temp = a; + a = b; + b = temp; +} + +#ifdef FI_MEGAMOE_IN_CUDA_COMPILATION + +CUTLASS_DEVICE float2 mul2(const float2& a, const float2& b) { +#if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) + return __fmul2_rn(a, b); +#else + return make_float2(__fmul_rn(a.x, b.x), __fmul_rn(a.y, b.y)); +#endif +} + +CUTLASS_HOST_DEVICE float fast_rcp(const float& x) { + float ret; + asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(ret) : "f"(x)); + return ret; +} + +/// Casting +template +CUTLASS_DEVICE int cast_into_bf16_and_pack(old_t& x, old_t& y) { + auto bf16x2 = + __float22bfloat162_rn({*reinterpret_cast(&x), *reinterpret_cast(&y)}); + return *reinterpret_cast(&bf16x2); +} + +CUTLASS_DEVICE float fast_pow2(const int& x) { + uint32_t bits_x = (x + 127) << 23; + return *reinterpret_cast(&bits_x); +} + +CUTLASS_DEVICE int fast_log2_ceil(float x) { + const auto bits = *reinterpret_cast(&x); + const auto exp = bits >> 23; + const auto man = bits & ((1 << 23) - 1); + return exp - 127 + (man != 0); +} + +template +CUTLASS_DEVICE void get_e4m3_sf_and_sf_inv(const float2& amax, float2& sf, float2& sf_inv) { + FI_MEGAMOE_STATIC_ASSERT(kUseUE8M0, "Must use UE8M0"); + const float2 finfo_factor = {1.0 / 448.0, 1.0 / 448.0}; + const auto scaled = mul2(amax, finfo_factor); + const auto exp_x = fast_log2_ceil(scaled.x); + const auto exp_y = fast_log2_ceil(scaled.y); + sf.x = fast_pow2(exp_x), sf_inv.x = fast_pow2(-exp_x); + sf.y = fast_pow2(exp_y), sf_inv.y = fast_pow2(-exp_y); +} + +/// Reduction +CUTLASS_DEVICE uint32_t warp_inclusive_sum(uint32_t value, const uint32_t& lane_idx) { +#pragma unroll + for (uint32_t offset = 1; offset < 32; offset <<= 1) { + const uint32_t synced = __shfl_up_sync(0xffffffff, value, offset); + if (lane_idx >= offset) value += synced; + } + return value; +} + +// Operation functors +template +struct ReduceSum { + CUTLASS_DEVICE T operator()(T a, T b) const { return a + b; } +}; +template +struct ReduceMax { + CUTLASS_DEVICE T operator()(T a, T b) const { return a > b ? a : b; } +}; +template +struct ReduceMin { + CUTLASS_DEVICE T operator()(T a, T b) const { return a < b ? a : b; } +}; +template +struct ReduceAnd { + CUTLASS_DEVICE T operator()(T a, T b) const { return a & b; } +}; +template +struct ReduceOr { + CUTLASS_DEVICE T operator()(T a, T b) const { return a | b; } +}; + +// Unified reduction function +template +CUTLASS_DEVICE T warp_reduce(T value, Op op) { + FI_MEGAMOE_STATIC_ASSERT(kNumLanesPerGroup == 32 or kNumLanesPerGroup == 16 or + kNumLanesPerGroup == 8 or kNumLanesPerGroup == 4 or + kNumLanesPerGroup == 2 or kNumLanesPerGroup == 1, + "Invalid number of lanes"); + constexpr uint32_t mask = 0xffffffff; + if constexpr (kIntergroupReduce) { + if constexpr (kNumLanesPerGroup <= 1) value = op(value, __shfl_xor_sync(mask, value, 1)); + if constexpr (kNumLanesPerGroup <= 2) value = op(value, __shfl_xor_sync(mask, value, 2)); + if constexpr (kNumLanesPerGroup <= 4) value = op(value, __shfl_xor_sync(mask, value, 4)); + if constexpr (kNumLanesPerGroup <= 8) value = op(value, __shfl_xor_sync(mask, value, 8)); + if constexpr (kNumLanesPerGroup <= 16) value = op(value, __shfl_xor_sync(mask, value, 16)); + } else { + if constexpr (kNumLanesPerGroup >= 32) value = op(value, __shfl_xor_sync(mask, value, 16)); + if constexpr (kNumLanesPerGroup >= 16) value = op(value, __shfl_xor_sync(mask, value, 8)); + if constexpr (kNumLanesPerGroup >= 8) value = op(value, __shfl_xor_sync(mask, value, 4)); + if constexpr (kNumLanesPerGroup >= 4) value = op(value, __shfl_xor_sync(mask, value, 2)); + if constexpr (kNumLanesPerGroup >= 2) value = op(value, __shfl_xor_sync(mask, value, 1)); + } + return value; +} + +// Convenience aliases + +#endif + +} // namespace flashinfer::megamoe::math diff --git a/include/flashinfer/megamoe/detail/dispatch_rank_selection.inl b/include/flashinfer/megamoe/detail/dispatch_rank_selection.inl new file mode 100644 index 00000000000..dec330444e5 --- /dev/null +++ b/include/flashinfer/megamoe/detail/dispatch_rank_selection.inl @@ -0,0 +1,43 @@ +// Warp-wide round-robin selection from per-rank expert counts. +uint32_t current_rank_in_expert_idx; +uint32_t remaining[kNumRanksPerLane]; +#pragma unroll +for (uint32_t i = 0; i < kNumRanksPerLane; ++i) remaining[i] = stored_rank_count[i]; +uint32_t offset = 0; +uint32_t token_idx_in_expert = token_idx - expert_start_idx; +uint32_t slot_idx = token_idx_in_expert; +uint32_t token_idx_in_rank; +while (true) { + uint32_t num_actives_in_lane = 0; + uint32_t min_in_lane = 0xffffffff; +#pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++i) { + num_actives_in_lane += remaining[i] > 0; + if (remaining[i] > 0) min_in_lane = cute::min(min_in_lane, remaining[i]); + } + const uint32_t num_active_ranks = __reduce_add_sync(0xffffffff, num_actives_in_lane); + const uint32_t length = __reduce_min_sync(0xffffffff, min_in_lane); + + const uint32_t num_round_tokens = length * num_active_ranks; + if (slot_idx < num_round_tokens) { + const uint32_t slot_idx_in_round = slot_idx % num_active_ranks; + uint32_t num_seen_ranks = 0; + current_rank_in_expert_idx = 0; +#pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++i) { + const uint32_t mask = __ballot_sync(0xffffffff, remaining[i] > 0); + const uint32_t num_active_lanes = __popc(mask); + if (slot_idx_in_round >= num_seen_ranks and + slot_idx_in_round < num_seen_ranks + num_active_lanes) + current_rank_in_expert_idx = + i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1); + num_seen_ranks += num_active_lanes; + } + token_idx_in_rank = offset + (slot_idx / num_active_ranks); + break; + } + slot_idx -= num_round_tokens; + offset += length; +#pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++i) remaining[i] -= cute::min(remaining[i], length); +} diff --git a/include/flashinfer/megamoe/detail/fp4_decode.cuh b/include/flashinfer/megamoe/detail/fp4_decode.cuh new file mode 100644 index 00000000000..3dd0b393e3e --- /dev/null +++ b/include/flashinfer/megamoe/detail/fp4_decode.cuh @@ -0,0 +1,38 @@ +// Shared E2M1 nibble/sign decoding and register-source fragment assembly. +#pragma once +#include + +namespace flashinfer::megamoe::detail { + +__device__ __forceinline__ uint32_t byte_perm_unchecked(uint32_t a, uint32_t b, uint32_t sel) { + uint32_t out; + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(out) : "r"(a), "r"(b), "r"(sel)); + return out; +} + +template +__device__ __forceinline__ uint2 decode_fp4_word(uint32_t packed, const Lut& lut) { + const uint32_t selectors = packed & 0x77777777u; + uint32_t hi = byte_perm_unchecked(lut.x, lut.y, selectors); + uint32_t lo = byte_perm_unchecked(lut.x, lut.y, selectors >> 16); + asm("lop3.b32 %0, %0, %1, 0x80808080, 0xf8;" : "+r"(hi) : "r"(packed)); + const uint32_t shifted = packed << 4; + asm("lop3.b32 %0, %0, %1, 0x80808080, 0xf8;" : "+r"(lo) : "r"(shifted)); + return make_uint2(hi, lo); +} + +__device__ __forceinline__ void assemble_rs_word_pair(uint2 lo, uint2 hi, bool keep_hi, + uint32_t (&frag)[4]) { + const uint32_t kept_lo = keep_hi ? lo.x : lo.y; + const uint32_t sent_lo = keep_hi ? lo.y : lo.x; + const uint32_t recv_lo = __shfl_xor_sync(0xffffffffu, sent_lo, 1); + frag[0] = keep_hi ? kept_lo : recv_lo; + frag[1] = keep_hi ? recv_lo : kept_lo; + const uint32_t kept_hi = keep_hi ? hi.x : hi.y; + const uint32_t sent_hi = keep_hi ? hi.y : hi.x; + const uint32_t recv_hi = __shfl_xor_sync(0xffffffffu, sent_hi, 1); + frag[2] = keep_hi ? kept_hi : recv_hi; + frag[3] = keep_hi ? recv_hi : kept_hi; +} + +} // namespace flashinfer::megamoe::detail diff --git a/include/flashinfer/megamoe/detail/fused_epilogue.inl b/include/flashinfer/megamoe/detail/fused_epilogue.inl new file mode 100644 index 00000000000..2ba52324485 --- /dev/null +++ b/include/flashinfer/megamoe/detail/fused_epilogue.inl @@ -0,0 +1,438 @@ +// Shared NVFP4/MXFP4 fused SwiGLU, FP8 quantization and BF16 scatter. +if constexpr (!kBlockIsL2) { + const float l1_global_scale = + l1_global_scales == nullptr ? 1.0f : __ldg(l1_global_scales + local_expert_idx); + if constexpr (kSwapABRequested) { + auto silu = [](float x) -> float { + const float e = kFastMath ? __expf(-x) : expf(-x); + const float sig = kFastMath ? math::fast_rcp(1.0f + e) : 1.0f / (1.0f + e); + return x * sig; + }; + auto clamp_gate = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(x, kActivationClamp); + }; + auto clamp_up = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(cute::max(x, -kActivationClamp), kActivationClamp); + }; + + constexpr uint32_t reduce_warp_start = 0; + constexpr uint32_t reduce_warp_count = kNumEpilogueWarps; + const uint32_t scale_token_thread = epilogue_thread_idx; + constexpr uint32_t scale_token_stride = kNumEpilogueThreads; + const uint32_t sf_base_k_idx = n_block_idx * L1_OUT_BLOCK_N / kL2ActsSFGranK; + float swap_v0[kSwapABWeightHalves][kSwapABTokenChunks] = {}; + float swap_v1[kSwapABWeightHalves][kSwapABTokenChunks] = {}; + + auto store_l1_swap_chunk = [&](const uint32_t& i) { + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + + float v0_amax = 0.0f; + float v1_amax = 0.0f; +#pragma unroll + for (uint32_t half = 0; half < kSwapABWeightHalves; ++half) { + const uint32_t accum_offset = half * kSwapABHalfAccumPerThread + i * 4; + float v0 = 0.0f; + if (token_0 < valid_m) { + float g0 = final_accum[accum_offset + 0] * l1_global_scale; + float u0 = final_accum[accum_offset + 2] * l1_global_scale; + clamp_gate(g0); + clamp_up(u0); + const float weight_0 = *l1_topk_weights_buffer.get_data_buffer(m_idx + token_0) + .template get_base_ptr(); + v0 = silu(g0) * u0 * weight_0; + swap_v0[half][i] = v0; + v0_amax = cute::max(v0_amax, cute::abs(v0)); + } + + float v1 = 0.0f; + if (token_1 < valid_m) { + float g1 = final_accum[accum_offset + 1] * l1_global_scale; + float u1 = final_accum[accum_offset + 3] * l1_global_scale; + clamp_gate(g1); + clamp_up(u1); + const float weight_1 = *l1_topk_weights_buffer.get_data_buffer(m_idx + token_1) + .template get_base_ptr(); + v1 = silu(g1) * u1 * weight_1; + swap_v1[half][i] = v1; + v1_amax = cute::max(v1_amax, cute::abs(v1)); + } + } + + const float amax0 = math::warp_reduce<4, true>(v0_amax, math::ReduceMax()); + const float amax1 = math::warp_reduce<4, true>(v1_amax, math::ReduceMax()); + if (row_idx == 0) { + if (token_0 < valid_m) + smem_cd_l1_shared_sf[token_0 * kNumEpilogueWarps + epilogue_warp_idx] = amax0; + if (token_1 < valid_m) + smem_cd_l1_shared_sf[token_1 * kNumEpilogueWarps + epilogue_warp_idx] = amax1; + } + }; + + const uint32_t num_swap_token_chunks = (valid_m + 7u) / 8u; + store_l1_swap_chunk(0); + if (valid_m > 8) { +#pragma unroll + for (uint32_t i = 1; i < kSwapABTokenChunks; ++i) { + if (i < num_swap_token_chunks) store_l1_swap_chunk(i); + } + } + + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + + for (uint32_t token = scale_token_thread; token < valid_m; token += scale_token_stride) { + float amax = 0.0f; +#pragma unroll + for (uint32_t w = 0; w < reduce_warp_count; ++w) + amax = cute::max(amax, + smem_cd_l1_shared_sf[token * kNumEpilogueWarps + reduce_warp_start + w]); + float2 amax_pair = {amax, amax}; + float2 sf_pair, sf_inv_pair; + math::get_e4m3_sf_and_sf_inv(amax_pair, sf_pair, sf_inv_pair); + + auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); + const uint32_t token_idx = m_idx + token; + sf_base_ptr[sf_base_k_idx * kNumPaddedSFPoolTokens + token_idx] = sf_pair.x; + smem_cd_l1_shared_sf[token * kNumEpilogueWarps + reduce_warp_start] = sf_inv_pair.x; + } + + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + +#pragma unroll + for (uint32_t i = 0; i < kSwapABTokenChunks; ++i) { + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; +#pragma unroll + for (uint32_t half = 0; half < kSwapABWeightHalves; ++half) { + const uint32_t out_col_base = wg_l1_out_n_idx + half * 32u + warp_idx_in_wg * 8 + row_idx; + if (token_0 < valid_m) { + const float sf_inv = + smem_cd_l1_shared_sf[token_0 * kNumEpilogueWarps + reduce_warp_start]; + const __nv_fp8_e4m3 q(swap_v0[half][i] * sf_inv); + reinterpret_cast(smem_cd_l1)[token_0 * L1_OUT_BLOCK_N + out_col_base] = + *reinterpret_cast(&q); + } + if (token_1 < valid_m) { + const float sf_inv = + smem_cd_l1_shared_sf[token_1 * kNumEpilogueWarps + reduce_warp_start]; + const __nv_fp8_e4m3 q(swap_v1[half][i] * sf_inv); + reinterpret_cast(smem_cd_l1)[token_1 * L1_OUT_BLOCK_N + out_col_base] = + *reinterpret_cast(&q); + } + } + } + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (epilogue_wg_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { + const uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N; + cute::tma_store_fence(); + cute::SM90_TMA_STORE_2D::copy(&tensor_map_l1_output, smem_cd_l1, out_n_idx, m_idx); + cute::tma_store_arrive(); + } + __syncwarp(); + ptx::tma_store_wait<0>(); + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + notify_l1_ready(pool_block_idx, n_block_idx); + } else { + // ---------------- L1 EPILOGUE: SwiGLU + FP8 quantize + TMA store ---------------- + const bool valid_r0 = row_offset_r0 < valid_m; + const bool valid_r1 = row_offset_r1 < valid_m; + // Layout in `final_accum`: + // 16 chunks of 8 N-cols, each chunk = 4 floats per thread = (r0c0, r0c1, r1c0, r1c1). + // Gate chunks: even (0, 2, ..., 14). Up chunks: odd (1, 3, ..., 15). + // Pair `p` ∈ [0, 8): gate chunk = 2p, up chunk = 2p+1. + // + // For each pair we produce 4 post-SwiGLU floats per thread, mapped to + // output cols (p*8 + col_idx*2 + {0,1}) for both r0 and r1. + + constexpr uint32_t kNumPairs = kAccumPerThread / 8; + constexpr uint32_t kNumSFGroups = 1; + float swiglu_r0[kNumPairs][2]; + float swiglu_r1[kNumPairs][2]; + + // Per-row amax, one scale for each 64-col L1 output group. + float amax_r0[kNumSFGroups] = {}; + float amax_r1[kNumSFGroups] = {}; + +// Compute SwiGLU + per-group amax. +#pragma unroll + for (uint32_t p = 0; p < kNumPairs; ++p) { + const uint32_t gate = 2 * p, up = 2 * p + 1; + const uint32_t sf_group = p / 8; + + auto clamp_gate = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(x, kActivationClamp); + }; + auto clamp_up = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(cute::max(x, -kActivationClamp), kActivationClamp); + }; + float g_r0_c0 = final_accum[gate * 4 + 0] * l1_global_scale; + clamp_gate(g_r0_c0); + float g_r0_c1 = final_accum[gate * 4 + 1] * l1_global_scale; + clamp_gate(g_r0_c1); + float g_r1_c0 = final_accum[gate * 4 + 2] * l1_global_scale; + clamp_gate(g_r1_c0); + float g_r1_c1 = final_accum[gate * 4 + 3] * l1_global_scale; + clamp_gate(g_r1_c1); + float u_r0_c0 = final_accum[up * 4 + 0] * l1_global_scale; + clamp_up(u_r0_c0); + float u_r0_c1 = final_accum[up * 4 + 1] * l1_global_scale; + clamp_up(u_r0_c1); + float u_r1_c0 = final_accum[up * 4 + 2] * l1_global_scale; + clamp_up(u_r1_c0); + float u_r1_c1 = final_accum[up * 4 + 3] * l1_global_scale; + clamp_up(u_r1_c1); + + auto silu = [](float x) -> float { + const float e = kFastMath ? __expf(-x) : expf(-x); + const float sig = kFastMath ? math::fast_rcp(1.0f + e) : 1.0f / (1.0f + e); + return x * sig; + }; + + if (valid_r0) { + swiglu_r0[p][0] = silu(g_r0_c0) * u_r0_c0; + swiglu_r0[p][1] = silu(g_r0_c1) * u_r0_c1; + amax_r0[sf_group] = cute::max( + amax_r0[sf_group], cute::max(cute::abs(swiglu_r0[p][0]), cute::abs(swiglu_r0[p][1]))); + } else { + swiglu_r0[p][0] = 0.0f; + swiglu_r0[p][1] = 0.0f; + } + if (valid_r1) { + swiglu_r1[p][0] = silu(g_r1_c0) * u_r1_c0; + swiglu_r1[p][1] = silu(g_r1_c1) * u_r1_c1; + amax_r1[sf_group] = cute::max( + amax_r1[sf_group], cute::max(cute::abs(swiglu_r1[p][0]), cute::abs(swiglu_r1[p][1]))); + } else { + swiglu_r1[p][0] = 0.0f; + swiglu_r1[p][1] = 0.0f; + } + } + + const float weight_r0 = valid_r0 + ? *l1_topk_weights_buffer.get_data_buffer(m_idx + row_offset_r0) + .template get_base_ptr() + : 0.0f; + const float weight_r1 = valid_r1 + ? *l1_topk_weights_buffer.get_data_buffer(m_idx + row_offset_r1) + .template get_base_ptr() + : 0.0f; +#pragma unroll + for (uint32_t p = 0; p < kNumPairs; ++p) { + swiglu_r0[p][0] *= weight_r0; + swiglu_r0[p][1] *= weight_r0; + swiglu_r1[p][0] *= weight_r1; + swiglu_r1[p][1] *= weight_r1; + } +#pragma unroll + for (uint32_t g = 0; g < kNumSFGroups; ++g) { + amax_r0[g] *= cute::abs(weight_r0); + amax_r1[g] *= cute::abs(weight_r1); + } +#pragma unroll + for (uint32_t g = 0; g < kNumSFGroups; ++g) { + amax_r0[g] = math::warp_reduce<4, false>(amax_r0[g], math::ReduceMax()); + amax_r1[g] = math::warp_reduce<4, false>(amax_r1[g], math::ReduceMax()); + } + + if (col_idx == 0) { + smem_cd_l1_shared_sf[epilogue_wg_idx * BLOCK_M + row_offset_r0] = amax_r0[0]; + smem_cd_l1_shared_sf[epilogue_wg_idx * BLOCK_M + row_offset_r1] = amax_r1[0]; + } + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if constexpr (kSplitMDecodedWeightReuse) { + amax_r0[0] = smem_cd_l1_shared_sf[epilogue_wg_idx * BLOCK_M + row_offset_r0]; + amax_r1[0] = smem_cd_l1_shared_sf[epilogue_wg_idx * BLOCK_M + row_offset_r1]; + } else { + amax_r0[0] = cute::max(smem_cd_l1_shared_sf[row_offset_r0], + smem_cd_l1_shared_sf[BLOCK_M + row_offset_r0]); + amax_r1[0] = cute::max(smem_cd_l1_shared_sf[row_offset_r1], + smem_cd_l1_shared_sf[BLOCK_M + row_offset_r1]); + } + + float sf_r0[kNumSFGroups], sf_inv_r0[kNumSFGroups]; + float sf_r1[kNumSFGroups], sf_inv_r1[kNumSFGroups]; +#pragma unroll + for (uint32_t g = 0; g < kNumSFGroups; ++g) { + float2 amax_pair = {amax_r0[g], amax_r1[g]}; + float2 sf_pair, sf_inv_pair; + math::get_e4m3_sf_and_sf_inv(amax_pair, sf_pair, sf_inv_pair); + sf_r0[g] = sf_pair.x; + sf_inv_r0[g] = sf_inv_pair.x; + sf_r1[g] = sf_pair.y; + sf_inv_r1[g] = sf_inv_pair.y; + } + +// Quantize and write to smem_cd_l1 (row-major, no swizzle). +#pragma unroll + for (uint32_t p = 0; p < kNumPairs; ++p) { + const uint32_t sf_group = p / 8; + const float v00 = swiglu_r0[p][0] * sf_inv_r0[sf_group]; + const float v01 = swiglu_r0[p][1] * sf_inv_r0[sf_group]; + const float v10 = swiglu_r1[p][0] * sf_inv_r1[sf_group]; + const float v11 = swiglu_r1[p][1] * sf_inv_r1[sf_group]; + + const __nv_fp8x2_e4m3 r0_pair(make_float2(v00, v01)); + const __nv_fp8x2_e4m3 r1_pair(make_float2(v10, v11)); + + const uint32_t col = p * 8 + col_idx * 2; + auto* p0 = reinterpret_cast(smem_cd_l1 + row_offset_r0 * L1_OUT_BLOCK_N + + wg_l1_out_n_idx + col); + auto* p1 = reinterpret_cast(smem_cd_l1 + row_offset_r1 * L1_OUT_BLOCK_N + + wg_l1_out_n_idx + col); + if (valid_r0) *p0 = r0_pair.__x; + if (valid_r1) *p1 = r1_pair.__x; + } + + // Write one physical L2-activation scale per 128 output columns. + if (col_idx == 0) { + auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); + const uint32_t token_r0 = m_idx + row_offset_r0; + const uint32_t token_r1 = m_idx + row_offset_r1; + const uint32_t base_k_sf_idx = + (n_block_idx * L1_OUT_BLOCK_N + wg_l1_out_n_idx) / kL2ActsSFGranK; +#pragma unroll + for (uint32_t g = 0; g < kNumSFGroups; ++g) { + const uint32_t sf_k_idx = base_k_sf_idx + g; + if ((kSplitMDecodedWeightReuse || epilogue_wg_idx == 0) && valid_r0) + sf_base_ptr[sf_k_idx * kNumPaddedSFPoolTokens + token_r0] = sf_r0[g]; + if ((kSplitMDecodedWeightReuse || epilogue_wg_idx == 0) && valid_r1) + sf_base_ptr[sf_k_idx * kNumPaddedSFPoolTokens + token_r1] = sf_r1[g]; + } + } + + // Issue TMA store of the entire tile. Padding rows beyond + // `valid_m` are written with stale/garbage FP8 to the L1-output + // pool buffer, but they are never consumed downstream: the L2 + // GEMM tile loads them, but its NVLink-scatter epilogue is + // gated by `m_idx_in_block >= valid_m`, and stale SF in the + // padding rows can produce NaN accumulators that simply stay + // in registers (only valid rows are converted to BF16 and + // STSM'd into smem). Using TMA for partial tiles is a large + // win for low-batch / decode where every tile is partial. + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + const uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N; + cute::tma_store_fence(); + cute::SM90_TMA_STORE_2D::copy(&tensor_map_l1_output, smem_cd_l1, out_n_idx, m_idx); + cute::tma_store_arrive(); + } + __syncwarp(); + ptx::tma_store_wait<0>(); + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + notify_l1_ready(pool_block_idx, n_block_idx); + } +} else { + // ---------------- L2 EPILOGUE: BF16 cast + NVLink scatter ---------------- + if constexpr (kSwapABRequested) { + // Each active warp scatters a contiguous group of up to 16 rows. + constexpr uint32_t kNumRowsPerWarp = BLOCK_M == 8 ? 4u : 8u; + auto store_swap_bf16 = [&](const uint32_t& token, const uint32_t& col, const float& value) { + if (token < valid_m) + smem_cd_l2[token * BLOCK_N + wg_n_idx + col] = __float2bfloat16_rn(value * l2_global_scale); + }; + + const uint32_t num_swap_token_chunks = (valid_m + 7u) / 8u; + auto store_l2_swap_chunk = [&](const uint32_t& i) { + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; +#pragma unroll + for (uint32_t half = 0; half < kSwapABWeightHalves; ++half) { + const uint32_t accum_offset = half * kSwapABHalfAccumPerThread + i * 4; + const uint32_t col_offset = half * 64u; + store_swap_bf16(token_0, col_offset + r_0, final_accum[accum_offset + 0]); + store_swap_bf16(token_0, col_offset + r_1, final_accum[accum_offset + 2]); + store_swap_bf16(token_1, col_offset + r_0, final_accum[accum_offset + 1]); + store_swap_bf16(token_1, col_offset + r_1, final_accum[accum_offset + 3]); + } + }; + + store_l2_swap_chunk(0); + if (valid_m > 8) { +#pragma unroll + for (uint32_t i = 1; i < kSwapABTokenChunks; ++i) { + if (i < num_swap_token_chunks) store_l2_swap_chunk(i); + } + } + + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + const uint32_t row_in_warp_block = lane_idx / 16; + const uint32_t lane_in_row = lane_idx % 16; + constexpr uint32_t kColsPerScatterLane = WG_BLOCK_N / 16; + FI_MEGAMOE_STATIC_ASSERT(WG_BLOCK_N % 16 == 0, + "SwapAB L2 scatter expects an even lane partition"); + FI_MEGAMOE_STATIC_ASSERT(kColsPerScatterLane == 4 || kColsPerScatterLane == 8, + "SwapAB L2 scatter supports WG_BLOCK_N=64 or 128"); + +#pragma unroll + for (uint32_t j = 0; j < kNumRowsPerWarp; ++j) { + const uint32_t token = warp_idx_in_wg * 16 + j * 2 + row_in_warp_block; + if (token >= valid_m) break; + + const auto src_metadata = + *workspace.get_token_src_metadata_ptr(pool_block_idx * BLOCK_M + token); + const uint32_t dst_rank_idx = src_metadata.rank_idx; + const uint32_t dst_token_idx = src_metadata.token_idx; + const uint32_t dst_topk_idx = src_metadata.topk_idx; + const auto dst_token = + combine_token_buffer.get_rank_buffer(dst_topk_idx).get_data_buffer(dst_token_idx); + auto smem_ptr = smem_cd_l2 + token * BLOCK_N + wg_n_idx + lane_in_row * kColsPerScatterLane; + if constexpr (kColsPerScatterLane == 8) { + const auto packed = *reinterpret_cast(smem_ptr); + auto dst_ptr = math::advance_ptr( + dst_token.get_base_ptr(), n_idx * sizeof(nv_bfloat16) + lane_in_row * sizeof(uint4)); + *sym_buffer.map(dst_ptr, dst_rank_idx) = packed; + } else { + const auto packed = *reinterpret_cast(smem_ptr); + auto dst_ptr = math::advance_ptr( + dst_token.get_base_ptr(), n_idx * sizeof(nv_bfloat16) + lane_in_row * sizeof(uint2)); + *sym_buffer.map(dst_ptr, dst_rank_idx) = packed; + } + } + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } else { + FI_MEGAMOE_STATIC_ASSERT(WG_BLOCK_N == 64 || WG_BLOCK_N == 128, + "Direct L2 scatter requires N64/N128"); + const bool valid_r0 = row_offset_r0 < valid_m; + const bool valid_r1 = row_offset_r1 < valid_m; + + auto scatter_direct_row = [&](const uint32_t& row_offset, const bool& valid_row, + const uint32_t& row_accum_offset) { + if (valid_row) { + const auto src_metadata = + *workspace.get_token_src_metadata_ptr(pool_block_idx * BLOCK_M + row_offset); + const uint32_t dst_rank_idx = src_metadata.rank_idx; + const uint32_t dst_token_idx = src_metadata.token_idx; + const uint32_t dst_topk_idx = src_metadata.topk_idx; + const auto dst_token = + combine_token_buffer.get_rank_buffer(dst_topk_idx).get_data_buffer(dst_token_idx); + auto dst_base = + math::advance_ptr(dst_token.get_base_ptr(), n_idx * sizeof(nv_bfloat16)); + auto mapped_dst_base = sym_buffer.map(dst_base, dst_rank_idx); + +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 8; ++i) { + const uint32_t chunk_lo = 2 * i, chunk_hi = 2 * i + 1; + const uint32_t col_lo = chunk_lo * 8 + col_idx * 2; + const uint32_t col_hi = chunk_hi * 8 + col_idx * 2; + const uint32_t packed_lo = + cast_l2_scaled_bf16_pair(final_accum[chunk_lo * 4 + row_accum_offset + 0], + final_accum[chunk_lo * 4 + row_accum_offset + 1]); + const uint32_t packed_hi = + cast_l2_scaled_bf16_pair(final_accum[chunk_hi * 4 + row_accum_offset + 0], + final_accum[chunk_hi * 4 + row_accum_offset + 1]); + *reinterpret_cast(mapped_dst_base + col_lo * sizeof(nv_bfloat16)) = packed_lo; + *reinterpret_cast(mapped_dst_base + col_hi * sizeof(nv_bfloat16)) = packed_hi; + } + } + }; + + scatter_direct_row(row_offset_r0, valid_r0, 0); + scatter_direct_row(row_offset_r1, valid_r1, 2); + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } +} diff --git a/include/flashinfer/megamoe/detail/mma_sm90.cuh b/include/flashinfer/megamoe/detail/mma_sm90.cuh new file mode 100644 index 00000000000..aba07f539ed --- /dev/null +++ b/include/flashinfer/megamoe/detail/mma_sm90.cuh @@ -0,0 +1,266 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace flashinfer::megamoe::mma::sm90 { + +/// MMA +template +struct FP8MMA { + template + CUTLASS_DEVICE static void call_fma_impl(uint64_t const& desc_a, uint64_t const& desc_b, float* d, + bool scale_d, cute::index_sequence) { + using namespace cute::SM90::GMMA; + MMA::fma(desc_a, desc_b, d[Idx]..., (scale_d ? ScaleOut::One : ScaleOut::Zero)); + } + + CUTLASS_DEVICE static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, + bool scale_d) { + call_fma_impl(desc_a, desc_b, d, scale_d, cute::make_index_sequence{}); + } + + static constexpr int M = 64; + static constexpr int N = N_; + static constexpr int K = 32; + static constexpr int kNumAccum = M * N / 128; +}; + +template +struct FP8MMASelector { + static constexpr auto select_mma() { + using namespace cute::SM90::GMMA; + if constexpr (N == 8) return MMA_64x8x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 16) return MMA_64x16x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 24) return MMA_64x24x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 32) return MMA_64x32x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 40) return MMA_64x40x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 48) return MMA_64x48x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 56) return MMA_64x56x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 64) return MMA_64x64x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 72) return MMA_64x72x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 80) return MMA_64x80x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 88) return MMA_64x88x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 96) return MMA_64x96x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 104) return MMA_64x104x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 112) return MMA_64x112x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 120) return MMA_64x120x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 128) return MMA_64x128x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 136) return MMA_64x136x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 144) return MMA_64x144x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 152) return MMA_64x152x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 160) return MMA_64x160x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 168) return MMA_64x168x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 176) return MMA_64x176x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 184) return MMA_64x184x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 192) return MMA_64x192x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 200) return MMA_64x200x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 208) return MMA_64x208x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 216) return MMA_64x216x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 224) return MMA_64x224x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 232) return MMA_64x232x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 240) return MMA_64x240x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 248) return MMA_64x248x32_F32E4M3E4M3_SS_TN(); + if constexpr (N == 256) return MMA_64x256x32_F32E4M3E4M3_SS_TN(); + } + + static constexpr auto select_type() { return FP8MMA(); } + + using type = decltype(select_type()); +}; + +template +struct FP8MMARS { + template + CUTLASS_DEVICE static void call_fma_impl(uint32_t const* a, uint64_t const& desc_b, float* d, + bool scale_d, cute::index_sequence) { + using namespace cute::SM90::GMMA; + MMA::fma(a[0], a[1], a[2], a[3], desc_b, d[Idx]..., (scale_d ? ScaleOut::One : ScaleOut::Zero)); + } + + CUTLASS_DEVICE static void wgmma(uint32_t const* a, uint64_t const& desc_b, float* d, + bool scale_d) { + call_fma_impl(a, desc_b, d, scale_d, cute::make_index_sequence{}); + } + + static constexpr int M = 64; + static constexpr int N = N_; + static constexpr int K = 32; + static constexpr int kNumAccum = M * N / 128; +}; + +template +struct FP8RSMMASelector { + static constexpr auto select_mma() { + using namespace cute::SM90::GMMA; + FI_MEGAMOE_STATIC_ASSERT(N == 8 || N == 16 || N == 24 || N == 32 || N == 64 || N == 128, + "Invalid FP8 RS WGMMA N"); + if constexpr (N == 8) return MMA_64x8x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 16) return MMA_64x16x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 24) return MMA_64x24x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 32) return MMA_64x32x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 64) return MMA_64x64x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 128) return MMA_64x128x32_F32E4M3E4M3_RS_TN(); + } + + static constexpr auto select_type() { return FP8MMARS(); } + + using type = decltype(select_type()); +}; + +// Register-operand fence for RS-form WGMMA A fragments (uint32 payload), +// mirroring `ptx::warpgroup_fence_operand(float&)` for accumulators. +CUTLASS_DEVICE void warpgroup_fence_operand(uint32_t& reg) { + asm volatile("" : "+r"(reg)::"memory"); +} + +template +struct BF16MMA { + template + CUTLASS_DEVICE static void call_fma_impl(uint64_t const& desc_a, uint64_t const& desc_b, float* d, + bool scale_d, cute::index_sequence) { + using namespace cute::SM90::GMMA; + MMA::fma(desc_a, desc_b, d[Idx]..., (scale_d ? ScaleOut::One : ScaleOut::Zero)); + } + + CUTLASS_DEVICE static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, + bool scale_d) { + call_fma_impl(desc_a, desc_b, d, scale_d, cute::make_index_sequence{}); + } + + static constexpr int M = 64; + static constexpr int N = N_; + static constexpr int K = 16; + static constexpr int kNumAccum = M * N / 128; +}; + +template +constexpr cute::SM90::GMMA::Major to_sm90_major() { + FI_MEGAMOE_STATIC_ASSERT(kMajor == cute::UMMA::Major::K or kMajor == cute::UMMA::Major::MN, + "Invalid major-ness"); + return kMajor == cute::UMMA::Major::K ? cute::SM90::GMMA::Major::K : cute::SM90::GMMA::Major::MN; +} + +template +struct BF16MMASelector { + static constexpr auto select_mma() { + using namespace cute::SM90::GMMA; + constexpr auto kGMMAMajorA = to_sm90_major(); + constexpr auto kGMMAMajorB = to_sm90_major(); + if constexpr (N == 8) return MMA_64x8x16_F32BF16BF16_SS(); + if constexpr (N == 16) return MMA_64x16x16_F32BF16BF16_SS(); + if constexpr (N == 24) return MMA_64x24x16_F32BF16BF16_SS(); + if constexpr (N == 32) return MMA_64x32x16_F32BF16BF16_SS(); + if constexpr (N == 40) return MMA_64x40x16_F32BF16BF16_SS(); + if constexpr (N == 48) return MMA_64x48x16_F32BF16BF16_SS(); + if constexpr (N == 56) return MMA_64x56x16_F32BF16BF16_SS(); + if constexpr (N == 64) return MMA_64x64x16_F32BF16BF16_SS(); + if constexpr (N == 72) return MMA_64x72x16_F32BF16BF16_SS(); + if constexpr (N == 80) return MMA_64x80x16_F32BF16BF16_SS(); + if constexpr (N == 88) return MMA_64x88x16_F32BF16BF16_SS(); + if constexpr (N == 96) return MMA_64x96x16_F32BF16BF16_SS(); + if constexpr (N == 104) return MMA_64x104x16_F32BF16BF16_SS(); + if constexpr (N == 112) return MMA_64x112x16_F32BF16BF16_SS(); + if constexpr (N == 120) return MMA_64x120x16_F32BF16BF16_SS(); + if constexpr (N == 128) return MMA_64x128x16_F32BF16BF16_SS(); + if constexpr (N == 136) return MMA_64x136x16_F32BF16BF16_SS(); + if constexpr (N == 144) return MMA_64x144x16_F32BF16BF16_SS(); + if constexpr (N == 152) return MMA_64x152x16_F32BF16BF16_SS(); + if constexpr (N == 160) return MMA_64x160x16_F32BF16BF16_SS(); + if constexpr (N == 168) return MMA_64x168x16_F32BF16BF16_SS(); + if constexpr (N == 176) return MMA_64x176x16_F32BF16BF16_SS(); + if constexpr (N == 184) return MMA_64x184x16_F32BF16BF16_SS(); + if constexpr (N == 192) return MMA_64x192x16_F32BF16BF16_SS(); + if constexpr (N == 200) return MMA_64x200x16_F32BF16BF16_SS(); + if constexpr (N == 208) return MMA_64x208x16_F32BF16BF16_SS(); + if constexpr (N == 216) return MMA_64x216x16_F32BF16BF16_SS(); + if constexpr (N == 224) return MMA_64x224x16_F32BF16BF16_SS(); + if constexpr (N == 232) return MMA_64x232x16_F32BF16BF16_SS(); + if constexpr (N == 240) return MMA_64x240x16_F32BF16BF16_SS(); + if constexpr (N == 248) return MMA_64x248x16_F32BF16BF16_SS(); + if constexpr (N == 256) return MMA_64x256x16_F32BF16BF16_SS(); + } + + static constexpr auto select_type() { return BF16MMA(); } + + using type = decltype(select_type()); +}; + +template +struct TF32MMARS { + template + CUTLASS_DEVICE static void call_fma_impl(uint32_t* a, uint64_t const& desc_b, float* d, + bool scale_d, cute::index_sequence) { + using namespace cute::SM90::GMMA; + MMA::fma(a[0], a[1], a[2], a[3], desc_b, d[Idx]..., (scale_d ? ScaleOut::One : ScaleOut::Zero)); + } + + CUTLASS_DEVICE static void wgmma(float* a, uint64_t const& desc_b, float* d, bool scale_d) { + call_fma_impl(reinterpret_cast(a), desc_b, d, scale_d, + cute::make_index_sequence{}); + } + + static constexpr int M = 64; + static constexpr int N = N_; + static constexpr int K = 8; + static constexpr int kNumAccum = M * N / 128; +}; + +template +struct TF32MMASelector { + static constexpr auto select_mma() { + using namespace cute::SM90::GMMA; + if constexpr (kUseRS) { + if constexpr (N == 8) return MMA_64x8x8_F32TF32TF32_RS_TN(); + if constexpr (N == 16) return MMA_64x16x8_F32TF32TF32_RS_TN(); + if constexpr (N == 32) return MMA_64x32x8_F32TF32TF32_RS_TN(); + if constexpr (N == 64) return MMA_64x64x8_F32TF32TF32_RS_TN(); + if constexpr (N == 128) return MMA_64x128x8_F32TF32TF32_RS_TN(); + if constexpr (N == 256) return MMA_64x256x8_F32TF32TF32_RS_TN(); + FI_MEGAMOE_STATIC_ASSERT(N == 8 or N == 16 or N == 32 or N == 64 or N == 128 or N == 256, + "Invalid N"); + } + } + + static constexpr auto select_type() { + if constexpr (kUseRS) { + return TF32MMARS(); + } else { + FI_MEGAMOE_STATIC_ASSERT(kUseRS, "SS mode is not supported for TF32MMASelector for now"); + } + } + + using type = decltype(select_type()); +}; + +/// Shared memory descriptor +template +CUTLASS_DEVICE cute::GmmaDescriptor make_smem_desc(PointerType smem_ptr, const int& layout_type, + const uint32_t& leading_byte_offset = 0, + const uint32_t& stride_byte_offset = 1024) { + // NOTES: the default LBO and SBO are for K-major types + cute::GmmaDescriptor desc; + const auto uint_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + desc.bitfield.start_address_ = uint_ptr >> 4; + desc.bitfield.layout_type_ = layout_type; + desc.bitfield.leading_byte_offset_ = leading_byte_offset >> 4; + desc.bitfield.stride_byte_offset_ = stride_byte_offset >> 4; + desc.bitfield.base_offset_ = 0; + return desc; +} + +template +constexpr uint32_t get_inner_block_atom_size() { + return kSwizzleMode == 0 ? BLOCK_INNER : kSwizzleMode / sizeof(dtype_t); +} + +// ReSharper disable once CppNotAllPathsReturnValue + +// ReSharper disable once CppNotAllPathsReturnValue + +} // namespace flashinfer::megamoe::mma::sm90 diff --git a/include/flashinfer/megamoe/detail/ptx.cuh b/include/flashinfer/megamoe/detail/ptx.cuh new file mode 100644 index 00000000000..cc062b60d91 --- /dev/null +++ b/include/flashinfer/megamoe/detail/ptx.cuh @@ -0,0 +1,411 @@ +#pragma once + +#include +#include + +#include +#include +#include + +namespace flashinfer::megamoe::ptx { + +// ---- from ptx_utils.cuh ---- +CUTLASS_DEVICE uint32_t get_lane_idx() { + uint32_t lane_id; + asm("mov.u32 %0, %%laneid;" : "=r"(lane_id)); + return lane_id; +} + +CUTLASS_DEVICE void sync_aligned(const uint32_t& num_threads, const uint32_t& barrier_idx) { + asm volatile("bar.sync %0, %1;" : : "r"(barrier_idx), "r"(num_threads)); +} + +CUTLASS_DEVICE void sync_unaligned(const uint32_t& num_threads, const uint32_t& barrier_idx) { + asm volatile("barrier.sync %0, %1;" : : "r"(barrier_idx), "r"(num_threads)); +} + +template +CUTLASS_DEVICE dtype_t exchange(dtype_t ptr, const uint32_t& src_lane_idx) { + FI_MEGAMOE_STATIC_ASSERT(sizeof(dtype_t) % sizeof(uint32_t) == 0, ""); + const auto send_int_values = reinterpret_cast(&ptr); + dtype_t recv_dtype; + auto recv_int_values = reinterpret_cast(&recv_dtype); +#pragma unroll + for (uint32_t i = 0; i < sizeof(dtype_t) / sizeof(uint32_t); ++i) + recv_int_values[i] = + __shfl_sync(0xffffffff, send_int_values[i], static_cast(src_lane_idx)); + return recv_dtype; +} + +CUTLASS_DEVICE void accumulate(float2& a, nv_bfloat162 b) { +#if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) + // Use `add.rn.f32.bf16` instruction to perform fused (cast + add) operation on SM100 + asm("add.rn.f32.bf16 %0, %1, %0;\n" : "+f"(a.x) : "h"(*reinterpret_cast(&b.x))); + asm("add.rn.f32.bf16 %0, %1, %0;\n" : "+f"(a.y) : "h"(*reinterpret_cast(&b.y))); +#else + const auto [x, y] = __bfloat1622float2(b); + a.x += x, a.y += y; +#endif +} + +// ---- from ptx_ld_st.cuh ---- +// Compatibility: 256 bits LD/ST instructions +#if defined(CUDART_VERSION) and CUDART_VERSION >= 13000 +using longlong4_t = longlong4_32a; +#define make_longlong4_t make_longlong4_32a +#else +struct alignas(32) longlong4_t { + long long x, y, z, w; +}; +CUTLASS_HOST_DEVICE longlong4_t make_longlong4_t(const long long& x, const long long& y, + const long long& z, const long long& w) { + return {x, y, z, w}; +} +#endif + +/// LD/ST matrix +// TODO: remove `struct` +struct SM90_U32x2_LDSM_N { + CUTLASS_DEVICE static void copy(uint32_t& dst_0, uint32_t& dst_1, void* smem_src) { + asm volatile("ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];\n" + : "=r"(dst_0), "=r"(dst_1) + : "l"(__cvta_generic_to_shared(smem_src))); + } +}; + +struct SM90_U32x4_LDSM_N { + CUTLASS_DEVICE static void copy(uint32_t& dst_0, uint32_t& dst_1, uint32_t& dst_2, + uint32_t& dst_3, void* smem_src) { + asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(dst_0), "=r"(dst_1), "=r"(dst_2), "=r"(dst_3) + : "l"(__cvta_generic_to_shared(smem_src))); + } +}; + +template +struct SM90_U32x2_STSM_N { + CUTLASS_DEVICE static void copy(dtype_t src_0, dtype_t src_1, void* smem_dst) { + FI_MEGAMOE_STATIC_ASSERT(sizeof(dtype_t) == sizeof(uint32_t), "Invalid dtype"); + const uint32_t src[2] = {*reinterpret_cast(&src_0), + *reinterpret_cast(&src_1)}; + asm volatile("stmatrix.sync.aligned.x2.m8n8.shared.b16 [%0], {%1, %2};\n" ::"l"( + __cvta_generic_to_shared(smem_dst)), + "r"(src[0]), "r"(src[1])); + } +}; + +template +struct SM90_U32x4_STSM_T { + CUTLASS_DEVICE static void copy(dtype_t src_0, dtype_t src_1, dtype_t src_2, dtype_t src_3, + void* smem_dst) { + FI_MEGAMOE_STATIC_ASSERT(sizeof(dtype_t) == sizeof(uint32_t), "Invalid dtype"); + const uint32_t src[4] = { + *reinterpret_cast(&src_0), *reinterpret_cast(&src_1), + *reinterpret_cast(&src_2), *reinterpret_cast(&src_3)}; + asm volatile("stmatrix.sync.aligned.x4.m8n8.shared.b16.trans [%0], {%1, %2, %3, %4};\n" ::"l"( + __cvta_generic_to_shared(smem_dst)), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3])); + } +}; + +template +struct SM100_U8x4_STSM_T { + __device__ __forceinline__ static void copy(dtype_t src_0, void* smem_dst) { + FI_MEGAMOE_STATIC_ASSERT(sizeof(dtype_t) == sizeof(uint32_t), "Invalid dtype"); + const uint32_t src = *reinterpret_cast(&src_0); + asm volatile("stmatrix.sync.aligned.m16n8.x1.trans.shared.b8 [%0], {%1};\n" ::"l"( + __cvta_generic_to_shared(smem_dst)), + "r"(src)); + } +}; + +template +struct SM100_U8x8_STSM_T { + __device__ __forceinline__ static void copy(dtype_t src_0, dtype_t src_1, void* smem_dst) { + FI_MEGAMOE_STATIC_ASSERT(sizeof(dtype_t) == sizeof(uint32_t), "Invalid dtype"); + const uint32_t src[2] = {*reinterpret_cast(&src_0), + *reinterpret_cast(&src_1)}; + asm volatile("stmatrix.sync.aligned.m16n8.x2.trans.shared.b8 [%0], {%1, %2};\n" ::"l"( + __cvta_generic_to_shared(smem_dst)), + "r"(src[0]), "r"(src[1])); + } +}; + +/// Shared memory +CUTLASS_DEVICE uint32_t ld_shared(const uint8_t* ptr) { + uint32_t ret; + asm volatile("ld.shared.u8 %0, [%1];" : "=r"(ret) : "l"(__cvta_generic_to_shared(ptr))); + return ret; +} + +CUTLASS_DEVICE uint32_t ld_shared(const uint32_t* ptr) { + uint32_t ret; + asm volatile("ld.shared.u32 %0, [%1];" : "=r"(ret) : "l"(__cvta_generic_to_shared(ptr))); + return ret; +} + +CUTLASS_DEVICE uint2 ld_shared(const uint2* ptr) { + uint2 ret; + asm volatile("ld.shared.v2.u32 {%0, %1}, [%2];" + : "=r"(ret.x), "=r"(ret.y) + : "l"(__cvta_generic_to_shared(ptr))); + return ret; +} + +CUTLASS_DEVICE float2 ld_shared(const float2* ptr) { + float2 ret; + asm volatile("ld.shared.v2.f32 {%0, %1}, [%2];" + : "=f"(ret.x), "=f"(ret.y) + : "l"(__cvta_generic_to_shared(ptr))); + return ret; +} + +CUTLASS_DEVICE float4 ld_shared(const float4* ptr) { + float4 ret; + asm volatile("ld.shared.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(ret.x), "=f"(ret.y), "=f"(ret.z), "=f"(ret.w) + : "l"(__cvta_generic_to_shared(ptr))); + return ret; +} + +CUTLASS_DEVICE uint4 ld_shared(const uint4* ptr) { + uint4 ret; + asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w) + : "l"(__cvta_generic_to_shared(ptr))); + return ret; +} + +CUTLASS_DEVICE float ld_shared(const float* ptr) { + float ret; + asm volatile("ld.shared.f32 %0, [%1];" : "=f"(ret) : "l"(__cvta_generic_to_shared(ptr))); + return ret; +} + +CUTLASS_DEVICE void st_shared(const float* ptr, float val) { + asm volatile("st.shared.f32 [%0], %1;" ::"l"(__cvta_generic_to_shared(ptr)), "f"(val)); +} + +CUTLASS_DEVICE void st_shared_cluster_async( + const float2* ptr, float2 val, const cutlass::arch::ClusterTransactionBarrier* barrier_ptr, + const uint32_t& cta_rank) { + const auto remote_addr = cute::set_block_rank(cute::cast_smem_ptr_to_uint(ptr), cta_rank); + const auto remote_barrier_addr = + cute::set_block_rank(cute::cast_smem_ptr_to_uint(barrier_ptr), cta_rank); + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 " + "[%0], {%1, %2}, [%3];" ::"r"(remote_addr), + "f"(val.x), "f"(val.y), "r"(remote_barrier_addr)); +} + +CUTLASS_DEVICE void st_shared(const float2* ptr, float2 val) { + asm volatile("st.shared.v2.f32 [%0], {%1, %2};" ::"l"(__cvta_generic_to_shared(ptr)), "f"(val.x), + "f"(val.y)); +} + +CUTLASS_DEVICE void st_shared(const uint32_t* ptr, uint32_t val) { + asm volatile("st.shared.u32 [%0], %1;" ::"l"(__cvta_generic_to_shared(ptr)), "r"(val)); +} + +CUTLASS_DEVICE void st_shared(const void* ptr, uint32_t x, uint32_t y) { + asm volatile("st.shared.v2.u32 [%0], {%1, %2};" ::"l"(__cvta_generic_to_shared(ptr)), "r"(x), + "r"(y)); +} + +CUTLASS_DEVICE void st_shared(const void* ptr, uint32_t x, uint32_t y, uint32_t z, uint32_t w) { + asm volatile("st.shared.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(__cvta_generic_to_shared(ptr)), + "r"(x), "r"(y), "r"(z), "r"(w)); +} + +CUTLASS_DEVICE void st_shared(const __int128_t* ptr, __int128_t val) { + asm volatile("st.shared.b128 [%0], %1;" ::"l"(__cvta_generic_to_shared(ptr)), "q"(val)); +} + +/// Global memory +CUTLASS_DEVICE uint32_t ld_volatile(const uint32_t* ptr) { + uint32_t ret; + asm volatile("ld.volatile.global.b32 %0, [%1];" : "=r"(ret) : "l"(ptr)); + return ret; +} + +CUTLASS_DEVICE uint64_t ld_volatile(const uint64_t* ptr) { + uint64_t ret; + asm volatile("ld.volatile.global.b64 %0, [%1];" : "=l"(ret) : "l"(ptr)); + return ret; +} + +CUTLASS_DEVICE uint32_t ld_acq(const uint32_t* ptr) { + uint32_t ret; + asm volatile("ld.acquire.gpu.global.b32 %0, [%1];" : "=r"(ret) : "l"(ptr)); + return ret; +} + +CUTLASS_DEVICE uint64_t ld_acq_sys(const uint64_t* ptr) { + uint64_t ret; + asm volatile("ld.acquire.sys.global.b64 %0, [%1];" : "=l"(ret) : "l"(ptr)); + return ret; +} + +/// Atomics +CUTLASS_DEVICE uint32_t atomic_add(const uint32_t* ptr, const uint32_t& value) { + uint32_t ret; + asm volatile("atom.global.add.u32 %0, [%1], %2;" : "=r"(ret) : "l"(ptr), "r"(value)); + return ret; +} + +CUTLASS_DEVICE uint64_t atomic_add(const uint64_t* ptr, const uint64_t& value) { + uint64_t ret; + asm volatile("atom.global.add.u64 %0, [%1], %2;" : "=l"(ret) : "l"(ptr), "l"(value)); + return ret; +} + +CUTLASS_DEVICE uint64_t atomic_add_sys(const uint64_t* ptr, const uint64_t& value) { + uint64_t ret; + asm volatile("atom.sys.global.add.u64 %0, [%1], %2;" : "=l"(ret) : "l"(ptr), "l"(value)); + return ret; +} + +CUTLASS_DEVICE uint32_t atomic_add_rel(const uint32_t* ptr, const uint32_t& value) { + uint32_t ret; + asm volatile("atom.release.gpu.global.add.u32 %0, [%1], %2;" : "=r"(ret) : "l"(ptr), "r"(value)); + return ret; +} + +CUTLASS_DEVICE void red_add(const int* ptr, const int& value) { + asm volatile("red.gpu.global.add.s32 [%0], %1;" ::"l"(ptr), "r"(value)); +} + +CUTLASS_DEVICE void red_add(const uint32_t* ptr, const uint32_t& value) { + asm volatile("red.gpu.global.add.u32 [%0], %1;" ::"l"(ptr), "r"(value)); +} + +CUTLASS_DEVICE void red_or_rel_gpu(uint64_t* ptr, const uint64_t& value) { + asm volatile("red.release.gpu.global.or.b64 [%0], %1;" ::"l"(ptr), "l"(value)); +} + +CUTLASS_DEVICE void red_add_rel(const uint32_t* ptr, const uint32_t& value) { + asm volatile("red.release.gpu.global.add.u32 [%0], %1;" ::"l"(ptr), "r"(value)); +} + +CUTLASS_DEVICE void red_add_rel_sys(const int* ptr, const int& value) { + asm volatile("red.release.sys.global.add.s32 [%0], %1;" ::"l"(ptr), "r"(value)); +} + +CUTLASS_DEVICE int ld_acq_sys(const int* ptr) { + int ret; + asm volatile("ld.acquire.sys.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr)); + return ret; +} + +CUTLASS_DEVICE uint32_t ld_acq_sys(const uint32_t* ptr) { + uint32_t ret; + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(ret) : "l"(ptr)); + return ret; +} + +CUTLASS_DEVICE uint64_t ld_acq_gpu(const uint64_t* ptr) { + uint64_t ret; + asm volatile("ld.acquire.gpu.global.u64 %0, [%1];" : "=l"(ret) : "l"(ptr)); + return ret; +} + +/// Predicated loads + +CUTLASS_DEVICE int ld_volatile(const int* ptr) { + int ret; + asm volatile("ld.volatile.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr)); + return ret; +} + +CUTLASS_DEVICE uint32_t atomic_add_acq_rel(const uint32_t* ptr, const uint32_t& value) { + uint32_t ret; + asm volatile("atom.acq_rel.gpu.global.add.u32 %0, [%1], %2;" + : "=r"(ret) + : "l"(ptr), "r"(value) + : "memory"); + return ret; +} + +/// Prefetch + +// ---- from ptx_tma.cuh ---- +// Tensor-map instructions + +/// TMA instructions + +CUTLASS_DEVICE void mbarrier_arrive_and_set_tx(cutlass::arch::ClusterTransactionBarrier* ptr, + const uint32_t& num_bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%1], %0; \n\t" ::"r"(num_bytes), + "r"(static_cast(__cvta_generic_to_shared(ptr)))); +} + +CUTLASS_DEVICE void mbarrier_wait(cutlass::arch::ClusterTransactionBarrier* ptr, + const uint32_t& phase) { + asm volatile( + "{\n\t" + ".reg .pred P1; \n\t" + "LAB_WAIT: \n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2; \n\t" + "@P1 bra DONE; \n\t" + "bra LAB_WAIT; \n\t" + "DONE: \n\t" + "}" ::"r"(static_cast(__cvta_generic_to_shared(ptr))), + "r"(phase), "r"(0x989680)); +} + +CUTLASS_DEVICE void mbarrier_wait_and_flip_phase(cutlass::arch::ClusterTransactionBarrier* ptr, + uint32_t& phase) { + mbarrier_wait(ptr, phase); + phase ^= 1; +} + +CUTLASS_DEVICE void tma_load_1d( + const void* dst_ptr, const void* src_ptr, + cutlass::arch::ClusterTransactionBarrier* mbarrier_ptr, const uint32_t& num_bytes, + const cute::TMA::CacheHintSm90& hint = cute::TMA::CacheHintSm90::EVICT_FIRST) { + // NOTES: normally, the loaded part will be evicted soon + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], " + "[%1], %2, [%3], %4;\n" ::"r"(static_cast(__cvta_generic_to_shared(dst_ptr))), + "l"(src_ptr), "r"(num_bytes), + "r"(static_cast(__cvta_generic_to_shared(mbarrier_ptr))), "l"(hint) + : "memory"); +} + +CUTLASS_DEVICE void tma_store_1d( + const void* dst_ptr, const void* src_ptr, const uint32_t& num_bytes, + const cute::TMA::CacheHintSm90& hint = cute::TMA::CacheHintSm90::EVICT_NORMAL) { + // NOTES: normally, the stored part will be used soon + asm volatile( + "cp.async.bulk.global.shared::cta.bulk_group.L2::cache_hint [%0], [%1], %2, %3;\n" ::"l"( + dst_ptr), + "r"(static_cast(__cvta_generic_to_shared(src_ptr))), "r"(num_bytes), "l"(hint) + : "memory"); +} + +template +__forceinline__ __device__ void tma_store_wait() { + // NOTES: this function does not have `.read` + asm volatile("cp.async.bulk.wait_group %0;" ::"n"(kNumRemainingWaits) : "memory"); +} + +// ---- from ptx_wgmma.cuh ---- +CUTLASS_DEVICE void warpgroup_arrive() { asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); } + +CUTLASS_DEVICE void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +CUTLASS_DEVICE void warpgroup_fence_operand(float& reg) { asm volatile("" : "+f"(reg)::"memory"); } + +CUTLASS_DEVICE void warpgroup_fence_operand(uint32_t& reg) { + asm volatile("" : "+r"(reg)::"memory"); +} + +template +CUTLASS_DEVICE void warpgroup_wait() { + FI_MEGAMOE_STATIC_ASSERT(N >= 0 and N <= 7, "WGMMA wait: N must be in range [0, 7]"); + asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory"); +} + +} // namespace flashinfer::megamoe::ptx diff --git a/include/flashinfer/megamoe/detail/tma_copy.cuh b/include/flashinfer/megamoe/detail/tma_copy.cuh new file mode 100644 index 00000000000..cd61e88fdb2 --- /dev/null +++ b/include/flashinfer/megamoe/detail/tma_copy.cuh @@ -0,0 +1,97 @@ +#pragma once + +#include + +#include +#include +#include + +namespace flashinfer::megamoe::tma { + +template +constexpr uint32_t get_inner_block_atom_size() { + return kSwizzleMode == 0 ? BLOCK_INNER : kSwizzleMode / sizeof(dtype_t); +} + +template +CUTLASS_DEVICE void copy(void const* desc_ptr, + cutlass::arch::ClusterTransactionBarrier* barrier_ptr, dtype_t* smem_ptr, + const uint32_t& inner_idx, const uint32_t& outer_idx, + const uint32_t& num_tma_multicast = 1, const uint32_t& batch_idx = 0) { + FI_MEGAMOE_STATIC_ASSERT(static_cast(cute::TMA::CacheHintSm90::EVICT_NORMAL) == + static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), + "Invalid cache hint"); + constexpr uint32_t BLOCK_INNER_ATOM = + get_inner_block_atom_size(); + + if constexpr (not kIs3DTMA) { + if (num_tma_multicast == 1) { +#pragma unroll + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++i) { + cute::SM90_TMA_LOAD_2D::copy(desc_ptr, reinterpret_cast(barrier_ptr), + static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, + inner_idx + i * BLOCK_INNER_ATOM, outer_idx); + } + } else { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) +// 2-CTA function will send signals to the leader CTA only +#pragma unroll + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++i) { + cute::SM100_TMA_2SM_LOAD_2D::copy( + desc_ptr, reinterpret_cast(barrier_ptr), + static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, inner_idx + i * BLOCK_INNER_ATOM, + outer_idx); + } +#elif (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) + if (cute::block_rank_in_cluster() == 0) { +#pragma unroll + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++i) { + cute::SM90_TMA_LOAD_MULTICAST_2D::copy( + desc_ptr, reinterpret_cast(barrier_ptr), (1 << num_tma_multicast) - 1, + static_cast(cute::TMA::CacheHintSm90::EVICT_NORMAL), + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, inner_idx + i * BLOCK_INNER_ATOM, + outer_idx); + } + } +#endif + } + } else { + if (num_tma_multicast == 1) { +#pragma unroll + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++i) { + cute::SM90_TMA_LOAD_3D::copy(desc_ptr, reinterpret_cast(barrier_ptr), + static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, + inner_idx + i * BLOCK_INNER_ATOM, outer_idx, batch_idx); + } + } else { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) +// 2-CTA function will send signals to the leader CTA only +#pragma unroll + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++i) { + cute::SM100_TMA_2SM_LOAD_3D::copy( + desc_ptr, reinterpret_cast(barrier_ptr), + static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, inner_idx + i * BLOCK_INNER_ATOM, + outer_idx, batch_idx); + } +#elif (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) + if (cute::block_rank_in_cluster() == 0) { +#pragma unroll + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++i) { + cute::SM90_TMA_LOAD_MULTICAST_3D::copy( + desc_ptr, reinterpret_cast(barrier_ptr), (1 << num_tma_multicast) - 1, + static_cast(cute::TMA::CacheHintSm90::EVICT_NORMAL), + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, inner_idx + i * BLOCK_INNER_ATOM, + outer_idx, batch_idx); + } + } +#endif + } + } +} + +} // namespace flashinfer::megamoe::tma diff --git a/include/flashinfer/megamoe/layout.cuh b/include/flashinfer/megamoe/layout.cuh new file mode 100644 index 00000000000..9f3c64d43cc --- /dev/null +++ b/include/flashinfer/megamoe/layout.cuh @@ -0,0 +1,371 @@ +#pragma once + +#include +#include + +// Bounded spin (trap after ~10 s at 2 GHz) used by the counter-based +// synchronisation paths (push dispatch DONE flags). +// +// Trap only: a device `printf` (vprintf ABI call) anywhere in the persistent +// kernel costs ~15 % on the math warps (measured: the ABI call changes the +// register allocation of the whole kernel), so no message is printed. +// +// `FI_MEGAMOE_SPIN_WHILE_DEBUG` can publish the spin's call-site tag and the +// observed/expected counter words to `Workspace::get_spin_debug_ptr()` right +// before trapping, so that a bare `cudaErrorLaunchFailure` can be attributed to +// a specific spin after the fact. That store is compiled in ONLY when +// `FI_MEGAMOE_SPIN_DEBUG_TAGS` is defined (host env `FI_MEGAMOE_NVFP4_SPIN_DEBUG=1`, which also +// renames the JIT kernel): on this kernel even a never-executed branch costs +// instruction-cache footprint and ptxas scheduling -- measured +160 SASS +// instructions per kernel when it was compiled in unconditionally -- so the +// default build must generate byte-identical code to the no-debug variant. +// +// `dbg` may be nullptr (the pull path never sets it); the store is predicated. +#ifndef FI_MEGAMOE_SPIN_WHILE_DEBUG +#if defined(FI_MEGAMOE_SPIN_DEBUG_TAGS) +#define FI_MEGAMOE_SPIN_WHILE_DEBUG(cond, tag, dbg, observed, expected) \ + for (long long __spin_t0 = clock64(); (cond);) { \ + if (clock64() - __spin_t0 > 20000000000ll) { \ + uint32_t* __spin_dbg = (dbg); \ + if (__spin_dbg != nullptr) { \ + __spin_dbg[0] = static_cast(tag); \ + __spin_dbg[1] = static_cast(observed); \ + __spin_dbg[2] = static_cast(expected); \ + __spin_dbg[3] = \ + (static_cast(blockIdx.x) << 16) | static_cast(threadIdx.x); \ + __threadfence_system(); \ + } \ + asm volatile("trap;"); \ + } \ + } +#else +#define FI_MEGAMOE_SPIN_WHILE_DEBUG(cond, tag, dbg, observed, expected) \ + for (long long __spin_t0 = clock64(); (cond);) { \ + if (clock64() - __spin_t0 > 20000000000ll) asm volatile("trap;"); \ + } +#endif +#endif + +#ifndef FI_MEGAMOE_SPIN_WHILE +#define FI_MEGAMOE_SPIN_WHILE(cond, tag) FI_MEGAMOE_SPIN_WHILE_DEBUG(cond, tag, nullptr, 0u, 0u) +#endif + +namespace flashinfer::megamoe::layout { + +static constexpr int kNumCandidateBlockMs = 7; +static constexpr int kCandidateBlockM[kNumCandidateBlockMs] = {8, 16, 32, 64, 96, 128, 192}; +static constexpr int kMaxCandidateBlockM = 192; +static constexpr int kMinCandidateBlockM = 8; +static constexpr int kLCMCandidateBlockM = 384; +// D40's dev-m dynamic scheduler owns two global task counters. Keep the +// complete 128-byte scheduler-state ABI used by every accepted D40 run; the +// big-M static split kernels simply leave these counters unused. +static constexpr int kSM90InterleavedSchedulerSMEMBytes = 96; + +// Pool capacity for shared expert token pool: worst-case total tokens + per-expert BLOCK_M +// alignment padding, among all possible BLOCK_M +template +CUTLASS_HOST_DEVICE constexpr T get_num_max_pool_tokens(T num_ranks, T num_max_tokens_per_rank, + T num_topk, T num_experts_per_rank) { + const auto num_max_recv_tokens = num_ranks * num_max_tokens_per_rank; + const auto num_max_experts_per_token = math::constexpr_min(num_topk, num_experts_per_rank); + return math::constexpr_align(num_max_recv_tokens * num_max_experts_per_token + + num_experts_per_rank * (static_cast(kMaxCandidateBlockM) - 1), + static_cast(kLCMCandidateBlockM)); +} + +// SF pool capacity: all experts share a contiguous SF region, sized by pool blocks × SF_BLOCK_M +template +CUTLASS_HOST_DEVICE constexpr T get_num_padded_sf_pool_tokens(T num_max_pool_tokens, T block_m) { + return (num_max_pool_tokens / block_m) * math::constexpr_align(block_m, static_cast(128)); +} + +// Per-token source metadata for combine write-back +struct TokenSrcMetadata { + uint32_t rank_idx; + uint32_t token_idx; + uint32_t topk_idx; +}; + +struct Workspace { + void* base; + uint32_t num_ranks, num_experts; + uint32_t num_experts_per_rank; + uint32_t num_max_tokens_per_rank; + uint32_t num_max_recv_tokens_per_expert; + + // Pool capacity: all local experts share a contiguous token pool + uint32_t num_max_pool_tokens; + uint32_t num_max_pool_blocks; + + // Grid/NVLink state plus the two dev-m dynamic task counters. Expert + // counters begin at the next 128-byte cache line so the hot task counters + // do not share an L2 sector with communication state. + static constexpr uint64_t kNumBarrierSignalBytes = 128; + + CUTLASS_HOST_DEVICE + Workspace(void* base, const uint32_t& num_ranks, const uint32_t& num_experts, + const uint32_t& num_max_tokens_per_rank, const uint32_t& num_topk) + : base(base), + num_ranks(num_ranks), + num_experts(num_experts), + num_max_tokens_per_rank(num_max_tokens_per_rank) { + num_experts_per_rank = num_experts / num_ranks; + num_max_recv_tokens_per_expert = num_ranks * num_max_tokens_per_rank; + num_max_pool_tokens = + get_num_max_pool_tokens(num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank); + num_max_pool_blocks = num_max_pool_tokens / kMinCandidateBlockM; + } + + CUTLASS_HOST_DEVICE + uint64_t get_num_bytes() const { + uint64_t num_bytes = 0; + + // Barrier and in-kernel scheduling counters + num_bytes += kNumBarrierSignalBytes; + + // Expert send/recv count + num_bytes += num_experts * sizeof(uint64_t) * 2; + + // Expert recv count sum (two sets: launch parity for the rotated push pool), + // each set padded to whole 128 B lines (remote sys-scope tickets land here) + num_bytes += math::align(num_experts_per_rank * sizeof(uint64_t), 128) * 2; + + // L1 arrival count (padded to even entry count for `uint64_t` alignment of L2 mask) + num_bytes += math::align(num_max_pool_blocks, 2u) * sizeof(uint32_t); + + // L2 block arrival mask + num_bytes += num_max_pool_blocks * sizeof(uint64_t); + + // Dispatch pulling source token-topk + num_bytes += num_experts_per_rank * num_ranks * num_max_recv_tokens_per_expert * sizeof(int); + + // Combine push source indices + num_bytes += num_max_pool_tokens * sizeof(TokenSrcMetadata); + + // Align to TMA descriptor requirements + num_bytes = math::align(num_bytes, 16); + return num_bytes; + } + + CUTLASS_HOST_DEVICE + void* get_end_ptr() const { return math::advance_ptr(base, get_num_bytes()); } + + // Grid sync counters: `kNumBarrierSignalBytes` layout + // [ 0..15]: 4 x `uint32_t` grid sync counters + // [16..20]: `uint32_t` NVLink barrier counter + // [20..27]: 2 x `int` NVLink barrier signals (phase 0 and 1) + static constexpr uint32_t kNumMaxGridSyncCounters = 4; + + template + CUTLASS_DEVICE uint32_t* get_grid_sync_count_ptr() const { + FI_MEGAMOE_STATIC_ASSERT(kIndex < kNumMaxGridSyncCounters, "Grid sync index out of bounds"); + return static_cast(base) + kIndex; + } + + CUTLASS_DEVICE + uint32_t* get_nvl_barrier_counter_ptr() const { + return static_cast(base) + kNumMaxGridSyncCounters; + } + + CUTLASS_DEVICE + int* get_nvl_barrier_signal_ptr(const uint32_t& phase) const { + // NOTES: the signal is signed, as we may minus + return math::advance_ptr( + base, (kNumMaxGridSyncCounters + 1) * sizeof(uint32_t) + phase * sizeof(int)); + } + + CUTLASS_DEVICE + uint32_t* get_l1_task_count_ptr() const { return math::advance_ptr(base, 28u); } + + CUTLASS_DEVICE + uint32_t* get_l2_task_count_ptr() const { return math::advance_ptr(base, 32u); } + + // Counter-based words (second 64 B sector of the signal line, away from the + // dev-m task counters that gpu-scope atomics hammer all kernel long): + // [64..67] int push DONE count (monotonic; remote red.release.sys, +1 per rank per launch) + // [68..71] push epoch (launches completed; SM0 bumps it in the cleanup; DONE target = + // ranks * (epoch + 1); also the pool parity slot) + // [72..75] push CTA arrival count (last CTA signals DONE and resets it) + // [76..91] spin-timeout debug word (u32[4] = {tag, observed, expected, + // blockIdx|threadIdx}; written once right before a spin traps) + CUTLASS_DEVICE + int* get_push_done_count_ptr() const { return math::advance_ptr(base, 64u); } + + CUTLASS_DEVICE + uint32_t* get_push_epoch_ptr() const { return math::advance_ptr(base, 68u); } + + CUTLASS_DEVICE + uint32_t* get_push_cta_arrival_ptr() const { return math::advance_ptr(base, 72u); } + + // Attribution for a `FI_MEGAMOE_SPIN_WHILE_DEBUG` timeout: a kernel that traps out of a + // counter spin leaves its call-site tag and the observed/expected counter values + // here, so the host can tell a protocol stall from any other launch failure. + CUTLASS_DEVICE + uint32_t* get_spin_debug_ptr() const { return math::advance_ptr(base, 76u); } + + CUTLASS_DEVICE + uint64_t* get_expert_send_count_ptr(const uint32_t& expert_idx = 0) const { + return math::advance_ptr(base, kNumBarrierSignalBytes) + expert_idx; + } + + CUTLASS_DEVICE + uint64_t* get_expert_recv_count_ptr(const uint32_t& rank_idx = 0, + const uint32_t& expert_idx = 0) const { + return get_expert_send_count_ptr(num_experts) + rank_idx * num_experts_per_rank + expert_idx; + } + + CUTLASS_DEVICE + uint32_t get_recv_count_sum_set_stride() const { + // uint64 words per parity set (whole 128 B lines) + return static_cast( + math::align(num_experts_per_rank * sizeof(uint64_t), 128) / sizeof(uint64_t)); + } + + CUTLASS_DEVICE + uint64_t* get_expert_recv_count_sum_ptr(const uint32_t& expert_idx = 0) const { + return get_expert_send_count_ptr(num_experts * 2) + expert_idx; + } + + // Second recv-count-sum set (push dispatch with the rotated pool: launch parity 1) + CUTLASS_DEVICE + uint64_t* get_expert_recv_count_sum_ptr(const uint32_t& expert_idx, + const uint32_t& parity) const { + return get_expert_send_count_ptr(num_experts * 2) + + (parity & 1u) * get_recv_count_sum_set_stride() + expert_idx; + } + + CUTLASS_DEVICE + uint32_t* get_l1_arrival_count_ptr(const uint32_t& pool_block_idx = 0) const { + const auto base = get_expert_recv_count_sum_ptr(get_recv_count_sum_set_stride() * 2); + return reinterpret_cast(base) + pool_block_idx; + } + + CUTLASS_DEVICE + uint64_t* get_l2_arrival_mask_ptr(const uint32_t& pool_block_idx = 0) const { + // Pad L1 entry count to even so that the `l2_arrival_mask` is 8-byte aligned + const auto base = get_l1_arrival_count_ptr(math::align(num_max_pool_blocks, 2u)); + return reinterpret_cast(base) + pool_block_idx; + } + + // For dispatch pulling + CUTLASS_DEVICE + uint32_t* get_src_token_topk_idx_ptr(const uint32_t& expert_idx = 0, const uint32_t& rank_idx = 0, + const uint32_t& token_idx = 0) const { + const auto base = get_l2_arrival_mask_ptr(num_max_pool_blocks); + return reinterpret_cast(base) + + expert_idx * (num_ranks * num_max_recv_tokens_per_expert) + + rank_idx * num_max_recv_tokens_per_expert + token_idx; + } + + // For combine usages + CUTLASS_DEVICE + TokenSrcMetadata* get_token_src_metadata_ptr(const uint32_t& pool_token_idx = 0) const { + const auto base = + reinterpret_cast(get_src_token_topk_idx_ptr(num_experts_per_rank)); + return base + pool_token_idx; + } +}; + +struct Data { + uint32_t num_bytes; + bool require_tma_alignment; + void* base; + + CUTLASS_HOST_DEVICE + constexpr explicit Data(const uint32_t& num_bytes, const bool& require_tma_alignment = true, + void* base = nullptr) + : num_bytes(num_bytes), require_tma_alignment(require_tma_alignment), base(base) { + FI_MEGAMOE_UNIFIED_ASSERT(num_bytes % 16 == 0 or not require_tma_alignment); + } + + template + CUTLASS_HOST_DEVICE constexpr dtype_t get_num_bytes() const { + return static_cast(num_bytes); + } + + template + CUTLASS_HOST_DEVICE dtype_t* get_base_ptr() const { + return static_cast(base); + } + + CUTLASS_HOST_DEVICE void set_base_ptr(void* ptr) { base = ptr; } +}; + +struct Buffer { + Data data_layout; + uint32_t num_ranks; + uint32_t num_max_tokens_per_rank; + + void* base; + + CUTLASS_HOST_DEVICE + Buffer(const Data& data_layout, const uint32_t& num_ranks, + const uint32_t& max_num_tokens_per_rank, void* base = nullptr) + : data_layout(data_layout), + num_ranks(num_ranks), + num_max_tokens_per_rank(max_num_tokens_per_rank), + base(base) {} + + CUTLASS_HOST_DEVICE + uint64_t get_num_bytes_per_rank() const { + return num_max_tokens_per_rank * data_layout.get_num_bytes(); + } + + CUTLASS_HOST_DEVICE + uint64_t get_num_bytes() const { return get_num_bytes_per_rank() * num_ranks; } + + template + CUTLASS_HOST_DEVICE dtype_t* get_base_ptr() const { + return static_cast(base); + } + + CUTLASS_HOST_DEVICE + void* get_end_ptr() const { return math::advance_ptr(base, get_num_bytes()); } + + CUTLASS_HOST_DEVICE + Buffer get_rank_buffer(const uint32_t& rank_idx) const { + return {data_layout, 1, num_max_tokens_per_rank, + math::advance_ptr(base, get_num_bytes_per_rank() * rank_idx)}; + } + + CUTLASS_HOST_DEVICE + Data get_data_buffer(const uint32_t& token_idx, const bool& global = false) const { + FI_MEGAMOE_DEVICE_ASSERT(num_ranks == 1 or global); + return Data(data_layout.num_bytes, data_layout.require_tma_alignment, + math::advance_ptr(base, data_layout.get_num_bytes() * token_idx)); + } +}; + +// One allocation ABI shared by NVFP4 fused/split and MXFP4 fused. Activation +// scales reserve per-64 capacity for L2; BN256 kernels use its dense per-128 +// prefix. The combine region is indexed by the original token and top-k slot. +struct MegaMoEBuffers { + Workspace workspace; + Buffer input_token_buffer, input_sf_buffer, input_topk_idx_buffer, input_topk_weights_buffer; + Buffer l1_token_buffer, l1_sf_buffer, l1_topk_weights_buffer; + Buffer l2_token_buffer, l2_sf_buffer, combine_token_buffer; + + CUTLASS_HOST_DEVICE + MegaMoEBuffers(void* base, uint32_t ranks, uint32_t experts, uint32_t capacity, uint32_t topk, + uint32_t hidden, uint32_t intermediate, uint32_t pool_tokens, + uint32_t sf_pool_tokens) + : workspace(base, ranks, experts, capacity, topk), + input_token_buffer(Data(hidden), 1, capacity, workspace.get_end_ptr()), + input_sf_buffer(Data(hidden / 32, false), 1, capacity, input_token_buffer.get_end_ptr()), + input_topk_idx_buffer(Data(topk * sizeof(int64_t), false), 1, capacity, + input_sf_buffer.get_end_ptr()), + input_topk_weights_buffer(Data(topk * sizeof(float), false), 1, capacity, + input_topk_idx_buffer.get_end_ptr()), + l1_token_buffer(Data(hidden), 1, pool_tokens, input_topk_weights_buffer.get_end_ptr()), + l1_sf_buffer(Data(hidden / 32, false), 1, sf_pool_tokens, l1_token_buffer.get_end_ptr()), + l1_topk_weights_buffer(Data(sizeof(float), false), 1, pool_tokens, + l1_sf_buffer.get_end_ptr()), + l2_token_buffer(Data(intermediate), 1, pool_tokens, l1_topk_weights_buffer.get_end_ptr()), + l2_sf_buffer(Data(intermediate / 16, false), 1, sf_pool_tokens, + l2_token_buffer.get_end_ptr()), + combine_token_buffer(Data(hidden * sizeof(uint16_t)), topk, capacity, + l2_sf_buffer.get_end_ptr()) {} +}; + +} // namespace flashinfer::megamoe::layout diff --git a/include/flashinfer/megamoe/mxfp4_dequant.cuh b/include/flashinfer/megamoe/mxfp4_dequant.cuh new file mode 100644 index 00000000000..02a245cfc9c --- /dev/null +++ b/include/flashinfer/megamoe/mxfp4_dequant.cuh @@ -0,0 +1,165 @@ +// SPDX-License-Identifier: MIT +// +// MXFP4 (E2M1 + E8M0) -> FP8 (E4M3) dequant for SM90 fused MegaMoE. +// +// E8M0 is a pure power of two, so scaling by 2^k is an exponent adjustment: +// adding k to an E4M3 value's 4-bit exponent field is adding (k << 3) to the +// byte, and the whole magnitude table for a scale group fits in two registers. +// +// Folding that shift arithmetically per group is nonetheless the wrong trade on +// SM90 -- the saturation/flush edge cases need byte-wise SIMD intrinsics Hopper +// does not implement. make_scaled_lut() below remains the definition of the +// numerics, but the mainloop reads the tables from a shared-memory copy built +// from it once per CTA; see kScaledLutSize. + +#pragma once + +#include +#include + +namespace flashinfer::megamoe { +namespace mxfp4 { + +#define FI_MEGAMOE_MXFP4_INLINE __device__ __forceinline__ + +// __byte_perm without the compiler's range check on the selector; the selectors +// below are masked to 0x7 lanes by construction. +FI_MEGAMOE_MXFP4_INLINE std::uint32_t byte_perm_unchecked(std::uint32_t a, std::uint32_t b, + std::uint32_t selector) { + return detail::byte_perm_unchecked(a, b, selector); +} + +// E2M1 magnitudes {0, .5, 1, 1.5, 2, 3, 4, 6} encoded as E4M3 at scale 2^0. +// Byte i of (LUT_X, LUT_Y) holds magnitude index i and i+4 respectively, the +// layout __byte_perm consumes. +static constexpr uint32_t kBaseLutX = 0x3c383000u; +static constexpr uint32_t kBaseLutY = 0x4c484440u; + +// Largest finite E4M3 (448.0). Saturating adds clamp to 0xff, so results are +// pulled back to this. +static constexpr uint32_t kE4M3MaxBytes = 0x7e7e7e7eu; +// Exponent field 0 encodes zero/subnormal; E4M3 subnormals are not produced by +// this path, so any byte below 0x08 must flush to zero. +static constexpr uint32_t kMinNormalBytes = 0x08080808u; + +// The 8-byte magnitude table for one E8M0 scale group, built in registers. +struct ScaledLut { + std::uint32_t x; + std::uint32_t y; +}; + +// Fold scale 2^(code - 127) into the base table -- branchless. +// +// k's sign and magnitude both come from the weight group's own scale byte, so +// they vary per element and are independent across the 32 lanes of a warp +// (each lane dequantizes a different B-tile row). An `if (k >= 0)` or +// `if (in fast range)` here is therefore not a rare/uniform branch: whenever +// two lanes in the same warp fall on different sides, the SM serializes them, +// and that scalar branch cost was found (via SASS inspection -- @!P0 BRA in +// the compiled make_scaled_lut) to dominate the two saturating vadds it was +// meant to save. So both signs are computed unconditionally and blended with +// an arithmetic mask (SEL/PRMT), never a control-flow branch. +// +// k >= 0: saturating byte add, then clamp to 448. Byte 0 of x is magnitude 0 +// and must stay zero, so it is masked off -- adding to it would +// manufacture a nonzero value. +// k < 0: saturating byte subtract, which already floors at zero, then flush +// any byte whose exponent field reached 0. That case is reachable with +// real weights: a group whose max is below 6 * 2^-5 gets k < -5, and +// its smallest magnitude underflows E4M3's normal range. +// +// |k| * 8 is clamped to 255 before broadcasting; beyond that every byte +// saturates anyway, and an unclamped broadcast would corrupt neighbouring +// bytes. +FI_MEGAMOE_MXFP4_INLINE ScaledLut make_scaled_lut(std::uint32_t scale_ue8m0) { + const int k = static_cast(scale_ue8m0 & 0xffu) - 127; + const std::uint32_t magnitude = static_cast(k >= 0 ? k : -k) * 8u; + const std::uint32_t delta = (magnitude > 255u ? 255u : magnitude) * 0x01010101u; + + const std::uint32_t pos_x = __vminu4(__vaddus4(kBaseLutX, delta), kE4M3MaxBytes) & 0xffffff00u; + const std::uint32_t pos_y = __vminu4(__vaddus4(kBaseLutY, delta), kE4M3MaxBytes); + + const std::uint32_t neg_x_raw = __vsubus4(kBaseLutX, delta); + const std::uint32_t neg_y_raw = __vsubus4(kBaseLutY, delta); + const std::uint32_t neg_x = neg_x_raw & __vcmpgeu4(neg_x_raw, kMinNormalBytes); + const std::uint32_t neg_y = neg_y_raw & __vcmpgeu4(neg_y_raw, kMinNormalBytes); + + // All-1s if k < 0, all-0s otherwise -- an arithmetic mask, not a predicate + // on control flow, so ptxas has nothing to branch on. + const std::uint32_t neg_mask = static_cast(-(k < 0)); + ScaledLut lut; + lut.x = (pos_x & ~neg_mask) | (neg_x & neg_mask); + lut.y = (pos_y & ~neg_mask) | (neg_y & neg_mask); + return lut; +} + +// Shared-memory window over the scaled tables. +// +// make_scaled_lut() is exact but expensive on SM90: its byte-wise saturating +// intrinsics (__vaddus4 / __vsubus4 / __vminu4 / __vcmpgeu4) have no native +// instruction on Hopper, so ptxas emulates them -- a single call compiles to +// 88 SASS instructions (39 LOP3 + 8 PRMT + ...) versus 40 for NVFP4's whole +// table-load-plus-decode. Calling it once per 32-element scale group, i.e. +// four times per BK128 weight row per thread, dominated the mainloop. +// +// The table is constant outside a narrow band of E8M0 codes: sweeping all 256 +// codes yields only 20 distinct tables, at codes 118..137 (k = -9..+10). Below +// that every magnitude underflows E4M3's normal range and flushes to zero, +// above it every magnitude saturates at 448. That makes a clamped 32-entry +// window bit-exact, and it is what this used to be -- but the clamp is not +// free: it costs two VIMNMX and a subtract per lookup, twelve instructions per +// decode iteration across four slices. Indexing all 256 codes directly costs +// 2 KB of shared memory and none of that arithmetic, so the table is stored +// whole. Still half as many lookups per row as NVFP4, which groups 16 elements +// where MXFP4 groups 32. +static constexpr std::uint32_t kScaledLutSize = 256; + +// A decode tile is `kBTileRows` fused rows stored 16-byte-chunk-major: the +// 16-byte chunk `c` of row `r` sits at `c * kBTileRows * 16 + r * 16`, and the +// row's four E8M0 bytes at `kBTileRows * 64 + r * 4`. Laying the chunks out +// this way is what lets the 12 bytes of per-row padding go -- 15 % of the +// weight stream -- while keeping every 16-byte decoder load aligned, which a +// flat 68-byte row would not. +static constexpr std::uint32_t kBTileRows = 128; +static constexpr std::uint32_t kBChunkStride = kBTileRows * 16; +static constexpr std::uint32_t kBTileFP4Bytes = kBTileRows * 64; +static constexpr std::uint32_t kBTileBytes = kBTileFP4Bytes + kBTileRows * 4; + +struct DecodeTileRow { + const std::uint8_t* chunk; // chunk c at `chunk + c * kBChunkStride` + const std::uint8_t* scale; +}; + +FI_MEGAMOE_MXFP4_INLINE DecodeTileRow decode_tile_row(const std::uint8_t* __restrict__ packed_b, + const std::uint32_t row) { + const std::uint8_t* tile = packed_b + (row / kBTileRows) * kBTileBytes; + const std::uint32_t sub = row % kBTileRows; + return {tile + sub * 16, tile + kBTileFP4Bytes + sub * 4}; +} + +// Fill the table once during initialization, striding so that a block with +// fewer than kScaledLutSize threads still fills all of it. The entries are +// built by the same make_scaled_lut() the decoders used to call inline, so the +// numerics are unchanged by construction. +FI_MEGAMOE_MXFP4_INLINE void init_scaled_lut(ScaledLut* __restrict__ smem_lut, + const std::uint32_t thread_idx, + const std::uint32_t num_threads) { + for (std::uint32_t code = thread_idx; code < kScaledLutSize; code += num_threads) + smem_lut[code] = make_scaled_lut(code); +} + +FI_MEGAMOE_MXFP4_INLINE ScaledLut load_scaled_lut(const ScaledLut* __restrict__ smem_lut, + const std::uint32_t scale_ue8m0) { + return smem_lut[scale_ue8m0 & 0xffu]; +} + +// Decode eight packed FP4 nibbles (one uint32) into eight FP8 bytes, using a +// table already scaled by make_scaled_lut. +FI_MEGAMOE_MXFP4_INLINE uint2 dequant_word(std::uint32_t packed, const ScaledLut& lut) { + return detail::decode_fp4_word(packed, lut); +} + +#undef FI_MEGAMOE_MXFP4_INLINE + +} // namespace mxfp4 +} // namespace flashinfer::megamoe diff --git a/include/flashinfer/megamoe/nvfp4_dequant.cuh b/include/flashinfer/megamoe/nvfp4_dequant.cuh new file mode 100644 index 00000000000..0e530324d23 --- /dev/null +++ b/include/flashinfer/megamoe/nvfp4_dequant.cuh @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: MIT +// +// NVFP4 (E2M1 + UE4M3 scale) -> FP8 (E4M3) dequant helper for SM90 +// fused MegaMoE. The packed FP4 layout matches the existing Marlin-style +// byte packing used by the predecessor FP4 path, with NVFP4 UE4M3 +// scale bytes applied per 16 K elements. + +#pragma once + +#include + +#include +#include + +namespace flashinfer::megamoe { +namespace nvfp4 { + +#define FI_MEGAMOE_NVFP4_INLINE __device__ __forceinline__ + +// Per-scale lookup table for positive E2M1 magnitudes 0..7 converted to FP8. +// Each row packs magnitudes 0..3 in .x and 4..7 in .y so __byte_perm can map +// four FP4 nibbles to four FP8 bytes with one instruction. Scale code 0x7f is +// treated as a saturating 448.0 fallback; the quantizer emits only 0x00..0x7e. +static __device__ __constant__ __align__(16) const uint2 kE2M1AndUe4m3ToFp8Lut[128] = { + {0x00000000u, 0x00000000u}, {0x02010000u, 0x06040302u}, {0x03020100u, 0x0c080604u}, + {0x04030200u, 0x110c0906u}, {0x06040200u, 0x14100c08u}, {0x08050200u, 0x17120f0au}, + {0x09060300u, 0x1914110cu}, {0x0a070400u, 0x1a16120eu}, {0x0c080400u, 0x1c181410u}, + {0x0e090400u, 0x1e191611u}, {0x0f0a0500u, 0x1f1a1712u}, {0x100b0600u, 0x201b1813u}, + {0x110c0600u, 0x211c1914u}, {0x120d0600u, 0x221d1a15u}, {0x120e0700u, 0x221e1a16u}, + {0x130f0800u, 0x231f1b17u}, {0x14100800u, 0x24201c18u}, {0x16110900u, 0x26211e19u}, + {0x17120a00u, 0x27221f1au}, {0x18130b00u, 0x2823201bu}, {0x19140c00u, 0x2924211cu}, + {0x1a150d00u, 0x2a25221du}, {0x1a160e00u, 0x2a26221eu}, {0x1b170f00u, 0x2b27231fu}, + {0x1c181000u, 0x2c282420u}, {0x1e191100u, 0x2e292621u}, {0x1f1a1200u, 0x2f2a2722u}, + {0x201b1300u, 0x302b2823u}, {0x211c1400u, 0x312c2924u}, {0x221d1500u, 0x322d2a25u}, + {0x221e1600u, 0x322e2a26u}, {0x231f1700u, 0x332f2b27u}, {0x24201800u, 0x34302c28u}, + {0x26211900u, 0x36312e29u}, {0x27221a00u, 0x37322f2au}, {0x28231b00u, 0x3833302bu}, + {0x29241c00u, 0x3934312cu}, {0x2a251d00u, 0x3a35322du}, {0x2a261e00u, 0x3a36322eu}, + {0x2b271f00u, 0x3b37332fu}, {0x2c282000u, 0x3c383430u}, {0x2e292100u, 0x3e393631u}, + {0x2f2a2200u, 0x3f3a3732u}, {0x302b2300u, 0x403b3833u}, {0x312c2400u, 0x413c3934u}, + {0x322d2500u, 0x423d3a35u}, {0x322e2600u, 0x423e3a36u}, {0x332f2700u, 0x433f3b37u}, + {0x34302800u, 0x44403c38u}, {0x36312900u, 0x46413e39u}, {0x37322a00u, 0x47423f3au}, + {0x38332b00u, 0x4843403bu}, {0x39342c00u, 0x4944413cu}, {0x3a352d00u, 0x4a45423du}, + {0x3a362e00u, 0x4a46423eu}, {0x3b372f00u, 0x4b47433fu}, {0x3c383000u, 0x4c484440u}, + {0x3e393100u, 0x4e494641u}, {0x3f3a3200u, 0x4f4a4742u}, {0x403b3300u, 0x504b4843u}, + {0x413c3400u, 0x514c4944u}, {0x423d3500u, 0x524d4a45u}, {0x423e3600u, 0x524e4a46u}, + {0x433f3700u, 0x534f4b47u}, {0x44403800u, 0x54504c48u}, {0x46413900u, 0x56514e49u}, + {0x47423a00u, 0x57524f4au}, {0x48433b00u, 0x5853504bu}, {0x49443c00u, 0x5954514cu}, + {0x4a453d00u, 0x5a55524du}, {0x4a463e00u, 0x5a56524eu}, {0x4b473f00u, 0x5b57534fu}, + {0x4c484000u, 0x5c585450u}, {0x4e494100u, 0x5e595651u}, {0x4f4a4200u, 0x5f5a5752u}, + {0x504b4300u, 0x605b5853u}, {0x514c4400u, 0x615c5954u}, {0x524d4500u, 0x625d5a55u}, + {0x524e4600u, 0x625e5a56u}, {0x534f4700u, 0x635f5b57u}, {0x54504800u, 0x64605c58u}, + {0x56514900u, 0x66615e59u}, {0x57524a00u, 0x67625f5au}, {0x58534b00u, 0x6863605bu}, + {0x59544c00u, 0x6964615cu}, {0x5a554d00u, 0x6a65625du}, {0x5a564e00u, 0x6a66625eu}, + {0x5b574f00u, 0x6b67635fu}, {0x5c585000u, 0x6c686460u}, {0x5e595100u, 0x6e696661u}, + {0x5f5a5200u, 0x6f6a6762u}, {0x605b5300u, 0x706b6863u}, {0x615c5400u, 0x716c6964u}, + {0x625d5500u, 0x726d6a65u}, {0x625e5600u, 0x726e6a66u}, {0x635f5700u, 0x736f6b67u}, + {0x64605800u, 0x74706c68u}, {0x66615900u, 0x76716e69u}, {0x67625a00u, 0x77726f6au}, + {0x68635b00u, 0x7873706bu}, {0x69645c00u, 0x7974716cu}, {0x6a655d00u, 0x7a75726du}, + {0x6a665e00u, 0x7a76726eu}, {0x6b675f00u, 0x7b77736fu}, {0x6c686000u, 0x7c787470u}, + {0x6e696100u, 0x7e797671u}, {0x6f6a6200u, 0x7f7a7772u}, {0x706b6300u, 0x7f7b7873u}, + {0x716c6400u, 0x7f7c7974u}, {0x726d6500u, 0x7f7d7a75u}, {0x726e6600u, 0x7f7e7a76u}, + {0x736f6700u, 0x7f7f7b77u}, {0x74706800u, 0x7f7f7c78u}, {0x76716900u, 0x7f7f7e79u}, + {0x77726a00u, 0x7f7f7f7au}, {0x78736b00u, 0x7f7f7f7bu}, {0x79746c00u, 0x7f7f7f7cu}, + {0x7a756d00u, 0x7f7f7f7du}, {0x7a766e00u, 0x7f7f7f7eu}, {0x7b776f00u, 0x7f7f7f7fu}, + {0x7c787000u, 0x7f7f7f7fu}, {0x7e797100u, 0x7f7f7f7fu}, {0x7f7a7200u, 0x7f7f7f7fu}, + {0x7f7b7300u, 0x7f7f7f7fu}, {0x7f7c7400u, 0x7f7f7f7fu}, {0x7f7d7500u, 0x7f7f7f7fu}, + {0x7f7e7600u, 0x7f7f7f7fu}, {0x7f7e7600u, 0x7f7f7f7fu}}; + +FI_MEGAMOE_NVFP4_INLINE uint2 load_e2m1_ue4m3_lut(std::uint32_t scale_ue4m3) { + return kE2M1AndUe4m3ToFp8Lut[scale_ue4m3 & 0x7fu]; +} + +template +FI_MEGAMOE_NVFP4_INLINE std::uint32_t pack_nvfp4_magnitude_selector(std::uint32_t byte_magnitudes) { + if constexpr (kUseDp4a) { + const std::uint32_t lo = __dp4a(byte_magnitudes, 0x00001001u, 0u); + const std::uint32_t hi = __dp4a(byte_magnitudes, 0x10010000u, 0u); + return lo + (hi << 8); + } else { + const std::uint32_t packed_pairs = byte_magnitudes + (byte_magnitudes >> 4); + return __byte_perm(packed_pairs, 0u, 0x4420u); + } +} + +FI_MEGAMOE_NVFP4_INLINE std::uint32_t byte_perm_unchecked(std::uint32_t a, std::uint32_t b, + std::uint32_t selector) { + return detail::byte_perm_unchecked(a, b, selector); +} + +template +FI_MEGAMOE_NVFP4_INLINE uint2 dequant_nvfp4_to_fp8_pair_with_lut(std::uint32_t uq, + const uint2& lut) { + const std::uint32_t sel_hi = pack_nvfp4_magnitude_selector((uq >> 4) & 0x07070707u); + const std::uint32_t sel_lo = pack_nvfp4_magnitude_selector(uq & 0x07070707u); + + std::uint32_t out_hi = byte_perm_unchecked(lut.x, lut.y, sel_hi); + std::uint32_t out_lo = byte_perm_unchecked(lut.x, lut.y, sel_lo); + out_hi |= uq & 0x80808080u; + out_lo |= (uq << 4) & 0x80808080u; + return make_uint2(out_hi, out_lo); +} + +#undef FI_MEGAMOE_NVFP4_INLINE + +} // namespace nvfp4 +} // namespace flashinfer::megamoe diff --git a/include/flashinfer/megamoe/scheduler.cuh b/include/flashinfer/megamoe/scheduler.cuh new file mode 100644 index 00000000000..a1b02f7df1b --- /dev/null +++ b/include/flashinfer/megamoe/scheduler.cuh @@ -0,0 +1,526 @@ +#pragma once + +#include + +#include +#include +#include + +namespace flashinfer::megamoe::sched { + +// Computation phase for the current block +enum class BlockPhase { None = 0, Linear1 = 1, Linear2 = 2 }; + +template < + uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K, uint32_t L1_SHAPE_N, uint32_t L1_SHAPE_K, + uint32_t L2_SHAPE_N, uint32_t L2_SHAPE_K, uint32_t kNumExpertsPerRank, + uint32_t kNumExpertsPerWave, uint32_t kNumSMs, uint32_t kNumRanks, uint32_t kClusterSize = 2, + uint32_t kPoolBlockM = BLOCK_M, + uint32_t kNumExpertsPerLane = math::constexpr_ceil_div(kNumExpertsPerRank, 32u), + uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N, uint32_t kNumL2BlockNs = L2_SHAPE_N / BLOCK_N, + uint32_t kNumL1BlockKs = L1_SHAPE_K / BLOCK_K, uint32_t kNumL2BlockKs = L2_SHAPE_K / BLOCK_K> +struct MegaMoEScheduler { + FI_MEGAMOE_STATIC_ASSERT(L1_SHAPE_N % BLOCK_N == 0, "Invalid shape"); + FI_MEGAMOE_STATIC_ASSERT(L2_SHAPE_N % BLOCK_N == 0, "Invalid shape"); + FI_MEGAMOE_STATIC_ASSERT(L1_SHAPE_K % BLOCK_K == 0, "Invalid shape"); + FI_MEGAMOE_STATIC_ASSERT(L2_SHAPE_K % BLOCK_K == 0, "Invalid shape"); + FI_MEGAMOE_STATIC_ASSERT(kNumExpertsPerRank % kNumExpertsPerWave == 0, "Invalid wave config"); + FI_MEGAMOE_STATIC_ASSERT(kPoolBlockM > 0, "Pool block M must be positive"); + + // For 2-CTA clusters, neighbour SMs share the same m_block_idx with adjacent + // n_block_idx; the asserts below guarantee that pairing is always possible. + // SM90 / single-CTA paths set kClusterSize = 1 and do not need this. + FI_MEGAMOE_STATIC_ASSERT(kClusterSize == 1 or kClusterSize == 2, "Invalid cluster size"); + FI_MEGAMOE_STATIC_ASSERT(kClusterSize == 1 or kNumSMs % 2 == 0, + "Number of SMs must be even for 2-CTA cluster"); + FI_MEGAMOE_STATIC_ASSERT(kClusterSize == 1 or kNumL1BlockNs % 2 == 0, + "L1 N block count must be even for 2-CTA cluster"); + FI_MEGAMOE_STATIC_ASSERT(kClusterSize == 1 or kNumL2BlockNs % 2 == 0, + "L2 N block count must be even for 2-CTA cluster"); + + // Arrival counts + const layout::Workspace& workspace; + + // Scheduler state + BlockPhase next_phase = BlockPhase::Linear1; + + // Current expert and block indices + uint32_t current_local_expert_idx = 0; + uint32_t current_num_tokens = 0; + uint32_t current_pool_block_offset = 0; + uint32_t block_idx = 0; + uint32_t m_block_idx = 0; + uint32_t n_block_idx = 0; + + // Pre-cached per-expert token counts (filled during `for_each_block` init) + // Layout: `stored_num_tokens_per_expert[i]` holds expert (i * 32 + lane_idx)'s count + uint32_t stored_num_tokens_per_expert[kNumExpertsPerLane] = {}; + + CUTLASS_DEVICE explicit MegaMoEScheduler(const layout::Workspace& workspace) + : workspace(workspace) { + block_idx = blockIdx.x; + } + + CUTLASS_DEVICE uint32_t get_wave_expert_end_idx() const { + return math::align(current_local_expert_idx + 1, kNumExpertsPerWave); + } + + CUTLASS_DEVICE uint32_t get_num_tokens(const uint32_t& expert_idx) const { + uint32_t valid_value; +#pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++i) { + valid_value = (expert_idx == i * 32 + ptx::get_lane_idx()) ? stored_num_tokens_per_expert[i] + : valid_value; + } + return ptx::exchange(valid_value, expert_idx % 32); + } + + // Get pool block offset for a given expert index from a per-lane token count array + CUTLASS_DEVICE uint32_t get_pool_block_offset(const uint32_t& expert_idx) { + uint32_t num_blocks = 0; +#pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++i) { + if (i * 32 + ptx::get_lane_idx() < expert_idx) + num_blocks += math::ceil_div(stored_num_tokens_per_expert[i], kPoolBlockM); + } + return __reduce_add_sync(0xffffffff, num_blocks); + } + + CUTLASS_DEVICE void advance_expert_idx() { + current_pool_block_offset += get_current_num_pool_blocks(); + current_local_expert_idx += 1; + current_num_tokens = get_num_tokens(current_local_expert_idx); + } + + CUTLASS_DEVICE void set_expert_idx(const uint32_t& expert_idx) { + current_local_expert_idx = expert_idx; + current_num_tokens = get_num_tokens(expert_idx); + current_pool_block_offset = get_pool_block_offset(expert_idx); + } + + CUTLASS_DEVICE uint32_t get_current_pool_block_offset() const { + return current_pool_block_offset; + } + + CUTLASS_DEVICE uint32_t get_current_num_m_blocks() const { + return math::ceil_div(current_num_tokens, BLOCK_M); + } + + CUTLASS_DEVICE uint32_t get_current_num_pool_blocks() const { + return math::ceil_div(current_num_tokens, kPoolBlockM); + } + + CUTLASS_DEVICE uint32_t get_current_pool_token_offset() const { + return current_pool_block_offset * kPoolBlockM; + } + + CUTLASS_DEVICE uint32_t get_current_block_pool_token_idx() const { + return get_current_pool_token_offset() + m_block_idx * BLOCK_M; + } + + template + CUTLASS_DEVICE uint32_t get_valid_m() const { + const auto m = cute::min(current_num_tokens - m_block_idx * BLOCK_M, BLOCK_M); + return kDoUMMAAligned ? math::align(m, 16u) : m; + } + + CUTLASS_DEVICE bool fetch_next_l1_block() { + const auto wave_end_expert_idx = get_wave_expert_end_idx(); + while (current_local_expert_idx < wave_end_expert_idx) { + const auto num_m_blocks = get_current_num_m_blocks(); + m_block_idx = block_idx / kNumL1BlockNs; + if (m_block_idx < num_m_blocks) return true; + + // Current expert is fully assigned, move to the next + block_idx -= num_m_blocks * kNumL1BlockNs; + advance_expert_idx(); + } + return false; + } + + CUTLASS_DEVICE bool fetch_next_l2_block() { + const auto wave_end_expert_idx = get_wave_expert_end_idx(); + while (current_local_expert_idx < wave_end_expert_idx) { + const auto num_m_blocks = get_current_num_m_blocks(); + if (block_idx < num_m_blocks * kNumL2BlockNs) { + m_block_idx = block_idx / kNumL2BlockNs; + return true; + } + + // Current expert is fully assigned, move to the next + block_idx -= num_m_blocks * kNumL2BlockNs; + advance_expert_idx(); + } + return false; + } + + // Core state machine: assigns the next block + CUTLASS_DEVICE cute::tuple get_next_block() { + while (true) { + if (current_local_expert_idx >= kNumExpertsPerRank) break; + + if (next_phase == BlockPhase::Linear1) { + if (fetch_next_l1_block()) { + // Found a new L1 block + n_block_idx = block_idx - m_block_idx * kNumL1BlockNs; + // Jump to next block + block_idx += kNumSMs; + return {BlockPhase::Linear1, current_local_expert_idx, m_block_idx, n_block_idx}; + } else { + // L1 for the current wave is complete, transition to L2 + next_phase = BlockPhase::Linear2; + set_expert_idx( + math::align(current_local_expert_idx - 1, kNumExpertsPerWave)); + } + } else { + if (fetch_next_l2_block()) { + // Found a new L2 block + n_block_idx = block_idx - m_block_idx * kNumL2BlockNs; + // Jump to next block + block_idx += kNumSMs; + return {BlockPhase::Linear2, current_local_expert_idx, m_block_idx, n_block_idx}; + } else { + // Move to L1 of the next wave + next_phase = BlockPhase::Linear1; + } + } + } + + // All waves and experts are fully processed + return {BlockPhase::None, 0, 0, 0}; + } + + CUTLASS_DEVICE void fetch_expert_recv_count() { +// NOTES: each lane caches experts at indices (i * 32 + lane_idx) +#pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++i) { + const auto expert_idx = i * 32 + ptx::get_lane_idx(); + uint64_t value = 0; + if (expert_idx < kNumExpertsPerRank) { + do { + value = ptx::ld_volatile(workspace.get_expert_recv_count_sum_ptr(expert_idx)); + } while (static_cast(value >> 32) != kNumSMs * kNumRanks); + } + stored_num_tokens_per_expert[i] = static_cast(value); + } + __syncwarp(); + } + + template + CUTLASS_DEVICE void for_each_block(Func&& func) { + // Wait for all expert counters to be finalized + fetch_expert_recv_count(); + + // Initialize current expert with 0 + set_expert_idx(0); + + // Iterate over all blocks + // TODO: add swizzle within expert waves for better L2 cache utilization + while (true) { + CUTE_TIE_DECL(get_next_block(), block_phase, current_local_expert_idx, m_block_idx, + n_block_idx); + if (block_phase == BlockPhase::None) break; + + func(block_phase, current_local_expert_idx, + block_phase == BlockPhase::Linear2 ? kNumL2BlockKs : kNumL1BlockKs, m_block_idx, + n_block_idx); + } + } + + template + CUTLASS_DEVICE void for_each_linear1_block(Func&& func) { + // Split-kernel mode: K1 owns only dispatch + Linear1. Unlike + // for_each_block(), do not burn scheduler iterations on Linear2 blocks. + fetch_expert_recv_count(); + set_expert_idx(0); + while (current_local_expert_idx < kNumExpertsPerRank) { + if (fetch_next_l1_block()) { + n_block_idx = block_idx - m_block_idx * kNumL1BlockNs; + block_idx += kNumSMs; + func(current_local_expert_idx, kNumL1BlockKs, m_block_idx, n_block_idx); + } else if (current_local_expert_idx >= kNumExpertsPerRank) { + break; + } + } + } + + template + CUTLASS_DEVICE void for_each_linear2_block(Func&& func) { + // Split-kernel mode: K2 starts after K1 globally completes, so all L2 + // ready masks are already final. Schedule Linear2 blocks directly. + fetch_expert_recv_count(); + set_expert_idx(0); + while (current_local_expert_idx < kNumExpertsPerRank) { + if (fetch_next_l2_block()) { + n_block_idx = block_idx - m_block_idx * kNumL2BlockNs; + block_idx += kNumSMs; + func(current_local_expert_idx, kNumL2BlockKs, m_block_idx, n_block_idx); + } else if (current_local_expert_idx >= kNumExpertsPerRank) { + break; + } + } + } +}; + +// Return the minimum all-L1 warm-up needed before the D40 scheduler starts +// interleaving L2 claims. This prevents an L2 readiness wait from forming a +// cycle with L1 task publication. +CUTLASS_HOST_DEVICE constexpr int get_num_l1_warmup_waves(const int& num_total_m_blocks, + const int& num_workers, + const int& num_l1_n_blocks, + const int& num_l2_n_blocks) { + if (num_total_m_blocks <= 0) return 0; + const int num_first_l2_wave_m_blocks = math::constexpr_ceil_div(num_workers, num_l2_n_blocks); + const int num_l1_warmup_waves_for_first_l2_wave = + math::constexpr_ceil_div(num_first_l2_wave_m_blocks * num_l1_n_blocks, num_workers); + + const int num_interleave_task_diff_per_m_block = + num_l1_n_blocks > num_l2_n_blocks ? num_l1_n_blocks - num_l2_n_blocks : 0; + const int num_warmup_waves_for_interleave_schedule = + math::constexpr_ceil_div( + num_l1_n_blocks + (num_total_m_blocks - 1) * num_interleave_task_diff_per_m_block, + num_workers) + + 1; + + return cute::max(num_l1_warmup_waves_for_first_l2_wave, num_warmup_waves_for_interleave_schedule); +} + +// Dynamic task payload shared by the dev-m producer and all consumers in one +// CTA. The layout is part of the validated two-stage mailbox protocol. +struct alignas(16) TaskInfo { + BlockPhase block_phase; + uint32_t local_expert_idx; + uint32_t m_block_idx; + uint32_t n_block_idx; + uint32_t pool_block_idx; + uint32_t valid_m; + uint32_t shape_n; + uint32_t shape_k; + + CUTLASS_HOST_DEVICE + TaskInfo() : TaskInfo(BlockPhase::None, 0, 0, 0, 0, 0, 0, 0) {} + + CUTLASS_HOST_DEVICE + TaskInfo(const BlockPhase& block_phase, const uint32_t& local_expert_idx, + const uint32_t& m_block_idx, const uint32_t& n_block_idx, const uint32_t& pool_block_idx, + const uint32_t& valid_m, const uint32_t& shape_n, const uint32_t& shape_k) + : block_phase(block_phase), + local_expert_idx(local_expert_idx), + m_block_idx(m_block_idx), + n_block_idx(n_block_idx), + pool_block_idx(pool_block_idx), + valid_m(valid_m), + shape_n(shape_n), + shape_k(shape_k) {} + + CUTLASS_HOST_DEVICE bool is_valid() const { return block_phase != BlockPhase::None; } +}; + +FI_MEGAMOE_STATIC_ASSERT(sizeof(TaskInfo) == 32, "Invalid task payload layout"); + +// SM90 one-CTA dynamic scheduler from Aichen dev-m. The weight-loader warp +// is the sole task producer and publishes work through a two-stage shared +// mailbox. The static scheduler above remains intact for bigM split and the +// D40 KF static-RS arm. +template +struct InterleavedMegaMoEScheduler { + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using task_info_t = TaskInfo; + + static constexpr uint32_t kNumScheduleStages = 2; + static constexpr uint32_t kNumL1WavesDone = 0xffffffffu; + + FI_MEGAMOE_STATIC_ASSERT(L1_SHAPE_N % BLOCK_N == 0, "Invalid L1 shape"); + FI_MEGAMOE_STATIC_ASSERT(L2_SHAPE_N % BLOCK_N == 0, "Invalid L2 shape"); + FI_MEGAMOE_STATIC_ASSERT(L1_SHAPE_K % BLOCK_K == 0, "Invalid L1 K shape"); + FI_MEGAMOE_STATIC_ASSERT(L2_SHAPE_K % BLOCK_K == 0, "Invalid L2 K shape"); + FI_MEGAMOE_STATIC_ASSERT(kNumL1BlockNs <= 64, "L1 readiness mask is too small"); + FI_MEGAMOE_STATIC_ASSERT(kNumL2BlockNs >= kNumL1BlockNs, + "Alternating scheduler requires at least as many L2 tasks as L1 tasks"); + + const layout::Workspace& workspace; + Barrier* task_info_full_barriers; + Barrier* task_info_empty_barriers; + task_info_t* task_infos; + + uint32_t sched_stage_idx = 0; + uint32_t sched_phase = 0; + uint32_t stored_num_tokens_per_expert[kNumExpertsPerLane] = {}; + uint32_t num_total_m_blocks = 0; + uint32_t num_l1_warmup_waves = 0; + + CUTLASS_DEVICE + InterleavedMegaMoEScheduler(const layout::Workspace& workspace, Barrier* task_info_full_barriers, + Barrier* task_info_empty_barriers, task_info_t* task_infos) + : workspace(workspace), + task_info_full_barriers(task_info_full_barriers), + task_info_empty_barriers(task_info_empty_barriers), + task_infos(task_infos) {} + + CUTLASS_DEVICE void advance_schedule_pipeline() { + sched_stage_idx ^= 1u; + sched_phase ^= sched_stage_idx == 0; + } + + CUTLASS_DEVICE uint32_t get_num_tokens(const uint32_t& expert_idx) const { + uint32_t valid_value = 0; +#pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++i) { + valid_value = (expert_idx == i * 32 + ptx::get_lane_idx()) ? stored_num_tokens_per_expert[i] + : valid_value; + } + return ptx::exchange(valid_value, expert_idx % 32); + } + + CUTLASS_DEVICE uint32_t get_pool_block_offset(const uint32_t& expert_idx) const { + uint32_t num_blocks = 0; +#pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++i) { + if (i * 32 + ptx::get_lane_idx() < expert_idx) + num_blocks += math::ceil_div(stored_num_tokens_per_expert[i], BLOCK_M); + } + return __reduce_add_sync(0xffffffff, num_blocks); + } + + // Push dispatch (`push_done_ptr != nullptr`): wait for every rank's DONE + // signal (lane 0 relaxed spin, then one acquire.sys load; bar.warp.sync + // extends the ordering to the other lanes), then read the low words of the + // recv-count sums of the given parity slot (the per-row remote tickets). + CUTLASS_DEVICE void fetch_expert_recv_count(const int* push_done_ptr = nullptr, + const int& push_done_target = 0, + const uint32_t& parity = 0) { + if (push_done_ptr != nullptr) { + if (ptx::get_lane_idx() == 0) { + auto* spin_dbg = workspace.get_spin_debug_ptr(); + FI_MEGAMOE_SPIN_WHILE_DEBUG(ptx::ld_volatile(push_done_ptr) - push_done_target < 0, 90010, + spin_dbg, ptx::ld_volatile(push_done_ptr), push_done_target); + FI_MEGAMOE_SPIN_WHILE_DEBUG(ptx::ld_acq_sys(push_done_ptr) - push_done_target < 0, 90020, + spin_dbg, ptx::ld_acq_sys(push_done_ptr), push_done_target); + } + __syncwarp(); + } +#pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++i) { + const auto expert_idx = i * 32 + ptx::get_lane_idx(); + uint64_t value = 0; + if (expert_idx < kNumExpertsPerRank) { + if (push_done_ptr != nullptr) { + value = ptx::ld_volatile(workspace.get_expert_recv_count_sum_ptr(expert_idx, parity)); + } else { + do { + value = ptx::ld_volatile(workspace.get_expert_recv_count_sum_ptr(expert_idx)); + } while (static_cast(value >> 32) != kNumSMs * kNumRanks); + } + } + stored_num_tokens_per_expert[i] = static_cast(value); + } + __syncwarp(); + + num_total_m_blocks = get_pool_block_offset(kNumExpertsPerRank); + const uint32_t num_total_l1_tasks = num_total_m_blocks * kNumL1BlockNs; + const uint32_t num_total_l1_waves = math::ceil_div(num_total_l1_tasks, kNumSMs); + const uint32_t min_l1_warmup_waves = + get_num_l1_warmup_waves(num_total_m_blocks, kNumSMs, kNumL1BlockNs, kNumL2BlockNs); + num_l1_warmup_waves = cute::min(min_l1_warmup_waves, num_total_l1_waves); + } + + CUTLASS_DEVICE task_info_t create_task(const BlockPhase& block_phase, const uint32_t& task_idx, + const uint32_t& num_n_blocks, const uint32_t& shape_n, + const uint32_t& shape_k) const { + const uint32_t lane_idx = ptx::get_lane_idx(); + const uint32_t pool_block_idx = task_idx / num_n_blocks; + const uint32_t n_block_idx = task_idx % num_n_blocks; + + task_info_t result(block_phase, 0, 0, n_block_idx, pool_block_idx, 0, shape_n, shape_k); + uint32_t block_offset = 0; +#pragma unroll + for (uint32_t i = 0; i < kNumExpertsPerLane; ++i) { + const uint32_t expert_idx = i * 32 + lane_idx; + const uint32_t num_tokens = stored_num_tokens_per_expert[i]; + const uint32_t num_m_blocks = math::ceil_div(num_tokens, BLOCK_M); + const uint32_t inclusive_num_m_blocks = math::warp_inclusive_sum(num_m_blocks, lane_idx); + const uint32_t lane_pool_block_offset = block_offset + inclusive_num_m_blocks - num_m_blocks; + const bool is_owner = expert_idx < kNumExpertsPerRank && + pool_block_idx >= lane_pool_block_offset && + pool_block_idx < lane_pool_block_offset + num_m_blocks; + const uint32_t owner_mask = __ballot_sync(0xffffffff, is_owner); + + if (owner_mask) { + const uint32_t owner_lane_idx = static_cast(__ffs(owner_mask) - 1); + const uint32_t owner_m_block_idx = pool_block_idx - lane_pool_block_offset; + const uint32_t owner_valid_m = cute::min(num_tokens - owner_m_block_idx * BLOCK_M, BLOCK_M); + result.local_expert_idx = ptx::exchange(expert_idx, owner_lane_idx); + result.m_block_idx = ptx::exchange(owner_m_block_idx, owner_lane_idx); + result.valid_m = ptx::exchange(owner_valid_m, owner_lane_idx); + } + block_offset += ptx::exchange(inclusive_num_m_blocks, 31); + } + return result; + } + + static CUTLASS_DEVICE uint32_t get_next_task_idx(uint32_t* task_count_ptr) { + uint32_t result = 0; + if (cute::elect_one_sync()) result = ptx::atomic_add(task_count_ptr, 1u); + return ptx::exchange(result, 0); + } + + CUTLASS_DEVICE task_info_t claim_next_task() { + while (true) { + if (num_l1_warmup_waves != kNumL1WavesDone && num_l1_warmup_waves > 0) { + --num_l1_warmup_waves; + const uint32_t task_idx = get_next_task_idx(workspace.get_l1_task_count_ptr()); + if (task_idx >= num_total_m_blocks * kNumL1BlockNs) { + num_l1_warmup_waves = kNumL1WavesDone; + continue; + } + return create_task(BlockPhase::Linear1, task_idx, kNumL1BlockNs, L1_SHAPE_N, L1_SHAPE_K); + } + + const uint32_t task_idx = get_next_task_idx(workspace.get_l2_task_count_ptr()); + if (task_idx >= num_total_m_blocks * kNumL2BlockNs) break; + + if (num_l1_warmup_waves != kNumL1WavesDone) num_l1_warmup_waves = 1; + + auto task_info = + create_task(BlockPhase::Linear2, task_idx, kNumL2BlockNs, L2_SHAPE_N, L2_SHAPE_K); + const uint32_t num_required_l1_tasks = (task_info.pool_block_idx + 1) * kNumL1BlockNs; + while (ptx::ld_volatile(workspace.get_l1_task_count_ptr()) < num_required_l1_tasks) { + } + return task_info; + } + return task_info_t(); + } + + CUTLASS_DEVICE void wait_task_slot_empty() const { + task_info_empty_barriers[sched_stage_idx].wait(sched_phase ^ 1u); + } + + CUTLASS_DEVICE void publish_task(const task_info_t& task_info) { + if (cute::elect_one_sync()) { + task_infos[sched_stage_idx] = task_info; + __threadfence_block(); + task_info_full_barriers[sched_stage_idx].arrive(); + } + __syncwarp(); + advance_schedule_pipeline(); + } + + CUTLASS_DEVICE bool get_published_task(task_info_t& task_info) { + task_info_full_barriers[sched_stage_idx].wait(sched_phase); + asm volatile("" ::: "memory"); + task_info = task_infos[sched_stage_idx]; + advance_schedule_pipeline(); + return task_info.is_valid(); + } + + CUTLASS_DEVICE void release_task_info(const uint32_t& lane_idx) const { + if (lane_idx == 0) task_info_empty_barriers[sched_stage_idx ^ 1u].arrive(); + } +}; + +} // namespace flashinfer::megamoe::sched diff --git a/include/flashinfer/megamoe/sm90_fp4_mega_moe.cuh b/include/flashinfer/megamoe/sm90_fp4_mega_moe.cuh new file mode 100644 index 00000000000..cf59fcf0278 --- /dev/null +++ b/include/flashinfer/megamoe/sm90_fp4_mega_moe.cuh @@ -0,0 +1,7 @@ +// SM90 FP4 MegaMoE kernel family. Common infrastructure lives in layout, +// scheduler, barrier and detail; each format retains its weight pipeline. +#pragma once + +#include +#include +#include diff --git a/include/flashinfer/megamoe/sm90_mxfp4_mega_moe_fused.cuh b/include/flashinfer/megamoe/sm90_mxfp4_mega_moe_fused.cuh new file mode 100644 index 00000000000..e33c176ace3 --- /dev/null +++ b/include/flashinfer/megamoe/sm90_mxfp4_mega_moe_fused.cuh @@ -0,0 +1,243 @@ +#pragma once + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace flashinfer::megamoe { +namespace mxfp4 { + +// MXFP4 decoders for the SM90 fused MegaMoE mainloop. +// +// Two differences from the NVFP4 decoders they replace, both consequences of +// the scale format: +// +// * Group width. NVFP4 groups 16 elements, so a 32-element quad spans two +// scales and its four packed words split across lut0/lut1. MXFP4 groups 32, +// so one quad is exactly one scale group and all four words share a table. +// A BK128 row therefore carries 4 E8M0 bytes (one uint32) instead of 8 +// UE4M3 bytes (a uint2). +// * Table provenance and size. Both formats look the scaled magnitudes up in +// shared memory. Combined with the wider group, MXFP4 needs half as many +// lookups per row -- see kScaledLutSize in mxfp4_fused_scale.cuh. + +template +__device__ __forceinline__ void dequant_mode2_nibble_row_regs( + uint8_t* __restrict__ fp8_dst, const uint4 (&fp4_quads)[4], const uint32_t scale_word, + const uint32_t row_swizzle, const ScaledLut* __restrict__ lut_smem) { +#pragma unroll + for (int quad_i = 0; quad_i < 4; ++quad_i) { + const uint4 q = fp4_quads[quad_i]; + const ScaledLut lut = load_scaled_lut(lut_smem, (scale_word >> (quad_i * 8)) & 0xffu); + + const uint2 w0 = dequant_word(q.x, lut); + const uint2 w1 = dequant_word(q.y, lut); + if constexpr (!kQuadILP) { + *reinterpret_cast(fp8_dst + (((quad_i * 2) * 16) ^ row_swizzle)) = + make_uint4(w0.x, w0.y, w1.x, w1.y); + } + const uint2 w2 = dequant_word(q.z, lut); + const uint2 w3 = dequant_word(q.w, lut); + if constexpr (kQuadILP) { + *reinterpret_cast(fp8_dst + (((quad_i * 2) * 16) ^ row_swizzle)) = + make_uint4(w0.x, w0.y, w1.x, w1.y); + } + *reinterpret_cast(fp8_dst + (((quad_i * 2 + 1) * 16) ^ row_swizzle)) = + make_uint4(w2.x, w2.y, w3.x, w3.y); + } +} + +template +__device__ __forceinline__ void dequant_smem_b_from_packed_mode2_nibble( + uint8_t* __restrict__ smem_b, const uint8_t* __restrict__ packed_b, const uint32_t row, + const ScaledLut* __restrict__ lut_smem) { + const DecodeTileRow src = decode_tile_row(packed_b, row); + uint4 fp4_quads[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + fp4_quads[i] = *reinterpret_cast(src.chunk + i * kBChunkStride); + const uint32_t scale_word = *reinterpret_cast(src.scale); + dequant_mode2_nibble_row_regs(smem_b + row * 128, fp4_quads, scale_word, + (row & 7u) << 4, lut_smem); +} + +// Threads 0-127 and 128-255 each decode one K64 half of the same N128 tile, +// letting the two M64 warpgroups reuse the decoded weights. A K64 half is two +// quads, i.e. two E8M0 bytes. +__device__ __forceinline__ void dequant_smem_b_from_packed_mode2_nibble_split_m( + uint8_t* __restrict__ smem_b, const uint8_t* __restrict__ packed_b, const uint32_t thread_idx, + const ScaledLut* __restrict__ lut_smem) { + const uint32_t row = thread_idx & 127u; + const uint32_t k_half_idx = thread_idx >> 7; + const DecodeTileRow src = decode_tile_row(packed_b, row); + const uint8_t* __restrict__ fp4_src = src.chunk + k_half_idx * 2u * kBChunkStride; + const uint32_t scale_pair = + *reinterpret_cast(src.scale + k_half_idx * sizeof(uint16_t)); + uint8_t* __restrict__ fp8_dst = smem_b + row * 128u; + const uint32_t row_swizzle = (row & 7u) << 4; + +#pragma unroll + for (uint32_t quad_i = 0; quad_i < 2; ++quad_i) { + const uint4 q = *reinterpret_cast(fp4_src + quad_i * kBChunkStride); + const ScaledLut lut = load_scaled_lut(lut_smem, (scale_pair >> (quad_i * 8u)) & 0xffu); + const uint2 w0 = dequant_word(q.x, lut); + const uint2 w1 = dequant_word(q.y, lut); + const uint2 w2 = dequant_word(q.z, lut); + const uint2 w3 = dequant_word(q.w, lut); + const uint32_t off0 = k_half_idx * 64u + (quad_i * 2u) * 16u; + const uint32_t off1 = k_half_idx * 64u + (quad_i * 2u + 1u) * 16u; + *reinterpret_cast(fp8_dst + (off0 ^ row_swizzle)) = make_uint4(w0.x, w0.y, w1.x, w1.y); + *reinterpret_cast(fp8_dst + (off1 ^ row_swizzle)) = make_uint4(w2.x, w2.y, w3.x, w3.y); + } +} + +// Decode the two packed words one lane pair owns directly into the four +// register-source WGMMA A operands, bypassing shared memory entirely. +// +// A K32 slice is 16 packed bytes. Lanes pair up: the even lane keeps the high +// half of each decoded word and ships the low half to its partner, the odd lane +// does the reverse, so one shuffle per word assembles both lanes' fragments. +// MXFP4's group of 32 is exactly the WGMMA's K, so both words of a slice share +// one table; NVFP4's group of 16 needs two. +__device__ __forceinline__ void dequant_rs_word_pair(const uint32_t w_lo, const uint32_t w_hi, + const ScaledLut& lut, const bool keep_hi, + uint32_t (&a_frag)[4]) { + detail::assemble_rs_word_pair(dequant_word(w_lo, lut), dequant_word(w_hi, lut), keep_hi, a_frag); +} + +__device__ __forceinline__ void dequant_braided_quad(uint8_t* __restrict__ fp8_dst, const uint4& q, + const ScaledLut& lut, const int quad_i, + const uint32_t row_swizzle) { + const uint2 w0 = dequant_word(q.x, lut); + const uint2 w1 = dequant_word(q.y, lut); + *reinterpret_cast(fp8_dst + (((quad_i * 2) * 16) ^ row_swizzle)) = + make_uint4(w0.x, w0.y, w1.x, w1.y); + const uint2 w2 = dequant_word(q.z, lut); + const uint2 w3 = dequant_word(q.w, lut); + *reinterpret_cast(fp8_dst + (((quad_i * 2 + 1) * 16) ^ row_swizzle)) = + make_uint4(w2.x, w2.y, w3.x, w3.y); +} + +} // namespace mxfp4 + +// Decode occupies its own warps only on a shared-memory swapAB tile tall +// enough for the WGMMA to hide it; the register-source path has no decoded-B +// ring to fill, and BLOCK_M 8 is faster decoding on the math warps. An +// unconditional 640 would also cut the compiler's per-thread register budget +// for the narrow tiers. Must stay in step with +// SM90MXFP4H200FusedConfig::num_threads(). +constexpr bool sm90_mxfp4_split_decode(const bool swap_ab, const bool use_rs, + const uint32_t block_m) { + return swap_ab and not use_rs and block_m >= 16; +} + +constexpr uint32_t sm90_mxfp4_num_threads(const bool swap_ab, const bool use_rs, + const uint32_t block_m) { + return 64 + 64 + 256 + (sm90_mxfp4_split_decode(swap_ab, use_rs, block_m) ? 256 : 0); +} + +template has one layout for every N, and the + // barriers, scheduler and workspace were already templated on it. + uint32_t kNumRanks, + // Shape is a template parameter, not a baked-in constant. The kernel body + // was already written against these names; only the wrapper hardcoded the + // H200 384-expert / 6144-hidden model, which locked out every other shape + // (e.g. DeepSeek-V4-Flash: 4096 hidden, 256 experts, topk 6). + uint32_t kHidden, uint32_t kIntermediateHidden, uint32_t kNumExperts, uint32_t kNumTopk, + uint32_t kNumMaxTokensPerRank, uint32_t kNumExpertsPerWave, uint32_t BLOCK_M, + uint32_t BLOCK_N, uint32_t kNumMaxPoolTokens, uint32_t kNumPaddedSFPoolTokens, + uint32_t kNumStages, float kActivationClamp, bool kFastMath, bool kSwapABRequested, + bool kSingleActiveDispatchWarp, bool kUseMode2RowDecoder, bool kUseInterleavedScheduler, + // Feed the WGMMA's A operand from registers instead of decoding the weight + // tile into shared memory first. swapAB only; removes the decoded-B ring. + bool kUseRSOperand, + // Phase timestamps are a compile-time variant, not a runtime flag. Checking + // a pointer per task is not free -- the same accumulators cost ~5% per task + // on the NVFP4 side when they were runtime-gated -- so the shipped kernel + // must not contain them at all. Defaulted, so every existing instantiation + // (and every compile-gate row) keeps its current meaning. + bool kPhaseStamps = false> +CUTLASS_GLOBAL __launch_bounds__( + sm90_mxfp4_num_threads(kSwapABRequested, kUseRSOperand, BLOCK_M), + 1) void sm90_mxfp4_mega_moe_fused_impl(void* y, int* cumulative_local_expert_recv_stats, + const uint32_t num_tokens, + const __grid_constant__ layout::SymBuffer + sym_buffer, + const __grid_constant__ cute::TmaDescriptor + tensor_map_l1_acts, + const __grid_constant__ cute::TmaDescriptor + tensor_map_l1_acts_sf, + const __grid_constant__ cute::TmaDescriptor + tensor_map_l1_output, + const __grid_constant__ cute::TmaDescriptor + tensor_map_l2_acts, + const __grid_constant__ cute::TmaDescriptor + tensor_map_l2_acts_sf, + // Weights are bulk-copied a tile at a time, not + // addressed through a tensor map: the tile is + // contiguous, so there is no shape to describe. + const uint8_t* __restrict__ l1_weights, + const uint8_t* __restrict__ l2_weights, + const float* __restrict__ l1_global_scales, + const float* __restrict__ l2_global_scales, + // Only written when kPhaseStamps; see the slot map in + // the body. + unsigned long long* __restrict__ phase_stamps) { + constexpr uint32_t BLOCK_K = 128; + constexpr uint32_t kNumDispatchThreads = 64; + constexpr uint32_t kNumNonEpilogueThreads = 64; + // Decode gets its own warps so it overlaps the WGMMA rather than + // serialising with it on the math warps. + constexpr bool kSplitDecodeWarps = + sm90_mxfp4_split_decode(kSwapABRequested, kUseRSOperand, BLOCK_M); + constexpr uint32_t kNumDecodeThreads = kSplitDecodeWarps ? 256 : 0; + constexpr uint32_t kNumEpilogueThreads = 256; + constexpr uint32_t kNumThreads = + kNumDispatchThreads + kNumNonEpilogueThreads + kNumDecodeThreads + kNumEpilogueThreads; + FI_MEGAMOE_STATIC_ASSERT( + kNumThreads == sm90_mxfp4_num_threads(kSwapABRequested, kUseRSOperand, BLOCK_M), + "Launch bounds and the role partition disagree"); + FI_MEGAMOE_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Experts must divide evenly across ranks"); + FI_MEGAMOE_STATIC_ASSERT(kHidden % 128 == 0, "Hidden must be a multiple of BLOCK_K"); + FI_MEGAMOE_STATIC_ASSERT(kIntermediateHidden % 128 == 0, + "Intermediate hidden must be a multiple of BLOCK_K"); + constexpr uint32_t L1_SHAPE_N = kIntermediateHidden * 2; + constexpr uint32_t L1_SHAPE_K = kHidden; + constexpr uint32_t L2_SHAPE_N = kHidden; + constexpr uint32_t L2_SHAPE_K = kIntermediateHidden; + constexpr uint32_t kNumDispatchWarps = kNumDispatchThreads / 32; + constexpr uint32_t kNumMMANonEpilogueWarps = kNumNonEpilogueThreads / 32; + constexpr uint32_t kNumDecodeWarps = kNumDecodeThreads / 32; + constexpr uint32_t kNumEpilogueWarps = kNumEpilogueThreads / 32; + // Warp map: [dispatch][A loader][B loader][decode][math] + constexpr uint32_t kFirstDecodeWarp = kNumDispatchWarps + kNumMMANonEpilogueWarps; + constexpr uint32_t kFirstMathWarp = kFirstDecodeWarp + kNumDecodeWarps; + constexpr uint32_t kNumEpilogueWarpgroups = kNumEpilogueWarps / 4; + constexpr uint32_t kNumTokensPerWarp = 32 / kNumTopk; + constexpr uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks; +#include +} + +} // namespace flashinfer::megamoe + +#pragma clang diagnostic pop diff --git a/include/flashinfer/megamoe/sm90_mxfp4_mega_moe_fused_body.inl b/include/flashinfer/megamoe/sm90_mxfp4_mega_moe_fused_body.inl new file mode 100644 index 00000000000..540891c73be --- /dev/null +++ b/include/flashinfer/megamoe/sm90_mxfp4_mega_moe_fused_body.inl @@ -0,0 +1,1424 @@ +// SM90 MXFP4 MegaMoE fused kernel body. +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900) and (__CUDA_ARCH__ < 1000)) or \ + defined(__CLION_IDE__) +using Barrier = cutlass::arch::ClusterTransactionBarrier; + +// ===================================================================== +// Template checks +// ===================================================================== +FI_MEGAMOE_STATIC_ASSERT(BLOCK_M == 8 || BLOCK_M == 16 || BLOCK_M == 24 || BLOCK_M == 32 || + BLOCK_M == 64 || BLOCK_M == 128, + "H200 fused kernel requires BM8/BM16/BM24/BM32/BM64/BM128"); +// Pipeline depth used to be pinned to one value per BLOCK_M, which made the +// shipped table the only reachable schedule. The binding constraint is +// shared memory, and the layout assertions below already enforce it +// exactly, so only the pipelining minimum is asserted here. +FI_MEGAMOE_STATIC_ASSERT(kNumStages >= 2, "Pipelining needs at least two stages"); +FI_MEGAMOE_STATIC_ASSERT(BLOCK_M != 128 || kNumStages == 6, + "BM128 split-M is only validated at stage depth 6"); +FI_MEGAMOE_STATIC_ASSERT((BLOCK_M == 128) == (BLOCK_N == 128), + "BM128 is paired with the BN128 split-M topology"); +// FP8MMASelector takes N=32, and the N_SWAP ladder below has a BM32 arm. +FI_MEGAMOE_STATIC_ASSERT(!kSwapABRequested || BLOCK_M <= 32, + "swap-AB packs tokens into the WGMMA N dimension"); + +// ===================================================================== +// Thread / warp identification +// ===================================================================== +const uint32_t sm_idx = blockIdx.x; +const uint32_t thread_idx = threadIdx.x; +const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); +const uint32_t lane_idx = ptx::get_lane_idx(); + +// Phase timestamps (globaltimer, ns), written only when kPhaseStamps. +// Compiled out entirely otherwise -- see the note on the template parameter. +// +// 0 min kernel entry 1 max dispatch barrier #1 done +// 2 max pool ready 3 min first math task +// 4 max last L1 task end 5 max last L2 task end +// 6 max combine barrier #2 done 7 max combine end +// 8 max cleanup barrier #3 done +// +// Accumulators, summed over every CTA (value, then count): +// 20/21 L1 task ns 22/23 L2 task ns +// 24/25 decode warps waiting on the weight TMA (decode starved by the load) +// 26/27 math warps waiting on the decoded ring (WGMMA starved by decode) +// 24 and 26 are the pair that says whether the mainloop is stream-bound or +// decode-bound, which is the question the end-to-end numbers cannot answer. +const auto stamp_min = [&](const uint32_t slot) { + if constexpr (kPhaseStamps) { + unsigned long long t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + atomicMin(phase_stamps + slot, t); + } +}; +const auto stamp_max = [&](const uint32_t slot) { + if constexpr (kPhaseStamps) { + unsigned long long t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + atomicMax(phase_stamps + slot, t); + } +}; +const auto stamp_now = [&]() -> unsigned long long { + unsigned long long t = 0; + if constexpr (kPhaseStamps) asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +}; +const auto stamp_add = [&](const uint32_t slot, const unsigned long long delta) { + if constexpr (kPhaseStamps) { + atomicAdd(phase_stamps + slot, delta); + atomicAdd(phase_stamps + slot + 1, 1ull); + } +}; +(void)stamp_now; +(void)stamp_add; + +if (warp_idx == 0 and cute::elect_one_sync()) { + stamp_min(0); + cute::prefetch_tma_descriptor(&tensor_map_l1_acts); + cute::prefetch_tma_descriptor(&tensor_map_l1_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l1_output); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts_sf); +} + +// ===================================================================== +// Workspaces and symmetric buffer slicing. The framework reserves +// per-64 SF capacity; this per-128 path uses its first half +// as a dense per-128 layout so no framework allocation change is needed. +// ===================================================================== +#include + +// ===================================================================== +// GEMM data types and shape constants +// ===================================================================== +using a_dtype_t = cutlass::float_e4m3_t; +using b_dtype_t = cutlass::float_e4m3_t; +using task_info_t = sched::TaskInfo; +using interleaved_scheduler_t = + sched::InterleavedMegaMoEScheduler; +constexpr uint32_t kNumRoutedL1BlockNs = L1_SHAPE_N / BLOCK_N; +constexpr bool kSplitMDecodedWeightReuse = + BLOCK_M == 128 && BLOCK_N == 128 && kNumEpilogueWarpgroups == 2; +constexpr uint32_t WG_BLOCK_M = kSplitMDecodedWeightReuse ? BLOCK_M / 2 : BLOCK_M; +constexpr uint32_t WG_BLOCK_N = kSplitMDecodedWeightReuse ? BLOCK_N : BLOCK_N / 2; +constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2; // post-SwiGLU tile N +constexpr uint32_t WG_L1_OUT_BLOCK_N = WG_BLOCK_N / 2; // post-SwiGLU per-WG N +constexpr uint32_t kSwapABTokenChunks = BLOCK_M / 8; +constexpr uint32_t kSwapABWeightHalves = WG_BLOCK_N / 64; +constexpr uint32_t kSwapABHalfAccumPerThread = 64 * 64 / 128; +FI_MEGAMOE_STATIC_ASSERT(!kSwapABRequested || WG_L1_OUT_BLOCK_N == 64, + "swapAB expects BN256 split-N with 64 L1 output columns per WG"); +// Both dispatch warps participate in CTA-wide barriers. Selected plans may +// use one warp for routing and token pulls, leaving the other warp's send +// buffer available for an additional GEMM stage. +constexpr uint32_t kNumActiveDispatchWarps = kSingleActiveDispatchWarp ? 1u : kNumDispatchWarps; +constexpr uint32_t kNumActiveDispatchThreads = kNumActiveDispatchWarps * 32; +constexpr bool kQuadDequantIlp = BLOCK_M == 8 && kNumStages == 4; +using L1WGMMA = typename mma::sm90::FP8MMASelector::type; +static_assert(L1WGMMA::M == 64 and L1WGMMA::N == WG_BLOCK_N and L1WGMMA::K == 32, + "Unexpected WGMMA shape"); +// A and B are CTA-local in the fixed cluster-size-one plan. +constexpr uint32_t LOAD_BLOCK_M = BLOCK_M; +constexpr uint32_t LOAD_BLOCK_N = BLOCK_N; +constexpr uint32_t kSwizzleAMode = BLOCK_K * sizeof(a_dtype_t); // 128 +constexpr uint32_t kL2ActsSFGranK = kSplitMDecodedWeightReuse ? 64u : 128u; +FI_MEGAMOE_STATIC_ASSERT(kSplitMDecodedWeightReuse || WG_L1_OUT_BLOCK_N < kL2ActsSFGranK, + "split-N warpgroups must share one L2 activation scale"); + +// ===================================================================== +// Shared memory layout +// ===================================================================== +constexpr uint32_t kSharedMemoryAlignment = 1024; +extern __shared__ __align__(kSharedMemoryAlignment) +uint8_t smem_buffer[]; + +constexpr uint32_t SMEM_EXPERT_COUNT_SIZE = + math::constexpr_align(kNumExperts * sizeof(uint32_t), kSharedMemoryAlignment); +constexpr uint32_t SMEM_SEND_BUFFER_SIZE = math::constexpr_align( + fp8_token_layout.get_num_bytes() * kNumActiveDispatchWarps, kSharedMemoryAlignment); +// Scaled-magnitude tables for all 256 E8M0 codes; see kScaledLutSize. +constexpr uint32_t SMEM_MXFP4_LUT_SIZE = math::constexpr_align( + mxfp4::kScaledLutSize * sizeof(mxfp4::ScaledLut), kSharedMemoryAlignment); +constexpr uint32_t SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K * sizeof(a_dtype_t); +constexpr uint32_t SMEM_B_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K * sizeof(b_dtype_t); +// BM128 split-M alternates two decoded-B slots. This lets one WG begin +// decoding K+1 after its K WGMMA completes without overwriting the slot +// that the paired WG may still be consuming. +// +// Every other plan needs exactly one. Decode and consume are serialized +// inside a K-block iteration -- the warpgroup drains its own WGMMA with +// warpgroup_wait<0> before it loops -- and the two warpgroups decode and +// read disjoint N halves of the same BN256 tile, so no slot is ever live +// across iterations. Sizing this ring by kNumStages spent +// (kNumStages - 1) * BN256 * BK128 bytes (64 KB at three stages) to hold +// tiles nobody could read, and that is what capped the TMA pipeline at +// three or four stages. +// One slot is enough outside split-M. Decode and consume are serialized +// inside a K-block iteration -- the warpgroup drains its own WGMMA with +// warpgroup_wait<0> before it loops -- and the two warpgroups decode and +// read disjoint N halves of the same BN256 tile, so no slot is live across +// iterations. +// +// Rotating two slots to decode K+1 under K's in-flight WGMMAs was measured +// and is worse: at MiMo/H20 M=64 it cost +5.9 / +7.1 / +10.8 % over three +// runs. The prefetch has to wait on stage K+1's TMA barrier, and that wait +// sits between the WGMMA issue and arrive_empty(K), so it delays the loader +// by exactly as much as it saves on the decode. +FI_MEGAMOE_STATIC_ASSERT(!kUseRSOperand || kSwapABRequested, + "The register-source A path is swapAB-only"); +constexpr uint32_t kNumDecodedBStages = + (kSplitMDecodedWeightReuse || kSplitDecodeWarps) ? 2u : (kUseRSOperand ? 0u : 1u); +// Only the split hands decoded tiles between warps, so only it needs the +// ring's barriers; every other plan decodes and consumes in one warp. +constexpr uint32_t kNumBRingStages = kSplitDecodeWarps ? 2u : 0u; +// Tile geometry (chunk-major, unpadded) lives with the decoders that read +// it; the loader only needs the tile size and a fixed row count, so the +// layout does not depend on the selected BLOCK_N. +// Local copies: `tma_load_1d` takes its byte count by reference, which +// odr-uses the operand, and a namespace-scope constant is host-side. +constexpr uint32_t kBTileRows = mxfp4::kBTileRows; +constexpr uint32_t kBTileBytes = mxfp4::kBTileBytes; +constexpr uint32_t SMEM_PACKED_B_SIZE_PER_STAGE = + (LOAD_BLOCK_N / kBTileRows) * kBTileBytes * sizeof(b_dtype_t); +FI_MEGAMOE_STATIC_ASSERT(LOAD_BLOCK_N % kBTileRows == 0, + "BLOCK_N must be a whole number of weight tiles"); +FI_MEGAMOE_STATIC_ASSERT(L1_SHAPE_N % kBTileRows == 0 && L2_SHAPE_N % kBTileRows == 0, + "Weight N must be a whole number of weight tiles"); +// L1 and L2 each consume one per-128 activation scale per row and K tile. +constexpr uint32_t kL2SFAHalfStride = + math::constexpr_align(BLOCK_M * sizeof(float), 128u) / sizeof(float); +constexpr uint32_t kNumL2SFAGroups = kSplitMDecodedWeightReuse ? 2u : 1u; +constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = kNumL2SFAGroups * kL2SFAHalfStride * sizeof(float); +// CD output: max of L1 FP8 (BLOCK_M * (BLOCK_N/2) * 1 byte * num_wg) and +// L2 BF16 (BLOCK_M * BLOCK_N * 2 bytes * num_wg). +constexpr uint32_t SMEM_CD_L1_SIZE = + kNumEpilogueWarpgroups * WG_BLOCK_M * WG_L1_OUT_BLOCK_N * sizeof(cutlass::float_e4m3_t); +constexpr uint32_t SMEM_CD_L2_SIZE = + kSwapABRequested ? BLOCK_M * BLOCK_N * sizeof(nv_bfloat16) : 0u; +constexpr uint32_t SMEM_CD_OUTPUT_BASE_SIZE = + SMEM_CD_L1_SIZE > SMEM_CD_L2_SIZE ? SMEM_CD_L1_SIZE : SMEM_CD_L2_SIZE; +constexpr uint32_t SMEM_CD_L1_SHARED_SF_SLOTS = kNumEpilogueWarpgroups * BLOCK_M; +constexpr uint32_t SMEM_CD_L1_SWAP_AMAX_SLOTS = kSwapABRequested ? BLOCK_M * kNumEpilogueWarps : 0u; +constexpr uint32_t SMEM_CD_L1_EXTRA_FLOAT_SLOTS = + SMEM_CD_L1_SHARED_SF_SLOTS > SMEM_CD_L1_SWAP_AMAX_SLOTS ? SMEM_CD_L1_SHARED_SF_SLOTS + : SMEM_CD_L1_SWAP_AMAX_SLOTS; +constexpr uint32_t SMEM_CD_L1_SHARED_SF_SIZE = SMEM_CD_L1_EXTRA_FLOAT_SLOTS * sizeof(float); +constexpr uint32_t SMEM_CD_OUTPUT_UNALIGNED_SIZE = + SMEM_CD_OUTPUT_BASE_SIZE + SMEM_CD_L1_SHARED_SF_SIZE; +constexpr uint32_t SMEM_CD_SIZE = + math::constexpr_align(SMEM_CD_OUTPUT_UNALIGNED_SIZE, kSharedMemoryAlignment); + +constexpr uint32_t SMEM_GEMM_END_SIZE = + SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_MXFP4_LUT_SIZE + SMEM_CD_SIZE + + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_PACKED_B_SIZE_PER_STAGE) + + kNumDecodedBStages * SMEM_B_SIZE_PER_STAGE + kNumStages * SMEM_SFA_SIZE_PER_STAGE; +// The combine phase reuses this whole region as its own ring, so it sets a +// floor on it that has nothing to do with the GEMM. Deriving the chunk +// count from kHidden alone keeps that floor non-circular; the static +// assertion further down re-derives it from the final region size and +// would fire if the two ever disagreed. +constexpr uint32_t kCombineNumChunks = kHidden <= 4096 ? 1u : 2u; +constexpr uint32_t SMEM_COMBINE_MIN_SIZE = + 3u * kNumEpilogueWarps * kHidden * sizeof(nv_bfloat16) / kCombineNumChunks; +constexpr uint32_t SMEM_BEFORE_BARRIER_SIZE = + SMEM_GEMM_END_SIZE > SMEM_COMBINE_MIN_SIZE ? SMEM_GEMM_END_SIZE : SMEM_COMBINE_MIN_SIZE; + +// SMEM pointers +auto smem_expert_count = reinterpret_cast(smem_buffer); +const auto smem_send_buffers = + layout::Buffer(fp8_token_layout, kNumActiveDispatchWarps, 1, + math::advance_ptr(smem_buffer, SMEM_EXPERT_COUNT_SIZE)); +auto smem_mxfp4_lut = reinterpret_cast( + math::advance_ptr(smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE)); +auto smem_gemm_base = math::advance_ptr( + smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_MXFP4_LUT_SIZE); + +auto smem_cd_base = smem_gemm_base; +// CD output is shared by L1 (FP8) and L2 (BF16); reinterpret-cast as needed. +auto smem_cd_l1 = reinterpret_cast(smem_cd_base); +auto smem_cd_l1_shared_sf = math::advance_ptr(smem_cd_base, SMEM_CD_OUTPUT_BASE_SIZE); +auto smem_cd_l2 = reinterpret_cast(smem_cd_base); + +auto smem_a = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE); +}); +auto smem_b = utils::PatternVisitor([=](const uint32_t& i) { + // Callers pass a decoded-slot index, not a pipeline stage. The RS path + // allocates no decoded tile at all, so the index degenerates to zero and + // the pointer is never dereferenced. + const uint32_t decoded_stage = kSplitMDecodedWeightReuse + ? (i & 1u) + : (kNumDecodedBStages <= 1u ? 0u : (i % kNumDecodedBStages)); + return math::advance_ptr( + smem_gemm_base, + SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + decoded_stage * SMEM_B_SIZE_PER_STAGE); +}); +auto smem_packed_b = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, + SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + + kNumDecodedBStages * SMEM_B_SIZE_PER_STAGE + + i * SMEM_PACKED_B_SIZE_PER_STAGE); +}); +auto sf_start_ptr = math::advance_ptr( + smem_gemm_base, SMEM_CD_SIZE + + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_PACKED_B_SIZE_PER_STAGE) + + kNumDecodedBStages * SMEM_B_SIZE_PER_STAGE); +auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast(sf_start_ptr + i * SMEM_SFA_SIZE_PER_STAGE); +}); +// Barriers live after SF. +auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_BEFORE_BARRIER_SIZE); +auto dispatch_barriers = + utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + i; }); +auto full_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + i; }); +auto empty_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + kNumStages + i; }); +auto combine_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + kNumStages * 2 + i; }); +// The decoded-B ring is a two-stage pipeline in its own right, so it gets +// the same barrier pair the stage pipeline uses. +constexpr uint32_t kNumRingBarriers = kNumDispatchWarps + kNumStages * 2 + kNumEpilogueWarps * 2; +auto decoded_full_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumRingBarriers + i; }); +auto decoded_empty_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumRingBarriers + kNumBRingStages + i; }); +constexpr uint32_t kNumBaseBarriers = kNumRingBarriers + kNumBRingStages * 2; +auto task_info_full_barriers = barrier_start_ptr + kNumBaseBarriers; +auto task_info_empty_barriers = + task_info_full_barriers + interleaved_scheduler_t::kNumScheduleStages; +auto task_infos = reinterpret_cast(task_info_empty_barriers + + interleaved_scheduler_t::kNumScheduleStages); +constexpr uint32_t kInterleavedSchedulerSMEMBytes = + 2 * interleaved_scheduler_t::kNumScheduleStages * sizeof(Barrier) + + interleaved_scheduler_t::kNumScheduleStages * sizeof(task_info_t); +FI_MEGAMOE_STATIC_ASSERT(kInterleavedSchedulerSMEMBytes == + layout::kSM90InterleavedSchedulerSMEMBytes, + "Host and device scheduler shared-memory layouts disagree"); +constexpr uint32_t kInterleavedSMEMEnd = + SMEM_BEFORE_BARRIER_SIZE + kNumBaseBarriers * sizeof(Barrier) + kInterleavedSchedulerSMEMBytes; +FI_MEGAMOE_STATIC_ASSERT(!kUseInterleavedScheduler || kInterleavedSMEMEnd <= 232448, + "Interleaved scheduler exceeds the SM90 shared-memory capacity"); + +// ===================================================================== +// Initialization +// ===================================================================== +mxfp4::init_scaled_lut(smem_mxfp4_lut, thread_idx, kNumThreads); + +if (warp_idx == 0) { +// Clean expert-count shared memory +#pragma unroll + for (uint32_t i = lane_idx; i < kNumExperts; i += 32) ptx::st_shared(smem_expert_count + i, 0u); +} else if (warp_idx == 1) { +// Init dispatch m-barriers +#pragma unroll + for (uint32_t i = lane_idx; i < kNumDispatchWarps; i += 32) dispatch_barriers[i]->init(1); + cutlass::arch::fence_barrier_init(); +} else if (warp_idx == 2) { + // Init GEMM full/empty barriers and combine barriers + if (cute::elect_one_sync()) { +#pragma unroll + for (uint32_t i = 0; i < kNumStages; ++i) { + // Producer arrivals: A(+SFA) + B(TMA+SFB). SFB is copied with + // cp.async.bulk and counted as B-loader transaction bytes, so + // it does not need a separate producer arrival. + full_barriers[i]->init(2); + // A stage holds both the packed weights and the activations. + // With split decode the two are read by different warps and + // both must release it, or the A loader refills the stage + // under a running WGMMA. + empty_barriers[i]->init(kNumEpilogueWarps + (kSplitDecodeWarps ? kNumDecodeWarps : 0u)); + } +#pragma unroll + for (uint32_t i = 0; i < kNumEpilogueWarps * 2; ++i) combine_barriers[i]->init(1); +#pragma unroll + for (uint32_t i = 0; i < kNumBRingStages; ++i) { + // One arrival per warp on each side: the tile is not readable + // until every decode warp has written its rows. + decoded_full_barriers[i]->init(kNumDecodeWarps); + decoded_empty_barriers[i]->init(kNumEpilogueWarps); + } + if constexpr (kUseInterleavedScheduler) { +#pragma unroll + for (uint32_t i = 0; i < interleaved_scheduler_t::kNumScheduleStages; ++i) { + task_info_full_barriers[i].init(1); + task_info_empty_barriers[i].init(kNumEpilogueWarps + + (kSplitDecodeWarps ? kNumDecodeWarps : 0u)); + } + } + } + cutlass::arch::fence_barrier_init(); +} +__syncthreads(); + +// ===================================================================== +// Scheduler (cluster=1) +// ===================================================================== +auto scheduler = sched::MegaMoEScheduler(workspace); +auto interleaved_scheduler = interleaved_scheduler_t(workspace, task_info_full_barriers, + task_info_empty_barriers, task_infos); + +// Pipeline state shared by TMA loaders and math warpgroups +uint32_t stage_idx = 0, phase = 0; +// The ring advances with the stage pipeline and, like it, runs continuously +// across blocks: its barriers carry state, so resetting the slot per block +// would leave the producer a slot ahead of its consumer. +uint32_t b_ring_idx = 0, b_ring_phase = 0; +auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++k_block_idx; + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + if constexpr (kSplitDecodeWarps) { + b_ring_idx = b_ring_idx == kNumBRingStages - 1 ? 0 : b_ring_idx + 1; + b_ring_phase ^= b_ring_idx == 0; + } +}; +// Intra-SM barrier indices (mirroring SM100) +constexpr uint32_t kDispatchBarrierIdx = 0; +constexpr uint32_t kDispatchWithEpilogueBarrierIdx = 1; +constexpr uint32_t kEpilogueFullBarrierIdx = 2; +constexpr uint32_t kEpilogueWGBarrierStartIdx = 3; +constexpr uint32_t kSplitMDecodeBarrierIdx = 8; + +// Cross-rank NVLink barrier tags +constexpr uint32_t kBeforeDispatchPullBarrierTag = 1; +constexpr uint32_t kBeforeCombineReduceBarrierTag = 2; +constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; + +// Register reconfiguration counts (chosen to fit in 64512 reg budget). +constexpr uint32_t kNumDispatchRegisters = 48; +constexpr uint32_t kNumNonEpilogueRegisters = kUseInterleavedScheduler ? 64 : 40; +constexpr uint32_t kNumDecodeRegisters = 48; +// The decode warps take their share out of the epilogue's budget, so the +// split runs the math warps leaner. +constexpr uint32_t kNumEpilogueRegisters = kSplitDecodeWarps ? 152 : 208; +// The register file is handed out per warp in 256-register granules, so +// what setmaxnreg can redistribute is not 64K but what the launch could +// actually reserve at this width -- 64512 at 384 threads, 61440 at 640. +constexpr uint32_t kNumWarps = kNumThreads / 32; +constexpr uint32_t kRegisterBudget = 65536 / kNumWarps / 256 * 256 * kNumWarps; +// setmaxnreg acts per warpgroup, so the budget has to be counted that way: +// dispatch and the two loaders share one warpgroup, which is held at the +// larger of their two requests. Counting per role instead understates the +// total, and an over-subscribed warpgroup_reg_alloc does not fail -- it +// blocks, so the kernel hangs rather than refusing to launch. +constexpr uint32_t kNumLoaderWarpgroupRegisters = kNumDispatchRegisters > kNumNonEpilogueRegisters + ? kNumDispatchRegisters + : kNumNonEpilogueRegisters; +FI_MEGAMOE_STATIC_ASSERT(kNumLoaderWarpgroupRegisters * 128 + + kNumDecodeRegisters * kNumDecodeThreads + + kNumEpilogueRegisters * kNumEpilogueThreads <= + kRegisterBudget, + "Too many registers"); + +constexpr uint32_t kDispatchGridSyncIndex = 0; +constexpr uint32_t kEpilogueGridSyncIndex = 1; + +const auto for_each_static_selected_block = [&](auto&& func) { + scheduler.fetch_expert_recv_count(); + scheduler.set_expert_idx(0); + while (true) { + CUTE_TIE_DECL(scheduler.get_next_block(), block_phase, local_expert_idx, m_block_idx, + n_block_idx); + if (block_phase == sched::BlockPhase::None) break; + if (block_phase == sched::BlockPhase::Linear1) { + func(std::integral_constant{}, + local_expert_idx, L1_SHAPE_K / BLOCK_K, m_block_idx, n_block_idx, + scheduler.get_current_pool_block_offset() + m_block_idx, + scheduler.template get_valid_m()); + } else { + func(std::integral_constant{}, + local_expert_idx, L2_SHAPE_K / BLOCK_K, m_block_idx, n_block_idx, + scheduler.get_current_pool_block_offset() + m_block_idx, + scheduler.template get_valid_m()); + } + } +}; + +const auto invoke_interleaved_task = [&](const task_info_t& task_info, auto&& func) { + if (task_info.block_phase == sched::BlockPhase::Linear1) { + func(std::integral_constant{}, + task_info.local_expert_idx, L1_SHAPE_K / BLOCK_K, task_info.m_block_idx, + task_info.n_block_idx, task_info.pool_block_idx, task_info.valid_m); + } else { + func(std::integral_constant{}, + task_info.local_expert_idx, L2_SHAPE_K / BLOCK_K, task_info.m_block_idx, + task_info.n_block_idx, task_info.pool_block_idx, task_info.valid_m); + } +}; + +const auto for_each_published_block = [&](auto&& func) { + task_info_t task_info; + while (interleaved_scheduler.get_published_task(task_info)) + invoke_interleaved_task(task_info, func); +}; + +const auto produce_interleaved_blocks = [&](auto&& func) { + interleaved_scheduler.fetch_expert_recv_count(); + while (true) { + interleaved_scheduler.wait_task_slot_empty(); + const auto task_info = interleaved_scheduler.claim_next_task(); + interleaved_scheduler.publish_task(task_info); + if (!task_info.is_valid()) break; + invoke_interleaved_task(task_info, func); + } +}; + +const auto cleanup_workspace = [&]() { + FI_MEGAMOE_STATIC_ASSERT(kNumSMs > 1, "Invalid SM count"); + if (sm_idx == 0) { +#pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) + *workspace.get_expert_send_count_ptr(i) = 0; + if constexpr (kUseInterleavedScheduler) { + if (thread_idx == 0) { + *workspace.get_l1_task_count_ptr() = 0; + *workspace.get_l2_task_count_ptr() = 0; + } + } + } else { + for (uint32_t i = sm_idx - 1; i < kNumExpertsPerRank; i += kNumSMs - 1) { + const auto num_recv_tokens = + static_cast(*workspace.get_expert_recv_count_sum_ptr(i)); + const auto num_recv_m_blocks = math::ceil_div(num_recv_tokens, BLOCK_M); + const auto cleanup_pool_block_offset = scheduler.get_pool_block_offset(i); + + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + FI_MEGAMOE_STATIC_ASSERT(kNumDispatchWarps >= 2, "Not enough dispatch warps"); + if (warp_idx == 0) { + *workspace.get_expert_recv_count_sum_ptr(i) = 0; + } else if (warp_idx == 1) { + if (cute::elect_one_sync() and cumulative_local_expert_recv_stats != nullptr) + ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(num_recv_tokens)); + __syncwarp(); + } + + for (uint32_t j = thread_idx; j < kNumRanks; j += kNumDispatchThreads) + *workspace.get_expert_recv_count_ptr(j, i) = 0; + __syncwarp(); + + for (uint32_t j = thread_idx; j < num_recv_m_blocks; j += kNumDispatchThreads) { + *workspace.get_l1_arrival_count_ptr(cleanup_pool_block_offset + j) = 0; + *workspace.get_l2_arrival_mask_ptr(cleanup_pool_block_offset + j) = 0; + } + __syncwarp(); + } + } +}; + +// ===================================================================== +// ROLE 1: DISPATCH WARPS +// Mirrors SM100 dispatch with two changes: +// * SF is per-128 channel float (no UTCCP transpose). We store the +// remote per-token SF directly into the local L1 SF buffer in +// MN-major layout: `local_sf[k_chunk * num_padded_sf_pool_tokens + token_idx]`. +// * The "token_idx_in_expert" → SF token index is now the simple +// per-block linear mapping (no 4×32 transpose). +// ===================================================================== +if (warp_idx < kNumDispatchWarps) { + cutlass::arch::warpgroup_reg_dealloc(); + + FI_MEGAMOE_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of topk"); + constexpr uint32_t kNumActivateLanes = kNumTokensPerWarp * kNumTopk; + const auto read_topk_idx = [&](const auto& process) { + if (warp_idx < kNumActiveDispatchWarps) { +#pragma unroll + for (uint32_t i = (sm_idx * kNumActiveDispatchWarps + warp_idx) * kNumTokensPerWarp; + i < num_tokens; i += kNumSMs * kNumActiveDispatchWarps * kNumTokensPerWarp) { + int expert_idx = -1; + if (i + (lane_idx / kNumTopk) < num_tokens and lane_idx < kNumActivateLanes) { + expert_idx = static_cast( + __ldg(input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); + if (expert_idx >= 0) process(i * kNumTopk + lane_idx, expert_idx); + } + __syncwarp(); + } + } + }; + + // Count tokens per expert + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + atomicAdd_block(smem_expert_count + expert_idx, 1); + }); + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + +// Stake out per-expert SM offsets via global atomic +#pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const uint64_t send_value = (1ull << 32) | static_cast(smem_expert_count[i]); + smem_expert_count[i] = + static_cast(ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value)); + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Write source token-topk indices to remote ranks + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + const auto dst_rank_idx = expert_idx / kNumExpertsPerRank; + const auto dst_slot_idx = atomicAdd_block(smem_expert_count + expert_idx, 1); + const auto dst_ptr = workspace.get_src_token_topk_idx_ptr(expert_idx % kNumExpertsPerRank, + sym_buffer.rank_idx, dst_slot_idx); + *sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx; + }); + + comm::grid_sync(workspace, sm_idx, thread_idx, [=]() { + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + }); + + if (sm_idx == 0 and thread_idx < kNumActiveDispatchThreads) { +#pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumActiveDispatchThreads) { + const auto dst_rank_idx = i / kNumExpertsPerRank; + const auto dst_local_expert_idx = i % kNumExpertsPerRank; + const auto expert_status = *workspace.get_expert_send_count_ptr(i); + *sym_buffer.map( + workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx), + dst_rank_idx) = expert_status & 0xffffffff; + ptx::atomic_add_sys( + sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), + dst_rank_idx), + expert_status); + } + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, false, true); + if (thread_idx == 0) stamp_max(1); + + // Sync with epilogue warps before pulling tokens + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Token / SF pull loop + if (warp_idx < kNumActiveDispatchWarps) { + uint32_t pull_mbarrier_phase = 0; + const auto pull_buffer = smem_send_buffers.get_rank_buffer(warp_idx).get_data_buffer(0); + const auto pull_mbarrier = dispatch_barriers[warp_idx]; + + scheduler.fetch_expert_recv_count(); + + constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u); + int current_expert_idx = -1; + uint32_t stored_rank_count[kNumRanksPerLane] = {}; + uint32_t expert_start_idx = 0, expert_end_idx = 0; + uint32_t expert_pool_block_offset = 0; + + constexpr uint32_t kNumGlobalWarps = kNumSMs * kNumActiveDispatchWarps; + for (uint32_t token_idx = sm_idx * kNumActiveDispatchWarps + warp_idx;; + token_idx += kNumGlobalWarps) { + int old_expert_idx = current_expert_idx; + while (token_idx >= expert_end_idx) { + if (++current_expert_idx >= kNumExpertsPerRank) break; + expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M); + expert_start_idx = expert_end_idx; + expert_end_idx += scheduler.get_num_tokens(current_expert_idx); + } + if (current_expert_idx >= kNumExpertsPerRank) break; + + if (old_expert_idx != current_expert_idx) { + old_expert_idx = current_expert_idx; +#pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++i) { + const uint32_t j = i * 32 + lane_idx; + stored_rank_count[i] = j < kNumRanks + ? static_cast(*workspace.get_expert_recv_count_ptr( + j, current_expert_idx)) + : 0; + } + } + + // Round-robin rank selection (identical to SM100) + +#include + const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr( + current_expert_idx, current_rank_in_expert_idx, token_idx_in_rank); + const uint32_t src_token_idx = src_token_topk_idx / kNumTopk; + const uint32_t src_topk_idx = src_token_topk_idx % kNumTopk; + + // TMA pull token data into SMEM + if (cute::elect_one_sync()) { + ptx::tma_load_1d( + pull_buffer.get_base_ptr(), + sym_buffer.map(input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx), + pull_mbarrier, kHidden); + } + __syncwarp(); + + // Copy SF: per-128 K floats, written linearly (no UTCCP transpose). + constexpr uint32_t kNumSFFloats = kHidden / 128; + FI_MEGAMOE_STATIC_ASSERT(kNumSFFloats > 0 and kHidden % 128 == 0, "Invalid SF"); + const auto remote_sf_ptr = + sym_buffer.map(input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx); + const auto local_sf_ptr = l1_sf_buffer.get_base_ptr(); + const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; +#pragma unroll + for (uint32_t i = 0; i < math::constexpr_ceil_div(kNumSFFloats, 32u); ++i) { + const uint32_t j = i * 32 + lane_idx; + if (j < kNumSFFloats) + local_sf_ptr[j * kNumPaddedSFPoolTokens + pool_token_idx] = remote_sf_ptr[j]; + } + __syncwarp(); + + if (cute::elect_one_sync()) { + const auto weight = + *sym_buffer.map(input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, + current_rank_in_expert_idx); + *l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr() = weight; + + ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kHidden); + ptx::mbarrier_wait_and_flip_phase(pull_mbarrier, pull_mbarrier_phase); + + ptx::tma_store_1d(l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr(), + pull_buffer.get_base_ptr(), pull_buffer.get_num_bytes()); + + *workspace.get_token_src_metadata_ptr(pool_token_idx) = {current_rank_in_expert_idx, + src_token_idx, src_topk_idx}; + + cute::tma_store_arrive(); + ptx::tma_store_wait<0>(); + ptx::red_add_rel(workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + + token_idx_in_expert / BLOCK_M), + 1u); + } + __syncwarp(); + } + } + + // Cleanup workspace, overlapping with combine + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + cleanup_workspace(); + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, true, false); + if (thread_idx == 0) stamp_max(8); +} else if (warp_idx == kNumDispatchWarps) { + // ===================================================================== + // ROLE 2: GEMM TMA LOAD warps (load A+SFA, B+SFB) + // The two warps inside `kNumNonEpilogueThreads` load A + SFA and + // B + SFB, respectively. + // ===================================================================== + cutlass::arch::warpgroup_reg_dealloc(); + + const auto load_a_task = [&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx, const uint32_t& pool_block_idx, + const uint32_t& valid_m) { + using BlockPhaseTag = std::remove_cv_t>; + constexpr bool kBlockIsL2 = BlockPhaseTag::value == sched::BlockPhase::Linear2; + const auto tensor_map_a_ptr = kBlockIsL2 ? &tensor_map_l2_acts : &tensor_map_l1_acts; + const auto tensor_map_sfa_ptr = kBlockIsL2 ? &tensor_map_l2_acts_sf : &tensor_map_l1_acts_sf; + + const bool has_valid_m = valid_m > 0; + + // Wait for the pool to be ready. + if (has_valid_m) { + if constexpr (!kBlockIsL2) { + const auto ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx); + while (ptx::ld_acq(ptr) != valid_m) { + } + } else { + const auto ptr = workspace.get_l2_arrival_mask_ptr(pool_block_idx); + const uint64_t expected = + (kNumRoutedL1BlockNs >= 64) ? ~0ull : ((1ull << kNumRoutedL1BlockNs) - 1ull); + while (ptx::ld_acq_gpu(ptr) != expected) { + } + } + } + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + if (cute::elect_one_sync()) { + if (has_valid_m) { + const uint32_t m_idx = pool_block_idx * BLOCK_M; + const uint32_t k_idx = k_block_idx * BLOCK_K; + + // TMA load A + tma::copy( + tensor_map_a_ptr, full_barriers[stage_idx], smem_a[stage_idx], k_idx, m_idx, 1); + + if constexpr (!kBlockIsL2 || !kSplitMDecodedWeightReuse) { + tma::copy(tensor_map_sfa_ptr, full_barriers[stage_idx], + smem_sfa[stage_idx], m_idx, k_block_idx, 1); + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + + BLOCK_M * sizeof(float)); + } else { + // BN128 L1 produces per-64 activation scales. L2 + // consumes both scale groups for each BK128 tile. + tma::copy(tensor_map_sfa_ptr, full_barriers[stage_idx], + smem_sfa[stage_idx], m_idx, k_block_idx * 2, 1); + tma::copy(tensor_map_sfa_ptr, full_barriers[stage_idx], + smem_sfa[stage_idx] + kL2SFAHalfStride, m_idx, + k_block_idx * 2 + 1, 1); + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + + 2 * BLOCK_M * sizeof(float)); + } + } else { + full_barriers[stage_idx]->arrive(); + } + } + __syncwarp(); + } + }; + if constexpr (kUseInterleavedScheduler) + for_each_published_block(load_a_task); + else + for_each_static_selected_block(load_a_task); + +} else if (warp_idx == kNumDispatchWarps + 1) { + cutlass::arch::warpgroup_reg_dealloc(); + + const auto load_b_task = [&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx, const uint32_t& pool_block_idx, + const uint32_t& valid_m) { + using BlockPhaseTag = std::remove_cv_t>; + constexpr bool kBlockIsL2 = BlockPhaseTag::value == sched::BlockPhase::Linear2; + const uint8_t* b_base = kBlockIsL2 ? l2_weights : l1_weights; + constexpr uint32_t shape_n = kBlockIsL2 ? L2_SHAPE_N : L1_SHAPE_N; + // One tile is `kBTileRows` fused rows laid out contiguously, so a + // stage is `BLOCK_N / kBTileRows` bulk copies of adjacent tiles. + constexpr uint32_t kNumTileLoads = LOAD_BLOCK_N / kBTileRows; + constexpr uint32_t kNumNTiles = shape_n / kBTileRows; + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + const uint32_t n_tile_idx = local_expert_idx * kNumNTiles + n_block_idx * kNumTileLoads; + if (cute::elect_one_sync()) { +#pragma unroll + for (uint32_t i = 0; i < kNumTileLoads; ++i) { + const uint64_t tile_idx = + static_cast(n_tile_idx + i) * num_k_blocks + k_block_idx; + // EVICT_NORMAL, not tma_load_1d's streaming default: + // above the swapAB tiers an expert spans several + // m-blocks, and those CTAs share these tiles via L2. + ptx::tma_load_1d(smem_packed_b[stage_idx] + i * kBTileBytes, + b_base + tile_idx * kBTileBytes, full_barriers[stage_idx], kBTileBytes, + cute::TMA::CacheHintSm90::EVICT_NORMAL); + } + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_PACKED_B_SIZE_PER_STAGE); + } + __syncwarp(); + } + }; + if constexpr (kUseInterleavedScheduler) + produce_interleaved_blocks(load_b_task); + else + for_each_static_selected_block(load_b_task); + +} else if (kSplitDecodeWarps and warp_idx < kFirstMathWarp) { + // ===================================================================== + // ROLE 3a: DECODE WARPS (FP4 -> FP8 into the decoded ring) + // ===================================================================== + cutlass::arch::warpgroup_reg_dealloc(); + const uint32_t decode_thread_idx = (warp_idx - kFirstDecodeWarp) * 32 + lane_idx; + + const auto run_decode_task = [&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx, const uint32_t& pool_block_idx, + const uint32_t& valid_m) { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + { + const unsigned long long w0 = stamp_now(); + full_barriers[stage_idx]->wait(phase); + if (decode_thread_idx == 0) stamp_add(24, stamp_now() - w0); + } + if constexpr (kUseInterleavedScheduler) { + if (k_block_idx == 0) interleaved_scheduler.release_task_info(lane_idx); + } + decoded_empty_barriers[b_ring_idx]->wait(b_ring_phase ^ 1); + flashinfer::megamoe::mxfp4::dequant_smem_b_from_packed_mode2_nibble( + reinterpret_cast(smem_b[b_ring_idx]), + reinterpret_cast(smem_packed_b[stage_idx]), decode_thread_idx, + smem_mxfp4_lut); + cutlass::arch::fence_view_async_shared(); + if (lane_idx == 0) { + // Publishing does not block, so this warp runs on into the + // other slot while the math warps consume this one. + decoded_full_barriers[b_ring_idx]->arrive(); + // The packed stage is ours to release once it is decoded. + empty_barriers[stage_idx]->arrive(); + } + } + }; + if constexpr (kUseInterleavedScheduler) + for_each_published_block(run_decode_task); + else + for_each_static_selected_block(run_decode_task); + +} else { + // ===================================================================== + // ROLE 3: MATH WARPGROUPS (WGMMA + epilogue + combine) + // ===================================================================== + cutlass::arch::warpgroup_reg_alloc(); + + const uint32_t epilogue_warp_idx = warp_idx - kFirstMathWarp; + const uint32_t epilogue_wg_idx = epilogue_warp_idx / 4; + const uint32_t epilogue_thread_idx = epilogue_warp_idx * 32 + lane_idx; + const uint32_t warp_idx_in_wg = epilogue_warp_idx % 4; + + const auto arrive_empty_barrier = [&](const uint32_t& s) { + if (lane_idx == 0) { + // Hand the decoded slot back, then release the stage itself: + // its activations were read by these warps. + if constexpr (kSplitDecodeWarps) decoded_empty_barriers[b_ring_idx]->arrive(); + empty_barriers[s]->arrive(); + } + }; + + const auto notify_l1_ready = [&](const uint32_t& ready_pool_block_idx, + const uint32_t& ready_n_block_idx) { + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + ptx::red_or_rel_gpu(workspace.get_l2_arrival_mask_ptr(ready_pool_block_idx), + 1ull << ready_n_block_idx); + } + __syncwarp(); + }; + + // WGMMA-output register layout helpers + const uint32_t row_idx = lane_idx / 4; + const uint32_t col_idx = lane_idx % 4; + const uint32_t r_0 = warp_idx_in_wg * 16 + row_idx; + const uint32_t r_1 = r_0 + 8; + + FI_MEGAMOE_STATIC_ASSERT( + kSwapABRequested || (WG_BLOCK_M == L1WGMMA::M and WG_BLOCK_N == L1WGMMA::N), + "Split-N WGs must each run one M64N128 WGMMA per K-block"); + + // Sync with dispatch + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + const auto run_math_task = [&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx, const uint32_t& pool_block_idx, + const uint32_t& valid_m) { + const unsigned long long task_t0 = stamp_now(); + if (epilogue_thread_idx == 0) stamp_min(3); + const uint32_t m_idx = pool_block_idx * BLOCK_M; + const uint32_t wg_n_idx = kSplitMDecodedWeightReuse ? 0u : epilogue_wg_idx * WG_BLOCK_N; + const uint32_t wg_l1_out_n_idx = + kSplitMDecodedWeightReuse ? 0u : epilogue_wg_idx * WG_L1_OUT_BLOCK_N; + const uint32_t n_idx = n_block_idx * BLOCK_N + wg_n_idx; + const uint32_t row_block_offset = kSplitMDecodedWeightReuse ? epilogue_wg_idx * WG_BLOCK_M : 0u; + const uint32_t row_offset_r0 = row_block_offset + r_0; + const uint32_t row_offset_r1 = row_block_offset + r_1; + using BlockPhaseTag = std::remove_cv_t>; + constexpr bool kBlockIsL2 = BlockPhaseTag::value == sched::BlockPhase::Linear2; + const float l2_global_scale = + l2_global_scales == nullptr ? 1.0f : __ldg(l2_global_scales + local_expert_idx); + const auto cast_l2_scaled_bf16_pair = [&](float x, float y) -> uint32_t { + x *= l2_global_scale; + y *= l2_global_scale; + return math::cast_into_bf16_and_pack(x, y); + }; + + // ---------------- GEMM ---------------- + using WGMMA = L1WGMMA; + constexpr uint32_t kAccumPerThread = WGMMA::kNumAccum; // 64 for M=64,N=128 + float final_accum[kAccumPerThread] = {}; + const auto decode_b_stage = [&](const uint32_t& packed_stage, const uint32_t& decoded_slot) { + if constexpr (kSplitMDecodedWeightReuse) { + // Both M64 consumers share the same decoded N128 tile. + // The two physical decoded slots remove the overwrite + // hazard; the trailing pair barrier publishes K+1 before + // either consumer issues its WGMMA. + FI_MEGAMOE_STATIC_ASSERT(kUseMode2RowDecoder, + "BM128 split-M uses the cooperative Mode2 decoder"); + flashinfer::megamoe::mxfp4::dequant_smem_b_from_packed_mode2_nibble_split_m( + reinterpret_cast(smem_b[decoded_slot]), + reinterpret_cast(smem_packed_b[packed_stage]), epilogue_thread_idx, + smem_mxfp4_lut); + cutlass::arch::fence_view_async_shared(); + asm volatile("bar.sync %0, %1;" : : "n"(kSplitMDecodeBarrierIdx), "n"(256) : "memory"); + } else { + flashinfer::megamoe::mxfp4::dequant_smem_b_from_packed_mode2_nibble( + reinterpret_cast(smem_b[decoded_slot]), + reinterpret_cast(smem_packed_b[packed_stage]), epilogue_thread_idx, + smem_mxfp4_lut); + cutlass::arch::fence_view_async_shared(); + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + } + }; + + // Register-source decode: one lane pair owns one weight row of the + // half-tile, and each slice of four A registers is produced from two + // packed words plus one shuffle. Nothing lands in shared memory, so + // this path allocates no decoded-B tile at all. + const auto decode_rs_half = [&](const uint32_t& rs_stage, const uint32_t& half, + uint32_t(&a_frag)[4][4]) { + const uint32_t frag_row0 = wg_n_idx + half * 64u + warp_idx_in_wg * 16u + row_idx; + const uint32_t decode_row = frag_row0 + ((lane_idx & 1u) << 3); + const uint32_t word_sel = (lane_idx >> 1) & 1u; + const bool keep_hi = (lane_idx & 1u) == 0; + const auto src = mxfp4::decode_tile_row( + reinterpret_cast(smem_packed_b[rs_stage]), decode_row); + const uint32_t scale_word = ptx::ld_shared(reinterpret_cast(src.scale)); + +#pragma unroll + for (uint32_t slice = 0; slice < 4; ++slice) { + const auto* __restrict__ chunk = + src.chunk + slice * mxfp4::kBChunkStride + word_sel * sizeof(uint32_t); + const uint32_t w_lo = ptx::ld_shared(reinterpret_cast(chunk)); + const uint32_t w_hi = ptx::ld_shared(reinterpret_cast(chunk + 8u)); + const auto lut = + mxfp4::load_scaled_lut(smem_mxfp4_lut, (scale_word >> (slice * 8u)) & 0xffu); + mxfp4::dequant_rs_word_pair(w_lo, w_hi, lut, keep_hi, a_frag[slice]); +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) mma::sm90::warpgroup_fence_operand(a_frag[slice][i]); + } + }; + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + // BM128 split-M alternates two decoded slots keyed by the + // pipeline stage, because its paired warpgroups share the same + // N range and one must not overwrite a tile the other is still + // reading. Every other plan decodes and consumes inside one + // iteration and has a single slot. + const uint32_t b_slot = + kSplitDecodeWarps ? b_ring_idx : (kSplitMDecodedWeightReuse ? stage_idx : 0u); + { + const unsigned long long w0 = stamp_now(); + if constexpr (kSplitDecodeWarps) { + // The decode warps own the packed stage and its release, + // so wait on the decoded tile rather than the TMA. + decoded_full_barriers[b_ring_idx]->wait(b_ring_phase); + } else { + full_barriers[stage_idx]->wait(phase); + } + if (epilogue_thread_idx == 0) stamp_add(26, stamp_now() - w0); + } + if constexpr (kUseInterleavedScheduler) { + if (k_block_idx == 0) interleaved_scheduler.release_task_info(lane_idx); + } + if constexpr (!kUseRSOperand and !kSplitDecodeWarps) decode_b_stage(stage_idx, b_slot); + + // Read SF (must precede warpgroup_arrive) + const float scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); + const float scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); + float scale_a_0_hi = 0.0f; + float scale_a_1_hi = 0.0f; + if constexpr (kBlockIsL2 && kSplitMDecodedWeightReuse) { + scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + kL2SFAHalfStride + row_offset_r0); + scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + kL2SFAHalfStride + row_offset_r1); + } + + // NVFP4 UE4M3 weight scales are applied during FP4 -> FP8 smem + // expansion, so the WGMMA accumulator only needs activation SF. + + if constexpr (!kBlockIsL2) { + if constexpr (kSwapABRequested) { + auto run_swap_ab_l1 = [&]() { + using SwapWGMMA = typename mma::sm90::FP8MMASelector::type; + constexpr uint32_t kSwapAccum = SwapWGMMA::kNumAccum; + FI_MEGAMOE_STATIC_ASSERT(kSwapABWeightHalves == 2, + "BN256 swapAB splits the weight tile into two halves"); + float swap_accum[2][kSwapAccum]; + + // The two halves are independent WGMMA groups, so issue both and only + // then wait. The window between the last issue and the first wait is + // where the next K-block's dequant hides. + using RSWGMMA = typename mma::sm90::FP8RSMMASelector::type; + // One fragment set per half. Under RS the decode sits + // between the stage barrier and the WGMMA, so giving + // each half its own registers lets half 1's decode run + // under half 0's MMA. Unlike prefetching the next + // K-block, this stays inside one stage and so never + // puts a TMA wait in front of arrive_empty. + uint32_t a_frags[kUseRSOperand ? 2 : 1][4][4]; + const auto issue_half = [&](const uint32_t half) { + auto& a_frag = a_frags[kUseRSOperand ? half : 0]; + if constexpr (kUseRSOperand) decode_rs_half(stage_idx, half, a_frag); +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(swap_accum[half][i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / SwapWGMMA::K; ++k) { + if constexpr (kUseRSOperand) { + auto desc_b = mma::sm90::make_smem_desc(smem_a[stage_idx] + k * RSWGMMA::K, 1); + RSWGMMA::wgmma(a_frag[k], desc_b, swap_accum[half], k); + } else { + auto desc_a = mma::sm90::make_smem_desc( + smem_b[b_slot] + (wg_n_idx + half * 64u) * BLOCK_K + k * SwapWGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc(smem_a[stage_idx] + k * SwapWGMMA::K, 1); + SwapWGMMA::wgmma(desc_a, desc_b, swap_accum[half], k); + } + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(swap_accum[half][i]); + }; + // The next issue_half overwrites a_frag, so the + // fragments have to be re-fenced once the WGMMA that + // reads them has drained. + const auto fence_frags = [&](const uint32_t half) { + if constexpr (kUseRSOperand) { + auto& a_frag = a_frags[half]; +#pragma unroll + for (uint32_t slice = 0; slice < 4; ++slice) +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) + mma::sm90::warpgroup_fence_operand(a_frag[slice][i]); + } + }; + const auto accumulate_half = [&](const uint32_t half) { +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum / 4; ++i) { + const uint32_t accum_offset = half * kSwapABHalfAccumPerThread + i * 4; + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + if (token_0 < valid_m) { + const float scale_0 = ptx::ld_shared(smem_sfa[stage_idx] + token_0); + final_accum[accum_offset + 0] += scale_0 * swap_accum[half][i * 4 + 0]; + final_accum[accum_offset + 2] += scale_0 * swap_accum[half][i * 4 + 2]; + } + if (token_1 < valid_m) { + const float scale_1 = ptx::ld_shared(smem_sfa[stage_idx] + token_1); + final_accum[accum_offset + 1] += scale_1 * swap_accum[half][i * 4 + 1]; + final_accum[accum_offset + 3] += scale_1 * swap_accum[half][i * 4 + 3]; + } + } + }; + + if constexpr (kUseRSOperand) { + issue_half(0); + issue_half(1); // decodes under half 0's MMA + ptx::warpgroup_wait<1>(); + fence_frags(0); + accumulate_half(0); + ptx::warpgroup_wait<0>(); + fence_frags(1); + accumulate_half(1); + } else { + issue_half(0); + ptx::warpgroup_wait<0>(); + accumulate_half(0); + issue_half(1); + ptx::warpgroup_wait<0>(); + accumulate_half(1); + } + + arrive_empty_barrier(stage_idx); + }; + + if constexpr (BLOCK_M == 8) { + run_swap_ab_l1.template operator()<8>(); + } else if constexpr (BLOCK_M == 16) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if (n_swap <= 8) { + run_swap_ab_l1.template operator()<8>(); + } else { + run_swap_ab_l1.template operator()<16>(); + } + } else if constexpr (BLOCK_M == 24) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if (n_swap <= 8) { + run_swap_ab_l1.template operator()<8>(); + } else if (n_swap <= 16) { + run_swap_ab_l1.template operator()<16>(); + } else { + run_swap_ab_l1.template operator()<24>(); + } + } else if constexpr (BLOCK_M == 32) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if (n_swap <= 8) { + run_swap_ab_l1.template operator()<8>(); + } else if (n_swap <= 16) { + run_swap_ab_l1.template operator()<16>(); + } else if (n_swap <= 24) { + run_swap_ab_l1.template operator()<24>(); + } else { + run_swap_ab_l1.template operator()<32>(); + } + } else { + // A tile with no arm here never runs the mainloop, + // so nothing releases the stage and the loader + // hangs. Fail at compile time instead. + FI_MEGAMOE_STATIC_ASSERT( + BLOCK_M == 8 || BLOCK_M == 16 || BLOCK_M == 24 || BLOCK_M == 32, + "swapAB tile has no N_SWAP arm"); + } + } else { + float accum[kAccumPerThread]; +// Single per-128 K-block WGMMA group +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + row_block_offset * BLOCK_K + k * WGMMA::K, 1); + auto desc_b = + mma::sm90::make_smem_desc(smem_b[b_slot] + wg_n_idx * BLOCK_K + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + arrive_empty_barrier(stage_idx); + +// L1: gate/up alternate at gran=8 along N; each `i` block +// of 8 cols belongs entirely to one of {gate, up}, so .x +// and .y share the same scalar. +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++i) { + final_accum[i * 4 + 0] += scale_a_0_lo * accum[i * 4 + 0]; + final_accum[i * 4 + 1] += scale_a_0_lo * accum[i * 4 + 1]; + final_accum[i * 4 + 2] += scale_a_1_lo * accum[i * 4 + 2]; + final_accum[i * 4 + 3] += scale_a_1_lo * accum[i * 4 + 3]; + } + } + } else { + if constexpr (kSwapABRequested) { + FI_MEGAMOE_STATIC_ASSERT(kL2ActsSFGranK == 128, + "L2 swap-AB requires per-128 activation scales"); + auto run_swap_ab_l2 = [&]() { + using SwapWGMMA = typename mma::sm90::FP8MMASelector::type; + constexpr uint32_t kSwapAccum = SwapWGMMA::kNumAccum; + FI_MEGAMOE_STATIC_ASSERT(kSwapABWeightHalves == 2, + "BN256 swapAB splits the weight tile into two halves"); + float swap_accum[2][kSwapAccum]; + + // The two halves are independent WGMMA groups, so issue both and only + // then wait. The window between the last issue and the first wait is + // where the next K-block's dequant hides. + using RSWGMMA = typename mma::sm90::FP8RSMMASelector::type; + // One fragment set per half. Under RS the decode sits + // between the stage barrier and the WGMMA, so giving + // each half its own registers lets half 1's decode run + // under half 0's MMA. Unlike prefetching the next + // K-block, this stays inside one stage and so never + // puts a TMA wait in front of arrive_empty. + uint32_t a_frags[kUseRSOperand ? 2 : 1][4][4]; + const auto issue_half = [&](const uint32_t half) { + auto& a_frag = a_frags[kUseRSOperand ? half : 0]; + if constexpr (kUseRSOperand) decode_rs_half(stage_idx, half, a_frag); +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(swap_accum[half][i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / SwapWGMMA::K; ++k) { + if constexpr (kUseRSOperand) { + auto desc_b = mma::sm90::make_smem_desc(smem_a[stage_idx] + k * RSWGMMA::K, 1); + RSWGMMA::wgmma(a_frag[k], desc_b, swap_accum[half], k); + } else { + auto desc_a = mma::sm90::make_smem_desc( + smem_b[b_slot] + (wg_n_idx + half * 64u) * BLOCK_K + k * SwapWGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc(smem_a[stage_idx] + k * SwapWGMMA::K, 1); + SwapWGMMA::wgmma(desc_a, desc_b, swap_accum[half], k); + } + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(swap_accum[half][i]); + }; + // The next issue_half overwrites a_frag, so the + // fragments have to be re-fenced once the WGMMA that + // reads them has drained. + const auto fence_frags = [&](const uint32_t half) { + if constexpr (kUseRSOperand) { + auto& a_frag = a_frags[half]; +#pragma unroll + for (uint32_t slice = 0; slice < 4; ++slice) +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) + mma::sm90::warpgroup_fence_operand(a_frag[slice][i]); + } + }; + const auto accumulate_half = [&](const uint32_t half) { +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum / 4; ++i) { + const uint32_t accum_offset = half * kSwapABHalfAccumPerThread + i * 4; + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + if (token_0 < valid_m) { + const float scale_0 = ptx::ld_shared(smem_sfa[stage_idx] + token_0); + final_accum[accum_offset + 0] += scale_0 * swap_accum[half][i * 4 + 0]; + final_accum[accum_offset + 2] += scale_0 * swap_accum[half][i * 4 + 2]; + } + if (token_1 < valid_m) { + const float scale_1 = ptx::ld_shared(smem_sfa[stage_idx] + token_1); + final_accum[accum_offset + 1] += scale_1 * swap_accum[half][i * 4 + 1]; + final_accum[accum_offset + 3] += scale_1 * swap_accum[half][i * 4 + 3]; + } + } + }; + + if constexpr (kUseRSOperand) { + issue_half(0); + issue_half(1); // decodes under half 0's MMA + ptx::warpgroup_wait<1>(); + fence_frags(0); + accumulate_half(0); + ptx::warpgroup_wait<0>(); + fence_frags(1); + accumulate_half(1); + } else { + issue_half(0); + ptx::warpgroup_wait<0>(); + accumulate_half(0); + issue_half(1); + ptx::warpgroup_wait<0>(); + accumulate_half(1); + } + + arrive_empty_barrier(stage_idx); + }; + + if constexpr (BLOCK_M == 8) { + run_swap_ab_l2.template operator()<8>(); + } else if constexpr (BLOCK_M == 16) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if (n_swap <= 8) { + run_swap_ab_l2.template operator()<8>(); + } else { + run_swap_ab_l2.template operator()<16>(); + } + } else if constexpr (BLOCK_M == 24) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if (n_swap <= 8) { + run_swap_ab_l2.template operator()<8>(); + } else if (n_swap <= 16) { + run_swap_ab_l2.template operator()<16>(); + } else { + run_swap_ab_l2.template operator()<24>(); + } + } else if constexpr (BLOCK_M == 32) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if (n_swap <= 8) { + run_swap_ab_l2.template operator()<8>(); + } else if (n_swap <= 16) { + run_swap_ab_l2.template operator()<16>(); + } else if (n_swap <= 24) { + run_swap_ab_l2.template operator()<24>(); + } else { + run_swap_ab_l2.template operator()<32>(); + } + } else { + // A tile with no arm here never runs the mainloop, + // so nothing releases the stage and the loader + // hangs. Fail at compile time instead. + FI_MEGAMOE_STATIC_ASSERT( + BLOCK_M == 8 || BLOCK_M == 16 || BLOCK_M == 24 || BLOCK_M == 32, + "swapAB tile has no N_SWAP arm"); + } + } else { + float accum[kAccumPerThread]; + const auto promote_l2_accum = [&](const float& scale_r0, const float& scale_r1) { +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++i) { + final_accum[i * 4 + 0] += scale_r0 * accum[i * 4 + 0]; + final_accum[i * 4 + 1] += scale_r0 * accum[i * 4 + 1]; + final_accum[i * 4 + 2] += scale_r1 * accum[i * 4 + 2]; + final_accum[i * 4 + 3] += scale_r1 * accum[i * 4 + 3]; + } + }; + if constexpr (kSplitMDecodedWeightReuse) { +#pragma unroll + for (uint32_t sf_group = 0; sf_group < 2; ++sf_group) { +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++k) { + const uint32_t k_off = sf_group * (BLOCK_K / 2) + k * WGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + row_block_offset * BLOCK_K + k_off, 1); + auto desc_b = + mma::sm90::make_smem_desc(smem_b[b_slot] + wg_n_idx * BLOCK_K + k_off, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + if (sf_group == 0) + promote_l2_accum(scale_a_0_lo, scale_a_1_lo); + else + promote_l2_accum(scale_a_0_hi, scale_a_1_hi); + } + arrive_empty_barrier(stage_idx); + } else { +// One per-128 scale permits a single four-instruction +// WGMMA group and one accumulator promotion per K tile. +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + row_block_offset * BLOCK_K + k * WGMMA::K, 1); + auto desc_b = + mma::sm90::make_smem_desc(smem_b[b_slot] + wg_n_idx * BLOCK_K + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + // All four warps must finish reading decoded B before the next K tile overwrites it. + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + arrive_empty_barrier(stage_idx); + promote_l2_accum(scale_a_0_lo, scale_a_1_lo); + } + } + } + } + + // Skip epilogue when block is past valid M (the GEMM loop already + // released its pipeline stages). Drain any prior L1 async store. + if (valid_m == 0) { + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + return; + } + +#include + + if (epilogue_thread_idx == 0) { + stamp_add(kBlockIsL2 ? 22u : 20u, stamp_now() - task_t0); + stamp_max(kBlockIsL2 ? 5u : 4u); + } + }; + if constexpr (kUseInterleavedScheduler) + for_each_published_block(run_math_task); + else + for_each_static_selected_block(run_math_task); + + // ---------------- COMBINE ---------------- + // NVLink barrier first: signals remote ranks that this rank's GEMM + // outputs (NVLink scatter targets) are fully written. + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, epilogue_thread_idx, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); }); + if (epilogue_thread_idx == 0) stamp_max(6); + + // Sync with dispatch (paired with dispatch's pre-cleanup sync) so that + // dispatch may now safely clean workspace state. + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + +// These are function fragments: sizing declarations must precede the combine body. +// clang-format off +#include +#include + // clang-format on +} +#else +if (blockIdx.x == 0 and threadIdx.x == 0) + FI_MEGAMOE_DEVICE_ASSERT(false and "This kernel only supports sm_90"); +#endif diff --git a/include/flashinfer/megamoe/sm90_nvfp4_mega_moe.cuh b/include/flashinfer/megamoe/sm90_nvfp4_mega_moe.cuh new file mode 100644 index 00000000000..428784066b2 --- /dev/null +++ b/include/flashinfer/megamoe/sm90_nvfp4_mega_moe.cuh @@ -0,0 +1,114 @@ +#pragma once + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace flashinfer::megamoe { + +// ============================================================================ +// SM90 (Hopper) NVFP4 MegaMoE kernels +// ---------------------------------------------------------------------------- +// Split pipeline: +// * L1 uses cluster=2 to pair adjacent 64-column output halves. +// * L2 uses independent cluster=1 CTAs. +// * Dispatch warps: pull tokens (FP8) and SF (per-128 channel float) from +// remote ranks via NVLink into the local L1 pool. +// * GEMM TMA-load warps (1 for A+SFA, 1 for B+SFB) feed the pipeline stages. +// * Math warpgroups (1 or 2, totalling kNumEpilogueThreads) consume each +// stage with WGMMA, accumulate into registers, then run the epilogue: +// - L1 (Linear1): SwiGLU with gate/up granularity-8 interleaved layout. +// The deployed split kernel pairs two 64-column N halves in one CTA, +// combines their row-wise maxima, quantizes the full 128-column output +// group with one E4M3 scale, and TMA-stores it to the local L1 buffer. +// - L2 (Linear2): BF16 cast of the GEMM output, STSM into SMEM, then +// NVLink scatter to remote combine buffers. +// * After all GEMM blocks, the L2 math warps run the COMBINE step (top-k +// reduction in BF16) — ported verbatim from the SM100 kernel. +// +// Public BN256 requests use sm90_nvfp4_mega_moe_small_m.cuh. The general host +// runtime instantiates only the split kernels below for BN128 deployment +// weights. +// ============================================================================ + +template +CUTLASS_GLOBAL __launch_bounds__(512, 1) void sm90_nvfp4_mega_moe_split_l1_impl( + int* cumulative_local_expert_recv_stats, const uint32_t num_tokens, + const __grid_constant__ layout::SymBuffer sym_buffer, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output, + const float* __restrict__ l1_global_scales) { + constexpr uint32_t BLOCK_M = 128; + constexpr uint32_t BLOCK_N = 128; + constexpr uint32_t BLOCK_K = 128; + constexpr uint32_t kClusterSize = 2; + constexpr uint32_t kNumDispatchThreads = 128; + constexpr uint32_t kNumNonEpilogueThreads = 128; + constexpr uint32_t kNumEpilogueThreads = 256; + constexpr uint32_t L1_SHAPE_N = kIntermediateHidden * 2; + constexpr uint32_t L1_SHAPE_K = kHidden; + constexpr uint32_t L2_SHAPE_N = kHidden; + constexpr uint32_t L2_SHAPE_K = kIntermediateHidden; + constexpr uint32_t kNumDispatchWarps = kNumDispatchThreads / 32; + constexpr uint32_t kNumMMANonEpilogueWarps = kNumNonEpilogueThreads / 32; + constexpr uint32_t kNumEpilogueWarps = kNumEpilogueThreads / 32; + constexpr uint32_t kNumEpilogueWarpgroups = kNumEpilogueWarps / 4; + constexpr uint32_t kNumTokensPerWarp = 32 / kNumTopk; + constexpr uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks; +#include +} + +template +CUTLASS_GLOBAL __launch_bounds__(384, 1) void sm90_nvfp4_mega_moe_split_l2_impl( + void* y, int* cumulative_local_expert_recv_stats, const uint32_t num_tokens, + const __grid_constant__ layout::SymBuffer sym_buffer, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, + const float* __restrict__ l2_global_scales) { + constexpr uint32_t BLOCK_M = 128; + constexpr uint32_t BLOCK_N = 128; + constexpr uint32_t BLOCK_K = 128; + constexpr uint32_t kNumNonEpilogueThreads = 128; + constexpr uint32_t kNumEpilogueThreads = 256; + constexpr uint32_t L1_SHAPE_N = kIntermediateHidden * 2; + constexpr uint32_t L1_SHAPE_K = kHidden; + constexpr uint32_t L2_SHAPE_N = kHidden; + constexpr uint32_t L2_SHAPE_K = kIntermediateHidden; + constexpr uint32_t kNumMMANonEpilogueWarps = kNumNonEpilogueThreads / 32; + constexpr uint32_t kNumEpilogueWarps = kNumEpilogueThreads / 32; + constexpr uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks; +#include +} + +} // namespace flashinfer::megamoe + +#pragma clang diagnostic pop diff --git a/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_fused.cuh b/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_fused.cuh new file mode 100644 index 00000000000..c1f3a9374c8 --- /dev/null +++ b/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_fused.cuh @@ -0,0 +1,404 @@ +#pragma once + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// The persistent fused kernel launches one CTA per SM and derives its +// dispatch/combine strides and grid barriers from the SM count. The JIT host +// defines MEGAMOE_NUM_SMS from the live device (132 on H200, 78 on H20-3e). +#ifndef MEGAMOE_NUM_SMS +#define MEGAMOE_NUM_SMS 132 +#endif + +namespace flashinfer::megamoe { +namespace nvfp4 { + +__device__ __forceinline__ uint2 dequant_mode2_nibble_word(const uint32_t packed, + const uint2& lut) { + return detail::decode_fp4_word(packed, lut); +} + +// Decode the two packed words owned by one lane pair directly into the four +// register-source WGMMA A operands. This is the archived mode-5 mapping: the +// exact shared LUT is retained and packed-row padding bytes 72..79 are never +// read. +__device__ __forceinline__ void dequant_mode2_rs_word_pair(const uint32_t w_lo, const uint32_t w_hi, + const uint2& lut_lo, const uint2& lut_hi, + const bool keep_hi, + uint32_t (&a_frag)[4]) { + detail::assemble_rs_word_pair(dequant_mode2_nibble_word(w_lo, lut_lo), + dequant_mode2_nibble_word(w_hi, lut_hi), keep_hi, a_frag); +} + +template +__device__ __forceinline__ void dequant_mode2_nibble_row_regs(uint8_t* __restrict__ fp8_dst, + const uint4 (&fp4_quads)[4], + const uint2& scale_words, + const uint32_t row_swizzle, + const uint2* __restrict__ lut_smem) { +#pragma unroll + for (int quad_i = 0; quad_i < 4; ++quad_i) { + const uint4 q = fp4_quads[quad_i]; + const uint32_t scale_word = quad_i < 2 ? scale_words.x : scale_words.y; + const int scale_i0 = quad_i * 2; + const int scale_i1 = scale_i0 + 1; + const uint32_t scale0 = (scale_word >> ((scale_i0 & 3) * 8)) & 0x7fu; + const uint32_t scale1 = (scale_word >> ((scale_i1 & 3) * 8)) & 0x7fu; + const uint2 lut0 = lut_smem[scale0]; + const uint2 lut1 = lut_smem[scale1]; + + const uint2 q0 = dequant_mode2_nibble_word(q.x, lut0); + const uint2 q1 = dequant_mode2_nibble_word(q.y, lut0); + if constexpr (!kQuadILP) { + *reinterpret_cast(fp8_dst + ((scale_i0 * 16) ^ row_swizzle)) = + make_uint4(q0.x, q0.y, q1.x, q1.y); + } + + const uint2 q2 = dequant_mode2_nibble_word(q.z, lut1); + const uint2 q3 = dequant_mode2_nibble_word(q.w, lut1); + if constexpr (kQuadILP) { + *reinterpret_cast(fp8_dst + ((scale_i0 * 16) ^ row_swizzle)) = + make_uint4(q0.x, q0.y, q1.x, q1.y); + } + *reinterpret_cast(fp8_dst + ((scale_i1 * 16) ^ row_swizzle)) = + make_uint4(q2.x, q2.y, q3.x, q3.y); + } +} + +template +__device__ __forceinline__ void dequant_smem_b_from_packed_mode2_nibble( + uint8_t* __restrict__ smem_b, const uint8_t* __restrict__ packed_b, const uint32_t row, + const uint2* __restrict__ lut_smem) { + const uint8_t* __restrict__ row_ptr = packed_b + row * 80; + const uint4* __restrict__ fp4_src = reinterpret_cast(row_ptr); + uint4 fp4_quads[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) fp4_quads[i] = fp4_src[i]; + const uint2 scale_words = *reinterpret_cast(row_ptr + 64); + dequant_mode2_nibble_row_regs(smem_b + row * 128, fp4_quads, scale_words, + (row & 7u) << 4, lut_smem); +} + +// Threads 0-127 and 128-255 each decode one K64 half of the same N128 tile, +// allowing the two M64 warpgroups to reuse the decoded weights. +__device__ __forceinline__ void dequant_smem_b_from_packed_mode2_nibble_split_m( + uint8_t* __restrict__ smem_b, const uint8_t* __restrict__ packed_b, const uint32_t thread_idx, + const uint2* __restrict__ lut_smem) { + const uint32_t row = thread_idx & 127u; + const uint32_t k_half_idx = thread_idx >> 7; + const uint8_t* __restrict__ row_ptr = packed_b + row * 80u; + const uint4* __restrict__ fp4_src = reinterpret_cast(row_ptr + k_half_idx * 32u); + const uint32_t scale_word = + *reinterpret_cast(row_ptr + 64u + k_half_idx * sizeof(uint32_t)); + uint8_t* __restrict__ fp8_dst = smem_b + row * 128u; + const uint32_t row_swizzle = (row & 7u) << 4; + +#pragma unroll + for (uint32_t quad_i = 0; quad_i < 2; ++quad_i) { + const uint4 q = fp4_src[quad_i]; + const uint32_t scale_i0 = quad_i * 2u; + const uint32_t scale_i1 = scale_i0 + 1u; + const uint32_t scale0 = (scale_word >> (scale_i0 * 8u)) & 0x7fu; + const uint32_t scale1 = (scale_word >> (scale_i1 * 8u)) & 0x7fu; + const uint2 lut0 = lut_smem[scale0]; + const uint2 lut1 = lut_smem[scale1]; + const uint2 q0 = dequant_mode2_nibble_word(q.x, lut0); + const uint2 q1 = dequant_mode2_nibble_word(q.y, lut0); + const uint2 q2 = dequant_mode2_nibble_word(q.z, lut1); + const uint2 q3 = dequant_mode2_nibble_word(q.w, lut1); + const uint32_t k_offset0 = k_half_idx * 64u + scale_i0 * 16u; + const uint32_t k_offset1 = k_half_idx * 64u + scale_i1 * 16u; + *reinterpret_cast(fp8_dst + (k_offset0 ^ row_swizzle)) = + make_uint4(q0.x, q0.y, q1.x, q1.y); + *reinterpret_cast(fp8_dst + (k_offset1 ^ row_swizzle)) = + make_uint4(q2.x, q2.y, q3.x, q3.y); + } +} + +__device__ __forceinline__ uint2 dequant_braided_selector_word(const uint32_t braided, + const uint2& lut) { + const uint32_t sel0 = braided & 0x00007777u; + const uint32_t sel1 = (braided >> 16) & 0x00007777u; + uint32_t out0 = byte_perm_unchecked(lut.x, lut.y, sel0); + uint32_t out1 = byte_perm_unchecked(lut.x, lut.y, sel1); + out0 |= braided & 0x80808080u; + out1 |= (braided << 4) & 0x80808080u; + return make_uint2(out0, out1); +} + +__device__ __forceinline__ void dequant_braided_quad(uint8_t* __restrict__ fp8_dst, const uint4& q, + const uint2& lut0, const uint2& lut1, + const int scale_i0, + const uint32_t row_swizzle) { + const uint2 q0 = dequant_braided_selector_word(q.x, lut0); + const uint2 q1 = dequant_braided_selector_word(q.y, lut0); + *reinterpret_cast(fp8_dst + ((scale_i0 * 16) ^ row_swizzle)) = + make_uint4(q0.x, q0.y, q1.x, q1.y); + + const uint2 q2 = dequant_braided_selector_word(q.z, lut1); + const uint2 q3 = dequant_braided_selector_word(q.w, lut1); + *reinterpret_cast(fp8_dst + (((scale_i0 + 1) * 16) ^ row_swizzle)) = + make_uint4(q2.x, q2.y, q3.x, q3.y); +} + +__device__ __forceinline__ void dequant_braided_quad_ilp(uint8_t* __restrict__ fp8_dst, + const uint4& q, const uint2& lut0, + const uint2& lut1, const int scale_i0, + const uint32_t row_swizzle) { + // Expose all four independent PRMT chains together so ptxas can overlap + // their selector/sign work before either 128-bit shared-memory store. + const uint32_t q0_sel0 = q.x & 0x00007777u; + const uint32_t q0_sel1 = (q.x >> 16) & 0x00007777u; + const uint32_t q1_sel0 = q.y & 0x00007777u; + const uint32_t q1_sel1 = (q.y >> 16) & 0x00007777u; + const uint32_t q2_sel0 = q.z & 0x00007777u; + const uint32_t q2_sel1 = (q.z >> 16) & 0x00007777u; + const uint32_t q3_sel0 = q.w & 0x00007777u; + const uint32_t q3_sel1 = (q.w >> 16) & 0x00007777u; + + uint32_t q0_out0 = byte_perm_unchecked(lut0.x, lut0.y, q0_sel0); + uint32_t q0_out1 = byte_perm_unchecked(lut0.x, lut0.y, q0_sel1); + uint32_t q1_out0 = byte_perm_unchecked(lut0.x, lut0.y, q1_sel0); + uint32_t q1_out1 = byte_perm_unchecked(lut0.x, lut0.y, q1_sel1); + uint32_t q2_out0 = byte_perm_unchecked(lut1.x, lut1.y, q2_sel0); + uint32_t q2_out1 = byte_perm_unchecked(lut1.x, lut1.y, q2_sel1); + uint32_t q3_out0 = byte_perm_unchecked(lut1.x, lut1.y, q3_sel0); + uint32_t q3_out1 = byte_perm_unchecked(lut1.x, lut1.y, q3_sel1); + + q0_out0 |= q.x & 0x80808080u; + q0_out1 |= (q.x << 4) & 0x80808080u; + q1_out0 |= q.y & 0x80808080u; + q1_out1 |= (q.y << 4) & 0x80808080u; + q2_out0 |= q.z & 0x80808080u; + q2_out1 |= (q.z << 4) & 0x80808080u; + q3_out0 |= q.w & 0x80808080u; + q3_out1 |= (q.w << 4) & 0x80808080u; + + *reinterpret_cast(fp8_dst + ((scale_i0 * 16) ^ row_swizzle)) = + make_uint4(q0_out0, q0_out1, q1_out0, q1_out1); + *reinterpret_cast(fp8_dst + (((scale_i0 + 1) * 16) ^ row_swizzle)) = + make_uint4(q2_out0, q2_out1, q3_out0, q3_out1); +} + +template +__device__ __forceinline__ void dequant_braided_quad_lut_window( + uint8_t* __restrict__ fp8_dst, const uint4 (&fp4_quads)[4], const uint32_t scale_word_lo, + const uint32_t scale_word_hi, const uint2* __restrict__ lut_smem, const uint2 lut0, + const uint2 lut1, const uint32_t row_swizzle) { + uint2 next_lut0; + uint2 next_lut1; + if constexpr (kQuad + 1 < 4) { + constexpr int kNextScaleI0 = (kQuad + 1) * 2; + constexpr int kNextScaleI1 = kNextScaleI0 + 1; + const uint32_t next_scale_word = kQuad + 1 < 2 ? scale_word_lo : scale_word_hi; + const uint32_t next_scale0 = (next_scale_word >> ((kNextScaleI0 & 3) * 8)) & 0x7fu; + const uint32_t next_scale1 = (next_scale_word >> ((kNextScaleI1 & 3) * 8)) & 0x7fu; + next_lut0 = lut_smem[next_scale0]; + next_lut1 = lut_smem[next_scale1]; + } + + if constexpr (kQuadIlp) { + dequant_braided_quad_ilp(fp8_dst, fp4_quads[kQuad], lut0, lut1, kQuad * 2, row_swizzle); + } else { + dequant_braided_quad(fp8_dst, fp4_quads[kQuad], lut0, lut1, kQuad * 2, row_swizzle); + } + + if constexpr (kQuad + 1 < 4) { + dequant_braided_quad_lut_window(fp8_dst, fp4_quads, scale_word_lo, + scale_word_hi, lut_smem, next_lut0, + next_lut1, row_swizzle); + } +} + +template +__device__ __forceinline__ void dequant_smem_b_from_packed_braided_lut_window( + uint8_t* __restrict__ smem_b, const uint8_t* __restrict__ packed_b, const uint32_t row, + const uint2* __restrict__ lut_smem) { + const uint8_t* __restrict__ row_ptr = packed_b + row * 80; + const uint4* __restrict__ fp4_src = reinterpret_cast(row_ptr); + uint4 fp4_quads[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) fp4_quads[i] = fp4_src[i]; + + const uint2 scale_words = *reinterpret_cast(row_ptr + 64); + const uint2 lut0 = lut_smem[scale_words.x & 0x7fu]; + const uint2 lut1 = lut_smem[(scale_words.x >> 8) & 0x7fu]; + dequant_braided_quad_lut_window<0, kQuadIlp>(smem_b + row * 128, fp4_quads, scale_words.x, + scale_words.y, lut_smem, lut0, lut1, + (row & 7u) << 4); +} + +} // namespace nvfp4 + +// Push dispatch (see the kernel body, `kPushDispatch`): the routed rows of this +// rank are written straight into the destination ranks' fixed-stride pools. +// Force-inlined: an ABI call (`__noinline__`) anywhere in the persistent kernel +// changes the register allocation of the whole kernel and costs ~15 % on the +// math warps (measured on H20). +template +__device__ __forceinline__ void sm90_nvfp4_push_dispatch_rows( + const layout::SymBuffer& sym_buffer, void* smem_row_buffer, + cutlass::arch::ClusterTransactionBarrier* row_mbarrier, + const int64_t* __restrict__ input_topk_idx, const uint8_t* __restrict__ input_tokens, + const float* __restrict__ input_sf, const float* __restrict__ input_topk_weights, + uint8_t* l1_tokens, float* l1_sf, float* l1_topk_weights, + layout::TokenSrcMetadata* token_src_metadata, uint64_t* recv_count_sum, + const uint32_t pool_block_base, const uint32_t num_tokens, const uint32_t global_warp_idx, + const uint32_t lane_idx) { + constexpr uint32_t kNumSFFloats = kHidden / 128; + constexpr uint32_t kNumTokenChunksPerLane = kHidden / (16 * 32); + constexpr uint32_t kChunkGroup = 2; + static_assert(kHidden % 512 == 0 && kNumTokenChunksPerLane % kChunkGroup == 0, + "Invalid token shape for push dispatch"); + // Rows (token, top-k slot) in contiguous chunks of R = ceil(rows / warps) per + // warp (<= 32 per ticket batch): every lane takes the remote ticket of one row + // of the batch at once (the NVLink round trips overlap), then the warp streams + // the batch's rows (16 B per lane per store) into the destination pools. + const uint32_t num_rows = num_tokens * kNumTopk; + const uint32_t rows_per_warp = (num_rows + kNumGlobalWarps - 1) / kNumGlobalWarps; + const uint32_t row_begin = global_warp_idx * rows_per_warp; + const uint32_t row_end = min(row_begin + rows_per_warp, num_rows); + uint32_t row_mbarrier_phase = 0; +#pragma unroll 1 + for (uint32_t batch_begin = row_begin; batch_begin < row_end; batch_begin += 32) { + const uint32_t batch_size = min(row_end - batch_begin, 32u); + const uint32_t r = batch_begin + lane_idx; + int expert_idx = -1; + if (lane_idx < batch_size) expert_idx = static_cast(__ldg(input_topk_idx + r)); + uint32_t lane_row_idx = 0; + if (expert_idx >= 0) { + const uint32_t dr = static_cast(expert_idx) / kNumExpertsPerRank; + const uint32_t de = static_cast(expert_idx) % kNumExpertsPerRank; + lane_row_idx = + static_cast(ptx::atomic_add_sys(sym_buffer.map(recv_count_sum + de, dr), 1ull)); + } + const uint32_t valid_mask = __ballot_sync(0xffffffff, expert_idx >= 0); +#pragma unroll 1 + for (uint32_t mask = valid_mask; mask != 0; mask &= mask - 1) { + const uint32_t src_lane = __ffs(mask) - 1; + const uint32_t row_r = batch_begin + src_lane; + const uint32_t e = static_cast(__shfl_sync(0xffffffff, expert_idx, src_lane)); + const uint32_t row_idx = __shfl_sync(0xffffffff, lane_row_idx, src_lane); + const uint32_t dr = e / kNumExpertsPerRank; + const uint32_t de = e % kNumExpertsPerRank; + const uint32_t src_token_idx = row_r / kNumTopk, src_topk_idx = row_r % kNumTopk; + FI_MEGAMOE_TRAP_ONLY_DEVICE_ASSERT(row_idx < kPushBlocksPerExpert * BLOCK_M); + const uint32_t pool_token_idx = + (pool_block_base + de * kPushBlocksPerExpert) * BLOCK_M + row_idx; + const auto* src_token = reinterpret_cast( + input_tokens + static_cast(src_token_idx) * kHidden); + auto* dst_token = sym_buffer.map( + reinterpret_cast(l1_tokens + static_cast(pool_token_idx) * kHidden), + dr); + if constexpr (kGenericRows) { +#pragma unroll + for (uint32_t g = 0; g < kNumTokenChunksPerLane; g += kChunkGroup) { + uint4 row[kChunkGroup]; +#pragma unroll + for (uint32_t c = 0; c < kChunkGroup; ++c) + row[c] = __ldg(src_token + (g + c) * 32 + lane_idx); +#pragma unroll + for (uint32_t c = 0; c < kChunkGroup; ++c) dst_token[(g + c) * 32 + lane_idx] = row[c]; + } + } else { + // Bulk copy (async proxy) local input row -> smem -> destination pool, + // the mirror image of the pull path's TMA pull. + if (cute::elect_one_sync()) { + ptx::tma_load_1d(smem_row_buffer, src_token, row_mbarrier, kHidden); + ptx::mbarrier_arrive_and_set_tx(row_mbarrier, kHidden); + ptx::mbarrier_wait_and_flip_phase(row_mbarrier, row_mbarrier_phase); + ptx::tma_store_1d(dst_token, smem_row_buffer, kHidden); + cute::tma_store_arrive(); + ptx::tma_store_wait<0>(); + } + __syncwarp(); + } + const float* src_sf = input_sf + static_cast(src_token_idx) * kNumSFFloats; + float* dst_sf = sym_buffer.map(l1_sf, dr); +#pragma unroll + for (uint32_t j = lane_idx; j < kNumSFFloats; j += 32) + dst_sf[static_cast(j) * kNumPaddedSFPoolTokens + pool_token_idx] = + __ldg(src_sf + j); + if (lane_idx == 0) { + *sym_buffer.map(l1_topk_weights + pool_token_idx, dr) = __ldg(input_topk_weights + row_r); + *sym_buffer.map(token_src_metadata + pool_token_idx, dr) = { + static_cast(sym_buffer.rank_idx), src_token_idx, src_topk_idx}; + } + } + __syncwarp(); + } +} + +template async proxy fence in the push A loader. + bool kStridedPoolDebug = false, bool kPushProxyFence = false, + // Push rows with 16 B generic stores instead of TMA bulk copies (diagnostic). + bool kPushGenericRows = false, + // Phase stamps (globaltimer) compiled in; host env FI_MEGAMOE_NVFP4_PHASE_STAMPS_PTR. + bool kPhaseStamps = false, uint32_t kNumRanks = 8> +CUTLASS_GLOBAL __launch_bounds__(384, 1) void sm90_nvfp4_mega_moe_fused_impl( + void* y, int* cumulative_local_expert_recv_stats, unsigned long long* phase_stamps, + const uint32_t num_tokens, const __grid_constant__ layout::SymBuffer sym_buffer, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, + const float* __restrict__ l1_global_scales, const float* __restrict__ l2_global_scales) { + constexpr uint32_t BLOCK_K = 128; + constexpr uint32_t kNumDispatchThreads = 64; + constexpr uint32_t kNumNonEpilogueThreads = 64; + constexpr uint32_t kNumEpilogueThreads = 256; + constexpr uint32_t kNumSMs = MEGAMOE_NUM_SMS; + static_assert(kNumRanks > 0 && kNumRanks <= 32, "Invalid EP width"); + static_assert(kNumExperts % kNumRanks == 0, "Experts must divide evenly across ranks"); + constexpr uint32_t L1_SHAPE_N = kIntermediateHidden * 2; + constexpr uint32_t L1_SHAPE_K = kHidden; + constexpr uint32_t L2_SHAPE_N = kHidden; + constexpr uint32_t L2_SHAPE_K = kIntermediateHidden; + constexpr uint32_t kNumDispatchWarps = kNumDispatchThreads / 32; + constexpr uint32_t kNumMMANonEpilogueWarps = kNumNonEpilogueThreads / 32; + constexpr uint32_t kNumEpilogueWarps = kNumEpilogueThreads / 32; + constexpr uint32_t kNumEpilogueWarpgroups = kNumEpilogueWarps / 4; + constexpr uint32_t kNumTokensPerWarp = 32 / kNumTopk; + constexpr uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks; +#include +} + +} // namespace flashinfer::megamoe + +#pragma clang diagnostic pop diff --git a/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_fused_body.inl b/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_fused_body.inl new file mode 100644 index 00000000000..1936744c914 --- /dev/null +++ b/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_fused_body.inl @@ -0,0 +1,1593 @@ +// Independent SM90 NVFP4 MegaMoE fused kernel body. +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900) and (__CUDA_ARCH__ < 1000)) or \ + defined(__CLION_IDE__) +using Barrier = cutlass::arch::ClusterTransactionBarrier; + +// ===================================================================== +// Template checks +// ===================================================================== +FI_MEGAMOE_STATIC_ASSERT(BLOCK_M == 8 || BLOCK_M == 16 || BLOCK_M == 24 || BLOCK_M == 64 || + BLOCK_M == 128, + "SM90 fused kernel requires BM8/BM16/BM24/BM64/BM128"); +FI_MEGAMOE_STATIC_ASSERT( + (BLOCK_M == 8 && (kNumStages == 3 || kNumStages == 4 || + (kNumSMs == 78 && kRSSwapABRequested && kNumStages == 6))) || + ((BLOCK_M == 16 || BLOCK_M == 24) && + (kNumStages == 3 || (kNumSMs == 78 && kRSSwapABRequested && kNumStages == 6))) || + (BLOCK_M == 64 && kNumStages == 3) || (BLOCK_M == 128 && kNumStages == 6), + "Unexpected SM90 pipeline depth"); +FI_MEGAMOE_STATIC_ASSERT((BLOCK_M == 128) == (BLOCK_N == 128), + "BM128 is paired with the BN128 split-M topology"); +FI_MEGAMOE_STATIC_ASSERT(!kSwapABRequested || BLOCK_M <= 24, + "swap-AB is only selected through the M64 bucket"); +FI_MEGAMOE_STATIC_ASSERT(!kRSSwapABRequested || kSwapABRequested, + "RS requires the existing transposed swap-AB epilogue"); +FI_MEGAMOE_STATIC_ASSERT(!kRSSwapABRequested || kUseMode2RowDecoder, + "RS requires the Mode2 braided register decoder ABI"); +FI_MEGAMOE_STATIC_ASSERT(!kRSSwapABRequested || BLOCK_M <= 24, + "minimum RS integration excludes BM64/BM128"); + +// ===================================================================== +// Thread / warp identification +// ===================================================================== +const uint32_t sm_idx = blockIdx.x; +const uint32_t thread_idx = threadIdx.x; +const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); +const uint32_t lane_idx = ptx::get_lane_idx(); + +// Optional phase timestamps (globaltimer, ns; `phase_stamps` != nullptr). Slots: +// 0 min kernel entry | 1 max dispatch barrier #1 / DONE wait done | 2 max pool ready +// 3 min first math task | 4 max last L1 task end | 5 max last L2 task end +// 6 max combine barrier #2 done (fine combine: first token ready) | 7 max combine end +// 8 max routing count done | 9 max routing writes / pushes issued | 10 max routing +// grid sync / DONE signalled | 11 max count broadcast done | 12 max barrier #3 done +const auto stamp_min = [&](const uint32_t slot) { + if constexpr (kPhaseStamps) { + unsigned long long t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + atomicMin(phase_stamps + slot, t); + } +}; +const auto stamp_max = [&](const uint32_t slot) { + if constexpr (kPhaseStamps) { + unsigned long long t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + atomicMax(phase_stamps + slot, t); + } +}; +// Accumulators (all SMs): 20/21 sum of L1 task ns / count, 22/23 L2 task ns / count, +// 24 sum of A-loader arrival-spin ns (L1), 25 sum of A-loader L2-mask-spin ns. +const auto stamp_now = [&]() { + unsigned long long t = 0; + if constexpr (kPhaseStamps) asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +}; +const auto stamp_accumulate = [&](const uint32_t slot, const unsigned long long& delta, + const unsigned long long& count) { + if constexpr (kPhaseStamps) { + atomicAdd(phase_stamps + slot, delta); + if (count) atomicAdd(phase_stamps + slot + 1, count); + } +}; + +if (warp_idx == 0 and cute::elect_one_sync()) { + stamp_min(0); + cute::prefetch_tma_descriptor(&tensor_map_l1_acts); + cute::prefetch_tma_descriptor(&tensor_map_l1_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l1_weights); + cute::prefetch_tma_descriptor(&tensor_map_l1_output); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l2_weights); +} + +// ===================================================================== +// Workspaces and symmetric buffer slicing. The framework reserves +// per-64 SF capacity; this per-128 path uses its first half +// as a dense per-128 layout so no framework allocation change is needed. +// ===================================================================== +#include + +// ===================================================================== +// GEMM data types and shape constants +// ===================================================================== +using a_dtype_t = cutlass::float_e4m3_t; +using b_dtype_t = cutlass::float_e4m3_t; +using task_info_t = sched::TaskInfo; +using interleaved_scheduler_t = + sched::InterleavedMegaMoEScheduler; +constexpr uint32_t kNumRoutedL1BlockNs = L1_SHAPE_N / BLOCK_N; +constexpr bool kSplitMDecodedWeightReuse = + BLOCK_M == 128 && BLOCK_N == 128 && kNumEpilogueWarpgroups == 2; +constexpr uint32_t WG_BLOCK_M = kSplitMDecodedWeightReuse ? BLOCK_M / 2 : BLOCK_M; +constexpr uint32_t WG_BLOCK_N = kSplitMDecodedWeightReuse ? BLOCK_N : BLOCK_N / 2; +constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2; // post-SwiGLU tile N +constexpr uint32_t WG_L1_OUT_BLOCK_N = WG_BLOCK_N / 2; // post-SwiGLU per-WG N +constexpr uint32_t kSwapABTokenChunks = BLOCK_M / 8; +constexpr uint32_t kSwapABWeightHalves = WG_BLOCK_N / 64; +constexpr uint32_t kSwapABHalfAccumPerThread = 64 * 64 / 128; +FI_MEGAMOE_STATIC_ASSERT(!kSwapABRequested || WG_L1_OUT_BLOCK_N == 64, + "swapAB expects BN256 split-N with 64 L1 output columns per WG"); +// Both dispatch warps participate in CTA-wide barriers. Selected plans may +// use one warp for routing and token pulls, leaving the other warp's send +// buffer available for an additional GEMM stage. +constexpr uint32_t kNumActiveDispatchWarps = kSingleActiveDispatchWarp ? 1u : kNumDispatchWarps; +constexpr uint32_t kNumActiveDispatchThreads = kNumActiveDispatchWarps * 32; +constexpr bool kQuadDequantIlp = BLOCK_M == 8 && kNumStages == 4; +using L1WGMMA = typename mma::sm90::FP8MMASelector::type; +static_assert(L1WGMMA::M == 64 and L1WGMMA::N == WG_BLOCK_N and L1WGMMA::K == 32, + "Unexpected WGMMA shape"); +// A and B are CTA-local in the fixed cluster-size-one plan. +constexpr uint32_t LOAD_BLOCK_M = BLOCK_M; +constexpr uint32_t LOAD_BLOCK_N = BLOCK_N; +constexpr uint32_t kSwizzleAMode = BLOCK_K * sizeof(a_dtype_t); // 128 +constexpr uint32_t kL2ActsSFGranK = kSplitMDecodedWeightReuse ? 64u : 128u; +FI_MEGAMOE_STATIC_ASSERT(kSplitMDecodedWeightReuse || WG_L1_OUT_BLOCK_N < kL2ActsSFGranK, + "split-N warpgroups must share one L2 activation scale"); + +// ===================================================================== +// Shared memory layout +// ===================================================================== +constexpr uint32_t kSharedMemoryAlignment = 1024; +extern __shared__ __align__(kSharedMemoryAlignment) +uint8_t smem_buffer[]; + +constexpr uint32_t SMEM_EXPERT_COUNT_SIZE = + math::constexpr_align(kNumExperts * sizeof(uint32_t), kSharedMemoryAlignment); +constexpr uint32_t SMEM_SEND_BUFFER_SIZE = math::constexpr_align( + fp8_token_layout.get_num_bytes() * kNumActiveDispatchWarps, kSharedMemoryAlignment); +constexpr uint32_t SMEM_NVFP4_LUT_SIZE = + math::constexpr_align(128u * sizeof(uint2), kSharedMemoryAlignment); +constexpr uint32_t SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K * sizeof(a_dtype_t); +constexpr uint32_t SMEM_B_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K * sizeof(b_dtype_t); +// BM128 split-M alternates two decoded-B slots. This lets one WG begin +// decoding K+1 after its K WGMMA completes without overwriting the slot +// that the paired WG may still be consuming. +constexpr bool kH20CompactRS = kNumSMs == 78 && kRSSwapABRequested; +constexpr uint32_t kNumDecodedBStages = + kH20CompactRS ? 0u : (kSplitMDecodedWeightReuse ? 2u : kNumStages); +constexpr uint32_t B_LOAD_BYTES_PER_ROW = 80u; +constexpr uint32_t SMEM_PACKED_B_SIZE_PER_STAGE = + LOAD_BLOCK_N * B_LOAD_BYTES_PER_ROW * sizeof(b_dtype_t); +// L1 and L2 each consume one per-128 activation scale per row and K tile. +constexpr uint32_t kL2SFAHalfStride = + math::constexpr_align(BLOCK_M * sizeof(float), 128u) / sizeof(float); +constexpr uint32_t kNumL2SFAGroups = kSplitMDecodedWeightReuse ? 2u : 1u; +constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = kNumL2SFAGroups * kL2SFAHalfStride * sizeof(float); +// CD output: max of L1 FP8 (BLOCK_M * (BLOCK_N/2) * 1 byte * num_wg) and +// L2 BF16 (BLOCK_M * BLOCK_N * 2 bytes * num_wg). +constexpr uint32_t SMEM_CD_L1_SIZE = + kNumEpilogueWarpgroups * WG_BLOCK_M * WG_L1_OUT_BLOCK_N * sizeof(cutlass::float_e4m3_t); +constexpr uint32_t SMEM_CD_L2_SIZE = + kSwapABRequested ? BLOCK_M * BLOCK_N * sizeof(nv_bfloat16) : 0u; +constexpr uint32_t SMEM_CD_OUTPUT_BASE_SIZE = + SMEM_CD_L1_SIZE > SMEM_CD_L2_SIZE ? SMEM_CD_L1_SIZE : SMEM_CD_L2_SIZE; +constexpr uint32_t SMEM_CD_L1_SHARED_SF_SLOTS = kNumEpilogueWarpgroups * BLOCK_M; +constexpr uint32_t SMEM_CD_L1_SWAP_AMAX_SLOTS = kSwapABRequested ? BLOCK_M * kNumEpilogueWarps : 0u; +constexpr uint32_t SMEM_CD_L1_EXTRA_FLOAT_SLOTS = + SMEM_CD_L1_SHARED_SF_SLOTS > SMEM_CD_L1_SWAP_AMAX_SLOTS ? SMEM_CD_L1_SHARED_SF_SLOTS + : SMEM_CD_L1_SWAP_AMAX_SLOTS; +constexpr uint32_t SMEM_CD_L1_SHARED_SF_SIZE = SMEM_CD_L1_EXTRA_FLOAT_SLOTS * sizeof(float); +constexpr uint32_t SMEM_CD_OUTPUT_UNALIGNED_SIZE = + SMEM_CD_OUTPUT_BASE_SIZE + SMEM_CD_L1_SHARED_SF_SIZE; +constexpr uint32_t SMEM_CD_SIZE = + math::constexpr_align(SMEM_CD_OUTPUT_UNALIGNED_SIZE, kSharedMemoryAlignment); + +constexpr uint32_t SMEM_GEMM_BEFORE_SFA_SIZE = + SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_NVFP4_LUT_SIZE + SMEM_CD_SIZE + + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_PACKED_B_SIZE_PER_STAGE) + + kNumDecodedBStages * SMEM_B_SIZE_PER_STAGE; +constexpr uint32_t SMEM_GEMM_END_SIZE = + SMEM_GEMM_BEFORE_SFA_SIZE + kNumStages * SMEM_SFA_SIZE_PER_STAGE; +constexpr uint32_t kCombineNumChunks = kHidden <= 4096 ? 1u : 2u; +constexpr uint32_t SMEM_COMBINE_MIN_SIZE = + 3u * kNumEpilogueWarps * kHidden * sizeof(nv_bfloat16) / kCombineNumChunks; +// Compact RS removes the unused decoded-B ring. Preserve the later +// combine scratch lifetime by padding only the pre-barrier region, exactly +// as the H20 Weave layout does; no extra decoded storage is reintroduced. +constexpr uint32_t SMEM_BEFORE_BARRIER_SIZE = + SMEM_GEMM_END_SIZE > SMEM_COMBINE_MIN_SIZE ? SMEM_GEMM_END_SIZE : SMEM_COMBINE_MIN_SIZE; + +// SMEM pointers +auto smem_expert_count = reinterpret_cast(smem_buffer); +const auto smem_send_buffers = + layout::Buffer(fp8_token_layout, kNumActiveDispatchWarps, 1, + math::advance_ptr(smem_buffer, SMEM_EXPERT_COUNT_SIZE)); +auto smem_nvfp4_lut = reinterpret_cast( + math::advance_ptr(smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE)); + +auto smem_gemm_base = math::advance_ptr( + smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_NVFP4_LUT_SIZE); + +auto smem_cd_base = smem_gemm_base; +// CD output is shared by L1 (FP8) and L2 (BF16); reinterpret-cast as needed. +auto smem_cd_l1 = reinterpret_cast(smem_cd_base); +auto smem_cd_l1_shared_sf = math::advance_ptr(smem_cd_base, SMEM_CD_OUTPUT_BASE_SIZE); +auto smem_cd_l2 = reinterpret_cast(smem_cd_base); + +auto smem_a = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE); +}); +auto smem_b = utils::PatternVisitor([=](const uint32_t& i) { + const uint32_t decoded_stage = kSplitMDecodedWeightReuse ? (i & 1u) : i; + return math::advance_ptr( + smem_gemm_base, + SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + decoded_stage * SMEM_B_SIZE_PER_STAGE); +}); +auto smem_packed_b = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, + SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + + kNumDecodedBStages * SMEM_B_SIZE_PER_STAGE + + i * SMEM_PACKED_B_SIZE_PER_STAGE); +}); +auto sf_start_ptr = math::advance_ptr( + smem_gemm_base, SMEM_CD_SIZE + + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_PACKED_B_SIZE_PER_STAGE) + + kNumDecodedBStages * SMEM_B_SIZE_PER_STAGE); +auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast(sf_start_ptr + i * SMEM_SFA_SIZE_PER_STAGE); +}); +// Barriers live after SF. +auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_BEFORE_BARRIER_SIZE); +auto dispatch_barriers = + utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + i; }); +auto full_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + i; }); +auto empty_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + kNumStages + i; }); +auto combine_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + kNumStages * 2 + i; }); +constexpr uint32_t kNumBaseBarriers = kNumDispatchWarps + kNumStages * 2 + kNumEpilogueWarps * 2; +auto task_info_full_barriers = barrier_start_ptr + kNumBaseBarriers; +auto task_info_empty_barriers = + task_info_full_barriers + interleaved_scheduler_t::kNumScheduleStages; +auto task_infos = reinterpret_cast(task_info_empty_barriers + + interleaved_scheduler_t::kNumScheduleStages); +constexpr uint32_t kInterleavedSchedulerSMEMBytes = + 2 * interleaved_scheduler_t::kNumScheduleStages * sizeof(Barrier) + + interleaved_scheduler_t::kNumScheduleStages * sizeof(task_info_t); +FI_MEGAMOE_STATIC_ASSERT(kInterleavedSchedulerSMEMBytes == + layout::kSM90InterleavedSchedulerSMEMBytes, + "Host and device scheduler shared-memory layouts disagree"); +constexpr uint32_t kInterleavedSMEMEnd = + SMEM_BEFORE_BARRIER_SIZE + kNumBaseBarriers * sizeof(Barrier) + kInterleavedSchedulerSMEMBytes; +FI_MEGAMOE_STATIC_ASSERT(!kUseInterleavedScheduler || kInterleavedSMEMEnd <= 232448, + "Interleaved scheduler exceeds the SM90 shared-memory capacity"); + +// ===================================================================== +// Initialization +// ===================================================================== +if (thread_idx < 64) { + reinterpret_cast(smem_nvfp4_lut)[thread_idx] = + reinterpret_cast(flashinfer::megamoe::nvfp4::kE2M1AndUe4m3ToFp8Lut)[thread_idx]; +} + +if (warp_idx == 0) { +// Clean expert-count shared memory +#pragma unroll + for (uint32_t i = lane_idx; i < kNumExperts; i += 32) ptx::st_shared(smem_expert_count + i, 0u); +} else if (warp_idx == 1) { +// Init dispatch m-barriers +#pragma unroll + for (uint32_t i = lane_idx; i < kNumDispatchWarps; i += 32) dispatch_barriers[i]->init(1); + cutlass::arch::fence_barrier_init(); +} else if (warp_idx == 2) { + // Init GEMM full/empty barriers and combine barriers + if (cute::elect_one_sync()) { +#pragma unroll + for (uint32_t i = 0; i < kNumStages; ++i) { + // Producer arrivals: A(+SFA) + B(TMA+SFB). SFB is copied with + // cp.async.bulk and counted as B-loader transaction bytes, so + // it does not need a separate producer arrival. + full_barriers[i]->init(2); + empty_barriers[i]->init(kNumEpilogueWarps); + } +#pragma unroll + for (uint32_t i = 0; i < kNumEpilogueWarps * 2; ++i) combine_barriers[i]->init(1); + if constexpr (kUseInterleavedScheduler) { +#pragma unroll + for (uint32_t i = 0; i < interleaved_scheduler_t::kNumScheduleStages; ++i) { + task_info_full_barriers[i].init(1); + task_info_empty_barriers[i].init(kNumEpilogueWarps); + } + } + } + cutlass::arch::fence_barrier_init(); +} +__syncthreads(); + +// ===================================================================== +// Scheduler (cluster=1) +// ===================================================================== +auto scheduler = sched::MegaMoEScheduler(workspace); +auto interleaved_scheduler = interleaved_scheduler_t(workspace, task_info_full_barriers, + task_info_empty_barriers, task_infos); + +// Pipeline state shared by TMA loaders and math warpgroups +uint32_t stage_idx = 0, phase = 0; +auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++k_block_idx; + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; +}; +// Intra-SM barrier indices (mirroring SM100) +constexpr uint32_t kDispatchBarrierIdx = 0; +constexpr uint32_t kDispatchWithEpilogueBarrierIdx = 1; +constexpr uint32_t kEpilogueFullBarrierIdx = 2; +constexpr uint32_t kEpilogueWGBarrierStartIdx = 3; +constexpr uint32_t kSplitMDecodeBarrierIdx = 8; + +// Cross-rank NVLink barrier tags +constexpr uint32_t kBeforeDispatchPullBarrierTag = 1; +constexpr uint32_t kBeforeCombineReduceBarrierTag = 2; +constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; + +// --------------------------------------------------------------------- +// Counter-based synchronisation knobs (host env FI_MEGAMOE_NVFP4_PUSH_DISPATCH / +// FI_MEGAMOE_NVFP4_FINE_COMBINE / FI_MEGAMOE_NVFP4_NO_CLEAN_BARRIER, see the host). +// +// Push dispatch (kPushDispatch): during routing the SOURCE rank takes a +// remote atomic ticket on the destination's per-expert recv-count word (low +// 32 bits grow by 1 per row) and writes the FP8 row + per-K128 SF + top-k +// weight + source metadata straight into the destination's pool with plain +// 16 B stores. The pool is addressed with a FIXED stride of +// kPushBlocksPerExpert blocks per local expert (rows inside an expert stay +// packed); the dynamic scheduler keeps its dense task indices and only the +// task's `pool_block_idx` is remapped (see `invoke_interleaved_task`). +// NVLink barrier #1 is replaced by per-rank DONE counts: after a CTA's +// dispatch warps issued their last pushed row, warp 0 takes a CTA arrival +// ticket (atom.acq_rel.gpu, after the CTA barrier) and the last CTA +// red.release.sys-adds 1 into every rank's DONE count. A launch's DONE target +// is kNumRanks * (epoch + 1), `epoch` = push launches completed on this rank +// (bumped by SM0 in the workspace cleanup). Waiters: every CTA's task producer +// (the B-loader warp) and SM e's dispatch warp 0, which then publishes expert +// e's L1 arrival counts (release.gpu) exactly like the pull path did. Lean +// routing is implied: no send counts, no cross-rank count broadcast, no +// top-k index writes. Pool reuse across launches rests on barrier #3 (a rank +// can only push launch N+1 rows after barrier #3 of launch N) or, with +// kNoCleanBarrier, on the parity-rotated pool (below). +// +// Rotated pool (kNoCleanBarrier, requires kPushDispatch): the strided push +// pool and the recv-count sums are double-buffered by launch parity, so no +// rank can push launch N+1 rows into a slot another rank still reads for +// launch N; barrier #3 (after the workspace cleanup) is dropped and the +// cleanup is ordered only by the local dispatch-warp grid sync. A rank pushes +// N+2 rows (slot of N) only after every rank signalled DONE for N+1, i.e. +// after every rank completed N (same stream) including its cleanup of N. +// --------------------------------------------------------------------- +constexpr bool kPushDispatch = kPushDispatchRequested && kUseInterleavedScheduler; +constexpr bool kNoCleanBarrier = kNoCleanBarrierRequested && kPushDispatch; +// Fixed-stride pool: implied by push dispatch; kStridedPoolDebug forces it +// under the pull protocol (diagnostic). +constexpr bool kStridedPool = kUseInterleavedScheduler && (kPushDispatch || kStridedPoolDebug); +constexpr uint32_t kNumPoolParitySlots = kNoCleanBarrier ? 2u : 1u; +constexpr uint32_t kNumPoolBlocksTotal = kNumMaxPoolTokens / BLOCK_M; +constexpr uint32_t kPushBlocksPerExpert = + kStridedPool ? kNumPoolBlocksTotal / (kNumExpertsPerRank * kNumPoolParitySlots) : 0u; +constexpr uint32_t kPushMaxRowsPerExpert = kPushBlocksPerExpert * BLOCK_M; +FI_MEGAMOE_STATIC_ASSERT(!kStridedPool || kPushBlocksPerExpert > 0, + "Push dispatch: empty pool stride"); +FI_MEGAMOE_STATIC_ASSERT(!kStridedPool || kNumPoolParitySlots * kNumExpertsPerRank * + kPushBlocksPerExpert * BLOCK_M <= + kNumMaxPoolTokens, + "Push dispatch: strided pool must fit the token pool"); +// Launch epoch (push launches completed on this rank): every role that needs it +// reads it at its own start, before any barrier of the launch, so SM0's bump in +// the cleanup (after every reader) cannot be observed by this launch. Selects +// the DONE target and the pool parity slot. Read inside the roles (not before +// the role split) so that no extra value is live across the split: ptxas +// re-schedules the math warps' code when the split-point pressure changes. +const auto read_launch_epoch = [&]() { + return kPushDispatch ? ptx::ld_volatile(workspace.get_push_epoch_ptr()) : 0u; +}; +const auto epoch_pool_parity = [&](const uint32_t& epoch) { + return kNoCleanBarrier ? (epoch & 1u) : 0u; +}; +const auto epoch_done_target = [&](const uint32_t& epoch) { + return static_cast(kNumRanks * (epoch + 1u)); +}; +const auto strided_pool_block_p = [&](const uint32_t& parity, const uint32_t& local_expert_idx, + const uint32_t& m_block_idx) { + return (parity * kNumExpertsPerRank + local_expert_idx) * kPushBlocksPerExpert + m_block_idx; +}; +const auto recv_count_sum_ptr_p = [&](const uint32_t& parity, const uint32_t& local_expert_idx) { + return kNoCleanBarrier ? workspace.get_expert_recv_count_sum_ptr(local_expert_idx, parity) + : workspace.get_expert_recv_count_sum_ptr(local_expert_idx); +}; + +// Register reconfiguration counts (chosen to fit in 64512 reg budget). +constexpr uint32_t kNumDispatchRegisters = 48; +constexpr uint32_t kNumNonEpilogueRegisters = kUseInterleavedScheduler ? 64 : 40; +constexpr uint32_t kNumEpilogueRegisters = 208; +FI_MEGAMOE_STATIC_ASSERT(kNumDispatchRegisters* kNumDispatchThreads + + kNumNonEpilogueRegisters * kNumNonEpilogueThreads + + kNumEpilogueRegisters * kNumEpilogueThreads <= + 64512, + "Too many registers"); + +constexpr uint32_t kDispatchGridSyncIndex = 0; +constexpr uint32_t kEpilogueGridSyncIndex = 1; + +const auto for_each_static_selected_block = [&](auto&& func) { + scheduler.fetch_expert_recv_count(); + scheduler.set_expert_idx(0); + while (true) { + CUTE_TIE_DECL(scheduler.get_next_block(), block_phase, local_expert_idx, m_block_idx, + n_block_idx); + if (block_phase == sched::BlockPhase::None) break; + if (block_phase == sched::BlockPhase::Linear1) { + func(std::integral_constant{}, + local_expert_idx, L1_SHAPE_K / BLOCK_K, m_block_idx, n_block_idx, + scheduler.get_current_pool_block_offset() + m_block_idx, + scheduler.template get_valid_m()); + } else { + func(std::integral_constant{}, + local_expert_idx, L2_SHAPE_K / BLOCK_K, m_block_idx, n_block_idx, + scheduler.get_current_pool_block_offset() + m_block_idx, + scheduler.template get_valid_m()); + } + } +}; + +// Pool parity of this launch for the task consumers (constant 0 unless the +// rotated pool is on; then one volatile load per role at first use). +uint32_t task_pool_parity = 0u; +bool task_pool_parity_read = !kNoCleanBarrier; +const auto invoke_interleaved_task = [&](const task_info_t& task_info, auto&& func) { + if (!task_pool_parity_read) { + task_pool_parity = epoch_pool_parity(read_launch_epoch()); + task_pool_parity_read = true; + } + // Push dispatch: the dense task index stays the scheduling key; the + // physical pool block is the fixed-stride one the source ranks wrote to. + const uint32_t pool_block_idx = + kStridedPool ? strided_pool_block_p(task_pool_parity, task_info.local_expert_idx, + task_info.m_block_idx) + : task_info.pool_block_idx; + if (task_info.block_phase == sched::BlockPhase::Linear1) { + func(std::integral_constant{}, + task_info.local_expert_idx, L1_SHAPE_K / BLOCK_K, task_info.m_block_idx, + task_info.n_block_idx, pool_block_idx, task_info.valid_m); + } else { + func(std::integral_constant{}, + task_info.local_expert_idx, L2_SHAPE_K / BLOCK_K, task_info.m_block_idx, + task_info.n_block_idx, pool_block_idx, task_info.valid_m); + } +}; + +const auto for_each_published_block = [&](auto&& func) { + task_info_t task_info; + while (interleaved_scheduler.get_published_task(task_info)) + invoke_interleaved_task(task_info, func); +}; + +const auto produce_interleaved_blocks = [&](auto&& func) { + if constexpr (kPushDispatch) { + const uint32_t epoch = read_launch_epoch(); + task_pool_parity = epoch_pool_parity(epoch); + task_pool_parity_read = true; + interleaved_scheduler.fetch_expert_recv_count(workspace.get_push_done_count_ptr(), + epoch_done_target(epoch), task_pool_parity); + } else { + interleaved_scheduler.fetch_expert_recv_count(); + } + while (true) { + interleaved_scheduler.wait_task_slot_empty(); + const auto task_info = interleaved_scheduler.claim_next_task(); + interleaved_scheduler.publish_task(task_info); + if (!task_info.is_valid()) break; + invoke_interleaved_task(task_info, func); + } +}; + +// Dispatch-role launch epoch / parity (read at the dispatch role's start) +uint32_t dispatch_epoch = 0u, dispatch_parity = 0u; +const auto cleanup_workspace = [&]() { + FI_MEGAMOE_STATIC_ASSERT(kNumSMs > 1, "Invalid SM count"); + if (sm_idx == 0) { + if constexpr (!kPushDispatch) { +#pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) + *workspace.get_expert_send_count_ptr(i) = 0; + } + if constexpr (kUseInterleavedScheduler) { + if (thread_idx == 0) { + *workspace.get_l1_task_count_ptr() = 0; + *workspace.get_l2_task_count_ptr() = 0; + } + } + if constexpr (kPushDispatch) { + // Next launch's DONE target / pool parity (see `kPushDispatch`) + if (thread_idx == 0) *workspace.get_push_epoch_ptr() = dispatch_epoch + 1u; + } + } else { + for (uint32_t i = sm_idx - 1; i < kNumExpertsPerRank; i += kNumSMs - 1) { + const auto num_recv_tokens = static_cast(*recv_count_sum_ptr_p(dispatch_parity, i)); + const auto num_recv_m_blocks = math::ceil_div(num_recv_tokens, BLOCK_M); + const auto cleanup_pool_block_offset = kStridedPool + ? strided_pool_block_p(dispatch_parity, i, 0) + : scheduler.get_pool_block_offset(i); + + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + FI_MEGAMOE_STATIC_ASSERT(kNumDispatchWarps >= 2, "Not enough dispatch warps"); + if (warp_idx == 0) { + *recv_count_sum_ptr_p(dispatch_parity, i) = 0; + } else if (warp_idx == 1) { + if (cute::elect_one_sync() and cumulative_local_expert_recv_stats != nullptr) + ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(num_recv_tokens)); + __syncwarp(); + } + + if constexpr (!kPushDispatch) { + for (uint32_t j = thread_idx; j < kNumRanks; j += kNumDispatchThreads) + *workspace.get_expert_recv_count_ptr(j, i) = 0; + } + __syncwarp(); + + for (uint32_t j = thread_idx; j < num_recv_m_blocks; j += kNumDispatchThreads) { + *workspace.get_l1_arrival_count_ptr(cleanup_pool_block_offset + j) = 0; + *workspace.get_l2_arrival_mask_ptr(cleanup_pool_block_offset + j) = 0; + } + __syncwarp(); + } + } +}; + +// ===================================================================== +// ROLE 1: DISPATCH WARPS +// Mirrors SM100 dispatch with two changes: +// * SF is per-128 channel float (no UTCCP transpose). We store the +// remote per-token SF directly into the local L1 SF buffer in +// MN-major layout: `local_sf[k_chunk * num_padded_sf_pool_tokens + token_idx]`. +// * The "token_idx_in_expert" → SF token index is now the simple +// per-block linear mapping (no 4×32 transpose). +// ===================================================================== +if (warp_idx < kNumDispatchWarps) { + cutlass::arch::warpgroup_reg_dealloc(); + if constexpr (kPushDispatch) { + dispatch_epoch = read_launch_epoch(); + dispatch_parity = epoch_pool_parity(dispatch_epoch); + } + + FI_MEGAMOE_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of topk"); + constexpr uint32_t kNumActivateLanes = kNumTokensPerWarp * kNumTopk; + const auto read_topk_idx = [&](const auto& process) { + if (warp_idx < kNumActiveDispatchWarps) { +#pragma unroll + for (uint32_t i = (sm_idx * kNumActiveDispatchWarps + warp_idx) * kNumTokensPerWarp; + i < num_tokens; i += kNumSMs * kNumActiveDispatchWarps * kNumTokensPerWarp) { + int expert_idx = -1; + if (i + (lane_idx / kNumTopk) < num_tokens and lane_idx < kNumActivateLanes) { + expert_idx = static_cast( + __ldg(input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); + if (expert_idx >= 0) process(i * kNumTopk + lane_idx, expert_idx); + } + __syncwarp(); + } + } + }; + + if constexpr (kPushDispatch) { + // Push: rows (token, top-k slot) are assigned to the global dispatch warps + // in contiguous chunks of R = ceil(rows / warps) (<= 32 per ticket batch). + // Every lane of the warp takes the remote ticket of one row of the batch + // at once (kNumRanks-way NVLink round trips overlap), then the warp + // streams the batch's rows (16 B per lane per store) into the + // destination pools. + if (warp_idx < kNumActiveDispatchWarps) { + constexpr uint32_t kNumGlobalWarps = kNumSMs * kNumActiveDispatchWarps; + sm90_nvfp4_push_dispatch_rows( + sym_buffer, smem_send_buffers.get_rank_buffer(warp_idx).get_data_buffer(0).get_base_ptr(), + dispatch_barriers[warp_idx], input_topk_idx_buffer.get_base_ptr(), + input_token_buffer.get_base_ptr(), input_sf_buffer.get_base_ptr(), + input_topk_weights_buffer.get_base_ptr(), l1_token_buffer.get_base_ptr(), + l1_sf_buffer.get_base_ptr(), l1_topk_weights_buffer.get_base_ptr(), + workspace.get_token_src_metadata_ptr(0), recv_count_sum_ptr_p(dispatch_parity, 0), + dispatch_parity * kNumExpertsPerRank * kPushBlocksPerExpert, num_tokens, + sm_idx * kNumActiveDispatchWarps + warp_idx, lane_idx); + } + if (thread_idx == 0) stamp_max(9); // pushes issued + + // CTA arrival ticket; the last CTA signals DONE to every rank (kNumRanks + // lanes of warp 0 at once: one warp-wide sys-scope release). + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + if (warp_idx == 0) { + FI_MEGAMOE_STATIC_ASSERT(kNumRanks <= 32, "Too many ranks for one signalling warp"); + uint32_t arrived = 0; + if (lane_idx == 0) + arrived = ptx::atomic_add_acq_rel(workspace.get_push_cta_arrival_ptr(), 1u); + arrived = __shfl_sync(0xffffffff, arrived, 0); + if (arrived == kNumSMs - 1) { + if (lane_idx == 0) *workspace.get_push_cta_arrival_ptr() = 0; + if (lane_idx < kNumRanks) + ptx::red_add_rel_sys(sym_buffer.map(workspace.get_push_done_count_ptr(), lane_idx), 1); + } + __syncwarp(); + } + if (thread_idx == 0) stamp_max(10); // DONE signalled (this CTA arrived) + + // Let the math warps go (they wait for published tasks anyway) + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // SM e (< experts per rank) publishes expert e's L1 arrival counts once + // every rank's rows are in the local pool (DONE acquired), one block per + // lane, with release.gpu: the loaders' acquire on the count then also + // covers the remotely written rows. + if (warp_idx == 0 and sm_idx < kNumExpertsPerRank) { + const int push_done_target = epoch_done_target(dispatch_epoch); + if (lane_idx == 0) { + auto* spin_dbg = workspace.get_spin_debug_ptr(); + FI_MEGAMOE_SPIN_WHILE_DEBUG( + ptx::ld_volatile(workspace.get_push_done_count_ptr()) - push_done_target < 0, 1094, + spin_dbg, ptx::ld_volatile(workspace.get_push_done_count_ptr()), push_done_target); + FI_MEGAMOE_SPIN_WHILE_DEBUG( + ptx::ld_acq_sys(workspace.get_push_done_count_ptr()) - push_done_target < 0, 1095, + spin_dbg, ptx::ld_acq_sys(workspace.get_push_done_count_ptr()), push_done_target); + stamp_max(1); + } + __syncwarp(); + const uint32_t num_recv_tokens = + static_cast(ptx::ld_volatile(recv_count_sum_ptr_p(dispatch_parity, sm_idx))); + const uint32_t num_blocks = math::ceil_div(num_recv_tokens, BLOCK_M); + for (uint32_t b = lane_idx; b < num_blocks; b += 32) + ptx::red_add_rel( + workspace.get_l1_arrival_count_ptr(strided_pool_block_p(dispatch_parity, sm_idx, b)), + cute::min(num_recv_tokens - b * BLOCK_M, BLOCK_M)); + __syncwarp(); + if (lane_idx == 0) stamp_max(2); // pool ready + } + } else { + // Count tokens per expert + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + atomicAdd_block(smem_expert_count + expert_idx, 1); + }); + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + if (thread_idx == 0) stamp_max(8); + +// Stake out per-expert SM offsets via global atomic +#pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const uint64_t send_value = (1ull << 32) | static_cast(smem_expert_count[i]); + smem_expert_count[i] = static_cast( + ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value)); + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Write source token-topk indices to remote ranks + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + const auto dst_rank_idx = expert_idx / kNumExpertsPerRank; + const auto dst_slot_idx = atomicAdd_block(smem_expert_count + expert_idx, 1); + const auto dst_ptr = workspace.get_src_token_topk_idx_ptr(expert_idx % kNumExpertsPerRank, + sym_buffer.rank_idx, dst_slot_idx); + *sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx; + }); + if (thread_idx == 0) stamp_max(9); + comm::grid_sync(workspace, sm_idx, thread_idx, [=]() { + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + }); + if (thread_idx == 0) stamp_max(10); + + if (sm_idx == 0 and thread_idx < kNumActiveDispatchThreads) { +#pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumActiveDispatchThreads) { + const auto dst_rank_idx = i / kNumExpertsPerRank; + const auto dst_local_expert_idx = i % kNumExpertsPerRank; + const auto expert_status = *workspace.get_expert_send_count_ptr(i); + *sym_buffer.map( + workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx), + dst_rank_idx) = expert_status & 0xffffffff; + ptx::atomic_add_sys( + sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), + dst_rank_idx), + expert_status); + } + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + if (thread_idx == 0) stamp_max(11); + + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, false, true); + if (thread_idx == 0) stamp_max(1); + + // Sync with epilogue warps before pulling tokens + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Token / SF pull loop + if (warp_idx < kNumActiveDispatchWarps) { + uint32_t pull_mbarrier_phase = 0; + const auto pull_buffer = smem_send_buffers.get_rank_buffer(warp_idx).get_data_buffer(0); + const auto pull_mbarrier = dispatch_barriers[warp_idx]; + + scheduler.fetch_expert_recv_count(); + + constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u); + int current_expert_idx = -1; + uint32_t stored_rank_count[kNumRanksPerLane] = {}; + uint32_t expert_start_idx = 0, expert_end_idx = 0; + uint32_t expert_pool_block_offset = 0; + + constexpr uint32_t kNumGlobalWarps = kNumSMs * kNumActiveDispatchWarps; + for (uint32_t token_idx = sm_idx * kNumActiveDispatchWarps + warp_idx;; + token_idx += kNumGlobalWarps) { + int old_expert_idx = current_expert_idx; + while (token_idx >= expert_end_idx) { + if (++current_expert_idx >= kNumExpertsPerRank) break; + expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M); + expert_start_idx = expert_end_idx; + expert_end_idx += scheduler.get_num_tokens(current_expert_idx); + } + if (current_expert_idx >= kNumExpertsPerRank) break; + if constexpr (kStridedPool) + expert_pool_block_offset = + strided_pool_block_p(dispatch_parity, static_cast(current_expert_idx), 0); + + if (old_expert_idx != current_expert_idx) { + old_expert_idx = current_expert_idx; +#pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++i) { + const uint32_t j = i * 32 + lane_idx; + stored_rank_count[i] = j < kNumRanks + ? static_cast(*workspace.get_expert_recv_count_ptr( + j, current_expert_idx)) + : 0; + } + } + + // Round-robin rank selection (identical to SM100) + +#include + const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr( + current_expert_idx, current_rank_in_expert_idx, token_idx_in_rank); + const uint32_t src_token_idx = src_token_topk_idx / kNumTopk; + const uint32_t src_topk_idx = src_token_topk_idx % kNumTopk; + + // TMA pull token data into SMEM + if (cute::elect_one_sync()) { + ptx::tma_load_1d( + pull_buffer.get_base_ptr(), + sym_buffer.map(input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx), + pull_mbarrier, kHidden); + } + __syncwarp(); + + // Copy SF: per-128 K floats, written linearly (no UTCCP transpose). + constexpr uint32_t kNumSFFloats = kHidden / 128; + FI_MEGAMOE_STATIC_ASSERT(kNumSFFloats > 0 and kHidden % 128 == 0, "Invalid SF"); + const auto remote_sf_ptr = + sym_buffer.map(input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx); + const auto local_sf_ptr = l1_sf_buffer.get_base_ptr(); + const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; +#pragma unroll + for (uint32_t i = 0; i < math::constexpr_ceil_div(kNumSFFloats, 32u); ++i) { + const uint32_t j = i * 32 + lane_idx; + if (j < kNumSFFloats) + local_sf_ptr[j * kNumPaddedSFPoolTokens + pool_token_idx] = remote_sf_ptr[j]; + } + __syncwarp(); + + if (cute::elect_one_sync()) { + const auto weight = + *sym_buffer.map(input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, + current_rank_in_expert_idx); + *l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr() = weight; + + ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kHidden); + ptx::mbarrier_wait_and_flip_phase(pull_mbarrier, pull_mbarrier_phase); + + ptx::tma_store_1d(l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr(), + pull_buffer.get_base_ptr(), pull_buffer.get_num_bytes()); + + *workspace.get_token_src_metadata_ptr(pool_token_idx) = {current_rank_in_expert_idx, + src_token_idx, src_topk_idx}; + + cute::tma_store_arrive(); + ptx::tma_store_wait<0>(); + ptx::red_add_rel(workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + + token_idx_in_expert / BLOCK_M), + 1u); + } + __syncwarp(); + } + } + if (thread_idx == 0) stamp_max(2); // pool ready (pull done) + } // !kPushDispatch + + // Cleanup workspace, overlapping with combine + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + if constexpr (kNoCleanBarrier) { + // Rotated pool: local grid sync (every CTA's math tasks are done), then + // clean; no cross-rank barrier (see `kNoCleanBarrier`). + comm::grid_sync(workspace, sm_idx, thread_idx, [=]() { + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + }); + cleanup_workspace(); + } else { + cleanup_workspace(); + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, true, false); + } + if (thread_idx == 0) stamp_max(12); +} else if (warp_idx == kNumDispatchWarps) { + // ===================================================================== + // ROLE 2: GEMM TMA LOAD warps (load A+SFA, B+SFB) + // The two warps inside `kNumNonEpilogueThreads` load A + SFA and + // B + SFB, respectively. + // ===================================================================== + cutlass::arch::warpgroup_reg_dealloc(); + + const auto load_a_task = [&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx, const uint32_t& pool_block_idx, + const uint32_t& valid_m) { + using BlockPhaseTag = std::remove_cv_t>; + constexpr bool kBlockIsL2 = BlockPhaseTag::value == sched::BlockPhase::Linear2; + const auto tensor_map_a_ptr = kBlockIsL2 ? &tensor_map_l2_acts : &tensor_map_l1_acts; + const auto tensor_map_sfa_ptr = kBlockIsL2 ? &tensor_map_l2_acts_sf : &tensor_map_l1_acts_sf; + + const bool has_valid_m = valid_m > 0; + + // Wait for the pool to be ready. + if (has_valid_m) { + const auto t_spin = stamp_now(); + if constexpr (!kBlockIsL2) { + const auto ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx); + while (ptx::ld_acq(ptr) != valid_m) { + } + if constexpr (kPushDispatch && kPushProxyFence) { + // The rows were written with generic stores (over NVLink); + // order the acquire before the async-proxy (TMA) loads. + asm volatile("fence.proxy.async.global;" ::: "memory"); + } + } else { + const auto ptr = workspace.get_l2_arrival_mask_ptr(pool_block_idx); + const uint64_t expected = + (kNumRoutedL1BlockNs >= 64) ? ~0ull : ((1ull << kNumRoutedL1BlockNs) - 1ull); + while (ptx::ld_acq_gpu(ptr) != expected) { + } + } + if (lane_idx == 0) stamp_accumulate(kBlockIsL2 ? 25 : 24, stamp_now() - t_spin, 0); + } + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + { + const auto t_w = stamp_now(); + empty_barriers[stage_idx]->wait(phase ^ 1); + if (lane_idx == 0) stamp_accumulate(29, stamp_now() - t_w, 0); + } + + if (cute::elect_one_sync()) { + if (has_valid_m) { + const uint32_t m_idx = pool_block_idx * BLOCK_M; + const uint32_t k_idx = k_block_idx * BLOCK_K; + + // TMA load A + tma::copy( + tensor_map_a_ptr, full_barriers[stage_idx], smem_a[stage_idx], k_idx, m_idx, 1); + + if constexpr (!kBlockIsL2 || !kSplitMDecodedWeightReuse) { + tma::copy(tensor_map_sfa_ptr, full_barriers[stage_idx], + smem_sfa[stage_idx], m_idx, k_block_idx, 1); + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + + BLOCK_M * sizeof(float)); + } else { + // BN128 L1 produces per-64 activation scales. L2 + // consumes both scale groups for each BK128 tile. + tma::copy(tensor_map_sfa_ptr, full_barriers[stage_idx], + smem_sfa[stage_idx], m_idx, k_block_idx * 2, 1); + tma::copy(tensor_map_sfa_ptr, full_barriers[stage_idx], + smem_sfa[stage_idx] + kL2SFAHalfStride, m_idx, + k_block_idx * 2 + 1, 1); + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + + 2 * BLOCK_M * sizeof(float)); + } + } else { + full_barriers[stage_idx]->arrive(); + } + } + __syncwarp(); + } + }; + if constexpr (kUseInterleavedScheduler) + for_each_published_block(load_a_task); + else + for_each_static_selected_block(load_a_task); + +} else if (warp_idx == kNumDispatchWarps + 1) { + cutlass::arch::warpgroup_reg_dealloc(); + + const auto load_b_task = [&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx, const uint32_t& pool_block_idx, + const uint32_t& valid_m) { + using BlockPhaseTag = std::remove_cv_t>; + constexpr bool kBlockIsL2 = BlockPhaseTag::value == sched::BlockPhase::Linear2; + const auto tensor_map_b_ptr = kBlockIsL2 ? &tensor_map_l2_weights : &tensor_map_l1_weights; + constexpr uint32_t shape_n = kBlockIsL2 ? L2_SHAPE_N : L1_SHAPE_N; + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + { + const auto t_w = stamp_now(); + empty_barriers[stage_idx]->wait(phase ^ 1); + if (lane_idx == 0) stamp_accumulate(27, stamp_now() - t_w, 0); + } + + const uint32_t n_idx = local_expert_idx * shape_n + n_block_idx * BLOCK_N; + // NVFP4 fused B+scale layout stores 64B packed FP4 + 8B + // UE4M3 scale + 8B zero padding per BK128 row. + const uint32_t k_idx = k_block_idx * B_LOAD_BYTES_PER_ROW; + if (cute::elect_one_sync()) { + tma::copy( + tensor_map_b_ptr, full_barriers[stage_idx], smem_packed_b[stage_idx], k_idx, n_idx, 1); + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_PACKED_B_SIZE_PER_STAGE); + } + __syncwarp(); + } + }; + if constexpr (kUseInterleavedScheduler) + produce_interleaved_blocks(load_b_task); + else + for_each_static_selected_block(load_b_task); + +} else { + // ===================================================================== + // ROLE 3: MATH WARPGROUPS (WGMMA + epilogue + combine) + // ===================================================================== + cutlass::arch::warpgroup_reg_alloc(); + + const uint32_t epilogue_warp_idx = warp_idx - (kNumDispatchWarps + kNumMMANonEpilogueWarps); + const uint32_t epilogue_wg_idx = epilogue_warp_idx / 4; + const uint32_t epilogue_thread_idx = epilogue_warp_idx * 32 + lane_idx; + const uint32_t warp_idx_in_wg = epilogue_warp_idx % 4; + + const auto arrive_empty_barrier = [&](const uint32_t& s) { + if (lane_idx == 0) empty_barriers[s]->arrive(); + }; + + const auto notify_l1_ready = [&](const uint32_t& ready_pool_block_idx, + const uint32_t& ready_n_block_idx) { + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + ptx::red_or_rel_gpu(workspace.get_l2_arrival_mask_ptr(ready_pool_block_idx), + 1ull << ready_n_block_idx); + } + __syncwarp(); + }; + + // WGMMA-output register layout helpers + const uint32_t row_idx = lane_idx / 4; + const uint32_t col_idx = lane_idx % 4; + const uint32_t r_0 = warp_idx_in_wg * 16 + row_idx; + const uint32_t r_1 = r_0 + 8; + + FI_MEGAMOE_STATIC_ASSERT( + kSwapABRequested || (WG_BLOCK_M == L1WGMMA::M and WG_BLOCK_N == L1WGMMA::N), + "Split-N WGs must each run one M64N128 WGMMA per K-block"); + + // Sync with dispatch + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + const auto run_math_task_impl = [&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx, const uint32_t& pool_block_idx, + const uint32_t& valid_m) { + const uint32_t m_idx = pool_block_idx * BLOCK_M; + const uint32_t wg_n_idx = kSplitMDecodedWeightReuse ? 0u : epilogue_wg_idx * WG_BLOCK_N; + const uint32_t wg_l1_out_n_idx = + kSplitMDecodedWeightReuse ? 0u : epilogue_wg_idx * WG_L1_OUT_BLOCK_N; + const uint32_t n_idx = n_block_idx * BLOCK_N + wg_n_idx; + const uint32_t row_block_offset = kSplitMDecodedWeightReuse ? epilogue_wg_idx * WG_BLOCK_M : 0u; + const uint32_t row_offset_r0 = row_block_offset + r_0; + const uint32_t row_offset_r1 = row_block_offset + r_1; + using BlockPhaseTag = std::remove_cv_t>; + constexpr bool kBlockIsL2 = BlockPhaseTag::value == sched::BlockPhase::Linear2; + const float l2_global_scale = + l2_global_scales == nullptr ? 1.0f : __ldg(l2_global_scales + local_expert_idx); + const auto cast_l2_scaled_bf16_pair = [&](float x, float y) -> uint32_t { + x *= l2_global_scale; + y *= l2_global_scale; + return math::cast_into_bf16_and_pack(x, y); + }; + + // ---------------- GEMM ---------------- + using WGMMA = L1WGMMA; + constexpr uint32_t kAccumPerThread = WGMMA::kNumAccum; // 64 for M=64,N=128 + float final_accum[kAccumPerThread] = {}; + const auto decode_b_stage = [&](const uint32_t& decoded_stage) { + if constexpr (kRSSwapABRequested) { + // Mode-5 RS consumes packed weights directly. The SS + // expanded-B allocation is intentionally retained in this + // first integration so RS and compact-stage tuning remain + // separate causal axes. + } else if constexpr (kSplitMDecodedWeightReuse) { + // Both M64 consumers share the same decoded N128 tile. + // The two physical decoded slots remove the overwrite + // hazard; the trailing pair barrier publishes K+1 before + // either consumer issues its WGMMA. + FI_MEGAMOE_STATIC_ASSERT(kUseMode2RowDecoder, + "BM128 split-M uses the cooperative Mode2 decoder"); + flashinfer::megamoe::nvfp4::dequant_smem_b_from_packed_mode2_nibble_split_m( + reinterpret_cast(smem_b[decoded_stage]), + reinterpret_cast(smem_packed_b[decoded_stage]), epilogue_thread_idx, + smem_nvfp4_lut); + cutlass::arch::fence_view_async_shared(); + asm volatile("bar.sync %0, %1;" : : "n"(kSplitMDecodeBarrierIdx), "n"(256) : "memory"); + } else { + if constexpr (kUseMode2RowDecoder) { + flashinfer::megamoe::nvfp4::dequant_smem_b_from_packed_mode2_nibble( + reinterpret_cast(smem_b[decoded_stage]), + reinterpret_cast(smem_packed_b[decoded_stage]), epilogue_thread_idx, + smem_nvfp4_lut); + } else { + flashinfer::megamoe::nvfp4::dequant_smem_b_from_packed_braided_lut_window< + kQuadDequantIlp>(reinterpret_cast(smem_b[decoded_stage]), + reinterpret_cast(smem_packed_b[decoded_stage]), + epilogue_thread_idx, smem_nvfp4_lut); + } + cutlass::arch::fence_view_async_shared(); + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + } + }; + + const auto decode_rs_swap_ab_half = [&](const uint32_t& rs_stage, + const uint32_t& half, + uint32_t(&a_frag)[4][4]) { + const uint32_t frag_row0 = wg_n_idx + half * 64u + warp_idx_in_wg * 16u + row_idx; + const uint32_t decode_row = frag_row0 + ((lane_idx & 1u) << 3); + const uint32_t word_sel = (lane_idx >> 1) & 1u; + const bool keep_hi = (lane_idx & 1u) == 0; + const auto* packed_row = reinterpret_cast(smem_packed_b[rs_stage]) + + decode_row * B_LOAD_BYTES_PER_ROW; + const uint2 scale_words = ptx::ld_shared(reinterpret_cast(packed_row + 64u)); + +#pragma unroll + for (uint32_t slice = 0; slice < 4; ++slice) { + const uint32_t slice_offset = slice * 16u; + const uint32_t word_offset = word_sel * sizeof(uint32_t); + const uint32_t w_lo = ptx::ld_shared( + reinterpret_cast(packed_row + slice_offset + word_offset)); + const uint32_t w_hi = ptx::ld_shared( + reinterpret_cast(packed_row + slice_offset + 8u + word_offset)); + const uint32_t scale_word = slice < 2u ? scale_words.x : scale_words.y; + const uint32_t scale_shift = (slice & 1u) * 16u; + const uint32_t scale_lo = (scale_word >> scale_shift) & 0x7fu; + const uint32_t scale_hi = (scale_word >> (scale_shift + 8u)) & 0x7fu; + flashinfer::megamoe::nvfp4::dequant_mode2_rs_word_pair( + w_lo, w_hi, smem_nvfp4_lut[scale_lo], smem_nvfp4_lut[scale_hi], keep_hi, a_frag[slice]); +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) mma::sm90::warpgroup_fence_operand(a_frag[slice][i]); + } + }; + + const auto issue_rs_swap_ab_half = [&](const uint32_t& rs_stage, + uint32_t(&a_frag)[4][4], + float* swap_accum) { + using RSWGMMA = typename mma::sm90::FP8RSMMASelector::type; + FI_MEGAMOE_STATIC_ASSERT(BLOCK_K / RSWGMMA::K == 4, "RS expects four K32 slices per BK128"); + +#pragma unroll + for (uint32_t i = 0; i < RSWGMMA::kNumAccum; ++i) ptx::warpgroup_fence_operand(swap_accum[i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t slice = 0; slice < 4; ++slice) { + const auto desc_b = mma::sm90::make_smem_desc(smem_a[rs_stage] + slice * RSWGMMA::K, 1); + RSWGMMA::wgmma(a_frag[slice], desc_b, swap_accum, slice > 0); + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < RSWGMMA::kNumAccum; ++i) ptx::warpgroup_fence_operand(swap_accum[i]); + }; + + const auto fence_rs_swap_ab_half = [&](uint32_t(&a_frag)[4][4]) { +#pragma unroll + for (uint32_t slice = 0; slice < 4; ++slice) { +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) mma::sm90::warpgroup_fence_operand(a_frag[slice][i]); + } + }; + + const auto finish_rs_swap_ab_half = [&](uint32_t(&a_frag)[4][4]) { + ptx::warpgroup_wait<0>(); + fence_rs_swap_ab_half.template operator()(a_frag); + }; + + // H200's accepted mode5 drains one half at a time. H20 reuses the + // same decode and MMA primitives but may keep both independent + // half-groups in flight before the waits below. + const auto run_rs_swap_ab_half = [&](const uint32_t& half, float* swap_accum) { + uint32_t a_frag[4][4]; + decode_rs_swap_ab_half.template operator()(stage_idx, half, a_frag); + issue_rs_swap_ab_half.template operator()(stage_idx, a_frag, swap_accum); + finish_rs_swap_ab_half.template operator()(a_frag); + }; + + if constexpr (kH20CompactRS && kSwapABRequested) { + // Keep half 1 of the current BK128 WGMMA in flight while the + // same warpgroup decodes half 0 of the next resident stage. + // Six stages guarantee that retaining current+next cannot + // block the two TMA producer warps at the ring boundary. + const auto run_h20_rotated_rs_task = [&]() { + using RSSwapWGMMA = typename mma::sm90::FP8RSMMASelector::type; + constexpr uint32_t kSwapAccum = RSSwapWGMMA::kNumAccum; + uint32_t a_frag[2][4][4]; + float swap_accum[2][kSwapAccum]; + + const auto accumulate_half = [&](const uint32_t rs_stage, const uint32_t half, + const float* half_accum) { +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum / 4; ++i) { + const uint32_t accum_offset = half * kSwapABHalfAccumPerThread + i * 4; + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + if (token_0 < valid_m) { + const float scale_0 = ptx::ld_shared(smem_sfa[rs_stage] + token_0); + final_accum[accum_offset + 0] += scale_0 * half_accum[i * 4 + 0]; + final_accum[accum_offset + 2] += scale_0 * half_accum[i * 4 + 2]; + } + if (token_1 < valid_m) { + const float scale_1 = ptx::ld_shared(smem_sfa[rs_stage] + token_1); + final_accum[accum_offset + 1] += scale_1 * half_accum[i * 4 + 1]; + final_accum[accum_offset + 3] += scale_1 * half_accum[i * 4 + 3]; + } + } + }; + + { + const auto t_w = stamp_now(); + full_barriers[stage_idx]->wait(phase); + if (epilogue_thread_idx == 0) stamp_accumulate(28, stamp_now() - t_w, 0); + } + if constexpr (kUseInterleavedScheduler) interleaved_scheduler.release_task_info(lane_idx); + decode_rs_swap_ab_half.template operator()(stage_idx, 0, a_frag[0]); + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + const uint32_t current_stage = stage_idx; + issue_rs_swap_ab_half.template operator()(current_stage, a_frag[0], + swap_accum[0]); + decode_rs_swap_ab_half.template operator()(current_stage, 1, a_frag[1]); + issue_rs_swap_ab_half.template operator()(current_stage, a_frag[1], + swap_accum[1]); + + ptx::warpgroup_wait<1>(); + fence_rs_swap_ab_half.template operator()(a_frag[0]); + accumulate_half(current_stage, 0, swap_accum[0]); + + if (k_block_idx + 1 < num_k_blocks) { + const uint32_t next_stage = current_stage == kNumStages - 1 ? 0 : current_stage + 1; + const uint32_t next_phase = phase ^ (next_stage == 0); + { + const auto t_w = stamp_now(); + full_barriers[next_stage]->wait(next_phase); + if (epilogue_thread_idx == 0) stamp_accumulate(28, stamp_now() - t_w, 0); + } + decode_rs_swap_ab_half.template operator()(next_stage, 0, a_frag[0]); + } + + ptx::warpgroup_wait<0>(); + fence_rs_swap_ab_half.template operator()(a_frag[1]); + accumulate_half(current_stage, 1, swap_accum[1]); + + // Each warp has completed its own asynchronous stage + // reads before arriving. No warpgroup-wide rendezvous + // is required on this H20-only rotated path. + arrive_empty_barrier(current_stage); + } + }; + + if constexpr (BLOCK_M == 8) { + run_h20_rotated_rs_task.template operator()<8>(); + } else if constexpr (BLOCK_M == 16) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if (n_swap <= 8) + run_h20_rotated_rs_task.template operator()<8>(); + else + run_h20_rotated_rs_task.template operator()<16>(); + } else if constexpr (BLOCK_M == 24) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if (n_swap <= 8) + run_h20_rotated_rs_task.template operator()<8>(); + else if (n_swap <= 16) + run_h20_rotated_rs_task.template operator()<16>(); + else + run_h20_rotated_rs_task.template operator()<24>(); + } + } else { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); + if constexpr (kUseInterleavedScheduler) { + if (k_block_idx == 0) interleaved_scheduler.release_task_info(lane_idx); + } + decode_b_stage(stage_idx); + + // Read SF (must precede warpgroup_arrive) + const float scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); + const float scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); + float scale_a_0_hi = 0.0f; + float scale_a_1_hi = 0.0f; + if constexpr (kBlockIsL2 && kSplitMDecodedWeightReuse) { + scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + kL2SFAHalfStride + row_offset_r0); + scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + kL2SFAHalfStride + row_offset_r1); + } + + // NVFP4 UE4M3 weight scales are applied during FP4 -> FP8 smem + // expansion, so the WGMMA accumulator only needs activation SF. + + if constexpr (!kBlockIsL2) { + if constexpr (kSwapABRequested) { + auto run_swap_ab_l1 = [&]() { + using SwapWGMMA = typename mma::sm90::FP8MMASelector::type; + constexpr uint32_t kSwapAccum = SwapWGMMA::kNumAccum; + const auto accumulate_half = [&](const uint32_t half, const float* swap_accum) { +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum / 4; ++i) { + const uint32_t accum_offset = half * kSwapABHalfAccumPerThread + i * 4; + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + if (token_0 < valid_m) { + const float scale_0 = ptx::ld_shared(smem_sfa[stage_idx] + token_0); + final_accum[accum_offset + 0] += scale_0 * swap_accum[i * 4 + 0]; + final_accum[accum_offset + 2] += scale_0 * swap_accum[i * 4 + 2]; + } + if (token_1 < valid_m) { + const float scale_1 = ptx::ld_shared(smem_sfa[stage_idx] + token_1); + final_accum[accum_offset + 1] += scale_1 * swap_accum[i * 4 + 1]; + final_accum[accum_offset + 3] += scale_1 * swap_accum[i * 4 + 3]; + } + } + }; + + if constexpr (kH20CompactRS) { + using RSSwapWGMMA = typename mma::sm90::FP8RSMMASelector::type; + FI_MEGAMOE_STATIC_ASSERT(RSSwapWGMMA::kNumAccum == kSwapAccum, + "RS/SS accumulator layouts differ"); + uint32_t a_frag[2][4][4]; + float swap_accum[2][kSwapAccum]; + + decode_rs_swap_ab_half.template operator()(stage_idx, 0, a_frag[0]); + issue_rs_swap_ab_half.template operator()(stage_idx, a_frag[0], + swap_accum[0]); + decode_rs_swap_ab_half.template operator()(stage_idx, 1, a_frag[1]); + issue_rs_swap_ab_half.template operator()(stage_idx, a_frag[1], + swap_accum[1]); + + ptx::warpgroup_wait<1>(); + fence_rs_swap_ab_half.template operator()(a_frag[0]); + accumulate_half(0, swap_accum[0]); + ptx::warpgroup_wait<0>(); + fence_rs_swap_ab_half.template operator()(a_frag[1]); + accumulate_half(1, swap_accum[1]); + } else { + float swap_accum[kSwapAccum]; +#pragma unroll + for (uint32_t half = 0; half < kSwapABWeightHalves; ++half) { + if constexpr (kRSSwapABRequested) { + using RSSwapWGMMA = typename mma::sm90::FP8RSMMASelector::type; + FI_MEGAMOE_STATIC_ASSERT(RSSwapWGMMA::kNumAccum == kSwapAccum, + "RS/SS accumulator layouts differ"); + run_rs_swap_ab_half.template operator()(half, swap_accum); + } else { +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(swap_accum[i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / SwapWGMMA::K; ++k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_b[stage_idx] + (wg_n_idx + half * 64u) * BLOCK_K + k * SwapWGMMA::K, + 1); + auto desc_b = + mma::sm90::make_smem_desc(smem_a[stage_idx] + k * SwapWGMMA::K, 1); + SwapWGMMA::wgmma(desc_a, desc_b, swap_accum, k); + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(swap_accum[i]); + ptx::warpgroup_wait<0>(); + } + accumulate_half(half, swap_accum); + } + } + + if constexpr (kRSSwapABRequested) { + // All 128 lanes read packed B and SFA. A + // leader-only mbarrier arrival cannot release + // the stage until those reads have retired. + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + } + + arrive_empty_barrier(stage_idx); + }; + + if constexpr (BLOCK_M == 8) { + run_swap_ab_l1.template operator()<8>(); + } else if constexpr (BLOCK_M == 16) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if (n_swap <= 8) { + run_swap_ab_l1.template operator()<8>(); + } else { + run_swap_ab_l1.template operator()<16>(); + } + } else if constexpr (BLOCK_M == 24) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if (n_swap <= 8) { + run_swap_ab_l1.template operator()<8>(); + } else if (n_swap <= 16) { + run_swap_ab_l1.template operator()<16>(); + } else { + run_swap_ab_l1.template operator()<24>(); + } + } + } else { + float accum[kAccumPerThread]; +// Single per-128 K-block WGMMA group +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + row_block_offset * BLOCK_K + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + wg_n_idx * BLOCK_K + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + arrive_empty_barrier(stage_idx); + +// L1: gate/up alternate at gran=8 along N; each `i` block +// of 8 cols belongs entirely to one of {gate, up}, so .x +// and .y share the same scalar. +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++i) { + final_accum[i * 4 + 0] += scale_a_0_lo * accum[i * 4 + 0]; + final_accum[i * 4 + 1] += scale_a_0_lo * accum[i * 4 + 1]; + final_accum[i * 4 + 2] += scale_a_1_lo * accum[i * 4 + 2]; + final_accum[i * 4 + 3] += scale_a_1_lo * accum[i * 4 + 3]; + } + } + } else { + if constexpr (kSwapABRequested) { + FI_MEGAMOE_STATIC_ASSERT(kL2ActsSFGranK == 128, + "L2 swap-AB requires per-128 activation scales"); + auto run_swap_ab_l2 = [&]() { + using SwapWGMMA = typename mma::sm90::FP8MMASelector::type; + constexpr uint32_t kSwapAccum = SwapWGMMA::kNumAccum; + const auto accumulate_half = [&](const uint32_t half, const float* swap_accum) { +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum / 4; ++i) { + const uint32_t accum_offset = half * kSwapABHalfAccumPerThread + i * 4; + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + if (token_0 < valid_m) { + const float scale_0 = ptx::ld_shared(smem_sfa[stage_idx] + token_0); + final_accum[accum_offset + 0] += scale_0 * swap_accum[i * 4 + 0]; + final_accum[accum_offset + 2] += scale_0 * swap_accum[i * 4 + 2]; + } + if (token_1 < valid_m) { + const float scale_1 = ptx::ld_shared(smem_sfa[stage_idx] + token_1); + final_accum[accum_offset + 1] += scale_1 * swap_accum[i * 4 + 1]; + final_accum[accum_offset + 3] += scale_1 * swap_accum[i * 4 + 3]; + } + } + }; + + if constexpr (kH20CompactRS) { + using RSSwapWGMMA = typename mma::sm90::FP8RSMMASelector::type; + FI_MEGAMOE_STATIC_ASSERT(RSSwapWGMMA::kNumAccum == kSwapAccum, + "RS/SS accumulator layouts differ"); + uint32_t a_frag[2][4][4]; + float swap_accum[2][kSwapAccum]; + + decode_rs_swap_ab_half.template operator()(stage_idx, 0, a_frag[0]); + issue_rs_swap_ab_half.template operator()(stage_idx, a_frag[0], + swap_accum[0]); + decode_rs_swap_ab_half.template operator()(stage_idx, 1, a_frag[1]); + issue_rs_swap_ab_half.template operator()(stage_idx, a_frag[1], + swap_accum[1]); + + ptx::warpgroup_wait<1>(); + fence_rs_swap_ab_half.template operator()(a_frag[0]); + accumulate_half(0, swap_accum[0]); + ptx::warpgroup_wait<0>(); + fence_rs_swap_ab_half.template operator()(a_frag[1]); + accumulate_half(1, swap_accum[1]); + } else { + float swap_accum[kSwapAccum]; +#pragma unroll + for (uint32_t half = 0; half < kSwapABWeightHalves; ++half) { + if constexpr (kRSSwapABRequested) { + using RSSwapWGMMA = typename mma::sm90::FP8RSMMASelector::type; + FI_MEGAMOE_STATIC_ASSERT(RSSwapWGMMA::kNumAccum == kSwapAccum, + "RS/SS accumulator layouts differ"); + run_rs_swap_ab_half.template operator()(half, swap_accum); + } else { +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(swap_accum[i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / SwapWGMMA::K; ++k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_b[stage_idx] + (wg_n_idx + half * 64u) * BLOCK_K + k * SwapWGMMA::K, + 1); + auto desc_b = + mma::sm90::make_smem_desc(smem_a[stage_idx] + k * SwapWGMMA::K, 1); + SwapWGMMA::wgmma(desc_a, desc_b, swap_accum, k); + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++i) + ptx::warpgroup_fence_operand(swap_accum[i]); + ptx::warpgroup_wait<0>(); + } + accumulate_half(half, swap_accum); + } + } + + if constexpr (kRSSwapABRequested) { + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + } + + arrive_empty_barrier(stage_idx); + }; + + if constexpr (BLOCK_M == 8) { + run_swap_ab_l2.template operator()<8>(); + } else if constexpr (BLOCK_M == 16) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if (n_swap <= 8) { + run_swap_ab_l2.template operator()<8>(); + } else { + run_swap_ab_l2.template operator()<16>(); + } + } else if constexpr (BLOCK_M == 24) { + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if (n_swap <= 8) { + run_swap_ab_l2.template operator()<8>(); + } else if (n_swap <= 16) { + run_swap_ab_l2.template operator()<16>(); + } else { + run_swap_ab_l2.template operator()<24>(); + } + } + } else { + float accum[kAccumPerThread]; + const auto promote_l2_accum = [&](const float& scale_r0, const float& scale_r1) { +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++i) { + final_accum[i * 4 + 0] += scale_r0 * accum[i * 4 + 0]; + final_accum[i * 4 + 1] += scale_r0 * accum[i * 4 + 1]; + final_accum[i * 4 + 2] += scale_r1 * accum[i * 4 + 2]; + final_accum[i * 4 + 3] += scale_r1 * accum[i * 4 + 3]; + } + }; + if constexpr (kSplitMDecodedWeightReuse) { +#pragma unroll + for (uint32_t sf_group = 0; sf_group < 2; ++sf_group) { +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++k) { + const uint32_t k_off = sf_group * (BLOCK_K / 2) + k * WGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + row_block_offset * BLOCK_K + k_off, 1); + auto desc_b = + mma::sm90::make_smem_desc(smem_b[stage_idx] + wg_n_idx * BLOCK_K + k_off, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + if (sf_group == 0) + promote_l2_accum(scale_a_0_lo, scale_a_1_lo); + else + promote_l2_accum(scale_a_0_hi, scale_a_1_hi); + } + arrive_empty_barrier(stage_idx); + } else { +// One per-128 scale permits a single four-instruction +// WGMMA group and one accumulator promotion per K tile. +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + row_block_offset * BLOCK_K + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + wg_n_idx * BLOCK_K + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + arrive_empty_barrier(stage_idx); + promote_l2_accum(scale_a_0_lo, scale_a_1_lo); + } + } + } + } + } + + // Skip epilogue when block is past valid M (the GEMM loop already + // released its pipeline stages). Drain any prior L1 async store. + if (valid_m == 0) { + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + return; + } + +#include + }; + const auto run_math_task = [&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx, const uint32_t& pool_block_idx, + const uint32_t& valid_m) { + using BlockPhaseTag = std::remove_cv_t>; + constexpr bool kBlockIsL2 = BlockPhaseTag::value == sched::BlockPhase::Linear2; + const auto t_task = stamp_now(); + const long long c_task = kPhaseStamps ? clock64() : 0ll; + if (epilogue_thread_idx == 0) stamp_min(3); + run_math_task_impl(block_phase, local_expert_idx, num_k_blocks, m_block_idx, n_block_idx, + pool_block_idx, valid_m); + if (epilogue_thread_idx == 0) { + stamp_accumulate(kBlockIsL2 ? 22 : 20, stamp_now() - t_task, 1); + // 26: SM clock cycles over L1 tasks (frequency = 26 / 20) + if constexpr (!kBlockIsL2) + stamp_accumulate(26, static_cast(clock64() - c_task), 0); + } + if (epilogue_thread_idx == 0) stamp_max(kBlockIsL2 ? 5 : 4); + }; + if constexpr (kUseInterleavedScheduler) + for_each_published_block(run_math_task); + else + for_each_static_selected_block(run_math_task); + + // ---------------- COMBINE ---------------- + // NVLink barrier first: signals remote ranks that this rank's GEMM + // outputs (NVLink scatter targets) are fully written. + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, epilogue_thread_idx, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); }); + if (epilogue_thread_idx == 0) stamp_max(6); + + // Sync with dispatch (paired with dispatch's pre-cleanup sync) so that + // dispatch may now safely clean workspace state. + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + +// These are function fragments: sizing declarations must precede the combine body. +// clang-format off +#include +#include + // clang-format on + + if (epilogue_thread_idx == 0) { + ptx::tma_store_wait<0>(); + stamp_max(7); + } +} +#else +if (blockIdx.x == 0 and threadIdx.x == 0) + FI_MEGAMOE_DEVICE_ASSERT(false and "This kernel only supports sm_90"); +#endif diff --git a/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_mode2_dequant.cuh b/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_mode2_dequant.cuh new file mode 100644 index 00000000000..c5bcc6df874 --- /dev/null +++ b/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_mode2_dequant.cuh @@ -0,0 +1,321 @@ +#pragma once + +#include + +#include + +namespace flashinfer::megamoe::nvfp4 { + +#ifndef FI_MEGAMOE_NVFP4_LUT_COMPACT +#define FI_MEGAMOE_NVFP4_LUT_COMPACT 0 +#endif + +// Scale-LUT gather. Compact mode (FI_MEGAMOE_NVFP4_LUT_COMPACT) halves the gathered entry +// from 8 B to 4 B: the E2M1 magnitudes hide two doubling chains (0.5->1->2->4 and +// 1.5->3->6), one doubling = one e4m3 exponent step = +0x08, so the high word is the +// low word's bytes (2,3,2,3) plus (0x08,0x08,0x10,0x10). The largest byte is +// 0x7f + 0x10 = 0x8f < 0x100, so the packed add cannot carry across byte lanes, and a +// per-byte clamp to 0x7f reproduces the table's saturation exactly. Bit-exact for every +// scale code >= 8 (verified exhaustively against all 128 rows); the eight +// subnormal-scale rows -- where e4m3 spacing is arithmetic, not geometric, and the +// chain breaks -- come verbatim from a 32-B side table stored right after lo[128], +// loaded only when some lane in the warp actually holds such a code. Callers are +// whole warps (the dequant teams), so the warp-wide vote is safe. +// Measured motivation: the 8-B gather over the 1024-B table costs ~64 base + ~20 +// conflict wavefronts per K-stage, 11.3% of the L1 SMEM pipe (dose-response probe). +__device__ __forceinline__ uint2 load_nvfp4_lut(const uint2* __restrict__ lut_smem, + const uint32_t idx) { +#if FI_MEGAMOE_NVFP4_LUT_COMPACT + const auto* lo_tbl = reinterpret_cast(lut_smem); + const uint32_t lo = lo_tbl[idx]; + uint32_t hi = __byte_perm(lo, 0u, 0x3232u) + 0x10100808u; + const uint32_t sat = hi & 0x80808080u; + hi = (hi | (sat - (sat >> 7u))) & ~sat; + if (__any_sync(0xffffffffu, idx < 8u)) { + const uint32_t exc = lo_tbl[128u + (idx & 7u)]; + hi = idx < 8u ? exc : hi; + } + return make_uint2(lo, hi); +#else + return lut_smem[idx]; +#endif +} + +// Reconstruct the straight NVFP4-scale -> E4M3 magnitude table in registers. +// For normal UE4M3 codes the low four entries are +// {0, round_even(scale / 2), scale, round_even(3 * scale / 2)} +// and the high four are the same two non-zero chains advanced by one/two E4M3 +// exponent steps. This is the direct-scale MegaMoE table, not Humming's +// deliberately /6-adjusted helper. Codes 0..7 are UE4M3 subnormals, where +// exponent stepping is not valid, so they retain the exact shared-LUT path. +// Code 0x7f is MegaMoE's documented saturating-448 fallback. + +__device__ __forceinline__ uint2 dequant_mode2_lop3_word(const uint32_t packed, const uint2& lut) { + return detail::decode_fp4_word(packed, lut); +} + +template +__device__ __forceinline__ void dequant_mode2_lop3_two_rows_lut_window( + uint8_t* __restrict__ fp8_dst0, uint8_t* __restrict__ fp8_dst1, const uint4 (&fp4_quads0)[4], + const uint4 (&fp4_quads1)[4], const uint2& scale_words0, const uint2& scale_words1, + const uint32_t row_swizzle, const uint2* __restrict__ lut_smem, const uint2 lut00, + const uint2 lut10, const uint2 lut01, const uint2 lut11, + const uint32_t half_ready_barrier_0 = 0, const uint32_t half_ready_barrier_1 = 0, + const uint32_t half_ready_threads = 0, const uint32_t producer_barrier_idx = 0, + const uint32_t producer_threads = 0, const uint32_t dequant_tid = 0, + const cutlass::arch::ClusterTransactionBarrier* half_ready_mbarrier_0 = nullptr, + const cutlass::arch::ClusterTransactionBarrier* half_ready_mbarrier_1 = nullptr) { + uint2 next_lut00; + uint2 next_lut10; + uint2 next_lut01; + uint2 next_lut11; + if constexpr (kQuad + 1 < 4 and not(kPostPublishLut and kQuad == 1)) { + constexpr int kNextScaleI0 = (kQuad + 1) * 2; + constexpr int kNextScaleI1 = kNextScaleI0 + 1; + const uint32_t next_scale_word0 = kQuad + 1 < 2 ? scale_words0.x : scale_words0.y; + const uint32_t next_scale_word1 = kQuad + 1 < 2 ? scale_words1.x : scale_words1.y; + const uint32_t next_scale00 = (next_scale_word0 >> ((kNextScaleI0 & 3) * 8)) & 0x7fu; + const uint32_t next_scale10 = (next_scale_word1 >> ((kNextScaleI0 & 3) * 8)) & 0x7fu; + const uint32_t next_scale01 = (next_scale_word0 >> ((kNextScaleI1 & 3) * 8)) & 0x7fu; + const uint32_t next_scale11 = (next_scale_word1 >> ((kNextScaleI1 & 3) * 8)) & 0x7fu; + next_lut00 = load_nvfp4_lut(lut_smem, next_scale00); + next_lut10 = load_nvfp4_lut(lut_smem, next_scale10); + next_lut01 = load_nvfp4_lut(lut_smem, next_scale01); + next_lut11 = load_nvfp4_lut(lut_smem, next_scale11); + } + + const uint4 q0 = fp4_quads0[kQuad]; + const uint4 q1 = fp4_quads1[kQuad]; + constexpr int kScaleI0 = kQuad * 2; + constexpr int kScaleI1 = kScaleI0 + 1; + const uint2 q0x = dequant_mode2_lop3_word(q0.x, lut00); + const uint2 q0y = dequant_mode2_lop3_word(q0.y, lut00); + const uint2 q1x = dequant_mode2_lop3_word(q1.x, lut10); + const uint2 q1y = dequant_mode2_lop3_word(q1.y, lut10); + *reinterpret_cast(fp8_dst0 + ((kScaleI0 * 16) ^ row_swizzle)) = + make_uint4(q0x.x, q0x.y, q0y.x, q0y.y); + *reinterpret_cast(fp8_dst1 + ((kScaleI0 * 16) ^ row_swizzle)) = + make_uint4(q1x.x, q1x.y, q1y.x, q1y.y); + + const uint2 q0z = dequant_mode2_lop3_word(q0.z, lut01); + const uint2 q0w = dequant_mode2_lop3_word(q0.w, lut01); + const uint2 q1z = dequant_mode2_lop3_word(q1.z, lut11); + const uint2 q1w = dequant_mode2_lop3_word(q1.w, lut11); + *reinterpret_cast(fp8_dst0 + ((kScaleI1 * 16) ^ row_swizzle)) = + make_uint4(q0z.x, q0z.y, q0w.x, q0w.y); + *reinterpret_cast(fp8_dst1 + ((kScaleI1 * 16) ^ row_swizzle)) = + make_uint4(q1z.x, q1z.y, q1w.x, q1w.y); + + if constexpr ((kHalfStream or kProducerMBarrier) and kQuad == 1) { + // The FP8 stores above use the generic proxy while WGMMA reads through + // the async proxy. Modes 2/3 rendezvous with all math threads; mode 4 + // rendezvous only the 64 writers and lets one writer release an + // mbarrier. Mode 5 replaces that CTA barrier with one warp barrier per + // producer warp, then each warp leader contributes one release arrival + // to an expected-count-2 mbarrier. Mode 6 publishes before loading the + // first four second-half LUT entries; the compiler memory clobber and + // the resulting LUT-to-PRMT dependency keep all second-half work after + // publication without the identity MOV gate used by modes 2--5. + cutlass::arch::fence_view_async_shared(); + if constexpr (kPostPublishLut) { + asm volatile("barrier.sync %0, %1;" + : + : "r"(producer_barrier_idx), "r"(producer_threads) + : "memory"); + if (dequant_tid == 0u) half_ready_mbarrier_0->arrive(); + + // kQuad==1 deferred exactly the kQuad==2 LUT window. Load it only + // after the first-half release so math can begin consuming K0:64. + constexpr int kNextScaleI0 = 4; + constexpr int kNextScaleI1 = 5; + const uint32_t next_scale_word0 = scale_words0.y; + const uint32_t next_scale_word1 = scale_words1.y; + const uint32_t next_scale00 = (next_scale_word0 >> ((kNextScaleI0 & 3) * 8)) & 0x7fu; + const uint32_t next_scale10 = (next_scale_word1 >> ((kNextScaleI0 & 3) * 8)) & 0x7fu; + const uint32_t next_scale01 = (next_scale_word0 >> ((kNextScaleI1 & 3) * 8)) & 0x7fu; + const uint32_t next_scale11 = (next_scale_word1 >> ((kNextScaleI1 & 3) * 8)) & 0x7fu; + next_lut00 = load_nvfp4_lut(lut_smem, next_scale00); + next_lut10 = load_nvfp4_lut(lut_smem, next_scale10); + next_lut01 = load_nvfp4_lut(lut_smem, next_scale01); + next_lut11 = load_nvfp4_lut(lut_smem, next_scale11); + } else { + uint2 gated_lut00, gated_lut10, gated_lut01, gated_lut11; + if constexpr (kWarpLeaderMBarrier) { + // bar.warp.sync orders every lane's preceding stores before + // its warp leader's release arrival. Completion after both + // leaders arrive is acquired by all math consumers. + asm volatile( + "bar.warp.sync 0xffffffff;\n" + "mov.b32 %0, %8;\n" + "mov.b32 %1, %9;\n" + "mov.b32 %2, %10;\n" + "mov.b32 %3, %11;\n" + "mov.b32 %4, %12;\n" + "mov.b32 %5, %13;\n" + "mov.b32 %6, %14;\n" + "mov.b32 %7, %15;" + : "=r"(gated_lut00.x), "=r"(gated_lut00.y), "=r"(gated_lut10.x), "=r"(gated_lut10.y), + "=r"(gated_lut01.x), "=r"(gated_lut01.y), "=r"(gated_lut11.x), "=r"(gated_lut11.y) + : "r"(next_lut00.x), "r"(next_lut00.y), "r"(next_lut10.x), "r"(next_lut10.y), + "r"(next_lut01.x), "r"(next_lut01.y), "r"(next_lut11.x), "r"(next_lut11.y) + : "memory"); + } else { + const uint32_t sync_barrier = + kProducerMBarrier ? producer_barrier_idx : half_ready_barrier_0; + const uint32_t sync_threads = kProducerMBarrier ? producer_threads : half_ready_threads; + asm volatile( + "barrier.sync %8, %9;\n" + "mov.b32 %0, %10;\n" + "mov.b32 %1, %11;\n" + "mov.b32 %2, %12;\n" + "mov.b32 %3, %13;\n" + "mov.b32 %4, %14;\n" + "mov.b32 %5, %15;\n" + "mov.b32 %6, %16;\n" + "mov.b32 %7, %17;" + : "=r"(gated_lut00.x), "=r"(gated_lut00.y), "=r"(gated_lut10.x), "=r"(gated_lut10.y), + "=r"(gated_lut01.x), "=r"(gated_lut01.y), "=r"(gated_lut11.x), "=r"(gated_lut11.y) + : "r"(sync_barrier), "r"(sync_threads), "r"(next_lut00.x), "r"(next_lut00.y), + "r"(next_lut10.x), "r"(next_lut10.y), "r"(next_lut01.x), "r"(next_lut01.y), + "r"(next_lut11.x), "r"(next_lut11.y)); + } + next_lut00 = gated_lut00; + next_lut10 = gated_lut10; + next_lut01 = gated_lut01; + next_lut11 = gated_lut11; + if constexpr (kProducerMBarrier) { + if constexpr (kWarpLeaderMBarrier) { + if ((dequant_tid & 31u) == 0u) half_ready_mbarrier_0->arrive(); + } else if (dequant_tid == 0u) { + half_ready_mbarrier_0->arrive(); + } + } + } + } else if constexpr ((kHalfStream or kProducerMBarrier) and kQuad == 3) { + cutlass::arch::fence_view_async_shared(); + if constexpr (kWarpLeaderMBarrier) { + asm volatile("bar.warp.sync 0xffffffff;" : : : "memory"); + } else { + const uint32_t sync_barrier = kProducerMBarrier ? producer_barrier_idx : half_ready_barrier_1; + const uint32_t sync_threads = kProducerMBarrier ? producer_threads : half_ready_threads; + asm volatile("barrier.sync %0, %1;" : : "r"(sync_barrier), "r"(sync_threads)); + } + if constexpr (kProducerMBarrier) { + if constexpr (kWarpLeaderMBarrier) { + if ((dequant_tid & 31u) == 0u) half_ready_mbarrier_1->arrive(); + } else if (dequant_tid == 0u) { + half_ready_mbarrier_1->arrive(); + } + } + } + + if constexpr (kQuad + 1 < 4) { + dequant_mode2_lop3_two_rows_lut_window( + fp8_dst0, fp8_dst1, fp4_quads0, fp4_quads1, scale_words0, scale_words1, row_swizzle, + lut_smem, next_lut00, next_lut10, next_lut01, next_lut11, half_ready_barrier_0, + half_ready_barrier_1, half_ready_threads, producer_barrier_idx, producer_threads, + dequant_tid, half_ready_mbarrier_0, half_ready_mbarrier_1); + } +} + +template +__device__ __forceinline__ void dequant_smem_b_inplace_two_rows_mode2_lop3( + uint8_t* __restrict__ smem_b, const uint32_t tid, const uint2* __restrict__ lut_smem, + const uint32_t half_ready_barrier_0 = 0, const uint32_t half_ready_barrier_1 = 0, + const uint32_t half_ready_threads = 0, + const cutlass::arch::ClusterTransactionBarrier* half_ready_mbarrier_0 = nullptr, + const cutlass::arch::ClusterTransactionBarrier* half_ready_mbarrier_1 = nullptr) { + static_assert(not(kSyncAfter and (kHalfStream or kProducerMBarrier)), + "half-ready handoff replaces the final writer rendezvous"); + static_assert(not(kHalfStream and kProducerMBarrier), "select exactly one half-ready protocol"); + static_assert(not kWarpLeaderMBarrier or kProducerMBarrier, + "warp-leader publication requires producer mbarriers"); + static_assert(not kPostPublishLut or (kProducerMBarrier and not kWarpLeaderMBarrier), + "post-publish LUT scheduling requires mode-4 mbarriers"); + const uint32_t row0 = tid; + const uint32_t row1 = tid + kNumDequantThreads; + const uint8_t* __restrict__ row_ptr0 = smem_b + row0 * kFusedRowBytes; + const uint8_t* __restrict__ row_ptr1 = smem_b + row1 * kFusedRowBytes; + const uint2 scale_words0 = *reinterpret_cast(row_ptr0 + 64); + const uint2 scale_words1 = *reinterpret_cast(row_ptr1 + 64); + const uint32_t scale00 = scale_words0.x & 0x7fu; + const uint32_t scale10 = scale_words1.x & 0x7fu; + const uint32_t scale01 = (scale_words0.x >> 8) & 0x7fu; + const uint32_t scale11 = (scale_words1.x >> 8) & 0x7fu; + const uint2 lut00 = load_nvfp4_lut(lut_smem, scale00); + const uint2 lut10 = load_nvfp4_lut(lut_smem, scale10); + const uint2 lut01 = load_nvfp4_lut(lut_smem, scale01); + const uint2 lut11 = load_nvfp4_lut(lut_smem, scale11); + const uint4* __restrict__ fp4_src0 = reinterpret_cast(row_ptr0); + const uint4* __restrict__ fp4_src1 = reinterpret_cast(row_ptr1); + uint4 fp4_quads0[4]; + uint4 fp4_quads1[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + fp4_quads0[i] = fp4_src0[i]; + fp4_quads1[i] = fp4_src1[i]; + } + + asm volatile("bar.sync %0, %1;" : : "n"(kBarIdx), "n"(kNumDequantThreads)); + + uint8_t* __restrict__ fp8_dst0 = smem_b + row0 * 128; + uint8_t* __restrict__ fp8_dst1 = smem_b + row1 * 128; + const uint32_t row_swizzle = (tid & 7u) << 4; + dequant_mode2_lop3_two_rows_lut_window<0, kHalfStream, kProducerMBarrier, kWarpLeaderMBarrier, + kPostPublishLut>( + fp8_dst0, fp8_dst1, fp4_quads0, fp4_quads1, scale_words0, scale_words1, row_swizzle, lut_smem, + lut00, lut10, lut01, lut11, half_ready_barrier_0, half_ready_barrier_1, half_ready_threads, + kBarIdx, kNumDequantThreads, tid, half_ready_mbarrier_0, half_ready_mbarrier_1); + if constexpr (kSyncAfter) { + // Every writer must publish its generic-proxy stores before the + // writer-set rendezvous allows one thread to signal the mbarrier. + cutlass::arch::fence_view_async_shared(); + asm volatile("bar.sync %0, %1;" : : "n"(kBarIdx), "n"(kNumDequantThreads)); + } +} + +// ============================================================================ +// RS-form register-target decode (L1 RS-swapAB arm) +// ---------------------------------------------------------------------------- +// Decodes the two packed FP4 words a math lane owns in one k32 slice straight +// into its four WGMMA A-fragment registers, with no SMEM write-back. For +// wgmma.m64nNk32 e4m3 the lane's fragment is (per k32 slice): +// a0 = row r, K nibble-group 4*(l%4) a2 = a0 at K offset +16 +// a1 = row r + 8, K nibble-group 4*(l%4) a3 = a1 at K offset +16 +// `dequant_mode2_lop3_word` returns out_hi = K-consecutive elements 0..3 and +// out_lo = 4..7, so each half IS one fragment register. Lane pairs (l ^ 1) +// need the same two packed words but complementary rows: the even lane +// decodes row r, the odd lane row r + 8, and the discarded halves are +// exchanged over one shfl per word, so every word is decoded exactly once. +// `keep_hi` must be true on even lanes (own K half = elements 0..3) and +// false on odd lanes (elements 4..7). +__device__ __forceinline__ void dequant_mode2_lop3_rs_word_pair( + const uint32_t w_lo, const uint32_t w_hi, const uint2& lut_lo, const uint2& lut_hi, + const bool keep_hi, uint32_t (&a_frag)[4]) { + detail::assemble_rs_word_pair(dequant_mode2_lop3_word(w_lo, lut_lo), + dequant_mode2_lop3_word(w_hi, lut_hi), keep_hi, a_frag); +} + +// Debug fallback (decode-and-discard): the lane decodes both of its fragment +// rows' words itself and discards the complementary halves: no shfl, ~2x +// decode ALU and LUT traffic. Kept for bisecting shfl/register-lifetime +// issues; never the production arm. +__device__ __forceinline__ void dequant_mode2_lop3_rs_word_pair_noshfl( + const uint32_t w_lo_r0, const uint32_t w_hi_r0, const uint32_t w_lo_r8, const uint32_t w_hi_r8, + const uint2& lut_lo_r0, const uint2& lut_hi_r0, const uint2& lut_lo_r8, const uint2& lut_hi_r8, + const bool keep_hi, uint32_t (&a_frag)[4]) { + const uint2 d_lo_r0 = dequant_mode2_lop3_word(w_lo_r0, lut_lo_r0); + const uint2 d_lo_r8 = dequant_mode2_lop3_word(w_lo_r8, lut_lo_r8); + a_frag[0] = keep_hi ? d_lo_r0.x : d_lo_r0.y; + a_frag[1] = keep_hi ? d_lo_r8.x : d_lo_r8.y; + const uint2 d_hi_r0 = dequant_mode2_lop3_word(w_hi_r0, lut_hi_r0); + const uint2 d_hi_r8 = dequant_mode2_lop3_word(w_hi_r8, lut_hi_r8); + a_frag[2] = keep_hi ? d_hi_r0.x : d_hi_r0.y; + a_frag[3] = keep_hi ? d_hi_r8.x : d_hi_r8.y; +} + +} // namespace flashinfer::megamoe::nvfp4 diff --git a/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_split_l1_body.inl b/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_split_l1_body.inl new file mode 100644 index 00000000000..244cf532a13 --- /dev/null +++ b/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_split_l1_body.inl @@ -0,0 +1,1508 @@ +// Independent SM90 NVFP4 MegaMoE l1 kernel body. +// RS-swapAB compile-time flag, emitted into the generated TU by the host +// codegen (csrc/jit_kernels/impls/sm90_nvfp4_mega_moe.hpp) from the +// FI_MEGAMOE_NVFP4_L1_RS_SWAPAB env var. Default off: the SS path below must stay +// byte-identical in behaviour when the flag is unset. +#ifndef FI_MEGAMOE_NVFP4_L1_RS_SWAPAB +#define FI_MEGAMOE_NVFP4_L1_RS_SWAPAB 0 +#endif +#ifndef FI_MEGAMOE_NVFP4_L1_SINGLE_ACCUM +#define FI_MEGAMOE_NVFP4_L1_SINGLE_ACCUM 0 +#endif +#ifndef FI_MEGAMOE_NVFP4_LUT_COMPACT +#define FI_MEGAMOE_NVFP4_LUT_COMPACT 0 +#endif +#ifndef FI_MEGAMOE_NVFP4_L1_SFA_LOOKAHEAD +#define FI_MEGAMOE_NVFP4_L1_SFA_LOOKAHEAD 0 +#endif +#ifndef FI_MEGAMOE_NVFP4_L1_DEQUANT_HALF_STREAM +#define FI_MEGAMOE_NVFP4_L1_DEQUANT_HALF_STREAM 0 +#endif +#ifndef FI_MEGAMOE_NVFP4_L1_DEQUANT_WARP_REMAP +#define FI_MEGAMOE_NVFP4_L1_DEQUANT_WARP_REMAP 0 +#endif +// Debug arm of the RS decoder: decode-and-discard (no shfl pairing) with +// fully serialized WGMMA groups, for bisecting shfl/register-lifetime issues. +#ifndef FI_MEGAMOE_NVFP4_L1_RS_DEBUG_DECODE +#define FI_MEGAMOE_NVFP4_L1_RS_DEBUG_DECODE 0 +#endif +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900) and (__CUDA_ARCH__ < 1000)) or \ + defined(__CLION_IDE__) +using Barrier = cutlass::arch::ClusterTransactionBarrier; + +// ===================================================================== +// Template checks +// ===================================================================== +FI_MEGAMOE_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks"); +FI_MEGAMOE_STATIC_ASSERT(kNumSMs % 2 == 0, "SM count must be divisible by the L1 cluster size"); + +// ===================================================================== +// Thread / warp identification +// ===================================================================== +const uint32_t sm_idx = blockIdx.x; +const uint32_t thread_idx = threadIdx.x; +const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); +const uint32_t lane_idx = ptx::get_lane_idx(); + +if (warp_idx == 0 and cute::elect_one_sync()) { + cute::prefetch_tma_descriptor(&tensor_map_l1_acts); + cute::prefetch_tma_descriptor(&tensor_map_l1_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l1_weights); + cute::prefetch_tma_descriptor(&tensor_map_l1_output); +} + +// ===================================================================== +// Workspaces and symmetric buffer slicing. The L2 activation SF allocation +// intentionally retains its old per-64 physical capacity in stage 1, while +// the active logical layout densely addresses only per-128 groups. +// ===================================================================== +#include + +// ===================================================================== +// GEMM data types and shape constants +// ===================================================================== +using a_dtype_t = cutlass::float_e4m3_t; +using b_dtype_t = cutlass::float_e4m3_t; +constexpr uint32_t WG_BLOCK_M = 64; +constexpr uint32_t WG_BLOCK_N = 128; +constexpr uint32_t L1_OUT_BLOCK_N = 64; +constexpr uint32_t WG_L1_OUT_BLOCK_N = 64; +constexpr bool kL2ArrivalCounter = kL2ArrivalCounterRequested; +// Use two active dispatch warps. The other two dispatch warps form the +// even-stage dequant team when dispatch-assisted dequant is on. +constexpr uint32_t kNumActiveDispatchWarps = 2; +constexpr uint32_t kNumActiveDispatchThreads = 64; +// RS-swapAB: swap the WGMMA A/B roles and decode the FP4 weights straight +// into A-register fragments (RS form). The FP8 SMEM expansion and both +// dequant teams compile out, removing the FP8 write-back and the WGMMA +// weight re-read from the SMEM pipe entirely. +constexpr bool kRSSwapAB = FI_MEGAMOE_NVFP4_L1_RS_SWAPAB != 0; +// Single-accumulator form. Microscaling forces a per-k-block flush -- the WGMMA +// output must be multiplied by the token's activation SF for THAT k range before it +// can join the running total -- which is why the default path carries two 64-register +// arrays. Keeping the running total in the WGMMA accumulator itself, expressed in +// units of the current block's SF and rescaled in place when the SF changes, needs +// only one. That frees 64 registers per math thread, which is the binding constraint +// on BLOCK_M * BLOCK_N (accum bytes = 2 * BLOCK_M * BLOCK_N / 256 threads). +constexpr bool kSingleAccum = FI_MEGAMOE_NVFP4_L1_SINGLE_ACCUM != 0 and not kRSSwapAB; +constexpr bool kLutCompact = FI_MEGAMOE_NVFP4_LUT_COMPACT != 0; +constexpr uint32_t kSFALookahead = FI_MEGAMOE_NVFP4_L1_SFA_LOOKAHEAD; +FI_MEGAMOE_STATIC_ASSERT(kSFALookahead <= 2, + "SFA lookahead modes: 0=OFF, 1=EMPTY full-wait, 2=REAL prefetch"); +// 0: original whole-stage mbarrier handoff. 1: whole-stage decode followed +// by the same two named ready barriers as the real arm (EMPTY/control). +// 2: publish K64 halves so the first two K32 WGMMAs overlap the second half. +// 3: producer is identical to mode 2, but math waits for both halves before +// issuing any WGMMA. This isolates split/fence cost from SMEM contention. +// 4: keep the mode-2 overlap, but rendezvous only the 64 dequant writers; +// one writer releases each half through a stage-local mbarrier, so math +// waits without forcing the producer to rendezvous with 256 consumers. +// 5: replace each mode-4 producer CTA barrier with per-warp rendezvous; +// both producer warp leaders release the expected-count-2 mbarrier. +// 6: retain mode-4 publication, but move the first four second-half LUT +// loads after the K0:64 release and remove its identity-MOV dependency gate. +constexpr uint32_t kDequantHalfStream = FI_MEGAMOE_NVFP4_L1_DEQUANT_HALF_STREAM; +FI_MEGAMOE_STATIC_ASSERT(kDequantHalfStream <= 6, + "dequant half stream modes: 0=OFF, 1=EMPTY, 2=NAMED, 3=NO_OVERLAP, " + "4=MBARRIER, 5=WARP_LEADER, 6=POST_PUBLISH_LUT"); +FI_MEGAMOE_STATIC_ASSERT(not(kDequantHalfStream != 0 and kSFALookahead != 0), + "half stream owns the full/dequant handoff ordering"); +FI_MEGAMOE_STATIC_ASSERT(not(kDequantHalfStream != 0 and kRSSwapAB), + "RS-swapAB has no shared-memory dequant producer"); +// The RS arm gathers full-width uint2 entries straight off smem_nvfp4_lut and +// never went through the shared loader; the two layouts cannot coexist. +FI_MEGAMOE_STATIC_ASSERT(not(kRSSwapAB and kLutCompact), + "RS-swapAB reads the full-width LUT layout"); +constexpr bool kDispatchDequant = kDispatchDequantRequested and not kRSSwapAB; +constexpr bool kDequantWarpRemap = FI_MEGAMOE_NVFP4_L1_DEQUANT_WARP_REMAP != 0; +FI_MEGAMOE_STATIC_ASSERT(not kDequantWarpRemap or (kDispatchDequant and kDequantHalfStream == 4), + "dequant warp remap is calibrated only for dispatch mode4"); +using L1WGMMA = typename mma::sm90::FP8MMASelector<128>::type; +// Same M64N128K32 shape, A sourced from per-thread fragments instead of +// an SMEM descriptor. Under swapAB each math WG's "M" is its 64 weight +// rows and "N" is the 128 activation tokens. +using L1WGMMARS = typename mma::sm90::FP8RSMMASelector<128>::type; +static_assert(L1WGMMA::M == 64 and L1WGMMA::N == WG_BLOCK_N and L1WGMMA::K == 32, + "Unexpected WGMMA shape"); + +// A is always CTA-local. When kClusterSize=2 the scheduler pairs adjacent +// M blocks with identical expert/N/K coordinates so the B TMA can multicast. +constexpr uint32_t LOAD_BLOCK_M = 128; +constexpr uint32_t LOAD_BLOCK_N = 128; +constexpr uint32_t kSwizzleAMode = 128; +constexpr uint32_t kLocalL1ActsSFGranK = 64; // each CTA's local half +constexpr uint32_t kL2ActsSFGranK = 128; // final L1 output / L2 input +FI_MEGAMOE_STATIC_ASSERT(L1_SHAPE_N / BLOCK_N <= 64, + "paired-N readiness must fit one 64-bit pool-block mask"); +FI_MEGAMOE_STATIC_ASSERT((L1_SHAPE_N / BLOCK_N) % 2 == 0, + "every L1 N block must have a cluster peer"); +FI_MEGAMOE_STATIC_ASSERT(kIntermediateHidden / kL2ActsSFGranK <= kIntermediateHidden / 64, + "logical per-128 SF groups must fit retained physical capacity"); + +// ===================================================================== +// Shared memory layout +// ===================================================================== +constexpr uint32_t kSharedMemoryAlignment = 1024; +extern __shared__ __align__(kSharedMemoryAlignment) +uint8_t smem_buffer[]; + +constexpr uint32_t SMEM_EXPERT_COUNT_SIZE = + math::constexpr_align(kNumExperts * sizeof(uint32_t), kSharedMemoryAlignment); +constexpr uint32_t SMEM_SEND_BUFFER_SIZE = math::constexpr_align( + fp8_token_layout.get_num_bytes() * kNumActiveDispatchWarps, kSharedMemoryAlignment); +constexpr uint32_t SMEM_NVFP4_LUT_SIZE = + math::constexpr_align(128u * sizeof(uint2), kSharedMemoryAlignment); +constexpr uint32_t SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K * sizeof(a_dtype_t); +constexpr uint32_t B_LOAD_BYTES_PER_ROW = 80u; +constexpr uint32_t SMEM_B_LOAD_SIZE_PER_STAGE = LOAD_BLOCK_N * B_LOAD_BYTES_PER_ROW; +// Under RS-swapAB nothing ever expands the packed rows to FP8 in SMEM, so +// each B stage shrinks to the 80-byte packed rows (10240 B, still a +// multiple of the 1024 B shared-memory alignment, so the SFA/barrier +// bases keep their alignment). This is safe for TMA: the B tensormap box +// is 80 B x 128 rows with no swizzle, written contiguously at the stage +// base; the 80-byte row stride lives in the tensormap, not in the stage +// stride. NOTE: the host smem-size heuristics are deliberately unchanged; +// the RS device layout uses less shared memory than the host allocates, +// which always fits and keeps num_stages/config selection identical +// between the two arms. +constexpr uint32_t SMEM_B_SIZE_PER_STAGE = + kRSSwapAB ? SMEM_B_LOAD_SIZE_PER_STAGE : LOAD_BLOCK_N * BLOCK_K * sizeof(b_dtype_t); +// Keep a two-slot SFA stage allocation so the six-stage ring retains its +// proven shared-memory/barrier placement. L1 reads only the first per-128 +// slot. +constexpr uint32_t kL2SFAHalfStride = + math::constexpr_align(BLOCK_M * sizeof(float), 128u) / sizeof(float); +constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = 2 * kL2SFAHalfStride * sizeof(float); +// CD output: max of L1 FP8 (BLOCK_M * (BLOCK_N/2) * 1 byte * num_wg) and +// L2 BF16 (BLOCK_M * BLOCK_N * 2 bytes * num_wg). +constexpr uint32_t SMEM_CD_L1_SIZE = + kNumEpilogueWarpgroups * WG_BLOCK_M * WG_L1_OUT_BLOCK_N * sizeof(cutlass::float_e4m3_t); +// RS-swapAB epilogue scratch: per-token amax slots (BLOCK_M tokens x +// kNumEpilogueWarps floats), appended to the CD tile. The DSM peer amax +// slots keep aliasing the CD tile itself, exactly as in the default path. +constexpr uint32_t SMEM_CD_RS_AMAX_SIZE = + kRSSwapAB ? BLOCK_M * kNumEpilogueWarps * static_cast(sizeof(float)) : 0u; +constexpr uint32_t SMEM_CD_OUTPUT_SIZE = + math::constexpr_align(SMEM_CD_L1_SIZE + SMEM_CD_RS_AMAX_SIZE, kSharedMemoryAlignment); +constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_OUTPUT_SIZE; + +// SMEM pointers +auto smem_expert_count = reinterpret_cast(smem_buffer); +const auto smem_send_buffers = + layout::Buffer(fp8_token_layout, kNumActiveDispatchWarps, 1, + math::advance_ptr(smem_buffer, SMEM_EXPERT_COUNT_SIZE)); +auto smem_nvfp4_lut = reinterpret_cast( + math::advance_ptr(smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE)); + +auto smem_gemm_base = math::advance_ptr( + smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_NVFP4_LUT_SIZE); + +auto smem_cd_base = smem_gemm_base; +auto smem_cd_l1 = reinterpret_cast(smem_cd_base); + +auto smem_a = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE); +}); +auto smem_b = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr( + smem_gemm_base, + SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); +}); +auto sf_start_ptr = math::advance_ptr( + smem_gemm_base, SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE)); +auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast(sf_start_ptr + i * SMEM_SFA_SIZE_PER_STAGE); +}); +constexpr uint32_t kNumDequantBarriers = kNumStages; +constexpr uint32_t kNumDequantHalf1Barriers = kDequantHalfStream >= 4 ? kNumStages : 0u; + +// Barriers live after SF. +auto barrier_start_ptr = + reinterpret_cast(sf_start_ptr + kNumStages * SMEM_SFA_SIZE_PER_STAGE); +auto dispatch_barriers = + utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + i; }); +auto full_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + i; }); +auto empty_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + kNumStages + i; }); +auto dequant_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + kNumStages * 2 + i; }); +auto dequant_half1_barriers = utils::PatternVisitor([=](const uint32_t& i) { + return barrier_start_ptr + kNumDispatchWarps + kNumStages * 2 + kNumDequantBarriers + i; +}); +auto combine_barriers = utils::PatternVisitor([=](const uint32_t& i) { + return barrier_start_ptr + kNumDispatchWarps + kNumStages * 2 + kNumDequantBarriers + + kNumDequantHalf1Barriers + i; +}); + +// ===================================================================== +// Initialization +// ===================================================================== +if constexpr (kLutCompact) { + // Compact layout: lo[128] (each row's low word), then the eight + // subnormal-scale rows' high words verbatim -- 544 B of the 1024-B region. + // The high words of rows >= 8 are reconstructed inline by load_nvfp4_lut. + if (thread_idx < 32) { + const auto* lut_src = flashinfer::megamoe::nvfp4::kE2M1AndUe4m3ToFp8Lut + thread_idx * 4; + reinterpret_cast(smem_nvfp4_lut)[thread_idx] = + make_uint4(lut_src[0].x, lut_src[1].x, lut_src[2].x, lut_src[3].x); + if (thread_idx < 8) + reinterpret_cast(smem_nvfp4_lut)[128 + thread_idx] = + flashinfer::megamoe::nvfp4::kE2M1AndUe4m3ToFp8Lut[thread_idx].y; + } +} else if (thread_idx < 64) { + reinterpret_cast(smem_nvfp4_lut)[thread_idx] = + reinterpret_cast(flashinfer::megamoe::nvfp4::kE2M1AndUe4m3ToFp8Lut)[thread_idx]; +} + +if (warp_idx == 0) { +// Clean expert-count shared memory +#pragma unroll + for (uint32_t i = lane_idx; i < kNumExperts; i += 32) ptx::st_shared(smem_expert_count + i, 0u); +} else if (warp_idx == 1) { +// Init dispatch m-barriers +#pragma unroll + for (uint32_t i = lane_idx; i < kNumDispatchWarps; i += 32) dispatch_barriers[i]->init(1); + cutlass::arch::fence_barrier_init(); +} else if (warp_idx == 2) { + // Init GEMM full/empty barriers and combine barriers + if (cute::elect_one_sync()) { +#pragma unroll + for (uint32_t i = 0; i < kNumStages; ++i) { + // Producer arrivals: A(+SFA) + B(TMA+SFB). SFB is copied with + // cp.async.bulk and counted as B-loader transaction bytes, so + // it does not need a separate producer arrival. + full_barriers[i]->init(2); + // With cluster multicast the leader CTA's TMA warp waits on peer + // empty barriers too, so every math warp releases both CTAs. + empty_barriers[i]->init(kClusterSize * kNumEpilogueWarpgroups); + } +#pragma unroll + for (uint32_t i = 0; i < kNumStages; ++i) + dequant_barriers[i]->init(kDequantHalfStream == 5 ? 2 : 1); + if constexpr (kDequantHalfStream >= 4) { +#pragma unroll + for (uint32_t i = 0; i < kNumStages; ++i) + dequant_half1_barriers[i]->init(kDequantHalfStream == 5 ? 2 : 1); + } +#pragma unroll + for (uint32_t i = 0; i < kNumEpilogueWarps * 2; ++i) combine_barriers[i]->init(1); + } + cutlass::arch::fence_barrier_init(); + if (cute::elect_one_sync()) { + // Arm the first DSM amax exchange before cluster peers leave the + // initialization rendezvous. Every tile completes one float + // transaction for each of its BLOCK_M rows. + combine_barriers[0]->arrive_and_expect_tx(BLOCK_M * sizeof(float)); + } +} +cute::cluster_sync(); + +// ===================================================================== +// Cluster-aware L1 scheduler +// ===================================================================== +auto scheduler = + sched::MegaMoEScheduler(workspace); + +// Pipeline state shared by TMA loaders and math warpgroups +uint32_t stage_idx = 0, phase = 0; +auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++k_block_idx; + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; +}; +// Separate named barriers let the two dequant teams publish K stages independently. +constexpr uint32_t kDequantBarrierIdx = 8; +constexpr uint32_t kAlternateDequantBarrierIdx = 9; +// With dispatch-assisted dequant, even and odd K stages have independent +// producer teams which may run concurrently. Give each team its own pair +// of K64 ready barriers. Indices 10..13 do not overlap dispatch/epilogue +// barriers (0..4) or the two producer-only in-place barriers (8..9). +constexpr uint32_t kDequantHalfReadyEven0 = 10; +constexpr uint32_t kDequantHalfReadyEven1 = 11; +constexpr uint32_t kDequantHalfReadyOdd0 = 12; +constexpr uint32_t kDequantHalfReadyOdd1 = 13; +constexpr uint32_t kDequantHalfReadyThreads = 64u + kNumEpilogueThreads; +const auto dequant_loaded_b_stage = [&](const uint32_t& s, const uint32_t& p, + const uint32_t& k_block_idx, + const uint32_t& non_epilogue_thread_idx) { + // RS-swapAB: math lanes decode the packed FP4 straight into A-register + // fragments, no SMEM FP8 expansion exists and no dequant team runs. + // The freed warps deliberately stay idle: the post-change stage floor + // still sits on the SMEM/MIO pipe, and any repurposing would add + // traffic to exactly that pipe (issue-active is only ~10%). + if constexpr (kRSSwapAB) return; + if constexpr (kDispatchDequant) { + if (non_epilogue_thread_idx >= 64u && (k_block_idx & 1u)) { + full_barriers[s]->wait(p); + const uint32_t dequant_tid = non_epilogue_thread_idx - 64u; + if constexpr (kDequantHalfStream == 0) { + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3< + 64u, kAlternateDequantBarrierIdx, true>(reinterpret_cast(smem_b[s]), + dequant_tid, smem_nvfp4_lut); + if (dequant_tid == 0) dequant_barriers[s]->arrive(); + } else if constexpr (kDequantHalfStream == 1) { + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3< + 64u, kAlternateDequantBarrierIdx>(reinterpret_cast(smem_b[s]), dequant_tid, + smem_nvfp4_lut); + cutlass::arch::fence_view_async_shared(); + ptx::sync_unaligned(kDequantHalfReadyThreads, kDequantHalfReadyOdd0); + ptx::sync_unaligned(kDequantHalfReadyThreads, kDequantHalfReadyOdd1); + } else if constexpr (kDequantHalfStream >= 4) { + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3< + 64u, kAlternateDequantBarrierIdx, false, 80u, false, true, kDequantHalfStream == 5, + kDequantHalfStream == 6>(reinterpret_cast(smem_b[s]), dequant_tid, + smem_nvfp4_lut, 0u, 0u, 0u, dequant_barriers[s], + dequant_half1_barriers[s]); + } else { + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3< + 64u, kAlternateDequantBarrierIdx, false, 80u, true>( + reinterpret_cast(smem_b[s]), dequant_tid, smem_nvfp4_lut, + kDequantHalfReadyOdd0, kDequantHalfReadyOdd1, kDequantHalfReadyThreads); + } + } + } else if (non_epilogue_thread_idx >= 64u) { + full_barriers[s]->wait(p); + const uint32_t dequant_tid = non_epilogue_thread_idx - 64u; + if constexpr (kDequantHalfStream == 0) { + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3<64u, + kDequantBarrierIdx>( + reinterpret_cast(smem_b[s]), dequant_tid, smem_nvfp4_lut); + cutlass::arch::fence_view_async_shared(); + ptx::sync_aligned(64, kDequantBarrierIdx); + if (dequant_tid == 0) dequant_barriers[s]->arrive(); + } else if constexpr (kDequantHalfStream == 1) { + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3<64u, + kDequantBarrierIdx>( + reinterpret_cast(smem_b[s]), dequant_tid, smem_nvfp4_lut); + cutlass::arch::fence_view_async_shared(); + ptx::sync_unaligned(kDequantHalfReadyThreads, kDequantHalfReadyEven0); + ptx::sync_unaligned(kDequantHalfReadyThreads, kDequantHalfReadyEven1); + } else if constexpr (kDequantHalfStream >= 4) { + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3< + 64u, kDequantBarrierIdx, false, 80u, false, true, kDequantHalfStream == 5, + kDequantHalfStream == 6>(reinterpret_cast(smem_b[s]), dequant_tid, + smem_nvfp4_lut, 0u, 0u, 0u, dequant_barriers[s], + dequant_half1_barriers[s]); + } else { + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3< + 64u, kDequantBarrierIdx, false, 80u, true>( + reinterpret_cast(smem_b[s]), dequant_tid, smem_nvfp4_lut, + kDequantHalfReadyEven0, kDequantHalfReadyEven1, kDequantHalfReadyThreads); + } + } +}; + +// Intra-SM barrier indices (mirroring SM100) +constexpr uint32_t kDispatchBarrierIdx = 0; +constexpr uint32_t kDispatchWithEpilogueBarrierIdx = 1; +constexpr uint32_t kEpilogueFullBarrierIdx = 2; +constexpr uint32_t kEpilogueWGBarrierStartIdx = 3; + +// Cross-rank NVLink barrier tags +constexpr uint32_t kBeforeDispatchPullBarrierTag = 1; + +// Register reconfiguration counts (chosen to fit in 64512 reg budget). +// Dispatch-assisted dequant keeps the same 63488-register CTA budget: +// 128*80 + 128*80 + 256*168. +constexpr uint32_t kNumDispatchRegisters = kDispatchDequant ? 80 : 48; +constexpr uint32_t kNumNonEpilogueRegisters = kDispatchDequant ? 80 : 64; +constexpr uint32_t kNumEpilogueRegisters = kDispatchDequant ? 168 : 192; +FI_MEGAMOE_STATIC_ASSERT(kNumDispatchRegisters* kNumDispatchThreads + + kNumNonEpilogueRegisters * kNumNonEpilogueThreads + + kNumEpilogueRegisters * kNumEpilogueThreads <= + 64512, + "Too many registers"); + +constexpr uint32_t kDispatchGridSyncIndex = 0; + +const auto for_each_selected_block = [&](auto&& func) { + scheduler.for_each_linear1_block([&](const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx) { + func(std::integral_constant{}, local_expert_idx, + num_k_blocks, m_block_idx, n_block_idx); + }); +}; + +// ===================================================================== +// ROLE 1: DISPATCH WARPS +// Mirrors SM100 dispatch with two changes: +// * SF is per-128 channel float (no UTCCP transpose). We store the +// remote per-token SF directly into the local L1 SF buffer in +// MN-major layout: `local_sf[k_chunk * num_padded_sf_pool_tokens + token_idx]`. +// * The "token_idx_in_expert" → SF token index is now the simple +// per-block linear mapping (no 4×32 transpose). +// ===================================================================== +if (warp_idx < kNumDispatchWarps) { + cutlass::arch::warpgroup_reg_dealloc(); + + FI_MEGAMOE_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of topk"); + constexpr uint32_t kNumActivateLanes = kNumTokensPerWarp * kNumTopk; + const auto read_topk_idx = [&](const auto& process) { + if (warp_idx < kNumActiveDispatchWarps) { +#pragma unroll + for (uint32_t i = (sm_idx * kNumActiveDispatchWarps + warp_idx) * kNumTokensPerWarp; + i < num_tokens; i += kNumSMs * kNumActiveDispatchWarps * kNumTokensPerWarp) { + int expert_idx = -1; + if (i + (lane_idx / kNumTopk) < num_tokens and lane_idx < kNumActivateLanes) { + expert_idx = static_cast( + __ldg(input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); + if (expert_idx >= 0) process(i * kNumTopk + lane_idx, expert_idx); + } + __syncwarp(); + } + } + }; + + // Count tokens per expert + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + atomicAdd_block(smem_expert_count + expert_idx, 1); + }); + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + +// Stake out per-expert SM offsets via global atomic +#pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const uint64_t send_value = (1ull << 32) | static_cast(smem_expert_count[i]); + smem_expert_count[i] = + static_cast(ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value)); + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Write source token-topk indices to remote ranks + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + const auto dst_rank_idx = expert_idx / kNumExpertsPerRank; + const auto dst_slot_idx = atomicAdd_block(smem_expert_count + expert_idx, 1); + const auto dst_ptr = workspace.get_src_token_topk_idx_ptr(expert_idx % kNumExpertsPerRank, + sym_buffer.rank_idx, dst_slot_idx); + *sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx; + }); + + comm::grid_sync(workspace, sm_idx, thread_idx, [=]() { + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + }); + + if (sm_idx == 0 and thread_idx < kNumActiveDispatchThreads) { +#pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumActiveDispatchThreads) { + const auto dst_rank_idx = i / kNumExpertsPerRank; + const auto dst_local_expert_idx = i % kNumExpertsPerRank; + const auto expert_status = *workspace.get_expert_send_count_ptr(i); + *sym_buffer.map( + workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx), + dst_rank_idx) = expert_status & 0xffffffff; + ptx::atomic_add_sys( + sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), + dst_rank_idx), + expert_status); + } + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, false, true); + + // Sync with epilogue warps before pulling tokens + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Token / SF pull loop + if (warp_idx < kNumActiveDispatchWarps) { + uint32_t pull_mbarrier_phase = 0; + const auto pull_buffer = smem_send_buffers.get_rank_buffer(warp_idx).get_data_buffer(0); + const auto pull_mbarrier = dispatch_barriers[warp_idx]; + + scheduler.fetch_expert_recv_count(); + + constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u); + int current_expert_idx = -1; + uint32_t stored_rank_count[kNumRanksPerLane] = {}; + uint32_t expert_start_idx = 0, expert_end_idx = 0; + uint32_t expert_pool_block_offset = 0; + + constexpr uint32_t kNumGlobalWarps = kNumSMs * kNumActiveDispatchWarps; + for (uint32_t token_idx = sm_idx * kNumActiveDispatchWarps + warp_idx;; + token_idx += kNumGlobalWarps) { + int old_expert_idx = current_expert_idx; + while (token_idx >= expert_end_idx) { + if (++current_expert_idx >= kNumExpertsPerRank) break; + expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M); + expert_start_idx = expert_end_idx; + expert_end_idx += scheduler.get_num_tokens(current_expert_idx); + } + if (current_expert_idx >= kNumExpertsPerRank) break; + + if (old_expert_idx != current_expert_idx) { + old_expert_idx = current_expert_idx; +#pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++i) { + const uint32_t j = i * 32 + lane_idx; + stored_rank_count[i] = j < kNumRanks + ? static_cast(*workspace.get_expert_recv_count_ptr( + j, current_expert_idx)) + : 0; + } + } + + // Round-robin rank selection (identical to SM100) + +#include + const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr( + current_expert_idx, current_rank_in_expert_idx, token_idx_in_rank); + const uint32_t src_token_idx = src_token_topk_idx / kNumTopk; + const uint32_t src_topk_idx = src_token_topk_idx % kNumTopk; + + // TMA pull token data into SMEM + if (cute::elect_one_sync()) { + ptx::tma_load_1d( + pull_buffer.get_base_ptr(), + sym_buffer.map(input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx), + pull_mbarrier, kHidden); + } + __syncwarp(); + + // Copy SF: per-128 K floats, written linearly (no UTCCP transpose). + constexpr uint32_t kNumSFFloats = kHidden / 128; + FI_MEGAMOE_STATIC_ASSERT(kNumSFFloats > 0 and kHidden % 128 == 0, "Invalid SF"); + const auto remote_sf_ptr = + sym_buffer.map(input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx); + const auto local_sf_ptr = l1_sf_buffer.get_base_ptr(); + const uint32_t sf_pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; +#pragma unroll + for (uint32_t i = 0; i < math::constexpr_ceil_div(kNumSFFloats, 32u); ++i) { + const uint32_t j = i * 32 + lane_idx; + if (j < kNumSFFloats) + local_sf_ptr[j * kNumPaddedSFPoolTokens + sf_pool_token_idx] = remote_sf_ptr[j]; + } + __syncwarp(); + + const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; + if (cute::elect_one_sync()) { + const auto weight = + *sym_buffer.map(input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, + current_rank_in_expert_idx); + *l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr() = weight; + + ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kHidden); + ptx::mbarrier_wait_and_flip_phase(pull_mbarrier, pull_mbarrier_phase); + + ptx::tma_store_1d(l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr(), + pull_buffer.get_base_ptr(), pull_buffer.get_num_bytes()); + + *workspace.get_token_src_metadata_ptr(pool_token_idx) = {current_rank_in_expert_idx, + src_token_idx, src_topk_idx}; + + cute::tma_store_arrive(); + ptx::tma_store_wait<0>(); + ptx::red_add_rel(workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + + token_idx_in_expert / BLOCK_M), + 1); + } + __syncwarp(); + } + + } else if constexpr (kDispatchDequant) { + const uint32_t dequant_tid = (warp_idx - kNumActiveDispatchWarps) * 32 + lane_idx; + for_each_selected_block([&](const auto&, const uint32_t&, const uint32_t& num_k_blocks, + const uint32_t&, const uint32_t&) { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + if ((k_block_idx & 1u) == 0) { + full_barriers[stage_idx]->wait(phase); + if constexpr (kDequantHalfStream == 0) { + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3< + 64u, kDequantBarrierIdx, true>(reinterpret_cast(smem_b[stage_idx]), + dequant_tid, smem_nvfp4_lut); + if (dequant_tid == 0) dequant_barriers[stage_idx]->arrive(); + } else if constexpr (kDequantHalfStream == 1) { + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3< + 64u, kDequantBarrierIdx>(reinterpret_cast(smem_b[stage_idx]), dequant_tid, + smem_nvfp4_lut); + cutlass::arch::fence_view_async_shared(); + ptx::sync_unaligned(kDequantHalfReadyThreads, kDequantHalfReadyEven0); + ptx::sync_unaligned(kDequantHalfReadyThreads, kDequantHalfReadyEven1); + } else if constexpr (kDequantHalfStream >= 4) { + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3< + 64u, kDequantBarrierIdx, false, 80u, false, true, kDequantHalfStream == 5, + kDequantHalfStream == 6>(reinterpret_cast(smem_b[stage_idx]), dequant_tid, + smem_nvfp4_lut, 0u, 0u, 0u, dequant_barriers[stage_idx], + dequant_half1_barriers[stage_idx]); + } else { + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3< + 64u, kDequantBarrierIdx, false, 80u, true>( + reinterpret_cast(smem_b[stage_idx]), dequant_tid, smem_nvfp4_lut, + kDequantHalfReadyEven0, kDequantHalfReadyEven1, kDequantHalfReadyThreads); + } + } + } + }); + } + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // ===================================================================== + // ROLE 2: GEMM TMA LOAD warps (load A+SFA, B+SFB) + // Default logical mapping is loaders 0/1 and odd-K dequant 2/3. + // The scheduler-remap experiment swaps those physical warp pairs while + // preserving the same logical TMA/dequant work and register budgets. + // ===================================================================== +} else if (warp_idx == kNumDispatchWarps + (kDequantWarpRemap ? 2u : 0u)) { + cutlass::arch::warpgroup_reg_dealloc(); + + for_each_selected_block([&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx) { + const auto tensor_map_a_ptr = [&]() { return &tensor_map_l1_acts; }(); + const auto tensor_map_sfa_ptr = [&]() { return &tensor_map_l1_acts_sf; }(); + + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + const uint32_t valid_m = scheduler.template get_valid_m(); + const bool has_valid_m = valid_m > 0; + + // Wait for the pool to be ready. Cluster peers can be dummy CTAs for + // the tail M unit when an expert has an odd number of M blocks. + if (has_valid_m) { + const auto ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx); + const auto expected = valid_m; + while (ptx::ld_acq(ptr) != expected); + } + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + if (cute::elect_one_sync()) { + if (has_valid_m) { + const uint32_t m_idx = pool_block_idx * BLOCK_M; + const uint32_t k_idx = k_block_idx * BLOCK_K; + + // Paired N CTAs share the same M/K coordinates, so the + // cluster leader multicasts A and SFA to both CTAs. + tma::copy( + tensor_map_a_ptr, full_barriers[stage_idx], smem_a[stage_idx], k_idx, m_idx, + kClusterSize); + + // TMA load SFA + + // L1 SFA per-128: load (BLOCK_M, 1) at K=k_block_idx + tma::copy(tensor_map_sfa_ptr, full_barriers[stage_idx], + smem_sfa[stage_idx], m_idx, k_block_idx, kClusterSize); + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + + BLOCK_M * sizeof(float)); + + } else { + full_barriers[stage_idx]->arrive(); + } + } + __syncwarp(); + dequant_loaded_b_stage(stage_idx, phase, k_block_idx, lane_idx); + } + }); + +} else if (warp_idx == kNumDispatchWarps + (kDequantWarpRemap ? 3u : 1u)) { + cutlass::arch::warpgroup_reg_dealloc(); + + for_each_selected_block([&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx) { + const auto tensor_map_b_ptr = [&]() { return &tensor_map_l1_weights; }(); + constexpr uint32_t shape_n = L1_SHAPE_N; + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + const uint32_t n_idx = local_expert_idx * shape_n + n_block_idx * BLOCK_N; + // NVFP4 fused B+scale layout stores 64B packed FP4 + 16B + // UE4M3 scale per BK128 row. + const uint32_t k_idx = k_block_idx * B_LOAD_BYTES_PER_ROW; + if (cute::elect_one_sync()) { + // The peer owns the adjacent N block, therefore B must be + // loaded independently rather than multicast. + tma::copy( + tensor_map_b_ptr, full_barriers[stage_idx], smem_b[stage_idx], k_idx, n_idx, 1); + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_B_LOAD_SIZE_PER_STAGE); + } + __syncwarp(); + dequant_loaded_b_stage(stage_idx, phase, k_block_idx, 32u + lane_idx); + } + }); + +} else if (warp_idx < kNumDispatchWarps + kNumMMANonEpilogueWarps) { + // The two remaining non-epilogue warps form the odd-K dequant team. + // Default: physical warps 6/7 (logical 2/3). Remap: physical warps + // 4/5 retain logical dequant IDs 2/3 while the TMA loaders move to 6/7. + // All four warps still participate in the warpgroup-collective + // `setmaxnreg.dec.sync.aligned` required by the math warpgroups. + cutlass::arch::warpgroup_reg_dealloc(); + + const uint32_t non_epilogue_warp_idx = warp_idx - kNumDispatchWarps; + const uint32_t dequant_warp_idx = non_epilogue_warp_idx + (kDequantWarpRemap ? 2u : 0u); + const uint32_t non_epilogue_thread_idx = dequant_warp_idx * 32 + lane_idx; + for_each_selected_block([&](const auto& block_phase, const uint32_t&, + const uint32_t& num_k_blocks, const uint32_t&, const uint32_t&) { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + dequant_loaded_b_stage(stage_idx, phase, k_block_idx, non_epilogue_thread_idx); + __syncwarp(); + } + }); +} else if (warp_idx >= kNumDispatchWarps + kNumMMANonEpilogueWarps) { + // ===================================================================== + // ROLE 3: MATH WARPGROUPS (WGMMA + epilogue + combine) + // ===================================================================== + cutlass::arch::warpgroup_reg_alloc(); + + const uint32_t epilogue_warp_idx = warp_idx - (kNumDispatchWarps + kNumMMANonEpilogueWarps); + const uint32_t epilogue_wg_idx = epilogue_warp_idx / 4; + const uint32_t epilogue_thread_idx = epilogue_warp_idx * 32 + lane_idx; + const uint32_t warp_idx_in_wg = epilogue_warp_idx % 4; + + const auto arrive_empty_barrier = [&](const uint32_t& s) { + if (warp_idx_in_wg != 0) return; + if (lane_idx < 2) empty_barriers[s]->arrive(lane_idx); + }; + + // WGMMA-output register layout helpers + const uint32_t row_idx = lane_idx / 4; + const uint32_t col_idx = lane_idx % 4; + const uint32_t r_0 = warp_idx_in_wg * 16 + row_idx; + const uint32_t r_1 = r_0 + 8; + + FI_MEGAMOE_STATIC_ASSERT(WG_BLOCK_M == L1WGMMA::M, + "Each warpgroup must run exactly one WGMMA per K-block"); + + // Sync with dispatch + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + uint32_t pair_scale_phase = 0; + + for_each_selected_block([&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx) { + const uint32_t valid_m = scheduler.template get_valid_m(); + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + const uint32_t m_idx = pool_block_idx * BLOCK_M; + constexpr uint32_t wg_n_idx = 0; + constexpr uint32_t wg_l1_out_n_idx = 0; + const uint32_t n_idx = n_block_idx * BLOCK_N + wg_n_idx; + const uint32_t row_block_offset = epilogue_wg_idx * WG_BLOCK_M; + const uint32_t row_offset_r0 = row_block_offset + r_0; + const uint32_t row_offset_r1 = row_block_offset + r_1; + const bool valid_r0 = row_offset_r0 < valid_m; + const bool valid_r1 = row_offset_r1 < valid_m; + // Under RS-swapAB each warpgroup owns 64 weight rows that are + // valid for every valid token, so a warpgroup may only drain when + // the whole CTA is a token-less dummy cluster peer (valid_m == 0), + // NOT when its per-WG row offset exceeds valid_m. + const bool inactive_math_wg = kRSSwapAB ? (valid_m == 0) : (row_block_offset >= valid_m); + + if (inactive_math_wg) { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + // RS-swapAB gates directly on the producer barrier: the + // dequant barriers stay allocated and initialized but are + // never armed in that arm. + if constexpr (kRSSwapAB) + full_barriers[stage_idx]->wait(phase); + else if constexpr (kDequantHalfStream == 0) + dequant_barriers[stage_idx]->wait(phase); + else { + full_barriers[stage_idx]->wait(phase); + if constexpr (kDequantHalfStream >= 4) { + dequant_barriers[stage_idx]->wait(phase); + dequant_half1_barriers[stage_idx]->wait(phase); + } else { + const bool use_odd_ready = kDispatchDequant and (k_block_idx & 1u); + ptx::sync_unaligned(kDequantHalfReadyThreads, + use_odd_ready ? kDequantHalfReadyOdd0 : kDequantHalfReadyEven0); + ptx::sync_unaligned(kDequantHalfReadyThreads, + use_odd_ready ? kDequantHalfReadyOdd1 : kDequantHalfReadyEven1); + } + } + arrive_empty_barrier(stage_idx); + __syncwarp(); + } + } + + // ---------------- GEMM ---------------- + using WGMMA = L1WGMMA; + constexpr uint32_t kAccumPerThread = WGMMA::kNumAccum; // 64 for M=64,N=128 + float final_accum[kAccumPerThread] = {}; + // In single-accumulator mode `final_accum` IS the WGMMA accumulator and + // `accum` is unused (and optimised away); the running total is held in units + // of `unit_scale_*`, the most recent NON-ZERO activation SF for each of the + // thread's two token rows. + float accum[kSingleAccum ? 1 : kAccumPerThread]; + float unit_scale_0 = 0.0f, unit_scale_1 = 0.0f; + + const auto run_default_gemm_loop = [&]() { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + const bool use_odd_ready = kDispatchDequant and (k_block_idx & 1u); + const uint32_t half_ready_barrier_0 = + use_odd_ready ? kDequantHalfReadyOdd0 : kDequantHalfReadyEven0; + const uint32_t half_ready_barrier_1 = + use_odd_ready ? kDequantHalfReadyOdd1 : kDequantHalfReadyEven1; + // The full barrier tracks TMA A + SFA + packed-B completion. + // Its successful wait is an acquire, so SFA is safe to read. + // The dequant barrier is a later acquire that publishes the + // generic-proxy FP8 stores into smem_b. Modes 2 and 4 use + // the gap between those dependency edges for useful work. + // Mode 1 executes the same extra full wait but keeps the + // original ordering, isolating wait/control perturbation. + if constexpr (kDequantHalfStream != 0) + full_barriers[stage_idx]->wait(phase); + else if constexpr (kSFALookahead != 0) + full_barriers[stage_idx]->wait(phase); + if constexpr (kDequantHalfStream == 0 and kSFALookahead != 2) + dequant_barriers[stage_idx]->wait(phase); + + // Read SF (must precede warpgroup_arrive). In mode 2 these + // loads can issue while the producer warps dequantize B. + const float scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); + const float scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); + + if constexpr (kDequantHalfStream == 0 and kSFALookahead == 2) { + // Do not issue WGMMA before this acquire: smem_b is + // overwritten through the generic proxy by all 64 + // dequant threads, and their writer rendezvous precedes + // the sole release arrival on dequant_barriers[s]. + dequant_barriers[stage_idx]->wait(phase); + } + + // NVFP4 UE4M3 weight scales are applied during FP4 -> FP8 + // smem expansion, so WGMMA only needs activation SF. + + if constexpr (kSingleAccum) { + // Rescale the running total from the previous block's units into this + // block's. A zero SF means the token's activations over this k range + // are all zero, so the block contributes nothing and the units must + // NOT change -- carrying the previous unit forward instead of dividing + // is what keeps a zero SF from producing inf. When the unit is still + // zero the running total is zero too, so a ratio of 1 is correct. + const float u_0 = scale_a_0 != 0.0f ? scale_a_0 : unit_scale_0; + const float u_1 = scale_a_1 != 0.0f ? scale_a_1 : unit_scale_1; + if (k_block_idx > 0) { + const float r_0 = u_0 != 0.0f ? unit_scale_0 / u_0 : 1.0f; + const float r_1 = u_1 != 0.0f ? unit_scale_1 / u_1 : 1.0f; +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++i) { + final_accum[i * 4 + 0] *= r_0; + final_accum[i * 4 + 1] *= r_0; + final_accum[i * 4 + 2] *= r_1; + final_accum[i * 4 + 3] *= r_1; + } + } + unit_scale_0 = u_0; + unit_scale_1 = u_1; + +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) + ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++k) { + if constexpr (kDequantHalfStream >= 4) { + if (k == 0) + dequant_barriers[stage_idx]->wait(phase); + else if (k == 2) + dequant_half1_barriers[stage_idx]->wait(phase); + } else if constexpr (kDequantHalfStream != 0) { + if (k == 0) { + ptx::sync_unaligned(kDequantHalfReadyThreads, half_ready_barrier_0); + if constexpr (kDequantHalfStream == 3) + ptx::sync_unaligned(kDequantHalfReadyThreads, half_ready_barrier_1); + } else if constexpr (kDequantHalfStream != 3) { + if (k == 2) ptx::sync_unaligned(kDequantHalfReadyThreads, half_ready_barrier_1); + } + } + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + row_block_offset * BLOCK_K + k * WGMMA::K, 1); + auto desc_b = + mma::sm90::make_smem_desc(smem_b[stage_idx] + wg_n_idx * BLOCK_K + k * WGMMA::K, 1); + // Accumulate in place across the whole K loop: only the very + // first WGMMA of the very first block may overwrite. + WGMMA::wgmma(desc_a, desc_b, final_accum, k > 0 or k_block_idx > 0); + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) + ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + arrive_empty_barrier(stage_idx); + } else { +// Single per-128 K-block WGMMA group +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++k) { + if constexpr (kDequantHalfStream >= 4) { + if (k == 0) + dequant_barriers[stage_idx]->wait(phase); + else if (k == 2) + dequant_half1_barriers[stage_idx]->wait(phase); + } else if constexpr (kDequantHalfStream != 0) { + if (k == 0) { + ptx::sync_unaligned(kDequantHalfReadyThreads, half_ready_barrier_0); + if constexpr (kDequantHalfStream == 3) + ptx::sync_unaligned(kDequantHalfReadyThreads, half_ready_barrier_1); + } else if constexpr (kDequantHalfStream != 3) { + if (k == 2) ptx::sync_unaligned(kDequantHalfReadyThreads, half_ready_barrier_1); + } + } + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + row_block_offset * BLOCK_K + k * WGMMA::K, 1); + auto desc_b = + mma::sm90::make_smem_desc(smem_b[stage_idx] + wg_n_idx * BLOCK_K + k * WGMMA::K, + 1); // NVFP4: no swizzle on B + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + arrive_empty_barrier(stage_idx); + +// L1: gate/up alternate at gran=8 along N; each `i` block of 8 +// cols belongs entirely to one of {gate, up}, so .x and .y +// share the same scalar. +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++i) { + final_accum[i * 4 + 0] += scale_a_0 * accum[i * 4 + 0]; + final_accum[i * 4 + 1] += scale_a_0 * accum[i * 4 + 1]; + final_accum[i * 4 + 2] += scale_a_1 * accum[i * 4 + 2]; + final_accum[i * 4 + 3] += scale_a_1 * accum[i * 4 + 3]; + } + } + } + }; + + // RS-swapAB main loop: math gates directly on `full_barriers` + // (packed-FP4 arrival AND activation/SFA arrival), decodes the + // packed weights into WGMMA A-register fragments per k32 slice, + // and runs weights-as-A / activations-as-B WGMMAs. No FP8 SMEM + // write-back, no WGMMA weight re-read. + const auto run_rs_swapab_gemm_loop = [&]() { + using WGMMARS = L1WGMMARS; + FI_MEGAMOE_STATIC_ASSERT(WGMMARS::kNumAccum == kAccumPerThread, + "RS and SS accumulator budgets must match"); + FI_MEGAMOE_STATIC_ASSERT(BLOCK_K / WGMMARS::K == 4, + "RS loop expects four k32 slices per stage"); + // Proven A-fragment map for wgmma.m64n128k32 e4m3: per k32 + // slice, thread (warp w, lane l) holds rows 16w + l/4 and + // 16w + l/4 + 8 of its WG's 64-row A tile at K nibble-group + // 4*(l%4) (+16 for a2/a3). Lane pairs (l ^ 1) share packed + // words: the even lane decodes row r, the odd lane row r + 8, + // and the discarded uint2 halves are exchanged over shfl so + // every word is decoded exactly once. + const uint32_t frag_row0 = epilogue_wg_idx * WG_BLOCK_M + warp_idx_in_wg * 16 + row_idx; + const uint32_t decode_row = frag_row0 + ((lane_idx & 1u) << 3); + const uint32_t word_sel = (lane_idx >> 1) & 1u; + const bool keep_hi = (lane_idx & 1u) == 0; + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + // Producer arrivals: TMA-B (packed FP4 + fused scales) and + // TMA-A (+SFA), both counted on the same barrier. + full_barriers[stage_idx]->wait(phase); + + const auto* packed_base = reinterpret_cast(smem_b[stage_idx]); +#if FI_MEGAMOE_NVFP4_L1_RS_DEBUG_DECODE + const auto* prow_r0 = packed_base + frag_row0 * B_LOAD_BYTES_PER_ROW; + const auto* prow_r8 = prow_r0 + 8 * B_LOAD_BYTES_PER_ROW; + const uint2 scale_words_r0 = *reinterpret_cast(prow_r0 + 64); + const uint2 scale_words_r8 = *reinterpret_cast(prow_r8 + 64); +#else + const auto* prow = packed_base + decode_row * B_LOAD_BYTES_PER_ROW; + const uint2 scale_words = *reinterpret_cast(prow + 64); +#endif + + // Slice `s` consumes packed words 4s..4s+3; this lane's + // words are 4s + word_sel (K-lo half, UE4M3 scale byte 2s) + // and 4s + 2 + word_sel (K-hi half, scale byte 2s + 1). + uint32_t a_frag[2][4]; + const auto decode_slice = [&](const uint32_t& s, uint32_t(&a)[4]) { +#if FI_MEGAMOE_NVFP4_L1_RS_DEBUG_DECODE + const uint4 q_r0 = reinterpret_cast(prow_r0)[s]; + const uint4 q_r8 = reinterpret_cast(prow_r8)[s]; + const uint32_t sw_r0 = s < 2 ? scale_words_r0.x : scale_words_r0.y; + const uint32_t sw_r8 = s < 2 ? scale_words_r8.x : scale_words_r8.y; + flashinfer::megamoe::nvfp4::dequant_mode2_lop3_rs_word_pair_noshfl( + word_sel ? q_r0.y : q_r0.x, word_sel ? q_r0.w : q_r0.z, word_sel ? q_r8.y : q_r8.x, + word_sel ? q_r8.w : q_r8.z, smem_nvfp4_lut[(sw_r0 >> ((s & 1u) * 16u)) & 0x7fu], + smem_nvfp4_lut[(sw_r0 >> ((s & 1u) * 16u + 8u)) & 0x7fu], + smem_nvfp4_lut[(sw_r8 >> ((s & 1u) * 16u)) & 0x7fu], + smem_nvfp4_lut[(sw_r8 >> ((s & 1u) * 16u + 8u)) & 0x7fu], keep_hi, a); +#else + const uint4 q = reinterpret_cast(prow)[s]; + const uint32_t sw = s < 2 ? scale_words.x : scale_words.y; + flashinfer::megamoe::nvfp4::dequant_mode2_lop3_rs_word_pair( + word_sel ? q.y : q.x, word_sel ? q.w : q.z, + smem_nvfp4_lut[(sw >> ((s & 1u) * 16u)) & 0x7fu], + smem_nvfp4_lut[(sw >> ((s & 1u) * 16u + 8u)) & 0x7fu], keep_hi, a); +#endif + }; + + decode_slice(0, a_frag[0]); +#pragma unroll + for (uint32_t s = 0; s < BLOCK_K / WGMMARS::K; ++s) { +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) mma::sm90::warpgroup_fence_operand(a_frag[s & 1u][i]); +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + // wgmma.fence is mandatory before EVERY wgmma here: + // the A registers are rewritten between groups (an + // accumulator-only chain would need only one). + ptx::warpgroup_arrive(); + // Swapped-role B descriptor: the full 128-token + // activation tile, B128 swizzle, no row offset (both + // WGs read the whole tile, the swapped-role cost). + const auto desc_b = mma::sm90::make_smem_desc(smem_a[stage_idx] + s * WGMMARS::K, 1); + WGMMARS::wgmma(a_frag[s & 1u], desc_b, accum, s > 0); + ptx::warpgroup_commit_batch(); + if (s + 1 < BLOCK_K / WGMMARS::K) { +#if FI_MEGAMOE_NVFP4_L1_RS_DEBUG_DECODE + // Debug arm: fully serialize the WGMMA groups. + ptx::warpgroup_wait<0>(); +#else + // Retire group s-1 before rewriting its A-fragment + // buffer (the load-bearing register-WAR gate); + // group s stays in flight under the next decode. + if (s >= 1) ptx::warpgroup_wait<1>(); +#endif +// Pin the retired buffer: the compiler must not sink +// the decode's register writes above the wait that +// retires the wgmma group still reading them. +#pragma unroll + for (uint32_t r = 0; r < 4; ++r) ptx::warpgroup_fence_operand(a_frag[(s + 1) & 1u][r]); + decode_slice(s + 1, a_frag[(s + 1) & 1u]); + } + } +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + // WAR-2: all groups retired, so the async smem_a + // descriptor reads are finished. WAR-1: group completion + // is warpgroup-collective, so every lane issued its wgmma + // and, by register data dependence through a_frag, every + // packed smem_b staging LDS has retired. + ptx::warpgroup_wait<0>(); + +// Per-token activation SF apply (the transposed axis: SFA +// now scales WGMMA N columns). Must complete before the +// empty release below: TMA-A rewrites smem_sfa together +// with smem_a under ring wraparound (the third WAR +// hazard). This mirrors the fused swapAB precedent. +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++i) { + const uint32_t token_0 = i * 8 + col_idx * 2; + const float2 sfa = + ptx::ld_shared(reinterpret_cast(smem_sfa[stage_idx] + token_0)); + final_accum[i * 4 + 0] += sfa.x * accum[i * 4 + 0]; + final_accum[i * 4 + 1] += sfa.y * accum[i * 4 + 1]; + final_accum[i * 4 + 2] += sfa.x * accum[i * 4 + 2]; + final_accum[i * 4 + 3] += sfa.y * accum[i * 4 + 3]; + } + + // The release is issued by warp 0 lanes 0-1 only, while the + // SFA loads above run per-thread across the whole warpgroup: + // without a rendezvous, warp 0 can release the stage while + // warps 1-3 are still reading smem_sfa. Both adversarial + // reviews found this independently. + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + // Release LAST, after the SFA reads (NOT at the default + // path's position right after the WGMMA wait). + arrive_empty_barrier(stage_idx); + } + }; + + if (!inactive_math_wg) { + if constexpr (kRSSwapAB) + run_rs_swapab_gemm_loop(); + else + run_default_gemm_loop(); + } + if constexpr (kSingleAccum) { +// The running total is in units of the last non-zero SF per row; restore +// absolute units. An inactive warpgroup never entered the loop, so its +// unit scales are still zero and its zero accumulators stay zero. +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++i) { + final_accum[i * 4 + 0] *= unit_scale_0; + final_accum[i * 4 + 1] *= unit_scale_0; + final_accum[i * 4 + 2] *= unit_scale_1; + final_accum[i * 4 + 3] *= unit_scale_1; + } + } + + // Even a fully invalid tail warpgroup must participate in the + // CTA-wide scale rendezvous. It carries zero accumulators and only + // writes padding rows, which L2 never scatters. + + const float l1_global_scale = + l1_global_scales == nullptr ? 1.0f : __ldg(l1_global_scales + local_expert_idx); + if constexpr (kRSSwapAB) { + // ---------------- RS-swapAB L1 EPILOGUE ---------------- + // D is transposed relative to the default path: WGMMA M is 64 + // weight rows, WGMMA N is 128 tokens. Gate/up alternate at + // granularity 8 along the weight rows, so each thread's row + // pair (r_0 = gate, r_1 = r_0 + 8 = up) is exactly one SwiGLU + // pair, and chunk `i` covers tokens i*8 + col_idx*2 + {0,1}. + // Weight rows 16g..16g+15 of this 128-row block produce + // output I-columns 8g..8g+7, hence: + const uint32_t out_col = epilogue_wg_idx * 32 + warp_idx_in_wg * 8 + row_idx; + constexpr uint32_t kNumTokenChunks = kAccumPerThread / 4; + auto* smem_rs_amax = math::advance_ptr(smem_cd_base, SMEM_CD_L1_SIZE); + auto* l2_sf_base = l2_sf_buffer.get_base_ptr(); + const uint32_t dense_sf_group_idx = n_block_idx / 2; + const uint32_t peer_cta_rank = cute::block_rank_in_cluster() ^ 1u; + + auto silu = [](float x) -> float { + const float e = kFastMath ? __expf(-x) : expf(-x); + const float sig = kFastMath ? math::fast_rcp(1.0f + e) : 1.0f / (1.0f + e); + return x * sig; + }; + auto clamp_gate = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(x, kActivationClamp); + }; + auto clamp_up = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(cute::max(x, -kActivationClamp), kActivationClamp); + }; + + // SwiGLU + per-token amax. Invalid tokens contribute exact + // zeros (their accumulators may hold NaN/Inf from garbage SFA + // in padding rows), so the DSM exchange below never ships + // garbage: a dummy peer CTA sends all-zero amax. + float v_0[kNumTokenChunks], v_1[kNumTokenChunks]; +#pragma unroll + for (uint32_t i = 0; i < kNumTokenChunks; ++i) { + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + v_0[i] = 0.0f; + v_1[i] = 0.0f; + if (token_0 < valid_m) { + float g = final_accum[i * 4 + 0] * l1_global_scale; + clamp_gate(g); + float u = final_accum[i * 4 + 2] * l1_global_scale; + clamp_up(u); + const float weight = *l1_topk_weights_buffer.get_data_buffer(m_idx + token_0) + .template get_base_ptr(); + v_0[i] = silu(g) * u * weight; + } + if (token_1 < valid_m) { + float g = final_accum[i * 4 + 1] * l1_global_scale; + clamp_gate(g); + float u = final_accum[i * 4 + 3] * l1_global_scale; + clamp_up(u); + const float weight = *l1_topk_weights_buffer.get_data_buffer(m_idx + token_1) + .template get_base_ptr(); + v_1[i] = silu(g) * u * weight; + } + // Reduce over the 8 lanes sharing col_idx (the 8 I-columns + // this warp owns for these two tokens). + const float amax_0 = + math::warp_reduce<4, true>(cute::abs(v_0[i]), math::ReduceMax()); + const float amax_1 = + math::warp_reduce<4, true>(cute::abs(v_1[i]), math::ReduceMax()); + if (row_idx == 0) { + smem_rs_amax[token_0 * kNumEpilogueWarps + epilogue_warp_idx] = amax_0; + smem_rs_amax[token_1 * kNumEpilogueWarps + epilogue_warp_idx] = amax_1; + } + } + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + + // Cross-warp per-token reduction and adjacent-N cluster DSM + // exchange. Token t is owned by epilogue thread t; even + // threads ship one float2 token pair (64 senders x 8 B = + // 512 B), matching the pre-armed BLOCK_M * sizeof(float) + // transaction. The peer slots alias the CD tile exactly as in + // the default path: the exchange completes before the + // quantize below overwrites the tile. + float token_amax = 0.0f; + if (epilogue_thread_idx < BLOCK_M) { +#pragma unroll + for (uint32_t w = 0; w < kNumEpilogueWarps; ++w) + token_amax = + cute::max(token_amax, smem_rs_amax[epilogue_thread_idx * kNumEpilogueWarps + w]); + const float token_amax_next = __shfl_down_sync(0xffffffff, token_amax, 1); + if ((epilogue_thread_idx & 1u) == 0) { + ptx::st_shared_cluster_async( + reinterpret_cast(smem_cd_l1) + epilogue_thread_idx / 2, + make_float2(token_amax, token_amax_next), combine_barriers[0], peer_cta_rank); + } + } + if (epilogue_thread_idx == 0) { + combine_barriers[0]->wait(pair_scale_phase); + pair_scale_phase ^= 1u; + // Pre-arm the next phase. The current quantize/store tail + // and the following GEMM separate this from peer stores. + combine_barriers[0]->arrive_and_expect_tx(BLOCK_M * sizeof(float)); + } + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + + // Both CTAs derive the same per-token scale; only the even-N + // owner publishes the dense logical per-128 SF entry. + if (epilogue_thread_idx < BLOCK_M) { + const uint32_t token = epilogue_thread_idx; + const float peer_amax = ptx::ld_shared(reinterpret_cast(smem_cd_l1) + token); + const float common_amax = cute::max(token_amax, peer_amax); + float2 sf_pair, sf_inv_pair; + math::get_e4m3_sf_and_sf_inv(make_float2(common_amax, common_amax), sf_pair, sf_inv_pair); + if ((n_block_idx & 1u) == 0 and token < valid_m) + l2_sf_base[dense_sf_group_idx * kNumPaddedSFPoolTokens + pool_block_idx * BLOCK_M + + token] = sf_pair.x; + smem_rs_amax[token * kNumEpilogueWarps] = sf_inv_pair.x; + } + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + +// Quantize and write to smem_cd_l1 ([token][I], row-major, no +// swizzle) so L2 consumes the pool buffer unchanged. +#pragma unroll + for (uint32_t i = 0; i < kNumTokenChunks; ++i) { + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + if (token_0 < valid_m) { + const float sf_inv = smem_rs_amax[token_0 * kNumEpilogueWarps]; + const __nv_fp8_e4m3 q(v_0[i] * sf_inv); + reinterpret_cast(smem_cd_l1)[token_0 * L1_OUT_BLOCK_N + out_col] = + *reinterpret_cast(&q); + } + if (token_1 < valid_m) { + const float sf_inv = smem_rs_amax[token_1 * kNumEpilogueWarps]; + const __nv_fp8_e4m3 q(v_1[i] * sf_inv); + reinterpret_cast(smem_cd_l1)[token_1 * L1_OUT_BLOCK_N + out_col] = + *reinterpret_cast(&q); + } + } + + // Both warpgroups wrote into both token halves, so the store + // must wait on the whole CTA rather than the default path's + // per-WG barrier. WG wg's CD sub-tile still holds token rows + // [64*wg, 64*wg + 64) x 64 I-columns, which is exactly what + // the unchanged per-WG TMA store below expects. + auto* smem_cd_l1_wg = smem_cd_l1 + epilogue_wg_idx * WG_BLOCK_M * L1_OUT_BLOCK_N; + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { + const uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N; + cute::tma_store_fence(); + cute::SM90_TMA_STORE_2D::copy(&tensor_map_l1_output, smem_cd_l1_wg, out_n_idx, + m_idx + row_block_offset); + cute::tma_store_arrive(); + } + __syncwarp(); + ptx::tma_store_wait<0>(); + if constexpr (!kL2ArrivalCounter) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } else { + // ---------------- L1 EPILOGUE: SwiGLU + FP8 quantize + TMA store ---------------- + // Layout in `final_accum`: + // 16 chunks of 8 N-cols, each chunk = 4 floats per thread = (r0c0, r0c1, r1c0, r1c1). + // Gate chunks: even (0, 2, ..., 14). Up chunks: odd (1, 3, ..., 15). + // Pair `p` ∈ [0, 8): gate chunk = 2p, up chunk = 2p+1. + // + // For each pair we produce 4 post-SwiGLU floats per thread, mapped to + // output cols (p*8 + col_idx*2 + {0,1}) for both r0 and r1. + + constexpr uint32_t kNumPairs = kAccumPerThread / 8; + constexpr uint32_t kNumSFGroups = WG_L1_OUT_BLOCK_N / kLocalL1ActsSFGranK; + FI_MEGAMOE_STATIC_ASSERT(kNumSFGroups == 1, + "each paired CTA must own one local 64-column SF group"); + float swiglu_r0[kNumPairs][2]; + float swiglu_r1[kNumPairs][2]; + + // Each CTA computes one local 64-column amax for all 128 rows. + // Its adjacent-N cluster peer owns the other half. + float amax_r0[kNumSFGroups] = {}; + float amax_r1[kNumSFGroups] = {}; + +// Compute SwiGLU + per-group amax. +#pragma unroll + for (uint32_t p = 0; p < kNumPairs; ++p) { + const uint32_t gate = 2 * p, up = 2 * p + 1; + const uint32_t sf_group = p / 8; + + auto clamp_gate = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(x, kActivationClamp); + }; + auto clamp_up = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(cute::max(x, -kActivationClamp), kActivationClamp); + }; + float g_r0_c0 = final_accum[gate * 4 + 0] * l1_global_scale; + clamp_gate(g_r0_c0); + float g_r0_c1 = final_accum[gate * 4 + 1] * l1_global_scale; + clamp_gate(g_r0_c1); + float g_r1_c0 = final_accum[gate * 4 + 2] * l1_global_scale; + clamp_gate(g_r1_c0); + float g_r1_c1 = final_accum[gate * 4 + 3] * l1_global_scale; + clamp_gate(g_r1_c1); + float u_r0_c0 = final_accum[up * 4 + 0] * l1_global_scale; + clamp_up(u_r0_c0); + float u_r0_c1 = final_accum[up * 4 + 1] * l1_global_scale; + clamp_up(u_r0_c1); + float u_r1_c0 = final_accum[up * 4 + 2] * l1_global_scale; + clamp_up(u_r1_c0); + float u_r1_c1 = final_accum[up * 4 + 3] * l1_global_scale; + clamp_up(u_r1_c1); + + auto silu = [](float x) -> float { + const float e = kFastMath ? __expf(-x) : expf(-x); + const float sig = kFastMath ? math::fast_rcp(1.0f + e) : 1.0f / (1.0f + e); + return x * sig; + }; + + if (valid_r0) { + swiglu_r0[p][0] = silu(g_r0_c0) * u_r0_c0; + swiglu_r0[p][1] = silu(g_r0_c1) * u_r0_c1; + amax_r0[sf_group] = cute::max( + amax_r0[sf_group], cute::max(cute::abs(swiglu_r0[p][0]), cute::abs(swiglu_r0[p][1]))); + } else { + swiglu_r0[p][0] = 0.0f; + swiglu_r0[p][1] = 0.0f; + } + if (valid_r1) { + swiglu_r1[p][0] = silu(g_r1_c0) * u_r1_c0; + swiglu_r1[p][1] = silu(g_r1_c1) * u_r1_c1; + amax_r1[sf_group] = cute::max( + amax_r1[sf_group], cute::max(cute::abs(swiglu_r1[p][0]), cute::abs(swiglu_r1[p][1]))); + } else { + swiglu_r1[p][0] = 0.0f; + swiglu_r1[p][1] = 0.0f; + } + } + + const float weight_r0 = valid_r0 + ? *l1_topk_weights_buffer.get_data_buffer(m_idx + row_offset_r0) + .template get_base_ptr() + : 0.0f; + const float weight_r1 = valid_r1 + ? *l1_topk_weights_buffer.get_data_buffer(m_idx + row_offset_r1) + .template get_base_ptr() + : 0.0f; +#pragma unroll + for (uint32_t p = 0; p < kNumPairs; ++p) { + swiglu_r0[p][0] *= weight_r0; + swiglu_r0[p][1] *= weight_r0; + swiglu_r1[p][0] *= weight_r1; + swiglu_r1[p][1] *= weight_r1; + } +#pragma unroll + for (uint32_t g = 0; g < kNumSFGroups; ++g) { + amax_r0[g] *= cute::abs(weight_r0); + amax_r1[g] *= cute::abs(weight_r1); + } +#pragma unroll + for (uint32_t g = 0; g < kNumSFGroups; ++g) { + amax_r0[g] = math::warp_reduce<4, false>(amax_r0[g], math::ReduceMax()); + amax_r1[g] = math::warp_reduce<4, false>(amax_r1[g], math::ReduceMax()); + } + + // Direct per-128 rendezvous across adjacent-N cluster CTAs. + // Pack the two output rows owned by each four-lane group into + // one contiguous float2 DSM slot. This keeps the same 512-byte + // transaction/barrier protocol while halving DSM store and + // shared-load instructions versus row-strided scalar slots. + auto* l2_sf_base = l2_sf_buffer.get_base_ptr(); + const uint32_t token_r0 = pool_block_idx * BLOCK_M + row_offset_r0; + const uint32_t token_r1 = pool_block_idx * BLOCK_M + row_offset_r1; + const uint32_t dense_sf_group_idx = n_block_idx / 2; + const uint32_t peer_cta_rank = cute::block_rank_in_cluster() ^ 1u; + auto* smem_cd_l1_wg = smem_cd_l1 + epilogue_wg_idx * WG_BLOCK_M * L1_OUT_BLOCK_N; + const uint32_t row_group_idx = epilogue_thread_idx / 4; + auto* peer_amax_pair_slot = reinterpret_cast(smem_cd_l1) + row_group_idx; + + if (col_idx == 0) { + ptx::st_shared_cluster_async(peer_amax_pair_slot, make_float2(amax_r0[0], amax_r1[0]), + combine_barriers[0], peer_cta_rank); + } + if (epilogue_thread_idx == 0) { + combine_barriers[0]->wait(pair_scale_phase); + pair_scale_phase ^= 1u; + // Pre-arm the next phase. The current quantize/store tail + // and the following GEMM separate this from peer stores. + combine_barriers[0]->arrive_and_expect_tx(BLOCK_M * sizeof(float)); + } + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + + float sf_r0 = 0.0f, sf_inv_r0 = 0.0f; + float sf_r1 = 0.0f, sf_inv_r1 = 0.0f; + if (col_idx == 0) { + const float2 peer_amax_pair = ptx::ld_shared(peer_amax_pair_slot); + const float2 common_amax_pair = {cute::max(amax_r0[0], peer_amax_pair.x), + cute::max(amax_r1[0], peer_amax_pair.y)}; + float2 sf_pair, sf_inv_pair; + math::get_e4m3_sf_and_sf_inv(common_amax_pair, sf_pair, sf_inv_pair); + sf_r0 = sf_pair.x; + sf_inv_r0 = sf_inv_pair.x; + sf_r1 = sf_pair.y; + sf_inv_r1 = sf_inv_pair.y; + } + const uint32_t row_group_leader = lane_idx & ~3u; + sf_inv_r0 = __shfl_sync(0xffffffff, sf_inv_r0, row_group_leader); + sf_inv_r1 = __shfl_sync(0xffffffff, sf_inv_r1, row_group_leader); + + // Both CTAs derive the same scale; only the even-N owner + // publishes the dense logical per-128 SF entry. + if ((n_block_idx & 1u) == 0 && col_idx == 0) { + if (valid_r0) l2_sf_base[dense_sf_group_idx * kNumPaddedSFPoolTokens + token_r0] = sf_r0; + if (valid_r1) l2_sf_base[dense_sf_group_idx * kNumPaddedSFPoolTokens + token_r1] = sf_r1; + } + + // Quantize and write to smem_cd_l1 (row-major, no swizzle). + constexpr uint32_t l1_store_stage = 0u; + smem_cd_l1_wg += l1_store_stage * kNumEpilogueWarpgroups * WG_BLOCK_M * L1_OUT_BLOCK_N; +#pragma unroll + for (uint32_t p = 0; p < kNumPairs; ++p) { + const float v00 = swiglu_r0[p][0] * sf_inv_r0; + const float v01 = swiglu_r0[p][1] * sf_inv_r0; + const float v10 = swiglu_r1[p][0] * sf_inv_r1; + const float v11 = swiglu_r1[p][1] * sf_inv_r1; + + const __nv_fp8x2_e4m3 r0_pair(make_float2(v00, v01)); + const __nv_fp8x2_e4m3 r1_pair(make_float2(v10, v11)); + + const uint32_t col = p * 8 + col_idx * 2; + auto* p0 = reinterpret_cast(smem_cd_l1_wg + r_0 * L1_OUT_BLOCK_N + + wg_l1_out_n_idx + col); + auto* p1 = reinterpret_cast(smem_cd_l1_wg + r_1 * L1_OUT_BLOCK_N + + wg_l1_out_n_idx + col); + if (valid_r0) *p0 = r0_pair.__x; + if (valid_r1) *p1 = r1_pair.__x; + } + + // Issue TMA store of the entire tile. Padding rows beyond + // `valid_m` are written with stale/garbage FP8 to the L1-output + // pool buffer, but they are never consumed downstream: the L2 + // GEMM tile loads them, but its NVLink-scatter epilogue is + // gated by `m_idx_in_block >= valid_m`, and stale SF in the + // padding rows can produce NaN accumulators that simply stay + // in registers (only valid rows are converted to BF16 and + // STSM'd into smem). Using TMA for partial tiles is a large + // win for low-batch / decode where every tile is partial. + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { + const uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N; + cute::tma_store_fence(); + cute::SM90_TMA_STORE_2D::copy(&tensor_map_l1_output, smem_cd_l1_wg, out_n_idx, + m_idx + row_block_offset); + cute::tma_store_arrive(); + } + __syncwarp(); + ptx::tma_store_wait<0>(); + if constexpr (!kL2ArrivalCounter) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } + }); + + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + return; +} +#else +if (blockIdx.x == 0 and threadIdx.x == 0) + FI_MEGAMOE_DEVICE_ASSERT(false and "This kernel only supports sm_90"); +#endif diff --git a/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_split_l2_body.inl b/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_split_l2_body.inl new file mode 100644 index 00000000000..b50930f4609 --- /dev/null +++ b/include/flashinfer/megamoe/sm90_nvfp4_mega_moe_split_l2_body.inl @@ -0,0 +1,620 @@ +#ifndef FI_MEGAMOE_NVFP4_LUT_COMPACT +#define FI_MEGAMOE_NVFP4_LUT_COMPACT 0 +#endif +// Independent SM90 NVFP4 MegaMoE l2 kernel body. +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900) and (__CUDA_ARCH__ < 1000)) or \ + defined(__CLION_IDE__) +using Barrier = cutlass::arch::ClusterTransactionBarrier; + +// ===================================================================== +// Template checks +// ===================================================================== +FI_MEGAMOE_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks"); + +// ===================================================================== +// Thread / warp identification +// ===================================================================== +const uint32_t sm_idx = blockIdx.x; +const uint32_t thread_idx = threadIdx.x; +const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); +const uint32_t lane_idx = ptx::get_lane_idx(); + +if (warp_idx == 0 and cute::elect_one_sync()) { + cute::prefetch_tma_descriptor(&tensor_map_l2_acts); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l2_weights); +} + +// ===================================================================== +// Workspaces and symmetric buffer slicing. The L2 activation SF allocation +// intentionally retains its old per-64 physical capacity in stage 1, while +// the active logical layout densely addresses only per-128 groups. +// ===================================================================== +#include + +// ===================================================================== +// GEMM data types and shape constants +// ===================================================================== +using a_dtype_t = cutlass::float_e4m3_t; +using b_dtype_t = cutlass::float_e4m3_t; +constexpr uint32_t WG_BLOCK_M = 64; +constexpr uint32_t WG_BLOCK_N = 128; +using L1WGMMA = typename mma::sm90::FP8MMASelector<128>::type; +static_assert(L1WGMMA::M == 64 and L1WGMMA::N == WG_BLOCK_N and L1WGMMA::K == 32, + "Unexpected WGMMA shape"); + +// Split L2 uses independent CTAs and CTA-local A/B tiles. +constexpr uint32_t LOAD_BLOCK_M = 128; +constexpr uint32_t LOAD_BLOCK_N = 128; +constexpr uint32_t kSwizzleAMode = 128; +constexpr uint32_t kL2ActsSFGranK = 128; // L1 output and L2 input SF granularity +// Split L1 and L2 share the same canonical BM128 physical pool-block unit, +// so scheduling and pool-prefix accounting use the same granularity. +constexpr uint32_t kPoolBlockM = 128; +FI_MEGAMOE_STATIC_ASSERT(kIntermediateHidden / kL2ActsSFGranK <= kIntermediateHidden / 64, + "logical per-128 SF groups must fit retained physical capacity"); + +// ===================================================================== +// Shared memory layout +// ===================================================================== +constexpr uint32_t kSharedMemoryAlignment = 1024; +extern __shared__ __align__(kSharedMemoryAlignment) +uint8_t smem_buffer[]; + +constexpr uint32_t SMEM_EXPERT_COUNT_SIZE = + math::constexpr_align(kNumExperts * sizeof(uint32_t), kSharedMemoryAlignment); +constexpr uint32_t SMEM_NVFP4_LUT_SIZE = + math::constexpr_align(128u * sizeof(uint2), kSharedMemoryAlignment); +constexpr uint32_t SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K * sizeof(a_dtype_t); +constexpr uint32_t SMEM_B_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K * sizeof(b_dtype_t); +constexpr uint32_t B_LOAD_BYTES_PER_ROW = 80u; +constexpr uint32_t SMEM_B_LOAD_SIZE_PER_STAGE = LOAD_BLOCK_N * B_LOAD_BYTES_PER_ROW; +// One per-128 activation scale is loaded for each row and BK128 tile. +constexpr uint32_t kL2SFAHalfStride = + math::constexpr_align(BLOCK_M * sizeof(float), 128u) / sizeof(float); +constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = kL2SFAHalfStride * sizeof(float); +constexpr uint32_t SMEM_BEFORE_BARRIER_SIZE = + SMEM_EXPERT_COUNT_SIZE + SMEM_NVFP4_LUT_SIZE + + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); + +// SMEM pointers +auto smem_expert_count = reinterpret_cast(smem_buffer); +auto smem_nvfp4_lut = + reinterpret_cast(math::advance_ptr(smem_buffer, SMEM_EXPERT_COUNT_SIZE)); + +// L2 remote-scatter staging. Epilogue warp `w` owns rows `[16w, 16w + 16)` +// of its warpgroup tile, so the tile is warp-private and only one 8-row +// half (`r_0` or `r_1`) needs to be resident at a time: a full 16-row tile +// would cost 32768 B and drop `max_num_stages` from 6 to 5. The row stride +// is padded by 8 BF16 so the 4-byte staging stores are bank-conflict free +// (`row * 68` words, all distinct modulo 32) while every staged row start +// stays 16-byte aligned for the `ld.shared.v4` read-back. +constexpr uint32_t kL2StageRows = 8; +constexpr uint32_t kL2StageRowPad = 8; +constexpr uint32_t kL2StageRowStride = WG_BLOCK_N + kL2StageRowPad; +constexpr uint32_t SMEM_CD_L2_PER_WARP = kL2StageRows * kL2StageRowStride * sizeof(nv_bfloat16); +constexpr uint32_t SMEM_CD_L2_SIZE = math::constexpr_align( + kNumEpilogueWarps * SMEM_CD_L2_PER_WARP, kSharedMemoryAlignment); +FI_MEGAMOE_STATIC_ASSERT(kL2StageRowStride % 8 == 0, "Staged rows must stay 16-byte aligned"); +FI_MEGAMOE_STATIC_ASSERT(SMEM_CD_L2_SIZE == kNumEpilogueWarps * SMEM_CD_L2_PER_WARP, + "Host and device `smem_cd_l2` must agree exactly"); + +auto smem_cd_l2_base = + math::advance_ptr(smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_NVFP4_LUT_SIZE); + +auto smem_gemm_base = + math::advance_ptr(smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_NVFP4_LUT_SIZE + SMEM_CD_L2_SIZE); + +auto smem_a = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, i * SMEM_A_SIZE_PER_STAGE); +}); +auto smem_b = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr( + smem_gemm_base, kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); +}); +auto sf_start_ptr = math::advance_ptr( + smem_gemm_base, kNumStages*(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE)); +auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast(sf_start_ptr + i * SMEM_SFA_SIZE_PER_STAGE); +}); +constexpr uint32_t kNumDequantBarriers = kNumStages; + +// Barriers live after SF. +auto barrier_start_ptr = + reinterpret_cast(sf_start_ptr + kNumStages * SMEM_SFA_SIZE_PER_STAGE); +auto full_barriers = + utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + i; }); +auto empty_barriers = + utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; }); +auto dequant_barriers = utils::PatternVisitor( + [=](const uint32_t& i) { return barrier_start_ptr + kNumStages * 2 + i; }); +auto combine_barriers = utils::PatternVisitor([=](const uint32_t& i) { + return barrier_start_ptr + kNumStages * 2 + kNumDequantBarriers + i; +}); + +// ===================================================================== +// Initialization +// ===================================================================== +if constexpr (FI_MEGAMOE_NVFP4_LUT_COMPACT != 0) { + // Compact layout: lo[128] (each row's low word), then the eight + // subnormal-scale rows' high words verbatim -- 544 B of the 1024-B region. + // The high words of rows >= 8 are reconstructed inline by load_nvfp4_lut. + if (thread_idx < 32) { + const auto* lut_src = flashinfer::megamoe::nvfp4::kE2M1AndUe4m3ToFp8Lut + thread_idx * 4; + reinterpret_cast(smem_nvfp4_lut)[thread_idx] = + make_uint4(lut_src[0].x, lut_src[1].x, lut_src[2].x, lut_src[3].x); + if (thread_idx < 8) + reinterpret_cast(smem_nvfp4_lut)[128 + thread_idx] = + flashinfer::megamoe::nvfp4::kE2M1AndUe4m3ToFp8Lut[thread_idx].y; + } +} else if (thread_idx < 64) { + reinterpret_cast(smem_nvfp4_lut)[thread_idx] = + reinterpret_cast(flashinfer::megamoe::nvfp4::kE2M1AndUe4m3ToFp8Lut)[thread_idx]; +} + +if (warp_idx == 0) { +// Clean expert-count shared memory +#pragma unroll + for (uint32_t i = lane_idx; i < kNumExperts; i += 32) ptx::st_shared(smem_expert_count + i, 0u); +} else if (warp_idx == 2) { + // Init GEMM full/empty barriers and combine barriers + if (cute::elect_one_sync()) { +#pragma unroll + for (uint32_t i = 0; i < kNumStages; ++i) { + // Producer arrivals: A(+SFA) + B(TMA+SFB). SFB is copied with + // cp.async.bulk and counted as B-loader transaction bytes, so + // it does not need a separate producer arrival. + full_barriers[i]->init(2); + empty_barriers[i]->init(kNumEpilogueWarps); + } +#pragma unroll + for (uint32_t i = 0; i < kNumStages; ++i) dequant_barriers[i]->init(1); +#pragma unroll + for (uint32_t i = 0; i < kNumEpilogueWarps * 2; ++i) combine_barriers[i]->init(1); + } + cutlass::arch::fence_barrier_init(); +} +__syncthreads(); + +// ===================================================================== +// Scheduler (cluster=1) +// ===================================================================== +auto scheduler = + sched::MegaMoEScheduler(workspace); + +// Pipeline state shared by TMA loaders and math warpgroups +uint32_t stage_idx = 0, phase = 0; +auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++k_block_idx; + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; +}; +const auto dequant_loaded_b_stage = [&](const uint32_t& s, const uint32_t& p, + const uint32_t& non_epilogue_thread_idx) { + if (non_epilogue_thread_idx >= 64u) { + full_barriers[s]->wait(p); + const uint32_t dequant_tid = non_epilogue_thread_idx - 64u; + flashinfer::megamoe::nvfp4::dequant_smem_b_inplace_two_rows_mode2_lop3<64u, 8u>( + reinterpret_cast(smem_b[s]), dequant_tid, smem_nvfp4_lut); + cutlass::arch::fence_view_async_shared(); + ptx::sync_aligned(64, 8); + if (dequant_tid == 0) dequant_barriers[s]->arrive(); + } +}; + +// Intra-SM barrier indices (mirroring SM100) +constexpr uint32_t kDispatchWithEpilogueBarrierIdx = 1; +constexpr uint32_t kEpilogueFullBarrierIdx = 2; +constexpr uint32_t kEpilogueWGBarrierStartIdx = 3; + +// Cross-rank NVLink barrier tags +constexpr uint32_t kBeforeCombineReduceBarrierTag = 2; +constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; + +// Register reconfiguration counts (chosen to fit in 64512 reg budget). +// For the 256-epilogue-thread loader-dequant case (block_m=128, 2 math WGs), +// give non-epilogue warps more room for two-row NVFP4 dequant while staying +// inside the 64K register file: 128*48 + 128*64 + 256*192 = 63488. +constexpr uint32_t kNumNonEpilogueRegisters = 112; +constexpr uint32_t kNumEpilogueRegisters = 192; +FI_MEGAMOE_STATIC_ASSERT(kNumNonEpilogueRegisters* kNumNonEpilogueThreads + + kNumEpilogueRegisters * kNumEpilogueThreads <= + 64512, + "Too many registers"); + +constexpr uint32_t kDispatchGridSyncIndex = 0; +constexpr uint32_t kEpilogueGridSyncIndex = 1; + +const auto for_each_selected_block = [&](auto&& func) { + scheduler.for_each_linear2_block([&](const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx) { + func(std::integral_constant{}, local_expert_idx, + num_k_blocks, m_block_idx, n_block_idx); + }); +}; + +// ===================================================================== +// ROLE 1: GEMM TMA LOAD warps +// ===================================================================== +if (warp_idx == 0) { + cutlass::arch::warpgroup_reg_dealloc(); + + for_each_selected_block([&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx) { + const auto tensor_map_a_ptr = [&]() { return &tensor_map_l2_acts; }(); + const auto tensor_map_sfa_ptr = [&]() { return &tensor_map_l2_acts_sf; }(); + + const uint32_t pool_token_idx = scheduler.get_current_block_pool_token_idx(); + const uint32_t valid_m = scheduler.template get_valid_m(); + const bool has_valid_m = valid_m > 0; + + // Wait for the pool to be ready. Cluster peers can be dummy CTAs for + // the tail M unit when an expert has an odd number of M blocks. + (void)has_valid_m; + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + if (cute::elect_one_sync()) { + if (has_valid_m) { + const uint32_t m_idx = pool_token_idx; + const uint32_t k_idx = k_block_idx * BLOCK_K; + + // TMA load A + tma::copy( + tensor_map_a_ptr, full_barriers[stage_idx], smem_a[stage_idx], k_idx, m_idx, 1); + + // One dense SFA group per BK128. The physical global + // allocation is still twice as large, but the descriptor + // and transaction address only its logical first half. + tma::copy(tensor_map_sfa_ptr, full_barriers[stage_idx], + smem_sfa[stage_idx], m_idx, k_block_idx, 1); + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + + BLOCK_M * sizeof(float)); + } else { + full_barriers[stage_idx]->arrive(); + } + } + __syncwarp(); + dequant_loaded_b_stage(stage_idx, phase, lane_idx); + } + }); + +} else if (warp_idx == 1) { + cutlass::arch::warpgroup_reg_dealloc(); + + for_each_selected_block([&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx) { + const auto tensor_map_b_ptr = [&]() { return &tensor_map_l2_weights; }(); + constexpr uint32_t shape_n = L2_SHAPE_N; + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + const uint32_t n_idx = local_expert_idx * shape_n + n_block_idx * BLOCK_N; + // NVFP4 fused B+scale layout stores 64B packed FP4 + 16B + // UE4M3 scale per BK128 row. + const uint32_t k_idx = k_block_idx * B_LOAD_BYTES_PER_ROW; + if (cute::elect_one_sync()) { + tma::copy( + tensor_map_b_ptr, full_barriers[stage_idx], smem_b[stage_idx], k_idx, n_idx, 1); + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_B_LOAD_SIZE_PER_STAGE); + } + __syncwarp(); + dequant_loaded_b_stage(stage_idx, phase, 32u + lane_idx); + } + }); + +} else if (warp_idx < kNumMMANonEpilogueWarps) { + // Idle non-epilogue warps (2 and 3). They must still + // participate in the warpgroup-collective `setmaxnreg.dec.sync.aligned` + // so that the math warpgroup's `warpgroup_reg_alloc` can succeed. + cutlass::arch::warpgroup_reg_dealloc(); + + const uint32_t non_epilogue_thread_idx = warp_idx * 32 + lane_idx; + for_each_selected_block([&](const auto& block_phase, const uint32_t&, + const uint32_t& num_k_blocks, const uint32_t&, const uint32_t&) { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + dequant_loaded_b_stage(stage_idx, phase, non_epilogue_thread_idx); + __syncwarp(); + } + }); +} else { + // ===================================================================== + // ROLE 3: MATH WARPGROUPS (WGMMA + epilogue + combine) + // ===================================================================== + cutlass::arch::warpgroup_reg_alloc(); + + const uint32_t epilogue_warp_idx = warp_idx - kNumMMANonEpilogueWarps; + const uint32_t epilogue_wg_idx = epilogue_warp_idx / 4; + const uint32_t epilogue_thread_idx = epilogue_warp_idx * 32 + lane_idx; + const uint32_t warp_idx_in_wg = epilogue_warp_idx % 4; + + const auto arrive_empty_barrier = [&](const uint32_t& s) { + if (lane_idx == 0) empty_barriers[s]->arrive(); + }; + + const auto cleanup_workspace_from_epilogue = [&]() { + FI_MEGAMOE_STATIC_ASSERT(kNumSMs > 1, "Invalid SM count"); + if (sm_idx == 0) { +#pragma unroll + for (uint32_t i = epilogue_thread_idx; i < kNumExperts; i += kNumEpilogueThreads) + *workspace.get_expert_send_count_ptr(i) = 0; + } else { + for (uint32_t i = sm_idx - 1; i < kNumExpertsPerRank; i += kNumSMs - 1) { + const auto num_recv_tokens = + static_cast(*workspace.get_expert_recv_count_sum_ptr(i)); + const auto num_recv_m_blocks = math::ceil_div(num_recv_tokens, kPoolBlockM); + const auto cleanup_pool_block_offset = scheduler.get_pool_block_offset(i); + + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (epilogue_thread_idx == 0) { + *workspace.get_expert_recv_count_sum_ptr(i) = 0; + if (cumulative_local_expert_recv_stats != nullptr) + ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(num_recv_tokens)); + } + + for (uint32_t j = epilogue_thread_idx; j < kNumRanks; j += kNumEpilogueThreads) + *workspace.get_expert_recv_count_ptr(j, i) = 0; + + for (uint32_t j = epilogue_thread_idx; j < num_recv_m_blocks; j += kNumEpilogueThreads) + *workspace.get_l1_arrival_count_ptr(cleanup_pool_block_offset + j) = 0; + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } + } + }; + + const auto finish_no_dispatch_cleanup = [&]() { + ptx::sync_unaligned(kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + cleanup_workspace_from_epilogue(); + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, epilogue_thread_idx, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); }, true, false); + }; + + // WGMMA-output register layout helpers + const uint32_t row_idx = lane_idx / 4; + const uint32_t col_idx = lane_idx % 4; + const uint32_t r_0 = warp_idx_in_wg * 16 + row_idx; + const uint32_t r_1 = r_0 + 8; + + FI_MEGAMOE_STATIC_ASSERT(WG_BLOCK_M == L1WGMMA::M, + "Each warpgroup must run exactly one WGMMA per K-block"); + + // The staged scatter writes 16 bytes at a time. The combine row pitch + // (`kHidden * 2`) is 16-byte aligned, so every destination row inherits + // the buffer base alignment: check it once, not once per row. + FI_MEGAMOE_DEVICE_ASSERT( + reinterpret_cast( + combine_token_buffer.get_rank_buffer(0u).get_data_buffer(0u).get_base_ptr()) % + 16 == + 0); + + ptx::sync_unaligned(kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + for_each_selected_block([&](const auto& block_phase, const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, const uint32_t& m_block_idx, + const uint32_t& n_block_idx) { + const uint32_t valid_m = scheduler.template get_valid_m(); + const uint32_t m_idx = scheduler.get_current_block_pool_token_idx(); + const uint32_t n_idx = n_block_idx * BLOCK_N; + const uint32_t row_block_offset = epilogue_wg_idx * WG_BLOCK_M; + const uint32_t row_offset_r0 = row_block_offset + r_0; + const uint32_t row_offset_r1 = row_block_offset + r_1; + const bool valid_r0 = row_offset_r0 < valid_m; + const bool valid_r1 = row_offset_r1 < valid_m; + const float l2_global_scale = + l2_global_scales == nullptr ? 1.0f : __ldg(l2_global_scales + local_expert_idx); + const auto cast_l2_scaled_bf16_pair = [&](float x, float y) -> uint32_t { + x *= l2_global_scale; + y *= l2_global_scale; + return math::cast_into_bf16_and_pack(x, y); + }; + + if (row_block_offset >= valid_m) { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + dequant_barriers[stage_idx]->wait(phase); + arrive_empty_barrier(stage_idx); + __syncwarp(); + } + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + return; + } + + // ---------------- GEMM ---------------- + using WGMMA = L1WGMMA; + constexpr uint32_t kAccumPerThread = WGMMA::kNumAccum; // 64 for M=64,N=128 + float final_accum[kAccumPerThread] = {}; + float accum[kAccumPerThread]; + + const auto run_default_gemm_loop = [&]() { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + dequant_barriers[stage_idx]->wait(phase); + + // One activation scale covers all four K32 WGMMAs in BK128. + // Read it before warpgroup_arrive as required by WGMMA. + const float scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); + const float scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); + + // NVFP4 UE4M3 weight scales are applied during FP4 -> FP8 smem + // expansion, so the WGMMA accumulator only needs activation SF. + +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + row_block_offset * BLOCK_K + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc(smem_b[stage_idx] + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + arrive_empty_barrier(stage_idx); + +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++i) { + final_accum[i * 4 + 0] += scale_a_0 * accum[i * 4 + 0]; + final_accum[i * 4 + 1] += scale_a_0 * accum[i * 4 + 1]; + final_accum[i * 4 + 2] += scale_a_1 * accum[i * 4 + 2]; + final_accum[i * 4 + 3] += scale_a_1 * accum[i * 4 + 3]; + } + } + }; + + run_default_gemm_loop(); + + // ---------------- L2 EPILOGUE: BF16 cast + staged NVLink scatter ---------------- + // The direct path stored 4 bytes per lane, so a warp's 8 row groups + // produced 8 scattered 16-byte requests per instruction, each billed as + // a 32-byte sector. Staging one 8-row half in warp-private shared memory + // lets every destination row leave as one 256-byte contiguous burst. + auto smem_cd_l2 = + math::advance_ptr(smem_cd_l2_base, epilogue_warp_idx * SMEM_CD_L2_PER_WARP); + + // Cast and pack one 8-row half into this warp's private tile. The + // arithmetic is unchanged, so the same values reach the same addresses. + const auto stage_rows = [&](const bool& valid_row, const uint32_t& row_accum_offset) { + // This guard must stay textually equivalent to the scatter guard + // below: invalid rows keep stale tile bytes and are skipped there. + if (not valid_row) return; +#pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 8; ++i) { + const uint32_t chunk_lo = 2 * i, chunk_hi = 2 * i + 1; + const uint32_t col_lo = chunk_lo * 8 + col_idx * 2; + const uint32_t col_hi = chunk_hi * 8 + col_idx * 2; + const uint32_t packed_lo = + cast_l2_scaled_bf16_pair(final_accum[chunk_lo * 4 + row_accum_offset + 0], + final_accum[chunk_lo * 4 + row_accum_offset + 1]); + const uint32_t packed_hi = + cast_l2_scaled_bf16_pair(final_accum[chunk_hi * 4 + row_accum_offset + 0], + final_accum[chunk_hi * 4 + row_accum_offset + 1]); + ptx::st_shared( + reinterpret_cast(smem_cd_l2 + row_idx * kL2StageRowStride + col_lo), + packed_lo); + ptx::st_shared( + reinterpret_cast(smem_cd_l2 + row_idx * kL2StageRowStride + col_hi), + packed_hi); + } + }; + + // Read the tile back with 16 lanes per destination row, 16 bytes each. + constexpr uint32_t kNumScatterLanesPerRow = 16; + constexpr uint32_t kScatterBytesPerLane = + WG_BLOCK_N * sizeof(nv_bfloat16) / kNumScatterLanesPerRow; + FI_MEGAMOE_STATIC_ASSERT(kScatterBytesPerLane == sizeof(uint4), "Expect one `uint4` per lane"); + const uint32_t scatter_row_in_pair = lane_idx / kNumScatterLanesPerRow; + const uint32_t lane_in_row = lane_idx % kNumScatterLanesPerRow; + const uint32_t scatter_group_leader = lane_idx - lane_in_row; + const uint32_t scatter_group_mask = 0xffffu << scatter_group_leader; + + const auto scatter_staged_rows = [&](const uint32_t& row_base) { +#pragma unroll + for (uint32_t j = 0; j < kL2StageRows / 2; ++j) { + const uint32_t stage_row = j * 2 + scatter_row_in_pair; + const uint32_t row_offset = row_base + stage_row; + // Uniform within each 16-lane group, so the shuffles stay converged. + if (row_offset >= valid_m) continue; + uint32_t dst_rank_idx = 0; + uint32_t dst_token_idx = 0; + uint32_t dst_topk_idx = 0; + if (lane_in_row == 0) { + const auto src_metadata = *workspace.get_token_src_metadata_ptr(m_idx + row_offset); + dst_rank_idx = src_metadata.rank_idx; + dst_token_idx = src_metadata.token_idx; + dst_topk_idx = src_metadata.topk_idx; + } + dst_rank_idx = __shfl_sync(scatter_group_mask, dst_rank_idx, scatter_group_leader); + dst_token_idx = __shfl_sync(scatter_group_mask, dst_token_idx, scatter_group_leader); + dst_topk_idx = __shfl_sync(scatter_group_mask, dst_topk_idx, scatter_group_leader); + const auto dst_token = + combine_token_buffer.get_rank_buffer(dst_topk_idx).get_data_buffer(dst_token_idx); + auto dst_ptr = math::advance_ptr( + dst_token.get_base_ptr(), + n_idx * sizeof(nv_bfloat16) + lane_in_row * kScatterBytesPerLane); + auto mapped_dst_ptr = sym_buffer.map(dst_ptr, dst_rank_idx); + const auto packed = ptx::ld_shared(reinterpret_cast( + smem_cd_l2 + stage_row * kL2StageRowStride + + lane_in_row * (kScatterBytesPerLane / sizeof(nv_bfloat16)))); + *reinterpret_cast(mapped_dst_ptr) = packed; + } + }; + + // `row_offset_r0 - row_idx`: the first of this warp's 16 rows. + const uint32_t warp_row_base = row_block_offset + warp_idx_in_wg * 16; + + // The staging tile is warp-private (producer and consumer are the same + // 32 lanes), so `__syncwarp` is a sufficient shared-memory ordering + // point and no named barrier index is consumed. The trailing + // `kEpilogueFullBarrierIdx` sync also covers the write-after-read on + // the tile across scheduled blocks. + stage_rows(valid_r0, 0); + __syncwarp(); + scatter_staged_rows(warp_row_base); + __syncwarp(); + stage_rows(valid_r1, 2); + __syncwarp(); + scatter_staged_rows(warp_row_base + kL2StageRows); + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + }); + + // ---------------- COMBINE ---------------- + // NVLink barrier first: signals remote ranks that this rank's GEMM + // outputs (NVLink scatter targets) are fully written. + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, epilogue_thread_idx, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); }); + + // Rendezvous before workspace cleanup. + ptx::sync_unaligned(kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + constexpr uint32_t kNumHiddenBytes = kHidden * sizeof(nv_bfloat16); + constexpr uint32_t kNumElemsPerUint4 = sizeof(uint4) / sizeof(nv_bfloat162); + + constexpr uint32_t kNumChunkSlots = 3; + constexpr uint32_t kNumMaxRegistersForBuffer = 128; + constexpr uint32_t kNumDefaultChunks = + (kNumChunkSlots * kNumEpilogueWarps * kNumHiddenBytes <= SMEM_BEFORE_BARRIER_SIZE and + kHidden <= 32 * kNumMaxRegistersForBuffer) + ? 1 + : 2; + constexpr uint32_t kNumChunks = kNumDefaultChunks; + constexpr uint32_t kNumChunkBytes = kNumHiddenBytes / kNumChunks; + constexpr uint32_t kNumChunkUint4 = kNumChunkBytes / sizeof(uint4); + constexpr uint32_t kNumUint4PerLane = kNumChunkUint4 / 32; + FI_MEGAMOE_STATIC_ASSERT(kHidden % kNumChunks == 0, + "Hidden must be divisible by number of chunks"); + FI_MEGAMOE_STATIC_ASSERT( + kNumChunkSlots * kNumEpilogueWarps * kNumHiddenBytes / kNumChunks <= SMEM_BEFORE_BARRIER_SIZE, + "Hidden is too large"); + FI_MEGAMOE_STATIC_ASSERT(kNumChunkBytes % 16 == 0, + "Combine chunk must be TMA-aligned (16 bytes)"); + FI_MEGAMOE_STATIC_ASSERT(kNumChunkBytes % sizeof(uint4) == 0, + "Combine chunk must be divisible by 16 bytes"); + FI_MEGAMOE_STATIC_ASSERT(kNumChunkUint4 % 32 == 0, + "Combine chunk must be a multiple of 32 16-byte elements"); + FI_MEGAMOE_STATIC_ASSERT(kNumTopk <= 32, "Top-k must fit in a single warp"); + + FI_MEGAMOE_DEVICE_ASSERT( + kNumChunkSlots * kNumEpilogueWarps * kNumChunkBytes <= + static_cast(reinterpret_cast(barrier_start_ptr) - smem_buffer)); + +#include + + finish_no_dispatch_cleanup(); +} +#else +if (blockIdx.x == 0 and threadIdx.x == 0) + FI_MEGAMOE_DEVICE_ASSERT(false and "This kernel only supports sm_90"); +#endif diff --git a/include/flashinfer/megamoe/sym_buffer.cuh b/include/flashinfer/megamoe/sym_buffer.cuh new file mode 100644 index 00000000000..f2a0a115fcb --- /dev/null +++ b/include/flashinfer/megamoe/sym_buffer.cuh @@ -0,0 +1,40 @@ +#pragma once + +#include + +namespace flashinfer::megamoe::layout { + +constexpr static uint32_t kNumMaxRanks = 72; + +template +struct SymBuffer { + int64_t base; + int64_t offsets[kNumMaxRanks]; + uint32_t rank_idx; + + FI_MEGAMOE_STATIC_ASSERT(kNumRanks <= kNumMaxRanks, "Too many ranks"); + + SymBuffer() = default; + + template + explicit SymBuffer(const Container& c, const uint32_t& rank_idx) : rank_idx(rank_idx) { + const auto size = static_cast(c.size()); + base = c[rank_idx]; + for (uint32_t i = 0; i < kNumMaxRanks; ++i) offsets[i] = i < size ? (c[i] - base) : 0; + } + +#if defined(__CUDA_ARCH__) or defined(__CLION_IDE__) + template + CUTLASS_DEVICE ptr_t get_base_ptr() const { + return reinterpret_cast(base); + } + + template + CUTLASS_DEVICE ptr_t map(const ptr_t& ptr, const uint32_t& dst_rank_idx) const { + int64_t mapped_ptr = offsets[dst_rank_idx] + reinterpret_cast(ptr); + return *reinterpret_cast(&mapped_ptr); + } +#endif +}; + +} // namespace flashinfer::megamoe::layout diff --git a/tests/moe/test_megamoe_fp4_ep.py b/tests/moe/test_megamoe_fp4_ep.py new file mode 100644 index 00000000000..c144accdd17 --- /dev/null +++ b/tests/moe/test_megamoe_fp4_ep.py @@ -0,0 +1,43 @@ +"""CPU checks for EP-specific configuration and JIT cache isolation.""" + +import pytest + +from flashinfer.jit.megamoe import get_megamoe_sm90_nvfp4_uri +from flashinfer.megamoe import select_fused_plan + + +@pytest.mark.parametrize("num_sms", [78, 132]) +@pytest.mark.parametrize( + "experts,hidden,intermediate,topk", + [(256, 4096, 2048, 6), (384, 6144, 2048, 8), (384, 7168, 3072, 6)], +) +@pytest.mark.parametrize("tokens", [1, 24, 64, 256]) +def test_ep4_fused_plan_uses_legal_generic_pipeline( + num_sms, experts, hidden, intermediate, topk, tokens +): + plan = select_fused_plan( + num_sms, 4, experts, topk, hidden, intermediate, 384, tokens + ) + assert not plan.rs_swap_ab + assert not plan.push_dispatch + assert not plan.no_clean_barrier + assert (experts // 4) % plan.num_experts_per_wave == 0 + if plan.block_m == 8: + assert plan.num_stages in (3, 4) + else: + assert plan.num_stages == 3 + + +def test_nvfp4_jit_cache_distinguishes_ep_width(): + plan = select_fused_plan(132, 8, 256, 6, 4096, 2048, 384, 1) + args = (plan, 132, 384, 256, 6, 4096, 2048, 10.0, True) + ep4 = get_megamoe_sm90_nvfp4_uri(*args, num_ranks=4) + ep8 = get_megamoe_sm90_nvfp4_uri(*args, num_ranks=8) + assert ep4 != ep8 + assert "_r4_" in ep4 and "_r8_" in ep8 + + +@pytest.mark.parametrize("ep", [1, 2, 3, 16]) +def test_nvfp4_planner_rejects_unvalidated_ep_width(ep): + with pytest.raises(ValueError, match="4 or 8 ranks"): + select_fused_plan(132, ep, 256, 6, 4096, 2048, 384, 1) diff --git a/tests/moe/test_megamoe_fp4_sm90.py b/tests/moe/test_megamoe_fp4_sm90.py new file mode 100644 index 00000000000..219c44e7d7e --- /dev/null +++ b/tests/moe/test_megamoe_fp4_sm90.py @@ -0,0 +1,327 @@ +"""Four/eight-rank SM90 checks with an independent dense W4A8 reference. + +Run with plain Python (workers are spawned here): + python tests/moe/test_megamoe_fp4_sm90.py --tokens 1 32 96 384 +No DeepGEMM extension or checkpoint is needed. +""" + +import argparse +import datetime +import os + +import pytest +import torch +import torch.distributed as dist + + +def _weights(e, n, k, group, generator): + codes = torch.randint( + 0, 16, (e, n, k), device="cuda", generator=generator, dtype=torch.uint8 + ) + chunks = codes.reshape(e, n, k // 8, 8) + packed = (chunks[..., 4:] | (chunks[..., :4] << 4)).reshape(e, n, k // 2) + level = torch.randint(0, 3, (e, n, k // group), device="cuda", generator=generator) + sf = (level + (16 if group == 16 else 122)).to(torch.uint8) + scales = ( + sf.view(torch.float8_e4m3fn).float() + if group == 16 + else torch.exp2(sf.float() - 127) + ) + lattice = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], device="cuda") + values = lattice[(codes & 7).long()] * torch.where(codes < 8, 1.0, -1.0) + values *= scales.repeat_interleave(group, dim=-1) + if group == 32: + values = torch.where(values.abs() < 2**-6, 0.0, values) + return (packed, sf), values.to(torch.float8_e4m3fn).float() + + +def _gather(t): + gathered = [torch.empty_like(t) for _ in range(dist.get_world_size())] + dist.all_gather(gathered, t) + return torch.cat(gathered) + + +def _reference(x, idx, weights, w1, w2, gs1, gs2, rank, local_experts): + all_x, all_idx, all_weights = _gather(x), _gather(idx), _gather(weights) + routes = torch.zeros( + (*all_idx.shape, w2.shape[1]), device="cuda", dtype=torch.float32 + ) + for expert in range(local_experts): + rows, slots = torch.where(all_idx == rank * local_experts + expert) + if rows.numel() == 0: + continue + intermediate = all_x[rows] @ w1[expert].T * gs1[expert] + gate, up = intermediate.chunk(2, dim=-1) + hidden = torch.nn.functional.silu(gate.clamp(max=10)) * up.clamp(-10, 10) + hidden *= all_weights[rows, slots, None] + grouped = hidden.reshape(hidden.shape[0], -1, 128) + sf = torch.exp2( + torch.ceil(torch.log2((grouped.abs().amax(-1) / 448).clamp_min(2**-126))) + ) + quantized = (grouped / sf[..., None]).to(torch.float8_e4m3fn).float() * sf[ + ..., None + ] + down = quantized.reshape_as(hidden) @ w2[expert].T * gs2[expert] + routes[rows, slots] = down.to(torch.bfloat16).float() + dist.all_reduce(routes) + result = routes.sum(1).to(torch.bfloat16) + m = x.shape[0] + counts = torch.stack( + [(all_idx == rank * local_experts + e).sum() for e in range(local_experts)] + ) + return result[rank * m : (rank + 1) * m], counts + + +def _worker(rank, args): + from flashinfer.megamoe import ( + MegaMoESymmBuffer, + nvfp4_mega_moe, + mxfp4_mega_moe, + transform_nvfp4_weights_for_mega_moe, + transform_mxfp4_weights_for_mega_moe, + ) + + ep = getattr(args, "ep", 8) + torch.cuda.set_device(rank) + torch.backends.cuda.matmul.allow_tf32 = False + dist.init_process_group( + "nccl", + rank=rank, + world_size=ep, + timeout=datetime.timedelta(minutes=10), + device_id=torch.device("cuda", rank), + ) + group = dist.group.WORLD + h, ih = args.hidden, args.intermediate + experts, topk = getattr(args, "experts", 256), getattr(args, "topk", 6) + epr = experts // ep + generator = torch.Generator(device="cuda").manual_seed(7919 + rank) + buffer = MegaMoESymmBuffer( + group, experts, args.capacity, topk, h, ih, fast_math=args.fast_math + ) + gs1 = torch.linspace(0.73, 1.37, epr, device="cuda") + gs2 = torch.linspace(1.31, 0.67, epr, device="cuda") + for fmt in args.formats: + gran = 16 if fmt == "nvfp4" else 32 + raw1, ref1 = _weights(epr, ih * 2, h, gran, generator) + raw2, ref2 = _weights(epr, h, ih, gran, generator) + transform = ( + transform_nvfp4_weights_for_mega_moe + if fmt == "nvfp4" + else transform_mxfp4_weights_for_mega_moe + ) + w1, w2 = transform(raw1, raw2) + for m in args.tokens: + # Binary input SFs isolate the CUDA kernels from input quantizer differences. + x = (torch.randn(m, h, device="cuda", generator=generator) * 32).to( + torch.float8_e4m3fn + ) + sf = torch.full((m, h // 128), 1 / 128, device="cuda") + scores = torch.randn(m, experts, device="cuda", generator=generator) + idx = scores.topk(topk, dim=-1).indices + routing = torch.randn(m, topk, device="cuda", generator=generator) * 0.25 + # Exercise masked slots, signed weights and uneven expert loading. + idx[::3, -1] = -1 + if args.routing == "hot": + idx[:] = torch.arange(topk, device="cuda") + idx[::3, -1] = -1 + elif args.routing == "empty": + idx.fill_(-1) + expected, counts = _reference( + x.float() / 128, idx, routing, ref1, ref2, gs1, gs2, rank, epr + ) + stats = torch.zeros(epr, dtype=torch.int32, device="cuda") + buffer.stage(x, sf, idx, routing) + y = torch.empty(m, h, device="cuda", dtype=torch.bfloat16) + family = "split" if fmt == "nvfp4" and m > 256 else "fused" + kwargs = dict( + cumulative_local_expert_recv_stats=stats, + l1_global_scales=gs1, + l2_global_scales=gs2, + ) + run = nvfp4_mega_moe if fmt == "nvfp4" else mxfp4_mega_moe + if fmt == "nvfp4": + kwargs["kernel_family"] = family + kwargs["force_push_dispatch"] = getattr( + args, "force_push_dispatch", None + ) + kwargs["force_no_clean_barrier"] = getattr( + args, "force_no_clean_barrier", None + ) + if getattr(args, "check_no_swap", False): + from flashinfer.megamoe import select_mxfp4_fused_plan + + plan = select_mxfp4_fused_plan( + buffer.num_sms, ep, experts, topk, h, ih, args.capacity, m + ) + assert not plan.swap_ab, "regression must exercise the SS L2 path" + direct_repeats = getattr(args, "repeats", 2) + for repetition in range(direct_repeats): + run(y, w1, w2, buffer, **kwargs) + torch.cuda.synchronize() + assert torch.isfinite(y).all() + error = (y.float() - expected.float()).abs() + relative_l2 = torch.linalg.vector_norm( + error + ) / torch.linalg.vector_norm(expected.float()).clamp_min(1e-12) + assert relative_l2 < 0.015, ( + fmt, + family, + m, + relative_l2.item(), + error.max().item(), + ) + assert torch.equal(stats.long(), counts * (repetition + 1)), ( + stats, + counts, + ) + if repetition == 0: + first = y.clone() + else: + assert torch.equal(y, first), ( + "workspace reuse changed the output", + fmt, + family, + m, + (y.float() - first.float()).abs().max().item(), + int((y != first).sum().item()), + ) + if getattr(args, "cuda_graph", False): + dist.barrier() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run(y, w1, w2, buffer, **kwargs) + for replay in range(direct_repeats): + graph.replay() + torch.cuda.synchronize() + assert torch.equal(y, first), "CUDA graph replay changed the output" + assert torch.equal( + stats.long(), counts * (direct_repeats + replay + 1) + ) + metrics = torch.tensor( + [relative_l2.item(), error.max().item()], device="cuda" + ) + dist.all_reduce(metrics, op=dist.ReduceOp.MAX) + if rank == 0: + print( + f"PASS {fmt} {family} M={m} fm={args.fast_math} routing={args.routing} " + f"relative_l2={metrics[0].item():.6g} max_abs={metrics[1].item():.6g} stats/reuse=exact", + flush=True, + ) + buffer.destroy() + dist.destroy_process_group() + + +def _run(args): + if getattr(args, "repeats", 2) < 1: + raise ValueError("repeats must be positive") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29648") + torch.multiprocessing.spawn(_worker, args=(args,), nprocs=getattr(args, "ep", 8)) + + +@pytest.mark.parametrize("ep", [4, 8]) +def test_megamoe_fp4_sm90(ep): + if torch.cuda.device_count() < ep or torch.cuda.get_device_capability()[0] != 9: + pytest.skip(f"requires {ep} Hopper GPUs") + _run( + argparse.Namespace( + ep=ep, + hidden=512, + intermediate=256, + capacity=384, + tokens=[1, 24, 32, 96, 384], + formats=["nvfp4", "mxfp4"], + fast_math=True, + routing="random", + cuda_graph=True, + ) + ) + + +def test_mxfp4_no_swap_decoded_b_reuse(): + if torch.cuda.device_count() < 8 or torch.cuda.get_device_capability()[0] != 9: + pytest.skip("requires 8 Hopper GPUs") + _run( + argparse.Namespace( + hidden=4096, + intermediate=2048, + capacity=384, + tokens=[384], + formats=["mxfp4"], + fast_math=True, + routing="random", + cuda_graph=True, + repeats=6, + check_no_swap=True, + ) + ) + + +@pytest.mark.parametrize( + "hidden,intermediate,experts,topk", + [(4096, 2048, 256, 6), (6144, 2048, 384, 8), (7168, 3072, 384, 6)], +) +def test_nvfp4_ep4_production_shape(hidden, intermediate, experts, topk): + if torch.cuda.device_count() < 4 or torch.cuda.get_device_capability()[0] != 9: + pytest.skip("requires 4 Hopper GPUs") + _run( + argparse.Namespace( + ep=4, + hidden=hidden, + intermediate=intermediate, + experts=experts, + topk=topk, + capacity=384, + tokens=[1, 96, 384], + formats=["nvfp4"], + fast_math=True, + routing="random", + cuda_graph=True, + ) + ) + + +@pytest.mark.parametrize("no_clean", [False, True]) +def test_nvfp4_ep4_counter_dispatch(no_clean): + if torch.cuda.device_count() < 4 or torch.cuda.get_device_capability()[0] != 9: + pytest.skip("requires 4 Hopper GPUs") + _run( + argparse.Namespace( + ep=4, + hidden=1024, + intermediate=256, + capacity=384, + tokens=[1, 24], + formats=["nvfp4"], + fast_math=True, + routing="random", + cuda_graph=True, + repeats=3, + force_push_dispatch=True, + force_no_clean_barrier=no_clean, + ) + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--ep", type=int, choices=[4, 8], default=8) + parser.add_argument("--tokens", type=int, nargs="+", default=[1, 24, 32, 96, 384]) + parser.add_argument("--hidden", type=int, default=512) + parser.add_argument("--intermediate", type=int, default=256) + parser.add_argument("--capacity", type=int, default=384) + parser.add_argument( + "--formats", nargs="+", choices=["nvfp4", "mxfp4"], default=["nvfp4", "mxfp4"] + ) + parser.add_argument( + "--routing", choices=["random", "hot", "empty"], default="random" + ) + parser.add_argument("--cuda-graph", action="store_true") + parser.add_argument("--repeats", type=int, default=2) + parser.add_argument("--check-no-swap", action="store_true") + parser.add_argument( + "--no-fast-math", dest="fast_math", action="store_false", default=True + ) + _run(parser.parse_args()) diff --git a/tests/moe/test_megamoe_fp4_weights.py b/tests/moe/test_megamoe_fp4_weights.py new file mode 100644 index 00000000000..af94aefbc18 --- /dev/null +++ b/tests/moe/test_megamoe_fp4_weights.py @@ -0,0 +1,72 @@ +"""Independent checks of the public transforms and decoder byte addresses.""" + +import pytest +import torch + +from flashinfer.megamoe.weights import ( + transform_mxfp4_weights_for_mega_moe, + transform_nvfp4_weights_for_mega_moe, +) + + +@pytest.mark.parametrize("weight_format", ["nvfp4", "mxfp4"]) +def test_decoder_addresses_preserve_nibbles_and_scales(weight_format): + group = 16 if weight_format == "nvfp4" else 32 + transform = ( + transform_nvfp4_weights_for_mega_moe + if group == 16 + else transform_mxfp4_weights_for_mega_moe + ) + generator = torch.Generator().manual_seed(19) + raw = [] + for n, k in ((512, 512), (512, 256)): + raw.append( + ( + torch.randint( + 0, 256, (2, n, k // 2), dtype=torch.uint8, generator=generator + ), + torch.randint( + 0, 255, (2, n, k // group), dtype=torch.uint8, generator=generator + ), + ) + ) + outputs = transform(*raw) + for layer, ((packed, scales), output) in enumerate(zip(raw, outputs, strict=True)): + e, n, half_k = packed.shape + kb = half_k // 64 + if group == 16: + rows = output.reshape(e, n, kb, 80) + payload, sf = rows[..., :64], rows[..., 64:72] + assert torch.count_nonzero(rows[..., 72:]) == 0 + else: + tiles = output.reshape(e, n // 128, kb, 8704) + payload = tiles[..., :8192].reshape(e, n // 128, kb, 4, 128, 16) + payload = payload.permute(0, 1, 4, 2, 3, 5).reshape(e, n, kb, 64) + sf = tiles[..., 8192:].reshape(e, n // 128, kb, 128, 4) + sf = sf.permute(0, 1, 3, 2, 4).reshape(e, n, kb, 4) + # Literal decoder mapping: PRMT selects low/high nibbles in order, + # while sign bits come from byte bit 7 and byte bit 3 respectively. + words = payload.reshape(e, n, kb, 16, 4) + magnitudes = torch.stack((words & 7, (words >> 4) & 7), dim=-1).flatten(-2) + signs = torch.cat((words >> 7, (words >> 3) & 1), dim=-1) + decoded = magnitudes | (signs << 3) + original_words = packed.reshape(e, n, kb, 16, 4) + expected = torch.cat((original_words >> 4, original_words & 15), dim=-1) + if layer == 0: + # Map each physical interleaved gate/up row back to its source row. + physical = torch.arange(n) + source = physical // 16 * 8 + physical % 8 + (physical % 16 >= 8) * (n // 2) + expected = expected[:, source] + scales = scales[:, source] + assert torch.equal(decoded, expected) + assert torch.equal(sf.reshape(e, n, -1), scales) + + +def test_weight_transform_rejects_mismatched_scales(): + p1 = torch.zeros(1, 512, 256, dtype=torch.uint8) + p2 = torch.zeros(1, 512, 128, dtype=torch.uint8) + with pytest.raises(ValueError, match="matching scales"): + transform_mxfp4_weights_for_mega_moe( + (p1, torch.zeros(1, 512, 8, dtype=torch.uint8)), + (p2, torch.zeros(1, 512, 8, dtype=torch.uint8)), + ) diff --git a/tests/moe/test_megamoe_nvfp4_sm90.py b/tests/moe/test_megamoe_nvfp4_sm90.py new file mode 100644 index 00000000000..571db096356 --- /dev/null +++ b/tests/moe/test_megamoe_nvfp4_sm90.py @@ -0,0 +1,208 @@ +"""SM90 NVFP4 MegaMoE: correctness against the reference implementation. + +Runs the FlashInfer kernel and the DeepGEMM reference on identical inputs and +weights and requires bit-identical output -- both drive the same CUDA kernel, so +any difference is a host-side (layout / descriptor / plan) bug. + +Requires 8 GPUs. The script spawns its own worker per GPU, so run it with plain +python, not torchrun: python tests/moe/test_megamoe_nvfp4_sm90.py +""" + +import argparse +import sys + +import torch +import torch.distributed as dist + +SHAPES = { + "flash": dict(hidden=4096, intermediate_hidden=2048, num_experts=256, num_topk=6), + "pro": dict(hidden=7168, intermediate_hidden=3072, num_experts=384, num_topk=6), + "mimo": dict(hidden=6144, intermediate_hidden=2048, num_experts=384, num_topk=8), +} + + +def make_routing(m, num_experts, num_topk, rank, seed): + torch.manual_seed(rank + 101 + seed * 1000) + scores = torch.randn((m, num_experts), dtype=torch.float, device="cuda") + w, idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False) + return idx, w.float() + + +def worker(local_rank, world_size, args): + import deep_gemm # reference + from deep_gemm.quantization_nvfp4 import quantize_to_nvfp4 + from deep_gemm.utils import per_token_cast_to_fp8 + from deep_gemm.utils.dist import init_dist + + from flashinfer.megamoe import ( + FAMILY_THRESHOLD, + MegaMoESymmBuffer, + nvfp4_mega_moe, + ) + + rank, num_ranks, group = init_dist(local_rank, world_size) + failures = 0 + + for shape_name in args.shapes: + s = SHAPES[shape_name] + hidden, ih = s["hidden"], s["intermediate_hidden"] + num_experts, num_topk = s["num_experts"], s["num_topk"] + experts_per_rank = num_experts // num_ranks + + g = torch.Generator(device="cuda") + g.manual_seed(7919 + rank) + l1_bf = ( + torch.randn( + (experts_per_rank, ih * 2, hidden), + dtype=torch.bfloat16, + device="cuda", + generator=g, + ) + * 0.05 + ) + l2_bf = ( + torch.randn( + (experts_per_rank, hidden, ih), + dtype=torch.bfloat16, + device="cuda", + generator=g, + ) + * 0.05 + ) + l1_packed, l1_scale = quantize_to_nvfp4(l1_bf, group_size=16) + l2_packed, l2_scale = quantize_to_nvfp4(l2_bf, group_size=16) + del l1_bf, l2_bf + torch.cuda.empty_cache() + + # The prepacked scales are tile-major (E, N/block_n, K/128, block_n, 8), + # so the weight layout is part of the family: block_n=256 for the fused + # kernel, block_n=128 for split L1/L2. Build each on first use. + transformed = {} + + ref_buf = deep_gemm.get_symm_buffer_for_mega_moe( + group, num_experts, args.cap, num_topk, hidden, ih + ) + fi_buf = MegaMoESymmBuffer( + group, num_experts, args.cap, num_topk, hidden, ih, fast_math=args.fast_math + ) + + for m in args.m: + x_bf = torch.randn( + (m, hidden), dtype=torch.bfloat16, device="cuda", generator=g + ) + x_fp8, x_sf = per_token_cast_to_fp8( + x_bf, use_ue8m0=False, gran_k=128, use_packed_ue8m0=False + ) + topk_idx, topk_w = make_routing(m, num_experts, num_topk, rank, 0) + + # Pin the family on both sides rather than letting each one's + # own "auto" rule decide -- a silent disagreement would compare + # two different kernels and still look like a pass. + family = "split" if m > FAMILY_THRESHOLD else "fused" + block_n = 128 if family == "split" else 256 + if block_n not in transformed: + transformed[block_n] = ( + deep_gemm.transform_nvfp4_weights_for_mega_moe_sm90( + (l1_packed, l1_scale), (l2_packed, l2_scale), block_n=block_n + ) + ) + l1w, l2w = transformed[block_n] + + y_ref = torch.zeros((m, hidden), dtype=torch.bfloat16, device="cuda") + ref_buf.x[:m].copy_(x_fp8) + ref_buf.x_sf[:m].copy_(x_sf) + ref_buf.topk_idx[:m].copy_(topk_idx) + ref_buf.topk_weights[:m].copy_(topk_w) + deep_gemm.nvfp4_mega_moe( + y_ref, + l1w, + l2w, + ref_buf, + recipe=(128, 128, 128), + activation="swiglu", + activation_clamp=10.0, + fast_math=args.fast_math, + kernel_family=family, + family_threshold=FAMILY_THRESHOLD, + ) + torch.cuda.synchronize() + dist.barrier(group=group) + + y_fi = torch.zeros((m, hidden), dtype=torch.bfloat16, device="cuda") + fi_buf.stage(x_fp8, x_sf, topk_idx, topk_w) + nvfp4_mega_moe( + y_fi, + l1w[0], + l2w[0], + fi_buf, + kernel_family=family, + family_threshold=FAMILY_THRESHOLD, + ) + torch.cuda.synchronize() + dist.barrier(group=group) + + # Reduce every reported field across ranks, so the line rank 0 + # prints describes all of them rather than its own slice. + d = (y_fi.float() - y_ref.float()).abs() + flags = torch.tensor( + [ + int(torch.equal(y_fi.view(torch.int16), y_ref.view(torch.int16))), + int(torch.isfinite(y_fi.float()).all().item()), + ], + device="cuda", + dtype=torch.int32, + ) + max_abs = d.max().reshape(1) + num_diff = (d > 0).sum().reshape(1) + dist.all_reduce(flags, op=dist.ReduceOp.MIN, group=group) + dist.all_reduce(max_abs, op=dist.ReduceOp.MAX, group=group) + dist.all_reduce(num_diff, op=dist.ReduceOp.SUM, group=group) + same, finite = bool(flags[0].item()), bool(flags[1].item()) + ok = same and finite + if rank == 0: + tag = "PASS" if ok else "FAIL" + extra = ( + "" + if same + else f" max_abs={max_abs.item():.3e} n={int(num_diff.item())}" + ) + print( + f"{tag} {shape_name} M={m} {family} fm={args.fast_math} " + f"bit-identical={same} finite={finite}{extra}", + flush=True, + ) + failures += 0 if ok else 1 + + ref_buf.destroy() + fi_buf.destroy() + del transformed, l1_packed, l2_packed, l1_scale, l2_scale + torch.cuda.empty_cache() + + dist.barrier(group=group) + if rank == 0: + print(f"DONE failures={failures}", flush=True) + dist.destroy_process_group() + if failures: + sys.exit(1) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--shapes", nargs="+", default=["flash"], choices=sorted(SHAPES)) + p.add_argument("--m", nargs="+", type=int, default=[1, 2, 8, 16, 32, 64]) + p.add_argument("--cap", type=int, default=8448) + p.add_argument("--num-processes", type=int, default=8) + # kFastMath is a template parameter of the kernel, so it must match on both + # sides. The False path additionally needs the JIT spec to override the + # -use_fast_math that gen_jit_spec() applies to every module. + p.add_argument( + "--no-fast-math", dest="fast_math", action="store_false", default=True + ) + args = p.parse_args() + torch.multiprocessing.spawn( + worker, args=(args.num_processes, args), nprocs=args.num_processes + ) + + +if __name__ == "__main__": + main()