diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6dabf552e..381f3ef21 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,5 +1,5 @@ -# Vendored third-party sources are kept byte-for-byte as upstream ships them. -exclude: '^include/cudnn_frontend/thirdparty/' +# Vendored third-party source bodies retain their upstream formatting. +exclude: '^(include/cudnn_frontend/thirdparty/|python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/)' repos: - repo: https://github.com/pre-commit/mirrors-clang-format diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md index f6150f3b3..6c768b089 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md @@ -92,6 +92,35 @@ operand pair. It preserves the pre-existing scale-factor contract: provide `global_scale_b` where the selected low-precision format requires them. BF16 does not reinterpret these controls; it rejects them instead. +Torch block-scaled callers in dense or discrete output mode that retain +operations for CUDA Graph replay may provide a caller-owned +`descriptor_workspace`. Allocate its size with +`get_grouped_gemm_wgrad_workspace_size_sm100`, keep it alive for as long as the +captured call site may replay, and do not share it between call sites that may +overlap. This lets multiple same-signature calls share one compiled kernel +without sharing mutable runtime TMA descriptors. Callers that omit this +argument retain the compatibility behavior that isolates cached API instances +by explicit dense output address; discrete callers retain the compiled +operation's internal workspace. + +```python +workspace = torch.empty( + cudnn.get_grouped_gemm_wgrad_workspace_size_sm100(num_experts), + dtype=torch.uint8, + device=a_tensor.device, +) +result = cudnn.grouped_gemm_wgrad_wrapper_sm100( + a_tensor=a_tensor, + b_tensor=b_tensor, + sfa_tensor=sfa_tensor, + sfb_tensor=sfb_tensor, + offsets_tensor=offsets_tensor, + wgrad_tensor=wgrad_tensor, + descriptor_workspace=workspace, + output_mode="dense", +) +``` + ## API usage ### BF16 diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md new file mode 100644 index 000000000..fe5ace08c --- /dev/null +++ b/docs/fe-oss-apis/moe_ep.md @@ -0,0 +1,349 @@ +# Mixture of Experts with Expert Parallelism + +`cudnn.moe_ep` provides Rubin SM107 fused SwiGLU MoE execution with optional +expert parallelism. Inference and training share one `MoeEp` object but use +separate call surfaces: + +- `MoeEp.__call__` and `warmup` for inference; +- `prepare_training`, `training_forward`, and `training_backward` for training. + +Training is stateless with respect to caller tensors. The operator retains +compiled kernels, runtime state, and private per-lane NVSHMEM scratch, but it +does not retain weights, forward state, output buffers, or fallback weight +staging. + +## Installation + +```bash +pip install "nvidia-cudnn-frontend[cutedsl,comm]" torch torch-c-dlpack-ext +``` + +Rubin MegaMoE requires `nvidia-cutlass-dsl>=4.8.0`. EP2+ also requires an +initialized NCCL process group and an NVSHMEM topology in which all +participating ranks are directly peer-addressable. + +## Constructing the operator + +```python +from cudnn import MoeEp + +op = MoeEp( + num_experts=num_experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, + ep_group=ep_group, + max_tokens_per_rank=max_tokens, + max_recv_size_per_rank=recv_capacity, + drop_on_overflow=False, + output_format="bf16", + combine_format="bf16", + apply_topk_in_fc1=True, + weight_interleave_size=32, +) +``` + +`max_recv_size_per_rank` is the physical receive-pool size in token rows, +including all per-expert padding. Padding is contained within this capacity. + +Native training requires `weight_interleave_size=32`. FC1 payloads then use +alternating 32-element gate/up strips. + +## Explicit sweep autotuning + +`MoeEp.autotune` measures inference forward. `MoeEp.autotune_training` +measures one training forward immediately followed by its matching backward: + +```python +result = op.autotune( + activation, fc1_weight, fc2_weight, topk_idx, topk_weights, + candidates=candidates, + warmup_iters=3, + timed_iters=10, +) + +training_result = op.autotune_training( + activation, grad_output, topk_idx, topk_weights, + forward_weights=native_fw, + backward_weights=native_bw, + candidates=candidates, +) +``` + +Both calls are collective over `ep_group`, must use the same ordered candidate +list on every rank, and must run outside CUDA Graph capture. +`autotune_training` must run before `prepare_training`. It accepts only native +weights; source packing and allocation are intentionally outside its measured +region. + +The current `MoeEpTuningConfig` is prepended as a baseline, duplicate values are +removed, and the normalized list is limited to 32 candidates. Autotuning keeps +`reduce_topk_in_kernel` fixed because that flag changes where top-k reduction +is performed. Each timed iteration is reduced with rank MAX and the candidate +score is the median of those slow-rank samples. Equal scores select the earlier +candidate. `MoeEpAutotuneResult` reports `winner`, per-candidate `latency_ms` +and `samples_ms`, and `evaluated_candidates`. + +The sweep is fail-fast. Any validation, allocation, compile, launch, timing, +synchronization, or teardown error ends the whole sweep. An error after +runtime/collective entry poisons the operator, and later execution is rejected; +close it and create a new instance. Compiled candidate kernels remain in the +process JIT cache. The production sweep does not compare candidate outputs at +runtime; supported candidates are covered by the separate correctness suite. + +Autotuning commits one active winner per instance. A later inference or +training sweep replaces it. Existing CUDA Graph executables are invalid after +the winner changes. Use these sequences: + +- inference: `autotune` → eager winner launch (performed by `autotune`) → + capture; +- training: `autotune_training` → `prepare_training` → allocate outputs → + eager forward/backward → rank synchronization → capture. + +## Stateless training preparation + +Preparation is collective over `ep_group` and must run outside CUDA Graph +capture: + +```python +requirements = op.prepare_training( + lane_count=1, + device=None, # current CUDA device; pass an explicit device for multi-GPU hosts +) +lane = op.training_lanes[0] +symmetric = op.training_symmetric_buffers(lane) +``` + +`prepare_training` does not accept or bind weights. It returns a plain mapping +whose values are: + +```text +(shape, stride, dtype, alignment_bytes) +``` + +The requirements mapping contains `output`, `fc1_preact`, `fc1_a`, `fc1_sfa`, +`valid_route_counts`, `expert_offsets`, `grad_activation`, `dprob`, `fc1_b`, +`fc1_sfb`, `fc2_a`, `fc2_sfa`, `fc2_b`, and `fc2_sfb`. +`training_symmetric_buffers(lane)` returns the cuDNN-allocated +`forward_input`, `forward_input_scale`, `backward_input`, +`backward_input_scale`, `output`, BF16 `grad_activation`, and FP32 `dprob` +buffers for that lane. TE quantizes directly into the input pairs and binds the +returned output buffers in the public output bundles. TE allocates the +remaining buffers from the requirements mapping. cuDNN validates exact shape, +stride, dtype, alignment, device, and non-aliasing before launch. + +`device=None` binds the current CUDA device. An explicit CUDA device takes +precedence. Every later training tensor must use the bound device. + +## Native weight ABI + +Forward and backward receive independent packs: + +```python +from cudnn import ( + MoeEpNativeForwardWeights, + MoeEpNativeBackwardWeights, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, +) +``` + +Each `MoeEpNativeWeight` contains: + +- `payload`: kernel-native E4M3 data; +- `scale`: contiguous Rubin-blocked E8M0 scales; +- `layout_id`: the exact versioned payload-and-scale layout. + +Execution validates the `layout_id` and passes payload and scale pointers to +the kernel without transformation or retention. Eager calls may use different +weight addresses. CUDA Graph capture pins every referenced address until the +graph executable is destroyed. + +Let `B(R, C) = round_up(R, 128) * round_up(C, 4)`. The native V1 contracts +are: + +- forward FC1: payload `(E_local, H, 2I)`, stride `(2HI, 1, H)`, with + gate/up 32-column strips; scale `(E_local, B(2I, H/32))`; +- forward FC2: payload `(E_local, I, H)`, stride `(IH, 1, I)`; scale + `(E_local, B(H, I/32))`; +- backward W2-transpose: contiguous payload `(E_local, H, I)`; scale + `(E_local, B(I, H/32))`; +- backward W1-transpose: contiguous payload `(E_local, 2I, H)`, with + gate/up 32-row strips; scale `(E_local, B(H, 2I/32))`. + +Every native scale tensor is contiguous E8M0. The corresponding +`MoeEpNativeWeightLayout` enum value is required; a compact or differently +swizzled scale tensor is rejected even when its element count matches. + +When upstream does not already produce native weights, use caller-owned +staging: + +```python +native_fw = op.pack_forward_weights(source_fw, out=forward_staging) +native_bw = op.pack_backward_weights(source_bw, out=backward_staging) +``` + +The equivalent standalone `pack_forward_weights` and `pack_backward_weights` +functions are also exported. Packing allocates nothing: every transformed +payload or scale is written to the supplied staging bundle. These fallback +packers consume logical gate-then-up `MoeEpForwardWeights` / +`MoeEpBackwardWeights` with compact axis-1 scales; already interleaved, +blocked producers should construct the native packs directly instead of +packing them again. + +## Forward + +```python +from cudnn import MoeEpTrainingForwardOutputs + +y = op.training_forward( + lane, + activation, + topk_idx, + topk_weights, + weights=native_fw, + out=MoeEpTrainingForwardOutputs( + output=y_out, + fc1_preact=fc1_preact, + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + valid_route_counts=valid_route_counts, + expert_offsets=expert_offsets, + ), +) +``` + +`activation` may be contiguous BF16/FP32 or an axis-1 MXFP8 +`BlockScaledTensor`. To avoid input staging, quantize directly into the +`forward_input` and `forward_input_scale` views returned by +`training_symmetric_buffers(lane)`. Construct the MXFP8 input from +`forward_input[:T]` and +`forward_input_scale[:T, :ceil_div(hidden_size, 32)]`. The logical scale view +has unit column stride and may retain the lane buffer's padded row stride; +training accepts this layout without copying. Routing metadata is still staged +privately. + +`fc1_preact` is required because the training forward kernel always runs with +`generate_c=True`; TE must provide its destination and retain it through the +matching backward. `fc1_a`, `fc1_sfa`, `valid_route_counts`, and +`expert_offsets` are also required caller-owned destinations after +`prepare_training()`. + +`output` is required and must be the lane's cuDNN-allocated symmetric +`output` buffer. The return is a logical `(T, H)` view of that buffer and the +forward kernel writes it directly. + +## Backward and WGrad + +```python +from cudnn import MoeEpTrainingBackwardOutputs + +dx, dprob, operands = op.training_backward( + lane, + grad_output, + topk_idx, + topk_weights, + weights=native_bw, + fc1_preact=fc1_preact, + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + valid_route_counts=valid_route_counts, + expert_offsets=expert_offsets, + out=MoeEpTrainingBackwardOutputs( + grad_activation=dx_out, + dprob=dprob_out, + fc1_b=fc1_b, + fc1_sfb=fc1_sfb, + fc2_a=fc2_a, + fc2_sfa=fc2_sfa, + fc2_b=fc2_b, + fc2_sfb=fc2_sfb, + ), +) +``` + +`grad_output` has the same BF16/FP32/MXFP8 input choices as forward. To avoid +staging, quantize it directly into the lane's `backward_input` and +`backward_input_scale` buffers and construct the same logical prefix views +described for `forward_input`. +`fc1_preact` and the four forward WGrad values are required and passed +explicitly because cuDNN does not retain the forward output bundle. + +`grad_activation` and `dprob` are required destinations. Both must be the +corresponding symmetric buffers returned by `training_symmetric_buffers(lane)` +and are written directly by the backward kernel. `grad_activation` is BF16 and +`dprob` is FP32. + +All six backward WGrad fields are required. `operands` is always a +`MoeEpTrainingWgradOperands` containing non-owning views of the exact caller +buffers. + +The producer-native ABI is directly consumable by the separately invoked +grouped WGrad kernel. Here `K_pool` is the fixed routed-token pool capacity, +not the model's top-k value: + +- `fc1_b` remains gate/up-interleaved with shape `(K_pool, 2I)` and stride + `(2I, 1)`; +- `fc1_a` and `fc2_a` use the advertised transpose-view layouts; +- all four scale tensors are written in the final grouped-WGrad 128x4 + interleaved layout; +- no public compact scale, deinterleave copy, physical transpose, slot export, + or scale-expansion kernel is used. + +## Ownership and lifetime + +- TE owns all native weights, output bundles, saved forward state, WGrad + operands, and optional pack staging. +- cuDNN borrows caller-allocated tensors for one call and does not cache their + Python objects or pointers. The per-lane symmetric buffers returned by + `training_symmetric_buffers` remain cuDNN-owned for the lane lifetime. +- TE must provide `fc1_preact` to forward and keep it live through the matching + backward; cuDNN has no private preactivation fallback or workspace alias. +- Forward WGrad outputs, segment metadata, and backward WGrad outputs remain + live until the independent grouped WGrad consumer completes. +- cuDNN owns per-lane local and NVSHMEM symmetric storage; only the documented + input and final-output views are exposed to the caller. +- One lane may be active on only one stream at a time. +- All EP ranks must submit distributed forward/backward calls in identical + order. +- The caller owns forward/backward weight-version consistency. +- `MoeEp.close()` releases only private runtime resources and never clears or + frees caller memory. + +## Overflow + +Overflow is private per-launch state. Each forward and backward applies the +configured policy before returning; there is no public overflow tensor or +`finalize_overflow` method. EP2+ retains the scalar MAX reduction required to +make the policy rank-consistent. + +## CUDA Graph capture + +1. Collectively call `prepare_training`. +2. Allocate every capture binding from the returned requirements. +3. Materialize or provide native weights at stable addresses. +4. Run ordinary forward/backward warmups for every captured specialization. +5. Capture calls using every caller-owned destination returned by + `prepare_training()`, including primary outputs, saved forward state, and + forward/backward WGrad tensors. +6. Keep all captured input, output, saved-state, staging, and native-pack + addresses stable until every referencing graph executable is destroyed. + +Dynamic contents may change at fixed addresses. Eager invocations may replace +addresses between calls. + +## Breaking migration + +Removed: + +- `MoeEpTrainingResources` +- `MoeEpTrainingSlot` +- `MoeEpTrainingWeights` +- `prepare_training_resources` +- `refresh_weights` +- `finalize_overflow` + +The old resource-owned forward/backward state is replaced by explicit +per-invocation native weight packs, caller-owned saved/WGrad buffers, and +cuDNN-owned symmetric input/final-output views. No compatibility shim is +retained. diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 9bc31f75f..da1902dc6 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -39,6 +39,9 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d - [RMSNorm + RHT + Amax](rmsnorm_rht_amax.md) - [SDPA Backward (SM120)](attention/sdpa_bwd_sm120.md) - [RMSNorm + SiLU](rmsnorm_silu.md) +- [MoE + Expert Parallel API](moe_ep.md) — Rubin SM107 fused SwiGLU with + stateless caller-owned training buffers and CUDA Graph support; see the + [MoeEP operation reference](../operations/moe_ep.md) for support details ## Installation and setup @@ -55,6 +58,16 @@ pip install --group jax # jax >= 0.5 (XLA entry points via cutlass.jax, ship ``` (For the published wheel, `pip install torch torch-c-dlpack-ext` or `pip install "jax>=0.5"` directly.) +MoE + Expert Parallel composes the reusable CuTeDSL and communication extras: +```bash +pip install "nvidia-cudnn-frontend[cutedsl,comm]" torch torch-c-dlpack-ext +``` + +MoeEP is currently CUDA/PyTorch-only and targets Rubin SM107. EP2+ execution +also requires NCCL, NVSHMEM, and a direct-P2P MNNVL peer-access domain. See the +[MoeEP support matrix and tensor contracts](../operations/moe_ep.md) +before integrating it. + After installation, you can import the APIs directly from the `cudnn` package, i.e. `from cudnn import {your_operation}` ## API Usage @@ -68,6 +81,13 @@ Most compiler-style operations expose the following two APIs. Functional autograd integrations such as BSA and Flex Attention instead document their own wrapper and reusable-plan lifecycle on their operation pages. +MoeEP is an exception to the generic wrapper/kernel pattern below. It exposes +an object API: `MoeEp.__call__` for inference and +`prepare_training` plus stateless `training_forward`/`training_backward` calls +for training. See +[MoE + Expert Parallel API](moe_ep.md) for its lifecycle and CUDA Graph +contract. + ### 1. High-level wrapper - Single pythonic function call diff --git a/docs/operations/moe_ep.md b/docs/operations/moe_ep.md new file mode 100644 index 000000000..e8f14cbd7 --- /dev/null +++ b/docs/operations/moe_ep.md @@ -0,0 +1,239 @@ +# Mixture of Experts with Expert Parallelism + +The MoeEP operation fuses token routing, expert SwiGLU computation, and +expert-parallel communication. Global experts are sharded contiguously across +the ranks of an expert-parallel process group. + +## Operation + +Let $x_t \in \mathbb{R}^{H}$ be token $t$, $e_{t,k}$ its $k$-th +selected global expert, and $p_{t,k}$ the corresponding routing weight. For +each valid route, split the FC1 result into gate and up projections: + +$$ +\left[g_{t,k}, u_{t,k}\right] + = x_t W^{\mathrm{fc1}}_{e_{t,k}}, +\qquad +h_{t,k} + = p_{t,k}\left(\mathrm{SiLU}(g_{t,k}) \odot u_{t,k}\right), +\qquad +z_{t,k} + = h_{t,k} W^{\mathrm{fc2}}_{e_{t,k}}. +$$ + +The final token output is the sum over its selected experts: + +$$ +y_t = \sum_{\substack{0 \le k < K \\ e_{t,k} \ne -1}} z_{t,k}. +$$ + +When `gate_up_clamp=C`, the operation uses +$\min(g_{t,k}, C)$ for the gate and +$\mathrm{clip}(u_{t,k}, -C, C)$ for the up projection. A route whose +expert ID is `-1` contributes zero. Because the executable backend requires +`apply_topk_in_fc1=True`, it applies $p_{t,k}$ to the SwiGLU result before +FC2. The backend also stages plain inputs to MXFP8 and requantizes the routed +intermediate before FC2, so the equations describe the mathematical operation +rather than its finite-precision rounding. + +With $E$ global experts and an expert-parallel group of size $P$, each rank +stores $E_{\mathrm{local}}=E/P$ consecutive experts. Global expert $e$ is +owned by group-relative rank + +$$ +\mathrm{owner}(e) + = \left\lfloor \frac{e}{E_{\mathrm{local}}} \right\rfloor. +$$ + +## Python API + +The operation is exposed by the frontend-only `cudnn.MoeEp` object API. Static +model, parallelism, capacity, and format choices are set in the constructor: + +```python +from cudnn import MoeEp + +op = MoeEp( + num_experts=E, + hidden_size=H, + intermediate_size=I, + top_k=K, + ep_group=ep_group, # None for EP1 + max_tokens_per_rank=max_tokens, + max_recv_size_per_rank=None, # Defaults to P * max_tokens * K + drop_on_overflow=False, + output_format="bf16", + combine_format="bf16", # "bf16" or "mxfp8" + apply_topk_in_fc1=True, + weight_interleave_size=None, # Or 32 for pre-interleaved MXFP8 W1 + gate_up_clamp=None, +) +``` + +`topk_idx` contains global expert IDs. Each rank passes its local tokens and +its contiguous shard of both expert-weight tensors: + +`weight_interleave_size=32` declares that MXFP8 FC1 values already use +alternating 32-element gate/up strips. The default `None` uses conventional +gate-then-up order. Plain BF16/FP16/FP32 weights remain conventional and reject +the interleaved contract because they must be quantized and staged internally. + +```python +output = op( + activation, # (T, H) + fc1_weight, # (E_local, H, 2I) + fc2_weight, # (E_local, I, H) + topk_idx, # (T, K), global expert IDs or -1 + topk_weights, # (T, K) +) # (T, H), BF16 +``` + +For inference CUDA Graph capture, call `op.warmup(...)` with the exact +bindings before capture. `MoeEp` supports `close()` and context-manager use. + +Explicit sweep autotuning is available before capture: + +```python +from cudnn import MoeEpTuningConfig + +result = op.autotune( + activation, fc1_weight, fc2_weight, topk_idx, topk_weights, + candidates=[ + MoeEpTuningConfig(token_in_flag_batch=2), + MoeEpTuningConfig(group_hint=256), + ], +) +``` + +The current tuning is always included as the baseline. Candidates are +de-duplicated and limited to 32 including that baseline. The winner is applied +only to this operator instance. + +Stateless training prepares only private execution lanes. Every invocation +receives independent native weights and caller-owned outputs: + +```python +requirements = op.prepare_training(lane_count=1, device=device) +lane = op.training_lanes[0] + +output = op.training_forward( + lane, activation, topk_idx, topk_weights, + weights=native_forward_weights, + out=forward_outputs, +) +grad_activation, dprob, wgrad_operands = op.training_backward( + lane, grad_output, topk_idx, topk_weights, + weights=native_backward_weights, + fc1_preact=forward_outputs.fc1_preact, + fc1_a=forward_outputs.fc1_a, + fc1_sfa=forward_outputs.fc1_sfa, + valid_route_counts=forward_outputs.valid_route_counts, + expert_offsets=forward_outputs.expert_offsets, + out=backward_outputs, +) +``` + +The WGrad result is a fixed-capacity grouped-GEMM operand bundle, not dense +optimizer-ready weight gradients. See the detailed +[MoE + Expert Parallel API](../fe-oss-apis/moe_ep.md) reference for +installation, all constructor arguments, native layouts, buffer ownership, +overflow handling, and CUDA Graph requirements. MoeEP is +distinct from the cuDNN graph [MoE Grouped Matmul](MoeGroupedMatmul.md) +operation. + +## Execution support + +- NVIDIA Rubin SM107 GPUs (compute capability 10.7). +- CUDA and PyTorch execution. +- `nvidia-cutlass-dsl>=4.8.0` for the Rubin kernels. The package-wide + `cutedsl` extra retains its 4.5.0 installation floor so other cuDNN Frontend + operations remain usable with older compatible DSL versions. +- Fused SwiGLU with contiguous expert sharding. +- `apply_topk_in_fc1=True`. +- `hidden_size` divisible by 128. +- `intermediate_size` divisible by 256. +- `top_k <= min(32, num_experts)`. +- `num_experts` divisible by the expert-parallel group size. +- An explicit positive `max_tokens_per_rank`. + +The stateless training CUDA Graph path has hardware acceptance through EP32 when +all ranks are in one direct-P2P MNNVL peer-access domain. The Python capability +layer does not impose an EP-size ceiling; cross-MNNVL execution is not part of +the validated support surface. + +## Data formats + +Inference activation and expert weights accept: + +- BF16, FP16, or FP32 plain tensors, staged internally to MXFP8; or +- MXFP8 `BlockScaledTensor` values with logical block axis 1. + +The current executable output format is BF16. The expert-combine path accepts +BF16 or MXFP8. NVFP4 types are represented by the public API but native NVFP4 +operands, combine, and output are not executable by this backend. + +Training accepts contiguous BF16/FP32 or MXFP8 block-scaled activation and +gradient inputs. Execution weights use versioned kernel-native E4M3 payload +and Rubin-blocked E8M0 scale layouts. + +## Tensor contracts + +Let: + +- $T$ be the local token count; +- $H$ be `hidden_size`; +- $I$ be `intermediate_size`; +- $K$ be `top_k`; +- $E_{\mathrm{local}}$ be the local expert count. + +Inference uses: + +- `activation`: `(T, H)`; +- `topk_idx`: `(T, K)`, Int32 or Int64, containing `-1` or a valid global + expert ID; +- `topk_weights`: `(T, K)`, floating point; +- FC1 weights: `(E_local, H, 2I)`; +- FC2 weights: `(E_local, I, H)`; +- output: `(T, H)`, BF16. + +All inference tensors must reside on one device, and the local token count must +satisfy `T <= max_tokens_per_rank`. + +Stateless training uses: + +- `activation` and `grad_output`: contiguous `(T, H)`, BF16, FP32, or MXFP8; +- `topk_idx`: contiguous `(T, K)`, Int32; +- `topk_weights`: contiguous `(T, K)`, FP32; +- independent forward and backward native weight packs with exact versioned + `layout_id` values; +- required caller-owned forward output: `(T, H)`, BF16; +- required caller-owned `fc1_preact`, produced by training forward with + `generate_c=True` and retained through matching backward; +- required caller-owned `grad_activation`: `(T, H)` view of a capacity buffer, + FP32; +- required caller-owned `dprob`: source-order `(T, K)`, FP32; +- required caller-owned WGrad saved state and a fixed-capacity + `MoeEpTrainingWgradOperands` bundle. + +All dynamic training tensors must reside on one device and satisfy +`T <= max_tokens_per_rank`. + +## Expert-parallel communication + +EP2+ execution requires: + +- an initialized NCCL process group; +- `nvshmem4py` and usable NVSHMEM libraries; +- direct peer access among every pair of participating ranks; and +- consistent rank ordering, buffer schemas, tuning, lane selection, and launch + ordering across the group. + +`max_recv_size_per_rank` is the physical receive-pool capacity in token rows, +including per-expert padding. + +```text +ep_size * max_tokens_per_rank * top_k +``` + +Private lane resources cannot grow during CUDA Graph replay. Capacity changes +require a new operator preparation; caller-address changes require recapture. diff --git a/llms.txt b/llms.txt index 17dcfc734..b47ba6304 100644 --- a/llms.txt +++ b/llms.txt @@ -18,6 +18,7 @@ Published documentation: https://docs.nvidia.com/deeplearning/cudnn/latest/devel - [Convolutions](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Convolutions.md) - [Normalizations (LayerNorm, RMSNorm, BatchNorm, InstanceNorm)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Normalizations.md) - [MoE Grouped Matmul](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/MoeGroupedMatmul.md) +- [MoE with Expert Parallelism](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/moe_ep.md) - [Pointwise](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Pointwise.md) - [Block Scaling (MXFP8/NVFP4 quantization)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/BlockScaling.md) - [RoPE](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/RoPE.md) @@ -32,6 +33,7 @@ Published documentation: https://docs.nvidia.com/deeplearning/cudnn/latest/devel - [FROST SDPA support-matrix tracker — what the FROST SDPA engines serve, per architecture](https://github.com/NVIDIA/cudnn-frontend/blob/main/python/cudnn/sdpa/frost/SUPPORT_MATRIX_TRACKER.md) - [Block-sparse attention (BSA)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/bsa.md), [DeepSeek sparse attention (DSA)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/dsa.md), [Native sparse attention (NSA)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/nsa.md) - [GEMM fusions (amax, SwiGLU, sReLU, grouped/discrete MoE variants)](https://github.com/NVIDIA/cudnn-frontend/tree/main/docs/fe-oss-apis/gemm_fusions) +- [MoE + Expert Parallel Python API](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/moe_ep.md) - [RMSNorm + RHT + Amax](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/rmsnorm_rht_amax.md), [RMSNorm + SiLU](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/rmsnorm_silu.md) ## How-to guides diff --git a/pyproject.toml b/pyproject.toml index 0a8bbcd19..734b859d4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "nvidia-cudnn-frontend" dynamic = ["version"] -description = "NVIDIA cuDNN Frontend — Python and C++ Graph API with SOTA attention (SDPA / Flash Attention), MoE grouped GEMM fusions, and FP8/MXFP8 kernels for Hopper and Blackwell GPUs." +description = "NVIDIA cuDNN Frontend — Python and C++ Graph API with SOTA attention (SDPA / Flash Attention), MoE grouped GEMM fusions, and FP8/MXFP8 kernels for Hopper, Blackwell, and Rubin GPUs." readme = "README.md" requires-python = ">=3.10" license = {text = "Apache-2.0 AND MIT"} @@ -21,6 +21,7 @@ keywords = [ "fp8", "mxfp8", "blackwell", + "rubin", "hopper", "pytorch", "kernel", @@ -72,6 +73,11 @@ cutedsl = [ # still writes it that way. "cuda-python", ] +comm = [ + # Communication runtimes are grouped by backend so future distributed + # operation graphs can reuse the same installation extra. + "nvshmem4py-cu13>=0.3.1", +] cutile = [ # The cuTile linear-attention engines. Base cuda-tile only -- its [tileiras] # extra pins cuda-toolkit>=13.2,<13.4, and that upper bound would cap the @@ -123,3 +129,7 @@ version = {attr = "cudnn.__version__"} [tool.setuptools.package-data] include = ["**/*"] +"cudnn.moe_ep._megamoe_backend.cutedsl_src" = [ + "LICENSE.Apache-2.0", + "VENDOR.md", +] diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 0c5126a25..ffb0df992 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -305,19 +305,81 @@ def _dlopen_cudnn(): ) __all__ = [*_EAGER_PUBLIC_NAMES, "Graph", "wrapper"] -_OPTIONAL_DEPENDENCY_INSTALL_HINT = "Install with 'pip install nvidia-cudnn-frontend[cutedsl]'" +_CUTEDSL_INSTALL_HINT = "Install with 'pip install nvidia-cudnn-frontend[cutedsl]'" +_MOE_EP_INSTALL_HINT = "Install with 'pip install " '"nvidia-cudnn-frontend[cutedsl,comm]" torch torch-c-dlpack-ext\'' +_MOE_EP_OPTIONAL_IMPORTS = { + "moe_ep", + "BlockScaledTensor", + "MoeEp", + "MoeEpAutotuneCandidateResult", + "MoeEpAutotuneResult", + "MoeEpBackwardWeightStaging", + "MoeEpBackwardWeights", + "MoeEpExecutionLane", + "MoeEpForwardWeightStaging", + "MoeEpForwardWeights", + "MoeEpNativeBackwardWeights", + "MoeEpNativeForwardWeights", + "MoeEpNativeWeight", + "MoeEpNativeWeightLayout", + "MoeEpTrainingBackwardOutputs", + "MoeEpTrainingForwardOutputs", + "MoeEpTrainingWgradOperands", + "MoeEpTuningConfig", + "MoeFormat", + "MoeTensor", + "pack_backward_weights", + "pack_forward_weights", +} _LAZY_OPTIONAL_IMPORTS = { "gnn": (".gnn", None), + "moe_ep": (".moe_ep", None), + "BlockScaledTensor": (".moe_ep", "BlockScaledTensor"), + "MoeEp": (".moe_ep", "MoeEp"), + "MoeEpAutotuneCandidateResult": ( + ".moe_ep", + "MoeEpAutotuneCandidateResult", + ), + "MoeEpAutotuneResult": (".moe_ep", "MoeEpAutotuneResult"), + "MoeEpBackwardWeightStaging": (".moe_ep", "MoeEpBackwardWeightStaging"), + "MoeEpBackwardWeights": (".moe_ep", "MoeEpBackwardWeights"), + "MoeEpExecutionLane": (".moe_ep", "MoeEpExecutionLane"), + "MoeEpForwardWeightStaging": (".moe_ep", "MoeEpForwardWeightStaging"), + "MoeEpForwardWeights": (".moe_ep", "MoeEpForwardWeights"), + "MoeEpNativeBackwardWeights": (".moe_ep", "MoeEpNativeBackwardWeights"), + "MoeEpNativeForwardWeights": (".moe_ep", "MoeEpNativeForwardWeights"), + "MoeEpNativeWeight": (".moe_ep", "MoeEpNativeWeight"), + "MoeEpNativeWeightLayout": (".moe_ep", "MoeEpNativeWeightLayout"), + "MoeEpTrainingBackwardOutputs": (".moe_ep", "MoeEpTrainingBackwardOutputs"), + "MoeEpTrainingForwardOutputs": (".moe_ep", "MoeEpTrainingForwardOutputs"), + "MoeEpTrainingWgradOperands": ( + ".moe_ep", + "MoeEpTrainingWgradOperands", + ), + "MoeEpTuningConfig": (".moe_ep", "MoeEpTuningConfig"), + "MoeFormat": (".moe_ep", "MoeFormat"), + "MoeTensor": (".moe_ep", "MoeTensor"), + "pack_backward_weights": (".moe_ep", "pack_backward_weights"), + "pack_forward_weights": (".moe_ep", "pack_forward_weights"), "FlexAttentionBwd": (".flex_attention", "FlexAttentionBwd"), "FlexAttentionFwd": (".flex_attention", "FlexAttentionFwd"), "create_mask_plan": (".flex_attention", "create_mask_plan"), "flex_attn_func": (".flex_attention", "flex_attn_func"), "sdpa_torch": (".sdpa.fwd.torch_op", "sdpa"), "BSA": (".block_sparse_attention", "BSA"), - "block_sparse_attention_forward": (".block_sparse_attention", "block_sparse_attention_forward"), - "block_sparse_attention_fp8_forward": (".block_sparse_attention", "block_sparse_attention_fp8_forward"), - "block_sparse_attention_backward": (".block_sparse_attention", "block_sparse_attention_backward"), + "block_sparse_attention_forward": ( + ".block_sparse_attention", + "block_sparse_attention_forward", + ), + "block_sparse_attention_fp8_forward": ( + ".block_sparse_attention", + "block_sparse_attention_fp8_forward", + ), + "block_sparse_attention_backward": ( + ".block_sparse_attention", + "block_sparse_attention_backward", + ), "DSA": (".deepseek_sparse_attention", "DSA"), "CSA": (".csa", "CSA"), "CSACompressorForward": (".csa", "CSACompressorForward"), @@ -326,62 +388,159 @@ def _dlopen_cudnn(): "csa_compressor_backward_wrapper": (".csa", "csa_compressor_backward_wrapper"), "NSA": (".native_sparse_attention", "NSA"), "GemmSwigluSm100": (".gemm.cutedsl.dense.swiglu", "GemmSwigluSm100"), - "gemm_swiglu_wrapper_sm100": (".gemm.cutedsl.dense.swiglu", "gemm_swiglu_wrapper_sm100"), + "gemm_swiglu_wrapper_sm100": ( + ".gemm.cutedsl.dense.swiglu", + "gemm_swiglu_wrapper_sm100", + ), "gemm_swiglu_jax_sm100": (".gemm.cutedsl.dense.swiglu", "gemm_swiglu_jax_sm100"), "gemm_srelu_jax_sm100": (".gemm.cutedsl.dense.srelu", "gemm_srelu_jax_sm100"), "gemm_dsrelu_jax_sm100": (".gemm.cutedsl.dense.dsrelu", "gemm_dsrelu_jax_sm100"), "GemmSreluSm100": (".gemm.cutedsl.dense.srelu", "GemmSreluSm100"), - "gemm_srelu_wrapper_sm100": (".gemm.cutedsl.dense.srelu", "gemm_srelu_wrapper_sm100"), + "gemm_srelu_wrapper_sm100": ( + ".gemm.cutedsl.dense.srelu", + "gemm_srelu_wrapper_sm100", + ), "GemmDsreluSm100": (".gemm.cutedsl.dense.dsrelu", "GemmDsreluSm100"), - "gemm_dsrelu_wrapper_sm100": (".gemm.cutedsl.dense.dsrelu", "gemm_dsrelu_wrapper_sm100"), + "gemm_dsrelu_wrapper_sm100": ( + ".gemm.cutedsl.dense.dsrelu", + "gemm_dsrelu_wrapper_sm100", + ), "GemmAmaxSm100": (".gemm.cutedsl.dense.amax", "GemmAmaxSm100"), "gemm_amax_wrapper_sm100": (".gemm.cutedsl.dense.amax", "gemm_amax_wrapper_sm100"), "gemm_amax_jax_sm100": (".gemm.cutedsl.dense.amax", "gemm_amax_jax_sm100"), - "GemmProjRopeMxfp8Bf16InSm100": (".gemm.cutedsl.dense.proj_rope_mxfp8", "GemmProjRopeMxfp8Bf16InSm100"), - "GemmProjRopeMxfp8Mxfp8InSm100": (".gemm.cutedsl.dense.proj_rope_mxfp8", "GemmProjRopeMxfp8Mxfp8InSm100"), - "gemm_proj_rope_mxfp8_wrapper_sm100": (".gemm.cutedsl.dense.proj_rope_mxfp8", "gemm_proj_rope_mxfp8_wrapper_sm100"), - "gemm_proj_rope_mxfp8_jax_sm100": (".gemm.cutedsl.dense.proj_rope_mxfp8", "gemm_proj_rope_mxfp8_jax_sm100"), + "GemmProjRopeMxfp8Bf16InSm100": ( + ".gemm.cutedsl.dense.proj_rope_mxfp8", + "GemmProjRopeMxfp8Bf16InSm100", + ), + "GemmProjRopeMxfp8Mxfp8InSm100": ( + ".gemm.cutedsl.dense.proj_rope_mxfp8", + "GemmProjRopeMxfp8Mxfp8InSm100", + ), + "gemm_proj_rope_mxfp8_wrapper_sm100": ( + ".gemm.cutedsl.dense.proj_rope_mxfp8", + "gemm_proj_rope_mxfp8_wrapper_sm100", + ), + "gemm_proj_rope_mxfp8_jax_sm100": ( + ".gemm.cutedsl.dense.proj_rope_mxfp8", + "gemm_proj_rope_mxfp8_jax_sm100", + ), "RmsNormRhtAmaxSm100": (".rmsnorm_rht_amax", "RmsNormRhtAmaxSm100"), - "rmsnorm_rht_amax_wrapper_sm100": (".rmsnorm_rht_amax", "rmsnorm_rht_amax_wrapper_sm100"), + "rmsnorm_rht_amax_wrapper_sm100": ( + ".rmsnorm_rht_amax", + "rmsnorm_rht_amax_wrapper_sm100", + ), "grouped_gemm": (".gemm.cutedsl.grouped", None), "GroupedGemmSm100": (".gemm.cutedsl.grouped", "GroupedGemmSm100"), - "grouped_gemm_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_wrapper_sm100"), + "grouped_gemm_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_wrapper_sm100", + ), "grouped_gemm_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_jax_sm100"), - "grouped_gemm_glu_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_glu_jax_sm100"), - "grouped_gemm_dglu_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dglu_jax_sm100"), - "grouped_gemm_dsrelu_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dsrelu_jax_sm100"), - "grouped_gemm_wgrad_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_wgrad_jax_sm100"), - "discrete_grouped_gemm_swiglu_jax_sm100": (".gemm.cutedsl.discrete_grouped", "discrete_grouped_gemm_swiglu_jax_sm100"), - "discrete_grouped_gemm_dswiglu_jax_sm100": (".gemm.cutedsl.discrete_grouped", "discrete_grouped_gemm_dswiglu_jax_sm100"), + "grouped_gemm_glu_jax_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_glu_jax_sm100", + ), + "grouped_gemm_dglu_jax_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_dglu_jax_sm100", + ), + "grouped_gemm_dsrelu_jax_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_dsrelu_jax_sm100", + ), + "grouped_gemm_wgrad_jax_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_wgrad_jax_sm100", + ), + "discrete_grouped_gemm_swiglu_jax_sm100": ( + ".gemm.cutedsl.discrete_grouped", + "discrete_grouped_gemm_swiglu_jax_sm100", + ), + "discrete_grouped_gemm_dswiglu_jax_sm100": ( + ".gemm.cutedsl.discrete_grouped", + "discrete_grouped_gemm_dswiglu_jax_sm100", + ), "GroupedGemmSwigluSm100": (".gemm.cutedsl.grouped", "GroupedGemmSwigluSm100"), - "grouped_gemm_swiglu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_swiglu_wrapper_sm100"), + "grouped_gemm_swiglu_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_swiglu_wrapper_sm100", + ), "GroupedGemmDswigluSm100": (".gemm.cutedsl.grouped", "GroupedGemmDswigluSm100"), - "grouped_gemm_dswiglu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dswiglu_wrapper_sm100"), + "grouped_gemm_dswiglu_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_dswiglu_wrapper_sm100", + ), "GroupedGemmSreluSm100": (".gemm.cutedsl.grouped", "GroupedGemmSreluSm100"), - "grouped_gemm_srelu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_srelu_wrapper_sm100"), + "grouped_gemm_srelu_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_srelu_wrapper_sm100", + ), "GroupedGemmDsreluSm100": (".gemm.cutedsl.grouped", "GroupedGemmDsreluSm100"), - "grouped_gemm_dsrelu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dsrelu_wrapper_sm100"), + "grouped_gemm_dsrelu_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_dsrelu_wrapper_sm100", + ), "HSTUFwdSm100": (".hstu_attention", "HSTUFwdSm100"), "HSTUBwdSm100": (".hstu_attention", "HSTUBwdSm100"), "hstu_attention_forward": (".hstu_attention", "hstu_attention_forward"), "hstu_attention_backward": (".hstu_attention", "hstu_attention_backward"), "GroupedGemmQuantSm100": (".gemm.cutedsl.grouped", "GroupedGemmQuantSm100"), - "grouped_gemm_quant_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_quant_wrapper_sm100"), + "grouped_gemm_quant_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_quant_wrapper_sm100", + ), "GroupedGemmGluSm100": (".gemm.cutedsl.grouped", "GroupedGemmGluSm100"), - "grouped_gemm_glu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_glu_wrapper_sm100"), - "GroupedGemmGluHadamardSm100": (".gemm.cutedsl.grouped", "GroupedGemmGluHadamardSm100"), - "grouped_gemm_glu_hadamard_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_glu_hadamard_wrapper_sm100"), - "GroupedGemmGluHadamardQuantSm100": (".gemm.cutedsl.grouped", "GroupedGemmGluHadamardQuantSm100"), - "grouped_gemm_glu_hadamard_quant_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_glu_hadamard_quant_wrapper_sm100"), + "grouped_gemm_glu_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_glu_wrapper_sm100", + ), + "GroupedGemmGluHadamardSm100": ( + ".gemm.cutedsl.grouped", + "GroupedGemmGluHadamardSm100", + ), + "grouped_gemm_glu_hadamard_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_glu_hadamard_wrapper_sm100", + ), + "GroupedGemmGluHadamardQuantSm100": ( + ".gemm.cutedsl.grouped", + "GroupedGemmGluHadamardQuantSm100", + ), + "grouped_gemm_glu_hadamard_quant_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_glu_hadamard_quant_wrapper_sm100", + ), "GroupedGemmDgluSm100": (".gemm.cutedsl.grouped", "GroupedGemmDgluSm100"), - "grouped_gemm_dglu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dglu_wrapper_sm100"), + "grouped_gemm_dglu_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_dglu_wrapper_sm100", + ), "GroupedGemmWgradSm100": (".gemm.cutedsl.grouped", "GroupedGemmWgradSm100"), - "grouped_gemm_wgrad_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_wgrad_wrapper_sm100"), + "get_grouped_gemm_wgrad_workspace_size_sm100": ( + ".gemm.cutedsl.grouped", + "get_grouped_gemm_wgrad_workspace_size_sm100", + ), + "grouped_gemm_wgrad_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_wgrad_wrapper_sm100", + ), "discrete_grouped_gemm": (".gemm.cutedsl.discrete_grouped", None), - "DiscreteGroupedGemmSwigluSm100": (".gemm.cutedsl.discrete_grouped", "DiscreteGroupedGemmSwigluSm100"), - "discrete_grouped_gemm_swiglu_wrapper_sm100": (".gemm.cutedsl.discrete_grouped", "discrete_grouped_gemm_swiglu_wrapper_sm100"), - "DiscreteGroupedGemmDswigluSm100": (".gemm.cutedsl.discrete_grouped", "DiscreteGroupedGemmDswigluSm100"), - "discrete_grouped_gemm_dswiglu_wrapper_sm100": (".gemm.cutedsl.discrete_grouped", "discrete_grouped_gemm_dswiglu_wrapper_sm100"), + "DiscreteGroupedGemmSwigluSm100": ( + ".gemm.cutedsl.discrete_grouped", + "DiscreteGroupedGemmSwigluSm100", + ), + "discrete_grouped_gemm_swiglu_wrapper_sm100": ( + ".gemm.cutedsl.discrete_grouped", + "discrete_grouped_gemm_swiglu_wrapper_sm100", + ), + "DiscreteGroupedGemmDswigluSm100": ( + ".gemm.cutedsl.discrete_grouped", + "DiscreteGroupedGemmDswigluSm100", + ), + "discrete_grouped_gemm_dswiglu_wrapper_sm100": ( + ".gemm.cutedsl.discrete_grouped", + "discrete_grouped_gemm_dswiglu_wrapper_sm100", + ), } @@ -408,7 +567,8 @@ def _optional_dependency_message(name: str, error: Exception) -> str: too_old = None if too_old is not None: return f"{too_old}: {error}" - return f"{name} requires optional dependencies. {_OPTIONAL_DEPENDENCY_INSTALL_HINT}: {error}" + install_hint = _MOE_EP_INSTALL_HINT if name in _MOE_EP_OPTIONAL_IMPORTS else _CUTEDSL_INSTALL_HINT + return f"{name} requires optional dependencies. {install_hint}: {error}" def __getattr__(name: str) -> Any: diff --git a/python/cudnn/gemm/cutedsl/grouped/__init__.py b/python/cudnn/gemm/cutedsl/grouped/__init__.py index b9b4ca937..e70a1d745 100644 --- a/python/cudnn/gemm/cutedsl/grouped/__init__.py +++ b/python/cudnn/gemm/cutedsl/grouped/__init__.py @@ -48,6 +48,7 @@ from .wgrad.api import ( GroupedGemmWgradSm100, + get_grouped_gemm_wgrad_workspace_size_sm100, grouped_gemm_wgrad_wrapper_sm100, ) @@ -76,6 +77,7 @@ "GroupedGemmDgluSm100", "grouped_gemm_dglu_wrapper_sm100", "GroupedGemmWgradSm100", + "get_grouped_gemm_wgrad_workspace_size_sm100", "grouped_gemm_wgrad_wrapper_sm100", "GroupedGemmSm100", "grouped_gemm_wrapper_sm100", diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/__init__.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/__init__.py index 85d717865..c81e20eec 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/__init__.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/__init__.py @@ -3,11 +3,13 @@ from .api import ( GroupedGemmWgradSm100, + get_grouped_gemm_wgrad_workspace_size_sm100, grouped_gemm_wgrad_wrapper_sm100, ) __all__ = [ "GroupedGemmWgradSm100", + "get_grouped_gemm_wgrad_workspace_size_sm100", "grouped_gemm_wgrad_wrapper_sm100", "grouped_gemm_wgrad_jax_sm100", ] diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py index 782fc3820..5aff1ecdd 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py @@ -155,6 +155,8 @@ def __init__( self.accumulate_on_output = accumulate_on_output self._kernel = _get_rubin_kernel() if self._is_rubin_kernel else BlockScaledMoEGroupedGemmWgradKernel self._workspace = None + self._workspace_bytes = None + self._workspace_arg = None def _validate_offsets(self, offsets_tensor: torch.Tensor, tokens_sum: int, name: str) -> Tuple[int, ...]: self._value_error_if(offsets_tensor.ndim != 1, f"{name} must be rank-1, got shape {tuple(offsets_tensor.shape)}") @@ -335,7 +337,8 @@ def compile(self) -> None: hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - self._workspace = torch.empty(max(kernel.get_workspace_bytes(), 1), dtype=torch.uint8, device=self.a_desc.device) + self._workspace_bytes = max(kernel.get_workspace_bytes(), 1) + self._workspace = torch.empty(self._workspace_bytes, dtype=torch.uint8, device=self.a_desc.device) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) if self.weight_mode == MoEWeightMode.DENSE: @@ -419,8 +422,11 @@ def _compile_dense(self, kernel, max_active_clusters, fake_stream) -> None: None, options="--enable-tvm-ffi", ) - - cached_workspace = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) + self._workspace_arg = from_dlpack( + self._workspace, + assumed_align=128, + enable_tvm_ffi=True, + ) def tensor_api( a_tensor: torch.Tensor, @@ -429,6 +435,7 @@ def tensor_api( sfb_tensor: torch.Tensor, wgrad_tensor: torch.Tensor, offsets_tensor: torch.Tensor, + workspace, stream: cuda.CUstream, global_scale_a: Optional[torch.Tensor], global_scale_b: Optional[torch.Tensor], @@ -440,7 +447,7 @@ def tensor_api( sfb_tensor, wgrad_tensor, offsets_tensor, - cached_workspace, + workspace, stream, global_scale_a, global_scale_b, @@ -530,7 +537,11 @@ def _compile_discrete(self, kernel, max_active_clusters, fake_stream) -> None: options="--enable-tvm-ffi", ) - cached_workspace = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) + self._workspace_arg = from_dlpack( + self._workspace, + assumed_align=128, + enable_tvm_ffi=True, + ) single_expert_placeholder = torch.empty_strided( self.single_expert_wgrad_desc.shape, self.single_expert_wgrad_desc.stride, @@ -550,6 +561,7 @@ def tensor_api( sfb_tensor: torch.Tensor, wgrad_ptrs: torch.Tensor, offsets_tensor: torch.Tensor, + workspace, stream: cuda.CUstream, global_scale_a: Optional[torch.Tensor], global_scale_b: Optional[torch.Tensor], @@ -561,7 +573,7 @@ def tensor_api( sfb_tensor, wgrad_ptrs.data_ptr(), offsets_tensor, - cached_workspace, + workspace, stream, global_scale_a, global_scale_b, @@ -579,6 +591,7 @@ def execute( offsets_tensor: torch.Tensor, wgrad_tensor: Optional[torch.Tensor] = None, wgrad_ptrs: Optional[torch.Tensor] = None, + descriptor_workspace: Optional[torch.Tensor] = None, global_scale_a: Optional[torch.Tensor] = None, global_scale_b: Optional[torch.Tensor] = None, current_stream: Optional[cuda.CUstream] = None, @@ -590,6 +603,30 @@ def execute( if self.weight_mode == MoEWeightMode.DENSE: self._value_error_if(wgrad_tensor is None, "wgrad_tensor is required in dense mode") + if descriptor_workspace is None: + workspace_arg = self._workspace_arg + else: + self._value_error_if( + descriptor_workspace.dtype != torch.uint8, + f"descriptor_workspace must have dtype uint8, got {descriptor_workspace.dtype}", + ) + self._value_error_if( + descriptor_workspace.device != wgrad_tensor.device, + "descriptor_workspace and wgrad_tensor must be on the same device", + ) + self._value_error_if( + not descriptor_workspace.is_contiguous(), + "descriptor_workspace must be contiguous", + ) + self._value_error_if( + descriptor_workspace.numel() < self._workspace_bytes, + f"descriptor_workspace requires at least {self._workspace_bytes} bytes, " f"got {descriptor_workspace.numel()}", + ) + workspace_arg = from_dlpack( + descriptor_workspace, + assumed_align=128, + enable_tvm_ffi=True, + ) self._compiled_kernel( a_tensor, b_tensor, @@ -597,6 +634,7 @@ def execute( sfb_tensor, wgrad_tensor, offsets_tensor, + workspace_arg, current_stream, global_scale_a, global_scale_b, @@ -614,6 +652,30 @@ def execute( ptrs = [wgrad_tensor.data_ptr() + i * expert_stride_bytes for i in range(wgrad_tensor.shape[0])] wgrad_ptrs = torch.tensor(ptrs, dtype=torch.int64, device=wgrad_tensor.device) _validate_pointer_tensor(wgrad_ptrs, "wgrad_ptrs", self.expert_cnt) + if descriptor_workspace is None: + workspace_arg = self._workspace_arg + else: + self._value_error_if( + descriptor_workspace.dtype != torch.uint8, + f"descriptor_workspace must have dtype uint8, got {descriptor_workspace.dtype}", + ) + self._value_error_if( + descriptor_workspace.device != a_tensor.device, + "descriptor_workspace and WGrad operands must be on the same device", + ) + self._value_error_if( + not descriptor_workspace.is_contiguous(), + "descriptor_workspace must be contiguous", + ) + self._value_error_if( + descriptor_workspace.numel() < self._workspace_bytes, + f"descriptor_workspace requires at least {self._workspace_bytes} bytes, " f"got {descriptor_workspace.numel()}", + ) + workspace_arg = from_dlpack( + descriptor_workspace, + assumed_align=128, + enable_tvm_ffi=True, + ) self._compiled_kernel( a_tensor, b_tensor, @@ -621,6 +683,7 @@ def execute( sfb_tensor, wgrad_ptrs, offsets_tensor, + workspace_arg, current_stream, global_scale_a, global_scale_b, diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py index 948d7bc39..38985dea7 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py @@ -30,7 +30,7 @@ backend_cache_key, select_grouped_gemm_backend, ) -from ..moe_utils import WGradInputOrder +from ..moe_utils import MoEWeightMode, WGradInputOrder, WgradSfTensormapConstructor def _block_scaled_dtype_pairs(): @@ -48,6 +48,33 @@ def _block_scaled_dtype_pairs(): _cache_of_GroupedGemmWgradSm100Objects = {} +def get_grouped_gemm_wgrad_workspace_size_sm100( + num_experts: int, + *, + output_mode: str = "dense", + input_order: WGradInputOrder | str = WGradInputOrder.Tensor2D, +) -> int: + """Return required runtime TMA-descriptor workspace bytes.""" + if num_experts <= 0: + raise ValueError(f"num_experts must be positive, got {num_experts}") + try: + weight_mode = MoEWeightMode(output_mode) + except ValueError as exc: + raise ValueError(f"unsupported output_mode {output_mode!r}") from exc + try: + normalized_input_order = WGradInputOrder(input_order) + except ValueError as exc: + raise ValueError(f"unsupported input_order {input_order!r}") from exc + return max( + WgradSfTensormapConstructor.get_workspace_size( + normalized_input_order, + weight_mode, + num_experts, + ), + 1, + ) + + from ._bf16_api import GroupedGemmWgradBf16API from ._blockscaled_api import ( GroupedGemmWgradBlockScaledAPI, @@ -173,6 +200,7 @@ def execute( offsets_tensor: torch.Tensor, wgrad_tensor: Optional[torch.Tensor] = None, wgrad_ptrs: Optional[torch.Tensor] = None, + descriptor_workspace: None = None, *, global_scale_a: None = None, global_scale_b: None = None, @@ -189,6 +217,7 @@ def execute( offsets_tensor: torch.Tensor, wgrad_tensor: Optional[torch.Tensor] = None, wgrad_ptrs: Optional[torch.Tensor] = None, + descriptor_workspace: Optional[torch.Tensor] = None, *, global_scale_a: Optional[torch.Tensor] = None, global_scale_b: Optional[torch.Tensor] = None, @@ -203,13 +232,19 @@ def execute( offsets_tensor: torch.Tensor, wgrad_tensor: Optional[torch.Tensor] = None, wgrad_ptrs: Optional[torch.Tensor] = None, + descriptor_workspace: Optional[torch.Tensor] = None, global_scale_a: Optional[torch.Tensor] = None, global_scale_b: Optional[torch.Tensor] = None, current_stream: Optional[cuda.CUstream] = None, ) -> None: if self._implementation is None: raise RuntimeError("Kernel not compiled; call compile() first") - self._implementation.execute( + if descriptor_workspace is not None and not isinstance( + self._implementation, + GroupedGemmWgradBlockScaledAPI, + ): + raise ValueError("descriptor_workspace requires the block-scaled WGrad backend") + execute_kwargs = dict( a_tensor=a_tensor, b_tensor=b_tensor, sfa_tensor=sfa_tensor, @@ -221,6 +256,9 @@ def execute( global_scale_b=global_scale_b, current_stream=current_stream, ) + if descriptor_workspace is not None: + execute_kwargs["descriptor_workspace"] = descriptor_workspace + self._implementation.execute(**execute_kwargs) def _wgrad_tensor_signature(tensor: Optional[torch.Tensor], *, dynamic_dims: tuple[int, ...] = (), exact_stride: bool): @@ -243,6 +281,7 @@ def grouped_gemm_wgrad_wrapper_sm100( output_mode: str = "dense", wgrad_tensor: Optional[torch.Tensor] = None, wgrad_ptrs: Optional[torch.Tensor] = None, + descriptor_workspace: Optional[torch.Tensor] = None, global_scale_a: Optional[torch.Tensor] = None, global_scale_b: Optional[torch.Tensor] = None, acc_dtype: Optional[torch.dtype] = None, @@ -294,6 +333,19 @@ def grouped_gemm_wgrad_wrapper_sm100( ) if framework == "jax" and backend is GroupedGemmBackend.BLOCK_SCALED: raise ValueError(_BLOCK_SCALED_JAX_ERROR) + if descriptor_workspace is not None and (backend is not GroupedGemmBackend.BLOCK_SCALED or framework != "torch"): + raise ValueError("descriptor_workspace is supported only for torch block-scaled WGrad") + explicit_dense_output_identity = None + if ( + backend is GroupedGemmBackend.BLOCK_SCALED + and framework == "torch" + and output_mode == "dense" + and wgrad_tensor is not None + and descriptor_workspace is None + ): + # Compatibility path: callers that do not own descriptor workspace keep + # the validated one-API-instance-per-output isolation. + explicit_dense_output_identity = int(wgrad_tensor.data_ptr()) if wgrad_tensor is None and wgrad_ptrs is None: wgrad_shape = (expert_cnt, hidden, intermediate) if framework == "torch": @@ -332,6 +384,7 @@ def grouped_gemm_wgrad_wrapper_sm100( accumulate_on_output, input_order, int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")), + explicit_dense_output_identity, ) op = _cache_of_GroupedGemmWgradSm100Objects.get(cache_key) if op is None: @@ -386,6 +439,7 @@ def _sample_wgrad_expert(): offsets_tensor=offsets_tensor, wgrad_tensor=wgrad_tensor, wgrad_ptrs=wgrad_ptrs, + descriptor_workspace=descriptor_workspace, global_scale_a=global_scale_a, global_scale_b=global_scale_b, current_stream=current_stream, diff --git a/python/cudnn/moe_ep/__init__.py b/python/cudnn/moe_ep/__init__.py new file mode 100644 index 000000000..c267cba51 --- /dev/null +++ b/python/cudnn/moe_ep/__init__.py @@ -0,0 +1,50 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +from ._tuning import ( + MoeEpAutotuneCandidateResult, + MoeEpAutotuneResult, + MoeEpTuningConfig, +) +from ._types import ( + BlockScaledTensor, + MoeEpBackwardWeightStaging, + MoeEpBackwardWeights, + MoeEpExecutionLane, + MoeEpForwardWeightStaging, + MoeEpForwardWeights, + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, + MoeEpTrainingBackwardOutputs, + MoeEpTrainingForwardOutputs, + MoeEpTrainingWgradOperands, + MoeFormat, + MoeTensor, +) +from .api import MoeEp, pack_backward_weights, pack_forward_weights + +__all__ = [ + "BlockScaledTensor", + "MoeEp", + "MoeEpAutotuneCandidateResult", + "MoeEpAutotuneResult", + "MoeEpBackwardWeightStaging", + "MoeEpBackwardWeights", + "MoeEpExecutionLane", + "MoeEpForwardWeightStaging", + "MoeEpForwardWeights", + "MoeEpNativeBackwardWeights", + "MoeEpNativeForwardWeights", + "MoeEpNativeWeight", + "MoeEpNativeWeightLayout", + "MoeEpTrainingBackwardOutputs", + "MoeEpTrainingForwardOutputs", + "MoeEpTrainingWgradOperands", + "MoeEpTuningConfig", + "MoeFormat", + "MoeTensor", + "pack_backward_weights", + "pack_forward_weights", +] diff --git a/python/cudnn/moe_ep/_autotune.py b/python/cudnn/moe_ep/_autotune.py new file mode 100644 index 000000000..7698d69c8 --- /dev/null +++ b/python/cudnn/moe_ep/_autotune.py @@ -0,0 +1,207 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Collective coordination and timing helpers for explicit MoeEP sweeps.""" + +from __future__ import annotations + +import math +import statistics +from collections.abc import Callable, Mapping, Sequence +from typing import TypeVar + +import torch +import torch.distributed as dist + +from ._tuning import MoeEpAutotuneCandidateResult, MoeEpTuningConfig +from ._types import MoeEpTrainingBackwardOutputs, MoeEpTrainingForwardOutputs + +_T = TypeVar("_T") +_MAX_AUTOTUNE_CANDIDATES = 32 + + +def normalize_candidates( + baseline: MoeEpTuningConfig, + candidates: Sequence[MoeEpTuningConfig], + *, + warmup_iters: int, + timed_iters: int, + max_candidates: int, +) -> tuple[MoeEpTuningConfig, ...]: + """Validate, de-duplicate, and prepend the current configuration.""" + + if isinstance(warmup_iters, bool) or not isinstance(warmup_iters, int) or warmup_iters < 0: + raise ValueError(f"warmup_iters must be a non-negative integer, got {warmup_iters!r}") + if isinstance(timed_iters, bool) or not isinstance(timed_iters, int) or timed_iters <= 0: + raise ValueError(f"timed_iters must be a positive integer, got {timed_iters!r}") + if isinstance(max_candidates, bool) or not isinstance(max_candidates, int) or not 1 <= max_candidates <= _MAX_AUTOTUNE_CANDIDATES: + raise ValueError(f"max_candidates must be an integer in [1, {_MAX_AUTOTUNE_CANDIDATES}], " f"got {max_candidates!r}") + if not isinstance(candidates, Sequence) or isinstance(candidates, (str, bytes)): + raise TypeError("candidates must be a sequence of MoeEpTuningConfig values") + if not candidates: + raise ValueError("candidates must not be empty") + + ordered: list[MoeEpTuningConfig] = [baseline] + seen = {baseline} + for index, candidate in enumerate(candidates): + if not isinstance(candidate, MoeEpTuningConfig): + raise TypeError("candidates must contain only MoeEpTuningConfig values; " f"candidates[{index}] is {type(candidate).__name__}") + if candidate.reduce_topk_in_kernel != baseline.reduce_topk_in_kernel: + raise ValueError( + "autotune does not sweep reduce_topk_in_kernel; " + f"candidate {index} has {candidate.reduce_topk_in_kernel}, " + f"baseline has {baseline.reduce_topk_in_kernel}" + ) + if candidate not in seen: + ordered.append(candidate) + seen.add(candidate) + + if len(ordered) > max_candidates: + raise ValueError(f"autotune has {len(ordered)} unique candidates including the baseline, " f"exceeding max_candidates={max_candidates}") + return tuple(ordered) + + +def verify_candidates_across_ranks( + candidates: tuple[MoeEpTuningConfig, ...], + group: dist.ProcessGroup | None, +) -> None: + """Fail before runtime allocation when EP ranks supplied different lists.""" + + if group is None: + return + gathered: list[object] = [None] * dist.get_world_size(group) + dist.all_gather_object(gathered, candidates, group=group) + if any(value != candidates for value in gathered): + raise RuntimeError(f"MoeEp autotune candidates must match on every EP rank; " f"rank candidate lists: {gathered}") + + +def verify_state_across_ranks( + state: tuple[object, ...], + group: dist.ProcessGroup | None, +) -> None: + """Require matching operator lifecycle state before collective teardown.""" + + if group is None: + return + gathered: list[object] = [None] * dist.get_world_size(group) + dist.all_gather_object(gathered, state, group=group) + if any(value != state for value in gathered): + raise RuntimeError(f"MoeEp autotune requires matching lifecycle state on every EP rank; " f"rank states: {gathered}") + + +def raise_preflight_errors( + error: BaseException | None, + *, + phase: str, + group: dist.ProcessGroup | None, +) -> None: + """Turn rank-local preflight failures into one collective failure.""" + + if group is None: + if error is not None: + raise error + return + local = None if error is None else (type(error).__name__, str(error)) + gathered: list[object] = [None] * dist.get_world_size(group) + dist.all_gather_object(gathered, local, group=group) + failures = [(rank, value) for rank, value in enumerate(gathered) if value is not None] + if failures: + raise RuntimeError(f"MoeEp autotune {phase} failed before runtime entry; rank errors: {failures}") from error + + +def synchronize_candidate( + device: torch.device, + group: dist.ProcessGroup | None, +) -> None: + """Drain device work and align ranks at a healthy candidate boundary.""" + + torch.cuda.synchronize(device) + if group is not None: + dist.barrier(group=group) + + +def benchmark_candidate( + run: Callable[[], _T], + *, + device: torch.device, + group: dist.ProcessGroup | None, + timed_iters: int, +) -> tuple[float, tuple[float, ...]]: + """Return median(per-iteration rank-MAX) in milliseconds.""" + + stream = torch.cuda.current_stream(device) + local_samples: list[float] = [] + for _ in range(timed_iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record(stream) + run() + end.record(stream) + end.synchronize() + local_samples.append(float(start.elapsed_time(end))) + + slow_rank_samples = torch.tensor(local_samples, dtype=torch.float64, device=device) + if group is not None: + dist.all_reduce(slow_rank_samples, op=dist.ReduceOp.MAX, group=group) + samples = tuple(float(value) for value in slow_rank_samples.cpu().tolist()) + latency_ms = float(statistics.median(samples)) + if not math.isfinite(latency_ms): + raise RuntimeError(f"MoeEp autotune produced a non-finite latency: {latency_ms}") + return latency_ms, samples + + +def select_winner( + results: Sequence[MoeEpAutotuneCandidateResult], +) -> MoeEpAutotuneCandidateResult: + """Choose the first minimum-latency candidate for stable tie-breaking.""" + + if not results: + raise ValueError("cannot select an autotune winner without results") + return min(results, key=lambda result: result.latency_ms) + + +def allocate_training_outputs( + requirements, + device: torch.device, + symmetric_buffers: Mapping[str, torch.Tensor], +) -> tuple[MoeEpTrainingForwardOutputs, MoeEpTrainingBackwardOutputs]: + """Bind symmetric outputs and allocate the remaining one-lane contracts.""" + + def allocate(name: str) -> torch.Tensor: + shape, stride, dtype, alignment = requirements[name] + tensor = torch.empty_strided(shape, stride, dtype=dtype, device=device) + if tensor.data_ptr() % alignment: + raise RuntimeError(f"autotune output {name} is not {alignment}-byte aligned") + return tensor + + forward = MoeEpTrainingForwardOutputs( + fc1_preact=allocate("fc1_preact"), + output=symmetric_buffers["output"], + fc1_a=allocate("fc1_a"), + fc1_sfa=allocate("fc1_sfa"), + valid_route_counts=allocate("valid_route_counts"), + expert_offsets=allocate("expert_offsets"), + ) + backward = MoeEpTrainingBackwardOutputs( + grad_activation=symmetric_buffers["grad_activation"], + dprob=symmetric_buffers["dprob"], + fc1_b=allocate("fc1_b"), + fc1_sfb=allocate("fc1_sfb"), + fc2_a=allocate("fc2_a"), + fc2_sfa=allocate("fc2_sfa"), + fc2_b=allocate("fc2_b"), + fc2_sfb=allocate("fc2_sfb"), + ) + return forward, backward + + +__all__ = [ + "allocate_training_outputs", + "benchmark_candidate", + "normalize_candidates", + "raise_preflight_errors", + "select_winner", + "synchronize_candidate", + "verify_candidates_across_ranks", + "verify_state_across_ranks", +] diff --git a/python/cudnn/moe_ep/_backend.py b/python/cudnn/moe_ep/_backend.py new file mode 100644 index 000000000..11e84a977 --- /dev/null +++ b/python/cudnn/moe_ep/_backend.py @@ -0,0 +1,68 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lightweight private backend seam for :class:`cudnn.moe_ep.MoeEp`. + +Capability policy and the executable factory are imported lazily through this +backend-neutral seam. Importing :mod:`cudnn` still does not load CuTeDSL or +initialize CUDA. +""" + +from __future__ import annotations + +from typing import Protocol + +import torch + +from ._contracts import ForwardConfig, ValidatedForwardRequest +from ._types import MoeTensor + + +class MoeEpBackend(Protocol): + """Instance-local backend created lazily for one static ``MoeEp`` config.""" + + def forward(self, request: ValidatedForwardRequest) -> MoeTensor: + """Execute one already-validated forward request.""" + + def close(self) -> None: + """Release backend-owned resources.""" + + +class BackendUnavailableError(RuntimeError): + """The requested supported path has no executable runtime backend yet.""" + + +def validate_config(config: ForwardConfig) -> None: + """Run the selected backend's static capability gate lazily.""" + + from ._megamoe_backend._capability import validate_config as validate + + validate(config) + + +def validate_request(request: ValidatedForwardRequest) -> None: + """Run the selected backend's per-request capability gate lazily.""" + + from ._megamoe_backend._capability import validate_request as validate + + validate(request) + + +def create_backend( + config: ForwardConfig, + device: torch.device, +) -> MoeEpBackend: + """Create the default backend without an allocation-only fallback.""" + + from ._megamoe_backend.mxfp8._backend import Mxfp8Backend + + return Mxfp8Backend(config, device) + + +__all__ = [ + "BackendUnavailableError", + "MoeEpBackend", + "create_backend", + "validate_config", + "validate_request", +] diff --git a/python/cudnn/moe_ep/_contracts.py b/python/cudnn/moe_ep/_contracts.py new file mode 100644 index 000000000..a88d31ed9 --- /dev/null +++ b/python/cudnn/moe_ep/_contracts.py @@ -0,0 +1,83 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Private data contracts shared by the MoE EP API, validation, and backend.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Any, Literal, Optional + +import torch + +from ._tuning import MoeEpTuningConfig +from ._types import MoeTensor + + +class Fc1WeightLayout(str, Enum): + """Logical gate/up ordering used by FC1 weights and their gradients.""" + + GATE_THEN_UP = "gate_then_up" + GATE_UP_INTERLEAVED_32 = "gate_up_interleaved_32" + + +def normalize_fc1_weight_layout(weight_interleave_size: Optional[int]) -> Fc1WeightLayout: + """Normalize the public compatibility flag into the internal layout ABI.""" + + if weight_interleave_size is None: + return Fc1WeightLayout.GATE_THEN_UP + if weight_interleave_size == 32: + return Fc1WeightLayout.GATE_UP_INTERLEAVED_32 + raise ValueError("weight_interleave_size must be None or 32") + + +@dataclass(frozen=True) +class ForwardConfig: + """Static configuration snapshot for one ``MoeEp`` instance.""" + + num_experts: int + hidden_size: int + intermediate_size: int + top_k: int + experts_per_rank: int + ep_size: int + ep_rank: int + ep_group: Any + ep_global_ranks: tuple[int, ...] + max_tokens_per_rank: Optional[int] + output_format: str + combine_format: str + apply_topk_in_fc1: bool + gate_up_clamp: Optional[float] + generate_c: bool + token_padding_size: int + sf_padding_size: int + tuning: MoeEpTuningConfig + backward_tuning: MoeEpTuningConfig | None = None + backward_wgrad_mode: Literal["none", "operands"] = "none" + max_recv_size_per_rank: Optional[int] = None + drop_on_overflow: bool = False + fc1_weight_layout: Fc1WeightLayout = Fc1WeightLayout.GATE_THEN_UP + + +@dataclass(frozen=True) +class ValidatedForwardRequest: + """Runtime inputs that have passed the public forward contract.""" + + config: ForwardConfig + activation: MoeTensor + fc1_weight: MoeTensor + fc2_weight: MoeTensor + topk_idx: torch.Tensor + topk_weights: torch.Tensor + token_count: int + device: torch.device + + +__all__ = [ + "Fc1WeightLayout", + "ForwardConfig", + "ValidatedForwardRequest", + "normalize_fc1_weight_layout", +] diff --git a/python/cudnn/moe_ep/_math.py b/python/cudnn/moe_ep/_math.py new file mode 100644 index 000000000..466649d19 --- /dev/null +++ b/python/cudnn/moe_ep/_math.py @@ -0,0 +1,17 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Small integer helpers shared by public contracts and private backends.""" + +from __future__ import annotations + + +def ceil_div(value: int, divisor: int) -> int: + return (value + divisor - 1) // divisor + + +def round_up(value: int, multiple: int) -> int: + return ceil_div(value, multiple) * multiple + + +__all__ = ["ceil_div", "round_up"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/README.md b/python/cudnn/moe_ep/_megamoe_backend/README.md new file mode 100644 index 000000000..f1bd9c9e8 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/README.md @@ -0,0 +1,84 @@ +# MegaMoE backend + +The private backend provides Rubin SM107 MXFP8 execution for `cudnn.moe_ep`. + +## Capability + +- CUDA Rubin SM107 +- BF16 output with BF16 or MXFP8 combine +- `hidden_size % 128 == 0` +- `intermediate_size % 256 == 0` +- `top_k <= min(32, num_experts)` +- positive `max_tokens_per_rank` +- `apply_topk_in_fc1=True` + +Inference accepts BF16/FP16/FP32 or MXFP8 operands. Training accepts +BF16/FP32 or MXFP8 activation and grad-output, contiguous Int32 routing +indices, contiguous FP32 routing weights, and independent native forward and +backward weight packs. + +## Execution state + +`MoeEp.prepare_training` creates one private `Mxfp8TrainingState`, which owns +only: + +- prepared forward/backward kernels and compile caches; +- NVSHMEM/runtime handles; +- one local and symmetric scratch slab per execution lane; +- private fixed-capacity transport and routing scratch used only during a call. + +It does not own or retain caller weights, output bundles, saved forward state, +WGrad operands, or weight-staging bundles. No slot is exposed by the public +API. + +## Native weights + +Training execution accepts only `MoeEpNativeForwardWeights` or +`MoeEpNativeBackwardWeights`. Validation checks the exact versioned +`layout_id`, shape, stride, dtype, alignment, and device. The launch adapter +creates aliases to payload and blocked E8M0 scale tensors without allocation, +copy, refresh, or persistent binding. + +`materialize_forward` and `materialize_backward` are allocation-free fallback +transforms. They write only caller-provided staging bundles and return native +packs that alias those destinations. + +## Inputs and outputs + +Plain training inputs use `Mxfp8TrainingStager`. MXFP8 +`BlockScaledTensor` inputs bypass quantization and copy their payload/scales +only into the symmetric transport plane required for peer addressing. + +Caller outputs are borrowed for one launch: + +- required FC1 preactivation is passed directly to forward and backward + kernels; +- all forward and backward WGrad payloads, scales, and route metadata are + required after `prepare_training()` and passed directly to the kernels; +- combine output and dprob first land in private symmetric buffers, then copy + to caller buffers because remote ranks address the symmetric plane; +- primary forward/backward outputs are required caller-owned destinations. + +The producing kernels already expose the final grouped-WGrad scale carriers +when token and scale-factor padding are both 128. Caller E8M0 matrices are +viewed through the producer's flat or matrix signature, so no scale expansion +kernel is launched. FC1-B remains gate/up-interleaved, and FC1-A/FC2-A use +legal transpose views without physical transpose copies. + +## Overflow and distributed ordering + +Each phase keeps overflow state private and applies the configured policy +before returning. EP2+ performs the scalar MAX needed for a rank-consistent +decision. There is no public `finalize_overflow`. + +One lane is exclusive to one active stream. Every EP rank must submit +distributed forward/backward launches in identical order. Distinct lanes do +not make unordered collective-kernel overlap valid. + +## CUDA Graph + +Preparation and first-time compilation happen before capture. Training calls +require every destination advertised by `prepare_training()` to be +caller-owned. Every input, output, saved-state, native weight, and staging +address referenced by a graph remains stable until that graph executable is +destroyed. Eager calls may change addresses between invocations. diff --git a/python/cudnn/moe_ep/_megamoe_backend/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/__init__.py new file mode 100644 index 000000000..8187472b9 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Private MegaMoE backend implementation for :mod:`cudnn.moe_ep`.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/_capability.py b/python/cudnn/moe_ep/_megamoe_backend/_capability.py new file mode 100644 index 000000000..993d3c621 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/_capability.py @@ -0,0 +1,106 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Capability policy for the private MegaMoE execution backend. + +Inputs reaching this module already satisfy the public :mod:`cudnn.moe_ep` +contract. These checks only describe the subset that the current MegaMoE +implementation can execute; they must run before runtime initialization, +allocation, compilation, or collectives. +""" + +from __future__ import annotations + +import torch + +from .._contracts import ( + ForwardConfig, + ValidatedForwardRequest, +) +from .._types import BlockScaledTensor, MoeFormat + + +def _validate_operand(name, tensor) -> None: + if isinstance(tensor, BlockScaledTensor): + if tensor.format is MoeFormat.MXFP8: + return + raise NotImplementedError("MoeEp training MegaMoE supports only MXFP8 BlockScaledTensor " f"inputs; {name} has format={tensor.format.value!r}") + if tensor.dtype not in { + torch.bfloat16, + torch.float16, + torch.float32, + }: + raise NotImplementedError(f"MoeEp MegaMoE {name} staging supports BF16, FP16, " f"or FP32 plain tensors, got {tensor.dtype}") + + +def _validate_device(device: torch.device) -> None: + if device.type != "cuda": + raise NotImplementedError(f"MoeEp MegaMoE backend requires a CUDA device, got {device}") + + major, minor = torch.cuda.get_device_capability(device) + if (major, minor) != (10, 7): + raise NotImplementedError("MoeEp MegaMoE backend requires Rubin SM107 " "(compute capability 10.7); " f"found compute capability {major}.{minor}") + + +def _is_cuda_stream_capturing(device: torch.device) -> bool: + """Return capture state for the request device.""" + + with torch.cuda.device(device): + return torch.cuda.is_current_stream_capturing() + + +def _validate_wgrad_config(config: ForwardConfig) -> None: + if config.backward_wgrad_mode not in ("none", "operands"): + raise ValueError("unsupported backward_wgrad_mode " f"{config.backward_wgrad_mode!r}") + if config.backward_wgrad_mode == "operands": + if not config.generate_c: + raise ValueError("backward_wgrad_mode='operands' requires generate_c=True") + if config.token_padding_size != 128: + raise ValueError("backward_wgrad_mode='operands' requires " "token_padding_size=128") + if config.sf_padding_size != 128: + raise ValueError("backward_wgrad_mode='operands' requires " "sf_padding_size=128") + + +def validate_config(config: ForwardConfig) -> None: + """Reject static configurations outside the current MegaMoE milestone.""" + + _validate_wgrad_config(config) + if config.output_format != MoeFormat.BF16.value: + raise NotImplementedError("MoeEp training MegaMoE supports output_format='bf16' only") + supported_combine_formats = { + MoeFormat.BF16.value, + MoeFormat.MXFP8.value, + } + if config.combine_format not in supported_combine_formats: + raise NotImplementedError("MoeEp training MegaMoE supports combine_format='bf16' " "or 'mxfp8'") + if config.max_tokens_per_rank is None: + raise NotImplementedError("MoeEp MegaMoE backend requires an explicit max_tokens_per_rank") + if config.max_tokens_per_rank == 0: + raise NotImplementedError("MoeEp SM107 MXFP8 execution requires " "max_tokens_per_rank to be positive") + if config.hidden_size % 128: + raise NotImplementedError("MoeEp SM107 MXFP8 kernel currently requires hidden_size " f"to be divisible by 128, got {config.hidden_size}") + if config.intermediate_size % 256: + raise NotImplementedError("MoeEp SM107 MXFP8 kernel currently requires intermediate_size " f"to be divisible by 256, got {config.intermediate_size}") + if config.top_k > 32: + raise NotImplementedError("MoeEp SM107 MXFP8 dispatch currently requires top_k <= 32") + if not config.apply_topk_in_fc1: + raise NotImplementedError("MoeEp Rubin training MegaMoE requires apply_topk_in_fc1=True") + + +def validate_request(request: ValidatedForwardRequest) -> None: + """Reject valid requests outside the current MegaMoE input/device family.""" + + for name, tensor in ( + ("activation", request.activation), + ("fc1_weight", request.fc1_weight), + ("fc2_weight", request.fc2_weight), + ): + _validate_operand(name, tensor) + + _validate_device(request.device) + + +__all__ = [ + "validate_config", + "validate_request", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/_comm.py b/python/cudnn/moe_ep/_megamoe_backend/_comm.py new file mode 100644 index 000000000..b4f67ad04 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/_comm.py @@ -0,0 +1,307 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Symmetric-memory ownership and peer-pointer descriptors for MegaMoE.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from typing import Optional, Protocol + +import torch + +from ._runtime import ( + RuntimeHandle, + RuntimeUnavailableError, + _runtime_debug, +) + + +class SymmetricMemoryProvider(Protocol): + """Injectable allocation boundary for a symmetric root slab.""" + + def allocate(self, nbytes: int, device: torch.device) -> torch.Tensor: ... + + def free(self, tensor: torch.Tensor) -> None: ... + + def peer_address(self, tensor: torch.Tensor, peer: int) -> int: ... + + +class _TorchMemoryProvider: + """CUDA tensor provider for local and single-rank symmetric memory.""" + + def allocate(self, nbytes: int, device: torch.device) -> torch.Tensor: + return torch.empty(nbytes, dtype=torch.uint8, device=device) + + def free(self, tensor: torch.Tensor) -> None: + del tensor + + def peer_address(self, tensor: torch.Tensor, peer: int) -> int: + if peer != 0: + raise ValueError(f"single-rank symmetric memory has no peer {peer}") + return int(tensor.data_ptr()) + + +class _NvshmemMemoryProvider: + """Lazy adapter over NVSHMEM symmetric tensor allocation.""" + + @staticmethod + def _core(): + try: + import nvshmem.core as core + except (ImportError, OSError) as exc: + raise RuntimeUnavailableError("symmetric workspace requires nvshmem4py and NVSHMEM libraries") from exc + return core + + def allocate(self, nbytes: int, device: torch.device) -> torch.Tensor: + del device # NVSHMEM allocates on the device bound during runtime init. + started_at = time.monotonic() + _runtime_debug("symmetric.allocate.begin", nbytes=nbytes) + try: + tensor = self._core().tensor( + (nbytes,), + dtype=torch.uint8, + release=False, + except_on_del=True, + ) + except Exception as exc: + _runtime_debug( + "symmetric.allocate.error", + nbytes=nbytes, + error_type=type(exc).__name__, + error=repr(exc), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) + raise RuntimeUnavailableError(f"failed to allocate {nbytes} bytes from the NVSHMEM symmetric heap") from exc + _runtime_debug( + "symmetric.allocate.end", + nbytes=nbytes, + data_ptr=hex(tensor.data_ptr()), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) + return tensor + + def free(self, tensor: torch.Tensor) -> None: + started_at = time.monotonic() + _runtime_debug( + "symmetric.free.begin", + nbytes=tensor.numel() * tensor.element_size(), + data_ptr=hex(tensor.data_ptr()), + ) + try: + self._core().free_tensor(tensor) + except Exception as exc: + _runtime_debug( + "symmetric.free.error", + error_type=type(exc).__name__, + error=repr(exc), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) + raise RuntimeUnavailableError("failed to free the NVSHMEM symmetric root slab") from exc + _runtime_debug( + "symmetric.free.end", + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) + + def peer_address(self, tensor: torch.Tensor, peer: int) -> int: + started_at = time.monotonic() + _runtime_debug( + "symmetric.peer-map.begin", + peer=peer, + data_ptr=hex(tensor.data_ptr()), + ) + try: + peer_tensor = self._core().get_peer_tensor(tensor, peer) + except Exception as exc: + _runtime_debug( + "symmetric.peer-map.error", + peer=peer, + error_type=type(exc).__name__, + error=repr(exc), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) + raise RuntimeUnavailableError(f"failed to map symmetric root slab for peer {peer}") from exc + peer_pointer = int(peer_tensor.data_ptr()) + _runtime_debug( + "symmetric.peer-map.end", + peer=peer, + peer_data_ptr=hex(peer_pointer), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) + return peer_pointer + + +@dataclass(frozen=True) +class PeerMapping: + """Dense EP-rank peer deltas packed into the vendored kernel mapper ABI.""" + + base_address: int + offsets: tuple[int, ...] + rank: int + + def __post_init__(self) -> None: + if len(self.offsets) == 0: + raise ValueError("peer mapping requires at least one rank") + if self.rank < 0 or self.rank >= len(self.offsets): + raise ValueError(f"peer mapping rank {self.rank} is outside {len(self.offsets)} ranks") + if self.offsets[self.rank] != 0: + raise ValueError(f"local peer offset must be zero, got {self.offsets[self.rank]}") + + @property + def world_size(self) -> int: + return len(self.offsets) + + def to_sym_buffer_host(self): + """Build the CuTeDSL host payload lazily at the launch boundary.""" + + from .cutedsl_src.communication.nvlink_domain.symmetric_buffer import ( + SymmetricBufferHost, + ) + + return SymmetricBufferHost( + base_address=self.base_address, + offsets=self.offsets, + rank=self.rank, + max_ranks=self.world_size, + ) + + +class SymmetricSlab: + """One stable root allocation shared by all peer-visible workspace views.""" + + def __init__( + self, + runtime: RuntimeHandle, + nbytes: int, + *, + provider: Optional[SymmetricMemoryProvider] = None, + ) -> None: + if nbytes <= 0: + raise ValueError(f"symmetric slab size must be positive, got {nbytes}") + runtime.ensure_open() + + self._runtime = runtime + self._nbytes = nbytes + self._provider = provider or (_NvshmemMemoryProvider() if runtime.nvshmem_enabled else _TorchMemoryProvider()) + self._root: Optional[torch.Tensor] = None + self._mapping: Optional[PeerMapping] = None + self._cleanup_required = False + + def ensure_allocated(self) -> None: + if self._cleanup_required: + raise RuntimeError("symmetric slab requires cleanup before allocation") + if self._root is not None and self._mapping is not None: + return + if self._root is not None: + raise RuntimeError("symmetric slab has an allocation pending cleanup") + + _runtime_debug( + "symmetric-slab.ensure.begin", + nbytes=self._nbytes, + world_size=self._runtime.world_size, + ep_rank=self._runtime.rank, + ) + root = self._provider.allocate(self._nbytes, self._runtime.device) + if not isinstance(root, torch.Tensor): + raise TypeError("symmetric memory provider must return a torch.Tensor") + self._root = root + try: + if root.dtype is not torch.uint8 or root.numel() < self._nbytes: + raise ValueError("symmetric root must be a uint8 tensor with at least " f"{self._nbytes} elements") + if root.device != self._runtime.device: + raise ValueError("symmetric root device does not match runtime device: " f"root={root.device}, runtime={self._runtime.device}") + if not root.is_contiguous(): + raise ValueError("symmetric root tensor must be contiguous") + except Exception: + self._cleanup_required = True + raise + + try: + _runtime_debug( + "symmetric-slab.zero.begin", + nbytes=self._nbytes, + data_ptr=hex(root.data_ptr()), + ) + root.zero_() + _runtime_debug("symmetric-slab.zero.enqueued") + + base_address = int(root.data_ptr()) + offsets = [] + for peer in range(self._runtime.world_size): + if peer == self._runtime.rank: + offsets.append(0) + continue + offsets.append(self._provider.peer_address(root, peer) - base_address) + mapping = PeerMapping( + base_address=base_address, + offsets=tuple(offsets), + rank=self._runtime.rank, + ) + if mapping.world_size != self._runtime.world_size or mapping.rank != self._runtime.rank: + raise RuntimeError("symmetric peer mapping does not match the EP subgroup") + except Exception: + self._cleanup_required = True + raise + + self._mapping = mapping + _runtime_debug( + "symmetric-slab.ensure.end", + nbytes=self._nbytes, + base_address=hex(mapping.base_address), + offsets=mapping.offsets, + ) + + @property + def nbytes(self) -> int: + return self._nbytes + + @property + def closed(self) -> bool: + return self._root is None + + @property + def allocated(self) -> bool: + return not self._cleanup_required and self._root is not None and self._mapping is not None + + @property + def mapping(self) -> PeerMapping: + if self._cleanup_required: + raise RuntimeError("symmetric slab requires cleanup") + if self._mapping is None: + raise RuntimeError("symmetric slab is closed") + return self._mapping + + @property + def root(self) -> torch.Tensor: + if self._cleanup_required: + raise RuntimeError("symmetric slab requires cleanup") + if self._root is None: + raise RuntimeError("symmetric slab is closed") + return self._root + + def byte_view(self, offset: int, nbytes: int) -> torch.Tensor: + if offset < 0 or nbytes < 0 or offset + nbytes > self._nbytes: + raise ValueError(f"byte view [{offset}, {offset + nbytes}) exceeds " f"symmetric slab size {self._nbytes}") + return self.root.narrow(0, offset, nbytes) + + def close(self) -> None: + root = self._root + if root is None: + self._cleanup_required = False + return + try: + self._provider.free(root) + except Exception: + self._cleanup_required = True + raise + self._root = None + self._mapping = None + self._cleanup_required = False + + +__all__ = [ + "PeerMapping", + "SymmetricMemoryProvider", + "SymmetricSlab", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/_plan.py b/python/cudnn/moe_ep/_megamoe_backend/_plan.py new file mode 100644 index 000000000..aad99cfc4 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/_plan.py @@ -0,0 +1,157 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy runtime/workspace owner for a compiled MegaMoE execution plan.""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from typing import Optional + +import torch + +from .._contracts import ForwardConfig, ValidatedForwardRequest +from ._comm import SymmetricMemoryProvider +from ._runtime import RuntimeHandle, RuntimeManager, get_runtime_manager +from ._workspace import ( + LocalMemoryProvider, + WorkspaceOwner, + WorkspaceRequirements, + WorkspaceViews, +) + + +@dataclass(frozen=True) +class PreparedResources: + """Resources prepared for staging, compile, and launch integration.""" + + runtime: RuntimeHandle + workspace: WorkspaceViews + + +class ExecutionPlanOwner: + """Own runtime and stable workspace without compiling or launching a kernel.""" + + def __init__( + self, + config: ForwardConfig, + device: torch.device, + requirements: WorkspaceRequirements, + *, + runtime_manager: Optional[RuntimeManager] = None, + symmetric_provider: Optional[SymmetricMemoryProvider] = None, + local_provider: Optional[LocalMemoryProvider] = None, + ) -> None: + if config.max_tokens_per_rank != requirements.max_tokens_per_rank: + raise ValueError("workspace capacity must match ForwardConfig.max_tokens_per_rank") + self.config = config + self.device = torch.device(device) + self.requirements = requirements + self._runtime_manager = runtime_manager or get_runtime_manager() + self._symmetric_provider = symmetric_provider + self._local_provider = local_provider + self._runtime: Optional[RuntimeHandle] = None + self._workspace: Optional[WorkspaceOwner] = None + self._closed = False + self._cleanup_required = False + self._lock = threading.RLock() + + @property + def prepared(self) -> bool: + return not self._cleanup_required and self._runtime is not None and self._workspace is not None and self._workspace.allocated + + @property + def cleanup_required(self) -> bool: + return self._cleanup_required + + @property + def closed(self) -> bool: + return self._closed + + def prepare( + self, + request: ValidatedForwardRequest, + ) -> PreparedResources: + with self._lock: + if self._closed: + raise RuntimeError("MegaMoE execution plan is closed") + if self._cleanup_required: + raise RuntimeError("MegaMoE execution plan requires cleanup before prepare") + if request.config is not self.config: + raise ValueError("request does not belong to this static plan") + if torch.device(request.device) != self.device: + raise ValueError(f"execution plan is bound to {self.device}, got {request.device}") + if request.token_count > self.requirements.max_tokens_per_rank: + raise ValueError(f"token count {request.token_count} exceeds " f"max_tokens_per_rank={self.requirements.max_tokens_per_rank}") + if not self.prepared and torch.cuda.is_current_stream_capturing(): + raise RuntimeError("MegaMoE runtime/workspace must be warmed up before " "CUDA graph capture") + + if not self.prepared: + runtime = self._runtime_manager.acquire(self.config, self.device) + self._runtime = runtime + try: + workspace = WorkspaceOwner( + self.requirements, + runtime, + symmetric_provider=self._symmetric_provider, + local_provider=self._local_provider, + ) + self._workspace = workspace + views = workspace.views(request.token_count) + except Exception: + try: + self._cleanup_failed_prepare() + except Exception: + self._cleanup_required = True + raise + raise + else: + assert self._runtime is not None + assert self._workspace is not None + views = self._workspace.views(request.token_count) + + return PreparedResources( + runtime=self._runtime, + workspace=views, + ) + + def _cleanup_failed_prepare(self) -> None: + if self._workspace is not None: + self._workspace.close() + self._workspace = None + if self._runtime is not None: + self._runtime.close() + self._runtime = None + + def close(self) -> None: + with self._lock: + if self._closed: + return + + try: + if self._workspace is not None: + self._workspace.close() + self._workspace = None + if self._runtime is not None: + self._runtime.close() + self._runtime = None + except Exception: + self._cleanup_required = True + raise + self._cleanup_required = False + self._closed = True + + def __enter__(self) -> "ExecutionPlanOwner": + with self._lock: + if self._closed: + raise RuntimeError("MegaMoE execution plan is closed") + return self + + def __exit__(self, exc_type, exc_value, traceback) -> bool: + del exc_type, exc_value, traceback + self.close() + return False + + +__all__ = ["ExecutionPlanOwner", "PreparedResources"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/_runtime.py b/python/cudnn/moe_ep/_megamoe_backend/_runtime.py new file mode 100644 index 000000000..434ba1d59 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/_runtime.py @@ -0,0 +1,768 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Process-level runtime ownership for the private MegaMoE backend. + +This module is import-light: importing it does not import CUDA Python or +NVSHMEM and does not initialize CUDA. Optional runtime modules are loaded only +when a distributed runtime is actually acquired. +""" + +from __future__ import annotations + +import faulthandler +import logging +import os +import socket +import sys +import threading +import time +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Callable, Optional, Protocol + +import torch +import torch.distributed as dist + +from .._contracts import ForwardConfig + +_logger = logging.getLogger(__name__) + + +def _runtime_debug_enabled() -> bool: + return os.environ.get("MOE_EP_DEBUG_RUNTIME", "0") == "1" + + +def _runtime_debug(event: str, **details: object) -> None: + if not _runtime_debug_enabled(): + return + fields = { + "time": f"{time.monotonic():.6f}", + "host": socket.gethostname(), + "pid": os.getpid(), + "rank": os.environ.get("RANK", "?"), + "local_rank": os.environ.get("LOCAL_RANK", "?"), + "event": event, + **details, + } + print( + "[moe-ep-runtime] " + " ".join(f"{name}={value}" for name, value in fields.items()), + file=sys.stderr, + flush=True, + ) + + +def _runtime_debug_init_status(core) -> object: + if not _runtime_debug_enabled(): + return "debug-disabled" + try: + status = core.init_status() + except (AttributeError, RuntimeError): + return "unavailable" + return getattr(status, "name", status) + + +class _RuntimeWatchdog: + """Emit Python stacks and kernel wait channels while NVSHMEM init is blocked.""" + + def __init__(self, event: str) -> None: + self._event = event + self._stopped = threading.Event() + try: + self._interval = float(os.environ.get("MOE_EP_RUNTIME_WATCHDOG_SECONDS", "30")) + except ValueError: + self._interval = 30.0 + self._thread: Optional[threading.Thread] = None + + def start(self) -> None: + if not _runtime_debug_enabled() or self._interval <= 0: + return + # faulthandler's timer is implemented outside the Python interpreter + # lock, so it still emits the main-thread stack if a native NVSHMEM + # call holds the GIL. The Python thread adds /proc wait-channel data + # whenever the GIL remains schedulable. + faulthandler.dump_traceback_later( + self._interval, + repeat=True, + file=sys.stderr, + ) + self._thread = threading.Thread( + target=self._run, + name="moe-ep-runtime-watchdog", + daemon=True, + ) + self._thread.start() + + def close(self) -> None: + self._stopped.set() + if _runtime_debug_enabled() and self._interval > 0: + faulthandler.cancel_dump_traceback_later() + if self._thread is not None: + self._thread.join(timeout=1) + + def _run(self) -> None: + sample = 0 + while not self._stopped.wait(self._interval): + sample += 1 + wait_channels: list[str] = [] + for task_dir in sorted(Path("/proc/self/task").glob("[0-9]*")): + try: + thread_name = (task_dir / "comm").read_text().strip() + wait_channel = (task_dir / "wchan").read_text().strip() + except OSError as exc: + wait_channels.append(f"{task_dir.name}:unavailable({exc.errno})") + else: + wait_channels.append(f"{task_dir.name}:{thread_name}:{wait_channel or '-'}") + _runtime_debug( + "watchdog", + blocked_event=self._event, + sample=sample, + threads=";".join(wait_channels), + ) + faulthandler.dump_traceback(file=sys.stderr, all_threads=True) + + +class RuntimeUnavailableError(RuntimeError): + """The requested runtime cannot be loaded or initialized.""" + + +class RuntimeInitState(Enum): + """Normalized NVSHMEM initialization state.""" + + NOT_INITIALIZED = "not_initialized" + INITIALIZED = "initialized" + PARTIAL = "partial" + + +@dataclass(frozen=True) +class RuntimeWorld: + """Group-relative geometry and ordered membership used for bootstrap.""" + + rank: int + size: int + group: object + global_ranks: tuple[int, ...] + + @property + def identity(self) -> tuple[int, int, tuple[int, ...]]: + """Stable process-group identity independent of ProcessGroup objects.""" + + return self.rank, self.size, self.global_ranks + + +class NvshmemRuntimeProvider(Protocol): + """Injectable NVSHMEM lifecycle boundary used by :class:`RuntimeManager`.""" + + def initialization_state(self) -> RuntimeInitState: ... + + def initialize(self, device: torch.device, world: RuntimeWorld) -> None: ... + + def rank(self) -> int: ... + + def world_size(self) -> int: ... + + def device(self) -> torch.device: ... + + def finalize(self) -> None: ... + + +def _resolve_world(config: ForwardConfig) -> RuntimeWorld: + if config.ep_group is None: + if config.ep_size != 1 or config.ep_rank != 0 or config.ep_global_ranks: + raise ValueError("ep_group=None requires ep_size=1, ep_rank=0, and no " "distributed rank membership") + return RuntimeWorld(rank=0, size=1, group=None, global_ranks=()) + + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("distributed MegaMoE runtime requires torch.distributed to be initialized") + + group = config.ep_group + rank = dist.get_rank(group) + size = dist.get_world_size(group) + global_ranks = tuple(dist.get_global_rank(group, group_rank) for group_rank in range(size)) + if (rank, size) != (config.ep_rank, config.ep_size): + raise RuntimeError( + "ForwardConfig EP geometry does not match its process group: " f"config=({config.ep_rank}, {config.ep_size}), " f"runtime=({rank}, {size})" + ) + if global_ranks != config.ep_global_ranks: + raise RuntimeError("ForwardConfig EP membership does not match its process group: " f"config={config.ep_global_ranks}, runtime={global_ranks}") + return RuntimeWorld( + rank=rank, + size=size, + group=group, + global_ranks=global_ranks, + ) + + +def _canonical_cuda_device(device: torch.device) -> torch.device: + device = torch.device(device) + if device.type != "cuda": + raise ValueError(f"MegaMoE runtime requires a CUDA device, got {device}") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + return device + + +def _spans_default_distributed_world(world: RuntimeWorld) -> bool: + """Whether ``world`` has the default world's complete ordered membership.""" + + if not dist.is_available() or not dist.is_initialized(): + return False + return world.global_ranks == tuple(range(dist.get_world_size())) + + +def _load_nvshmem_core(): + try: + import nvshmem.core as core + except (ImportError, OSError) as exc: + raise RuntimeUnavailableError("MegaMoE distributed runtime requires nvshmem4py and NVSHMEM libraries") from exc + return core + + +def _normalize_nvshmem_init_state(status) -> RuntimeInitState: + """Normalize enum and integer forms used across nvshmem4py releases.""" + + name = getattr(status, "name", "") + if name.endswith("NOT_INITIALIZED"): + return RuntimeInitState.NOT_INITIALIZED + if name.endswith("IS_INITIALIZED") or name.endswith(("LIMITED_MPG", "FULL_MPG")): + return RuntimeInitState.INITIALIZED + if name.endswith("IS_BOOTSTRAPPED"): + return RuntimeInitState.PARTIAL + + try: + value = int(getattr(status, "value", status)) + except (TypeError, ValueError): + return RuntimeInitState.PARTIAL + if value == 0: + return RuntimeInitState.NOT_INITIALIZED + if value in {2, 3, 4}: + return RuntimeInitState.INITIALIZED + return RuntimeInitState.PARTIAL + + +class _DefaultNvshmemRuntimeProvider: + """Lazy adapter over the installed ``nvshmem.core`` API.""" + + def initialization_state(self) -> RuntimeInitState: + core = _load_nvshmem_core() + try: + status = core.init_status() + except Exception as exc: + raise RuntimeUnavailableError("failed to query NVSHMEM initialization status") from exc + return _normalize_nvshmem_init_state(status) + + def initialize(self, device: torch.device, world: RuntimeWorld) -> None: + if world.size <= 1: + raise ValueError("NVSHMEM initialization requires a distributed subgroup") + if world.group is None: + raise ValueError("NVSHMEM initialization requires a process group") + + core = _load_nvshmem_core() + started_at = time.monotonic() + _runtime_debug( + "initialize.begin", + device=device, + ep_rank=world.rank, + ep_size=world.size, + global_ranks=world.global_ranks, + ) + try: + import numpy as np + + try: + from cuda.core.experimental import Device + except ImportError: + from cuda.core import Device + + torch.cuda.set_device(device) + cuda_device = Device(device.index) + cuda_device.set_current() + _runtime_debug("initialize.cuda-current", device=device) + + uid = core.get_unique_id(empty=(world.rank != 0)) + _runtime_debug("initialize.uid-created") + uid_bytes = uid._data.view(np.uint8).copy() + uid_tensor = torch.from_numpy(uid_bytes) + group_backend = dist.get_backend(world.group) + if group_backend == dist.Backend.NCCL or str(group_backend).lower() == "nccl": + uid_tensor = uid_tensor.to(device=device) + root_global_rank = dist.get_global_rank(world.group, 0) + if root_global_rank != world.global_ranks[0]: + raise RuntimeError("EP subgroup root changed during NVSHMEM bootstrap") + _runtime_debug( + "initialize.uid-broadcast.begin", + backend=group_backend, + root_global_rank=root_global_rank, + tensor_device=uid_tensor.device, + tensor_bytes=uid_tensor.numel() * uid_tensor.element_size(), + ) + dist.broadcast( + uid_tensor, + src=root_global_rank, + group=world.group, + ) + _runtime_debug("initialize.uid-broadcast.end") + _runtime_debug("initialize.torch-barrier.begin") + dist.barrier(group=world.group) + _runtime_debug("initialize.torch-barrier.end") + uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype) + + watchdog = _RuntimeWatchdog("core.init") + _runtime_debug("initialize.core-init.begin") + watchdog.start() + try: + core.init( + device=cuda_device, + uid=uid, + rank=world.rank, + nranks=world.size, + initializer_method="uid", + ) + finally: + watchdog.close() + _runtime_debug( + "initialize.core-init.end", + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + init_status=_runtime_debug_init_status(core), + ) + except RuntimeUnavailableError: + raise + except Exception as exc: + _runtime_debug( + "initialize.error", + error_type=type(exc).__name__, + error=repr(exc), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) + raise RuntimeUnavailableError("failed to initialize the NVSHMEM EP subgroup runtime") from exc + + def rank(self) -> int: + try: + return int(_load_nvshmem_core().my_pe()) + except Exception as exc: + raise RuntimeUnavailableError("failed to query the NVSHMEM PE rank") from exc + + def world_size(self) -> int: + try: + return int(_load_nvshmem_core().n_pes()) + except Exception as exc: + raise RuntimeUnavailableError("failed to query the NVSHMEM PE world size") from exc + + def device(self) -> torch.device: + try: + from nvshmem.core.memory import _cached_device + + cached = _cached_device["device"] + if cached is None: + raise RuntimeError("NVSHMEM cached device is empty") + return torch.device("cuda", int(cached.device_id)) + except Exception as exc: + raise RuntimeUnavailableError("failed to query the NVSHMEM initialization device") from exc + + def finalize(self) -> None: + core = _load_nvshmem_core() + started_at = time.monotonic() + _runtime_debug( + "finalize.begin", + init_status=_runtime_debug_init_status(core), + ) + watchdog = _RuntimeWatchdog("core.finalize") + watchdog.start() + try: + core.finalize() + except Exception as exc: + _runtime_debug( + "finalize.error", + error_type=type(exc).__name__, + error=repr(exc), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) + raise RuntimeUnavailableError("failed to finalize NVSHMEM") from exc + finally: + watchdog.close() + _runtime_debug( + "finalize.end", + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + init_status=_runtime_debug_init_status(core), + ) + + +@dataclass +class _ActiveRuntime: + token: object + device: torch.device + world: RuntimeWorld + provider: Optional[NvshmemRuntimeProvider] + owns_runtime: bool + ref_count: int = 1 + cleanup_required: bool = False + + +class _ProcessRuntimeRegistry: + """State shared by every RuntimeManager instance in this process.""" + + def __init__(self) -> None: + self.lock = threading.RLock() + self.active: Optional[_ActiveRuntime] = None + + +_PROCESS_RUNTIME_REGISTRY = _ProcessRuntimeRegistry() + + +class RuntimeHandle: + """Per-backend lease on process-level runtime state.""" + + def __init__( + self, + manager: "RuntimeManager", + token: object, + device: torch.device, + world: RuntimeWorld, + owns_runtime: bool, + ) -> None: + self._manager = manager + self._token = token + self.device = device + self.rank = world.rank + self.world_size = world.size + self.group = world.group + self.global_ranks = world.global_ranks + self.owns_runtime = owns_runtime + self._closed = False + self._close_lock = threading.Lock() + + @property + def nvshmem_enabled(self) -> bool: + return self.world_size > 1 + + @property + def closed(self) -> bool: + return self._closed + + def ensure_open(self) -> None: + if self._closed: + raise RuntimeError("MegaMoE runtime handle is closed") + + def current_stream(self) -> torch.cuda.Stream: + self.ensure_open() + return torch.cuda.current_stream(self.device) + + def close(self) -> None: + with self._close_lock: + if self._closed: + return + self._manager._release(self._token) + self._closed = True + + +class RuntimeManager: + """Reference-counted owner for one process-global runtime subgroup.""" + + def __init__( + self, + *, + provider_factory: Callable[[], NvshmemRuntimeProvider] = (_DefaultNvshmemRuntimeProvider), + world_resolver: Callable[[ForwardConfig], RuntimeWorld] = _resolve_world, + keep_alive: bool = False, + ) -> None: + self._provider_factory = provider_factory + self._world_resolver = world_resolver + self._keep_alive = keep_alive + + @property + def ref_count(self) -> int: + with _PROCESS_RUNTIME_REGISTRY.lock: + active = _PROCESS_RUNTIME_REGISTRY.active + return 0 if active is None else active.ref_count + + @property + def active_device(self) -> Optional[torch.device]: + with _PROCESS_RUNTIME_REGISTRY.lock: + active = _PROCESS_RUNTIME_REGISTRY.active + return None if active is None else active.device + + def acquire( + self, + config: ForwardConfig, + device: torch.device, + ) -> RuntimeHandle: + device = _canonical_cuda_device(device) + world = self._world_resolver(config) + + with _PROCESS_RUNTIME_REGISTRY.lock: + if _PROCESS_RUNTIME_REGISTRY.active is not None: + active = _PROCESS_RUNTIME_REGISTRY.active + _runtime_debug( + "manager.acquire-reuse.begin", + ref_count=active.ref_count, + owns_runtime=active.owns_runtime, + cleanup_required=active.cleanup_required, + ) + if active.cleanup_required: + raise RuntimeError("MegaMoE process runtime requires cleanup before reacquire") + if active.device != device: + raise ValueError(f"MegaMoE process runtime is bound to {active.device}; " f"cannot acquire it for {device}") + if active.world.identity != world.identity: + raise RuntimeError("MegaMoE process runtime is already bound to a different " "EP subgroup") + active.ref_count += 1 + _runtime_debug( + "manager.acquire-reuse.end", + ref_count=active.ref_count, + ) + return RuntimeHandle( + self, + active.token, + active.device, + active.world, + active.owns_runtime, + ) + + provider: Optional[NvshmemRuntimeProvider] = None + owns_runtime = False + _runtime_debug( + "manager.acquire-new.begin", + device=device, + ep_rank=world.rank, + ep_size=world.size, + ) + if world.size > 1: + provider = self._provider_factory() + status = provider.initialization_state() + _runtime_debug( + "manager.acquire-new.state", + init_status=status.value, + ) + if status is RuntimeInitState.PARTIAL: + raise RuntimeError("cannot attach to a partially initialized NVSHMEM runtime") + if status is RuntimeInitState.INITIALIZED and not _spans_default_distributed_world(world): + raise RuntimeError( + "cannot safely attach an externally initialized NVSHMEM " + "runtime to a non-WORLD EP subgroup because its ordered " + "membership cannot be verified" + ) + if status is RuntimeInitState.NOT_INITIALIZED: + try: + provider.initialize(device, world) + except Exception as initialization_error: + self._rollback_failed_initialization( + provider, + device, + world, + initialization_error, + ) + raise + owns_runtime = True + + try: + provider_device = provider.device() + provider_rank = provider.rank() + provider_size = provider.world_size() + except Exception as validation_error: + if owns_runtime: + self._cleanup_owned_runtime_after_error( + provider, + device, + world, + validation_error, + ) + raise + + if provider_device != device: + if owns_runtime: + self._cleanup_owned_runtime_after_error( + provider, + device, + world, + RuntimeError( + "NVSHMEM initialization device does not match " f"the requested device: nvshmem={provider_device}, " f"requested={device}" + ), + ) + ownership = "owned" if owns_runtime else "externally initialized" + raise RuntimeError(f"{ownership} NVSHMEM runtime is bound to " f"{provider_device}, not the requested device {device}") + if (provider_rank, provider_size) != (world.rank, world.size): + if owns_runtime: + self._cleanup_owned_runtime_after_error( + provider, + device, + world, + RuntimeError("NVSHMEM PE geometry mismatch"), + ) + raise RuntimeError( + "NVSHMEM PE geometry does not match the EP subgroup: " + f"nvshmem=({provider_rank}, {provider_size}), " + f"torch=({world.rank}, {world.size})" + ) + + token = object() + _PROCESS_RUNTIME_REGISTRY.active = _ActiveRuntime( + token=token, + device=device, + world=world, + provider=provider, + owns_runtime=owns_runtime, + ) + _runtime_debug( + "manager.acquire-new.end", + owns_runtime=owns_runtime, + ref_count=1, + ) + return RuntimeHandle(self, token, device, world, owns_runtime) + + @staticmethod + def _mark_cleanup_required( + provider: NvshmemRuntimeProvider, + device: torch.device, + world: RuntimeWorld, + ) -> None: + _PROCESS_RUNTIME_REGISTRY.active = _ActiveRuntime( + token=object(), + device=device, + world=world, + provider=provider, + owns_runtime=True, + ref_count=0, + cleanup_required=True, + ) + + @classmethod + def _rollback_failed_initialization( + cls, + provider: NvshmemRuntimeProvider, + device: torch.device, + world: RuntimeWorld, + initialization_error: Exception, + ) -> None: + try: + state = provider.initialization_state() + if state is not RuntimeInitState.NOT_INITIALIZED: + provider.finalize() + except Exception as cleanup_error: + cls._mark_cleanup_required(provider, device, world) + raise RuntimeError("NVSHMEM initialization failed and rollback requires retry") from cleanup_error + _logger.debug( + "rolled back failed NVSHMEM initialization: %s", + initialization_error, + ) + + @classmethod + def _cleanup_owned_runtime_after_error( + cls, + provider: NvshmemRuntimeProvider, + device: torch.device, + world: RuntimeWorld, + original_error: Exception, + ) -> None: + try: + provider.finalize() + except Exception as cleanup_error: + cls._mark_cleanup_required(provider, device, world) + raise RuntimeError("NVSHMEM validation failed and cleanup requires retry") from cleanup_error + _logger.debug( + "finalized owned NVSHMEM after validation failure: %s", + original_error, + ) + + def retry_cleanup(self) -> None: + """Retry cleanup after an acquire-time rollback failure.""" + + with _PROCESS_RUNTIME_REGISTRY.lock: + active = _PROCESS_RUNTIME_REGISTRY.active + if active is None: + return + if not active.cleanup_required or active.ref_count != 0: + raise RuntimeError("MegaMoE process runtime does not have retryable cleanup") + if active.provider is None: + raise RuntimeError("retryable MegaMoE runtime cleanup has no provider") + active.provider.finalize() + _PROCESS_RUNTIME_REGISTRY.active = None + + def shutdown(self) -> None: + """Finalize an idle process runtime at a caller-controlled collective point.""" + + with _PROCESS_RUNTIME_REGISTRY.lock: + active = _PROCESS_RUNTIME_REGISTRY.active + if active is None: + return + if active.ref_count != 0: + raise RuntimeError("cannot shut down the MegaMoE process runtime while " f"{active.ref_count} handles remain active") + if active.provider is not None and (active.owns_runtime or active.cleanup_required): + _runtime_debug( + "manager.shutdown-finalize.begin", + cleanup_required=active.cleanup_required, + ) + active.provider.finalize() + _runtime_debug("manager.shutdown-finalize.end") + _PROCESS_RUNTIME_REGISTRY.active = None + _runtime_debug("manager.shutdown.end") + + def _release(self, token: object) -> None: + with _PROCESS_RUNTIME_REGISTRY.lock: + active = _PROCESS_RUNTIME_REGISTRY.active + if active is None or active.token is not token: + _runtime_debug("manager.release-stale") + return + if active.ref_count <= 0: + raise RuntimeError("MegaMoE process runtime has invalid release state") + + _runtime_debug( + "manager.release.begin", + ref_count=active.ref_count, + owns_runtime=active.owns_runtime, + cleanup_required=active.cleanup_required, + ) + if active.cleanup_required: + if active.ref_count != 1 or active.provider is None: + raise RuntimeError("MegaMoE process runtime has invalid retry state") + active.provider.finalize() + _PROCESS_RUNTIME_REGISTRY.active = None + _runtime_debug("manager.release.cleanup-retry.end") + return + + if active.ref_count > 1: + active.ref_count -= 1 + _runtime_debug( + "manager.release-retained", + ref_count=active.ref_count, + ) + return + + if self._keep_alive and not active.cleanup_required: + active.ref_count = 0 + _runtime_debug( + "manager.release-idle", + owns_runtime=active.owns_runtime, + ) + return + + if active.owns_runtime and active.provider is not None: + _runtime_debug("manager.release-finalize.begin") + try: + active.provider.finalize() + except Exception: + active.cleanup_required = True + _runtime_debug("manager.release-finalize.error") + raise + _runtime_debug("manager.release-finalize.end") + _PROCESS_RUNTIME_REGISTRY.active = None + _runtime_debug("manager.release.end", ref_count=0) + + +_DEFAULT_RUNTIME_MANAGER = RuntimeManager(keep_alive=True) + + +def get_runtime_manager() -> RuntimeManager: + """Return the process-level manager used by the default MegaMoE backend.""" + + return _DEFAULT_RUNTIME_MANAGER + + +__all__ = [ + "NvshmemRuntimeProvider", + "RuntimeHandle", + "RuntimeInitState", + "RuntimeManager", + "RuntimeUnavailableError", + "RuntimeWorld", + "get_runtime_manager", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/_workspace.py new file mode 100644 index 000000000..4446881c6 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/_workspace.py @@ -0,0 +1,433 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Stable local and symmetric workspace ownership for MegaMoE.""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from types import MappingProxyType +from typing import Mapping, Optional, Protocol, Sequence + +import torch + +from .._contracts import ForwardConfig +from .._math import round_up +from ._comm import ( + PeerMapping, + SymmetricMemoryProvider, + SymmetricSlab, + _TorchMemoryProvider, +) +from ._runtime import RuntimeHandle, _runtime_debug + + +def padded_mxfp8_scale_columns(hidden: int) -> int: + """Return the E8M0 row width required by Rubin's 16-byte token-in copy.""" + + logical_columns = (hidden + 31) // 32 + return round_up(logical_columns, 16) + + +@dataclass(frozen=True) +class BufferRegion: + """One named byte region within a stable root allocation.""" + + name: str + nbytes: int + alignment: int = 256 + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("workspace region name must not be empty") + if self.nbytes < 0: + raise ValueError(f"workspace region {self.name!r} has negative size {self.nbytes}") + if self.alignment <= 0 or self.alignment & (self.alignment - 1): + raise ValueError(f"workspace region {self.name!r} alignment must be a power of two") + + +@dataclass(frozen=True) +class BufferPlacement: + """Resolved byte offset for one region.""" + + name: str + offset: int + nbytes: int + + +@dataclass(frozen=True) +class BufferLayout: + """Deterministic aligned layout for one root byte allocation.""" + + placements: tuple[BufferPlacement, ...] + total_bytes: int + + @classmethod + def build(cls, regions: Sequence[BufferRegion]) -> "BufferLayout": + names: set[str] = set() + placements = [] + offset = 0 + max_alignment = 1 + for region in regions: + if region.name in names: + raise ValueError(f"duplicate workspace region {region.name!r}") + names.add(region.name) + offset = round_up(offset, region.alignment) + placements.append( + BufferPlacement( + name=region.name, + offset=offset, + nbytes=region.nbytes, + ) + ) + offset += region.nbytes + max_alignment = max(max_alignment, region.alignment) + return cls( + placements=tuple(placements), + total_bytes=round_up(offset, max_alignment), + ) + + def placement(self, name: str) -> BufferPlacement: + for placement in self.placements: + if placement.name == name: + return placement + raise KeyError(name) + + +@dataclass(frozen=True) +class WorkspaceRequirements: + """Capacity-driven regions supplied before runtime allocation. + + The executable backend obtains exact Rubin kernel workspace sizes, then + passes them here without making the runtime owner import or instantiate + CuTeDSL kernels. + """ + + max_tokens_per_rank: int + symmetric_regions: tuple[BufferRegion, ...] + local_regions: tuple[BufferRegion, ...] + + def __post_init__(self) -> None: + if self.max_tokens_per_rank < 0: + raise ValueError("max_tokens_per_rank must be non-negative") + symmetric_names = {region.name for region in self.symmetric_regions} + local_names = {region.name for region in self.local_regions} + duplicates = symmetric_names & local_names + if duplicates: + raise ValueError("workspace region names must be unique across roots: " f"{sorted(duplicates)}") + + @classmethod + def for_mxfp8( + cls, + config: ForwardConfig, + *, + kernel_local_workspace_bytes: int, + kernel_shared_workspace_bytes: int, + col_quant_data_bytes: int = 0, + col_quant_sf_bytes: int = 0, + backward_dprob_bytes: int = 0, + backward_aux_data_bytes: int = 0, + backward_aux_scale_bytes: int = 0, + ) -> "WorkspaceRequirements": + if config.max_tokens_per_rank is None: + raise ValueError("MXFP8 workspace requires max_tokens_per_rank") + for name, value in ( + ("kernel_local_workspace_bytes", kernel_local_workspace_bytes), + ("kernel_shared_workspace_bytes", kernel_shared_workspace_bytes), + ("col_quant_data_bytes", col_quant_data_bytes), + ("col_quant_sf_bytes", col_quant_sf_bytes), + ("backward_dprob_bytes", backward_dprob_bytes), + ("backward_aux_data_bytes", backward_aux_data_bytes), + ("backward_aux_scale_bytes", backward_aux_scale_bytes), + ): + if value < 0: + raise ValueError(f"{name} must be non-negative, got {value}") + if bool(col_quant_data_bytes) != bool(col_quant_sf_bytes): + raise ValueError("column requant data and scale workspace must be enabled together") + backward_sizes = ( + backward_dprob_bytes, + backward_aux_data_bytes, + backward_aux_scale_bytes, + ) + if any(backward_sizes) and not all(backward_sizes): + raise ValueError("backward dprob, data, and scale workspace must be enabled together") + + tokens = config.max_tokens_per_rank + hidden = config.hidden_size + top_k = config.top_k + kernel_sf_columns = padded_mxfp8_scale_columns(hidden) + + backward_symmetric_regions = (BufferRegion("backward_dprob", backward_dprob_bytes),) if backward_dprob_bytes else () + symmetric_regions = ( + BufferRegion("activation_data", tokens * hidden), + BufferRegion("activation_scale", round_up(tokens, 128) * kernel_sf_columns), + BufferRegion("topk_weights", tokens * top_k * 4), + BufferRegion("output_data", tokens * hidden * 2), + *backward_symmetric_regions, + BufferRegion( + "kernel_shared_workspace", + kernel_shared_workspace_bytes, + ), + # Test-visible tail canary placed immediately after the opaque + # peer-visible kernel workspace. It does not enter the kernel ABI. + BufferRegion("symmetric_guard", 256, alignment=1), + ) + col_quant_regions = ( + ( + BufferRegion("col_quant_data", col_quant_data_bytes), + BufferRegion("col_quant_sf", col_quant_sf_bytes), + ) + if col_quant_data_bytes + else () + ) + backward_local_regions = ( + ( + BufferRegion( + "backward_aux_data", + backward_aux_data_bytes, + alignment=128, + ), + BufferRegion( + "backward_aux_scale", + backward_aux_scale_bytes, + alignment=128, + ), + ) + if backward_aux_data_bytes + else () + ) + local_regions = ( + BufferRegion("topk_idx", tokens * top_k * 4), + BufferRegion("overflow_flag", 4), + *col_quant_regions, + *backward_local_regions, + BufferRegion("kernel_local_workspace", kernel_local_workspace_bytes), + BufferRegion("local_guard", 256, alignment=1), + ) + return cls( + max_tokens_per_rank=tokens, + symmetric_regions=symmetric_regions, + local_regions=local_regions, + ) + + +class LocalMemoryProvider(Protocol): + """Injectable local allocation boundary.""" + + def allocate(self, nbytes: int, device: torch.device) -> torch.Tensor: ... + + def free(self, tensor: torch.Tensor) -> None: ... + + +class _LocalSlab: + def __init__( + self, + nbytes: int, + device: torch.device, + provider: LocalMemoryProvider, + ) -> None: + if nbytes <= 0: + raise ValueError(f"local slab size must be positive, got {nbytes}") + self._provider = provider + self._nbytes = nbytes + _runtime_debug("local-slab.allocate.begin", nbytes=nbytes, device=device) + root = provider.allocate(nbytes, device) + _runtime_debug( + "local-slab.allocate.end", + nbytes=nbytes, + data_ptr=hex(root.data_ptr()) if isinstance(root, torch.Tensor) else "?", + ) + self._root: Optional[torch.Tensor] = None + try: + if not isinstance(root, torch.Tensor): + raise TypeError("local memory provider must return a torch.Tensor") + if root.dtype is not torch.uint8 or root.numel() < nbytes: + raise ValueError("local root must be a uint8 tensor with at least " f"{nbytes} elements") + if root.device != device: + raise ValueError("local root device does not match runtime device: " f"root={root.device}, runtime={device}") + if not root.is_contiguous(): + raise ValueError("local root tensor must be contiguous") + _runtime_debug("local-slab.zero.begin", nbytes=nbytes) + root.zero_() + _runtime_debug("local-slab.zero.enqueued", nbytes=nbytes) + except Exception: + if isinstance(root, torch.Tensor): + provider.free(root) + raise + self._root = root + + @property + def root(self) -> torch.Tensor: + if self._root is None: + raise RuntimeError("local workspace slab is closed") + return self._root + + def byte_view(self, offset: int, nbytes: int) -> torch.Tensor: + if offset < 0 or nbytes < 0 or offset + nbytes > self._nbytes: + raise ValueError(f"byte view [{offset}, {offset + nbytes}) exceeds " f"local slab size {self._nbytes}") + return self.root.narrow(0, offset, nbytes) + + def close(self) -> None: + root = self._root + if root is None: + return + self._provider.free(root) + self._root = None + + +@dataclass(frozen=True) +class WorkspaceViews: + """Stable full-capacity byte views for one prepared request.""" + + token_count: int + symmetric: Mapping[str, torch.Tensor] + local: Mapping[str, torch.Tensor] + peer_mapping: PeerMapping + + +class WorkspaceOwner: + """Own local and symmetric slabs for one static execution plan.""" + + def __init__( + self, + requirements: WorkspaceRequirements, + runtime: RuntimeHandle, + *, + symmetric_provider: Optional[SymmetricMemoryProvider] = None, + local_provider: Optional[LocalMemoryProvider] = None, + ) -> None: + self.requirements = requirements + self.runtime = runtime + self.symmetric_layout = BufferLayout.build(requirements.symmetric_regions) + self.local_layout = BufferLayout.build(requirements.local_regions) + if self.symmetric_layout.total_bytes <= 0: + raise ValueError("workspace requires at least one symmetric byte") + if self.local_layout.total_bytes <= 0: + raise ValueError("workspace requires at least one local byte") + + self._symmetric_provider = symmetric_provider + self._local_provider = local_provider or _TorchMemoryProvider() + self._symmetric: Optional[SymmetricSlab] = None + self._local: Optional[_LocalSlab] = None + self._closed = False + self._cleanup_required = False + self._lock = threading.RLock() + + @property + def allocated(self) -> bool: + return not self._cleanup_required and self._symmetric is not None and self._symmetric.allocated and self._local is not None + + @property + def cleanup_required(self) -> bool: + return self._cleanup_required + + @property + def closed(self) -> bool: + return self._closed + + def ensure_allocated(self) -> None: + with self._lock: + if self._closed: + raise RuntimeError("workspace owner is closed") + if self._cleanup_required: + raise RuntimeError("workspace owner requires cleanup before allocation") + if self.allocated: + return + self.runtime.ensure_open() + + _runtime_debug( + "workspace.allocate.begin", + local_bytes=self.local_layout.total_bytes, + symmetric_bytes=self.symmetric_layout.total_bytes, + ) + local = _LocalSlab( + self.local_layout.total_bytes, + self.runtime.device, + self._local_provider, + ) + self._local = local + try: + _runtime_debug("workspace.symmetric-slab.create.begin") + symmetric = SymmetricSlab( + self.runtime, + self.symmetric_layout.total_bytes, + provider=self._symmetric_provider, + ) + self._symmetric = symmetric + _runtime_debug("workspace.symmetric-slab.ensure.begin") + symmetric.ensure_allocated() + _runtime_debug("workspace.symmetric-slab.ensure.end") + except Exception: + try: + if self._symmetric is not None: + self._symmetric.close() + self._symmetric = None + if self._local is not None: + self._local.close() + self._local = None + except Exception: + self._cleanup_required = True + raise + raise + _runtime_debug("workspace.allocate.end") + + def views(self, token_count: int) -> WorkspaceViews: + with self._lock: + if token_count < 0: + raise ValueError(f"token_count must be non-negative, got {token_count}") + if token_count > self.requirements.max_tokens_per_rank: + raise ValueError(f"token count {token_count} exceeds " f"max_tokens_per_rank={self.requirements.max_tokens_per_rank}") + self.ensure_allocated() + assert self._symmetric is not None + assert self._local is not None + + symmetric_views = { + placement.name: self._symmetric.byte_view( + placement.offset, + placement.nbytes, + ) + for placement in self.symmetric_layout.placements + } + local_views = { + placement.name: self._local.byte_view( + placement.offset, + placement.nbytes, + ) + for placement in self.local_layout.placements + } + return WorkspaceViews( + token_count=token_count, + symmetric=MappingProxyType(symmetric_views), + local=MappingProxyType(local_views), + peer_mapping=self._symmetric.mapping, + ) + + def close(self) -> None: + with self._lock: + if self._closed: + return + try: + if self._symmetric is not None: + self._symmetric.close() + self._symmetric = None + if self._local is not None: + self._local.close() + self._local = None + except Exception: + self._cleanup_required = True + raise + self._cleanup_required = False + self._closed = True + + +__all__ = [ + "BufferLayout", + "BufferPlacement", + "BufferRegion", + "LocalMemoryProvider", + "WorkspaceOwner", + "WorkspaceRequirements", + "WorkspaceViews", + "padded_mxfp8_scale_columns", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/LICENSE.Apache-2.0 b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/LICENSE.Apache-2.0 new file mode 100644 index 000000000..261eeb9e9 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/LICENSE.Apache-2.0 @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR.md b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR.md new file mode 100644 index 000000000..db8fc6f97 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR.md @@ -0,0 +1,72 @@ +# Vendoring record: cutedsl_megamoe + +This file records provenance and synchronization state for the CuTeDSL +MegaMoE source snapshot. Runtime behavior and integration details are +documented in the parent backend `README.md`. + +## Upstream + +- **Project**: `cutedsl_megamoe` (NVIDIA-internal repository; URL omitted). +- **Source tree**: `cutedsl_megamoe/next/sources`. +- **Current synchronized commit**: + `9b15c450e2d19472bdfaae37489317f029beb01c`. +- **Earlier import points**: + - base forward: + `882c83e2ce4086c3cd4211fc5a2296143c5e2aea`; + - selected forward updates and backward dGLU: + `92dd334af2eeedb36087834354b58ace08e880c6`; + - forward column-quantization output-layout updates: + `5b89819cb16069dfe20a1a0ba0778d35cb428352`. +- **Last synced**: 2026-09-02. Earlier imports occurred on 2026-08-11, + 2026-08-17, 2026-08-20, 2026-08-24, and 2026-08-28. +- **Vendored subset**: the recursive Python import closure required by Rubin + SM107 training MegaMoE forward GLU, optional forward MXFP8 column + requantization, and backward dGLU. + +Complete kernel products, runners, tests, repository scaffolding, and +unrelated Blackwell and Rubin inference sources are excluded. Three +architecture-neutral Blackwell donor modules are retained because the Rubin +`topk_reduce.py` and `tmem_transpose.py` source-copy shims import them. + +## Policy + +- Vendored Python source bodies track the corresponding upstream paths at the + synchronized commit. +- Repository-required copyright and BSD-3-Clause SPDX headers may be added + where the upstream snapshot did not carry them. +- Integration behavior belongs in the parent `_megamoe_backend` package, not + in the vendored source bodies. +- Local kernel fixes should go upstream first and then be synchronized here. + Any unavoidable local source difference must be listed below. +- Snapshot updates must preserve the minimal recursive import closure and + verify source-body equality while ignoring repository-added header lines. + +The synchronized Python sources use BSD-3-Clause SPDX identifiers. +`LICENSE.Apache-2.0` is retained as historical snapshot metadata. + +## Local differences from upstream + +- `kernel_src/rubin/training/__init__.py` is reduced to a package marker. This + avoids importing the unused traditional-wgrad product. +- Repository-required copyright and BSD-3-Clause SPDX headers are added to + source files that lacked explicit headers. + +No other vendored Python source-body differences are expected. + +## Integration boundary + +Public API validation, symmetric-workspace ownership, overflow reporting, +input and weight staging, CUDA Graph handling, dprob materialization, and +grouped-WGrad layout conversion live in the parent `_megamoe_backend` +package. + +The vendored Rubin sources require a CUTLASS DSL distribution that provides +`cutlass.utils.rubin_helpers`. The executable backend enforces +`nvidia-cutlass-dsl>=4.8.0` before importing these kernels. + +## Consumers + +- `_megamoe_backend/mxfp8/_compile.py`: Rubin MXFP8 forward preparation and + compilation. +- `_megamoe_backend/mxfp8/_backward_compile.py`: Rubin MXFP8 backward dGLU + preparation and compilation. diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/__init__.py new file mode 100644 index 000000000..6d10a7e6e --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Integration-ready CuTeDSL MegaMoE sources.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/api.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/api.py new file mode 100644 index 000000000..1935befbb --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/api.py @@ -0,0 +1,184 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Stable construction API for composable kernel implementations.""" + +import types +from abc import ABC, ABCMeta, abstractmethod +from collections.abc import Iterator, Mapping +from types import MappingProxyType +from typing import Dict, Union, get_args, get_origin + +from cutlass.cute.typing import SymInt + + +RuntimeIntegerType = SymInt +StaticIntegerType = int +StaticOrRuntimeIntegerType = Union[int, SymInt] + + +class OptionalRequirement: + """Descriptor field that a component may consume conditionally.""" + + __slots__ = ("expected_type",) + + def __init__(self, expected_type) -> None: + self.expected_type = expected_type + + +Requirement = Union[type, OptionalRequirement] + + +def _required_type(requirement: Requirement): + return requirement.expected_type if isinstance(requirement, OptionalRequirement) else requirement + + +def _matches_type(value, expected_type) -> bool: + origin = get_origin(expected_type) + if origin in (Union, types.UnionType): + return any( + _matches_type(value, candidate) + for candidate in get_args(expected_type) + ) + return isinstance(value, expected_type) + + +class Desc(Mapping[str, object]): + """Immutable descriptor mapping validated against component schemas.""" + + def __init__(self, values: Mapping[str, object]) -> None: + self._values = MappingProxyType(dict(values)) + + def __getitem__(self, key: str): + return self._values[key] + + def __iter__(self) -> Iterator[str]: + return iter(self._values) + + def __len__(self) -> int: + return len(self._values) + + def validate( + self, + requirements: Dict[str, Requirement], + *, + component_name: str, + ) -> None: + for name, requirement in requirements.items(): + if name not in self._values: + if isinstance(requirement, OptionalRequirement): + continue + raise KeyError( + f"{component_name} requires descriptor field {name!r}." + ) + expected_type = _required_type(requirement) + value = self._values[name] + if not _matches_type(value, expected_type): + raise TypeError( + f"{component_name} requires {name!r} to have type " + f"{expected_type}, got {type(value)}." + ) + + +class ProblemDesc(Desc): + """Problem semantics shared by every component in one kernel.""" + + +class ImplDesc(Desc): + """Fully static implementation choices shared by kernel components.""" + + +class _KernelComponentMeta(ABCMeta): + def __call__(cls, *args, **kwargs): + instance = super().__call__(*args, **kwargs) + instance._validate_component_init() + return instance + + +class KernelComponent(ABC, metaclass=_KernelComponentMeta): + """Component that consumes descriptors only during construction.""" + + @classmethod + @abstractmethod + def problem_desc_require(cls) -> Dict[str, Requirement]: + ... + + @classmethod + @abstractmethod + def impl_desc_require(cls) -> Dict[str, Requirement]: + ... + + def _validate_desc_inputs( + self, + problem_desc: ProblemDesc, + impl_desc: ImplDesc, + ) -> None: + component_name = type(self).__name__ + overlap = ( + self.problem_desc_require().keys() + & self.impl_desc_require().keys() + ) + if overlap: + raise ValueError( + f"{component_name} requires fields from both descriptors: " + f"{sorted(overlap)}." + ) + problem_desc.validate( + self.problem_desc_require(), + component_name=component_name, + ) + impl_desc.validate( + self.impl_desc_require(), + component_name=component_name, + ) + + def _validate_component_init(self) -> None: + component_name = type(self).__name__ + requirements = { + **self.problem_desc_require(), + **self.impl_desc_require(), + } + for name, requirement in requirements.items(): + if not hasattr(self, name): + if isinstance(requirement, OptionalRequirement): + continue + raise RuntimeError( + f"{component_name} did not bind required field {name!r}." + ) + expected_type = _required_type(requirement) + value = getattr(self, name) + if not _matches_type(value, expected_type): + raise TypeError( + f"{component_name}.{name} must have type " + f"{expected_type}, got {type(value)}." + ) + for name, value in vars(self).items(): + if isinstance(value, Desc): + raise RuntimeError( + f"{component_name}.{name} retains a descriptor." + ) + + +class KernelClass(KernelComponent): + """Top-level host wrapper for one composable kernel implementation.""" + + @abstractmethod + def name(self) -> str: + ... + + @abstractmethod + def aot_compile(self): + ... + + +__all__ = [ + "Desc", + "ImplDesc", + "KernelClass", + "KernelComponent", + "OptionalRequirement", + "ProblemDesc", + "RuntimeIntegerType", + "StaticIntegerType", + "StaticOrRuntimeIntegerType", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/__init__.py new file mode 100644 index 000000000..a3fa301fa --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/__init__.py @@ -0,0 +1,25 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Cross-rank communication protocols and implementations.""" + +from ..quant_def import CombineFormat, QuantKind +from .nvlink_domain.token_comm import ( + TokenBackScheduleMode, + TokenBackMode, + TokenCommArgs, + TokenCommNonDeterministic, +) +from .nvlink_domain.token_comm_deterministic import TokenCommDeterministic +from .token_protocol import TokenSrcMetadata + +__all__ = [ + "CombineFormat", + "QuantKind", + "TokenBackScheduleMode", + "TokenBackMode", + "TokenCommArgs", + "TokenCommDeterministic", + "TokenCommNonDeterministic", + "TokenSrcMetadata", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/__init__.py new file mode 100644 index 000000000..962122907 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""NVLink-domain pointer mapping and token communication.""" + +from ...helpers.software_sync import NvlinkBarrier, SoftwareGridSync +from ...quant_def import QuantKind +from .symmetric_buffer import SymmetricBufferDevice, SymmetricBufferHost +from .token_comm import ( + TokenBackMode, + TokenBackScheduleMode, + TokenCommArgs, + TokenCommNonDeterministic, +) +from .token_comm_deterministic import TokenCommDeterministic + +__all__ = [ + "NvlinkBarrier", + "QuantKind", + "SoftwareGridSync", + "SymmetricBufferDevice", + "SymmetricBufferHost", + "TokenBackMode", + "TokenBackScheduleMode", + "TokenCommArgs", + "TokenCommDeterministic", + "TokenCommNonDeterministic", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py new file mode 100644 index 000000000..58186c7fd --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py @@ -0,0 +1,182 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Symmetric-heap peer pointer mapping carried in kernel arguments.""" + +from dataclasses import dataclass +from typing import Any, Optional + +from packaging.version import Version + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith, llvm +from cutlass.base_dsl.dsl import extract_mlir_values, get_mlir_types, new_from_mlir_values +from cutlass.base_dsl.runtime.jit_arg_adapters import JitArgAdapterRegistry +from cutlass.base_dsl.typing import get_c_pointers +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64, dsl_user_op + + +try: + from cutlass.base_dsl.typing import MLIR_DYNAMIC_INDEX +except ImportError: + MLIR_DYNAMIC_INDEX = -(2**31) + + +_dsl_release = Version(Version(cutlass.__version__).base_version) +_grid_constant_width_is_free = _dsl_release < Version("4.0.0") or _dsl_release >= Version("4.7.0") +_byval_rank_limit = 2**31 if _grid_constant_width_is_free else 16 + + +def _byval_struct_type(rank_count: int) -> Any: + words = rank_count if _grid_constant_width_is_free else _byval_rank_limit + return ir.Type.parse(f"!llvm.struct<(array<{words} x i64>)>") + + +@dataclass(frozen=True) +class SymmetricBufferDevice: + """Device-side peer offset table in constant/by-value kernel arguments.""" + + value: Any + max_ranks: cutlass.Constexpr[int] + + def __extract_mlir_values__(self) -> list: + return [self.value] + + def __new_from_mlir_values__(self, values: list) -> "SymmetricBufferDevice": + return SymmetricBufferDevice(values[0], self.max_ranks) + + def __get_mlir_types__(self) -> list: + if self.max_ranks <= _byval_rank_limit: + return [ir.Type.parse("!llvm.ptr")] + return [ir.Type.parse(f"vector<{self.max_ranks}xi64>")] + + def __extract_mlir_attributes__(self) -> list: + if self.max_ranks <= _byval_rank_limit: + return [ + ir.DictAttr.get( + { + "cute_nvgpu.grid_constant": ir.UnitAttr.get(), + "llvm.byval": ir.TypeAttr.get(_byval_struct_type(self.max_ranks)), + } + ) + ] + return [ir.DictAttr.get({})] + + @cute.jit + def map(self, local_address: Int64, destination_rank: Int32, byte_offset: Int64 = Int64(0)) -> Int64: + if cutlass.const_expr(self.max_ranks <= _byval_rank_limit): + i64_type = ir.Type.parse("i64") + offset_pointer = llvm.getelementptr( + ir.Type.parse("!llvm.ptr"), + self.value, + [destination_rank.ir_value()], + [MLIR_DYNAMIC_INDEX], + i64_type, + no_wrap_flags="None", + ) + peer_offset = Int64(llvm.load(i64_type, offset_pointer)) + else: + peer_offset = Int64(llvm.extractelement(self.value, destination_rank.ir_value())) + return local_address + peer_offset + byte_offset + + @cute.jit + def map_pointer(self, pointer, destination_rank: Int32, byte_alignment: Optional[int] = None): + if cutlass.const_expr(pointer.memspace != AddressSpace.gmem): + raise ValueError("Only GMEM pointers can be mapped to a symmetric peer.") + if cutlass.const_expr(byte_alignment is None): + byte_alignment = pointer.max_alignment + return cute.make_ptr( + pointer.dtype, self.map(pointer.toint(), destination_rank), pointer.memspace, assumed_align=byte_alignment + ) + + +@dataclass(frozen=True) +class SymmetricBufferHost: + """Host launch payload used to construct a SymmetricBufferDevice.""" + + base_address: Int64 + offsets: tuple + rank: Int32 + max_ranks: cutlass.Constexpr[int] + + @staticmethod + def _as_int64(value) -> Int64: + return value if isinstance(value, Int64) else Int64(int(value)) + + @dsl_user_op + def make_device_object(self, *, loc=None, ip=None) -> SymmetricBufferDevice: + offsets = tuple(self.offsets) + if len(offsets) != self.max_ranks: + raise ValueError(f"Expected {self.max_ranks} peer offsets, got {len(offsets)}.") + + if self.max_ranks <= _byval_rank_limit: + pointer_type = ir.Type.parse("!llvm.ptr") + struct_type = _byval_struct_type(self.max_ranks) + i64_type = ir.Type.parse("i64") + one = arith.constant(value=ir.IntegerAttr.get(i64_type, 1), result=i64_type, loc=loc, ip=ip) + buffer = llvm.alloca(res=pointer_type, elem_type=struct_type, array_size=one, alignment=64, loc=loc, ip=ip) + for index, offset in enumerate(offsets): + slot = llvm.getelementptr( + pointer_type, buffer, [], [index], i64_type, no_wrap_flags="None", loc=loc, ip=ip + ) + llvm.store(self._as_int64(offset).ir_value(), slot, loc=loc, ip=ip) + return SymmetricBufferDevice(buffer, self.max_ranks) + + i32_type = ir.Type.parse("i32") + vector_type = ir.Type.parse(f"vector<{self.max_ranks}xi64>") + vector = llvm.mlir_zero(vector_type, loc=loc, ip=ip) + for index, offset in enumerate(offsets): + element_index = arith.constant(value=ir.IntegerAttr.get(i32_type, index), result=i32_type, loc=loc, ip=ip) + vector = llvm.insertelement(vector, self._as_int64(offset).ir_value(), element_index, loc=loc, ip=ip) + return SymmetricBufferDevice(vector, self.max_ranks) + + +@JitArgAdapterRegistry.register_jit_arg_adapter(SymmetricBufferHost) +class _SymmetricBufferHostAdapter: + def __init__(self, argument: SymmetricBufferHost) -> None: + self._argument = argument + offsets = tuple(argument.offsets) + if len(offsets) != int(argument.max_ranks): + raise ValueError(f"Expected {int(argument.max_ranks)} peer offsets, got {len(offsets)}.") + self._fields = (Int64(argument.base_address), *(Int64(offset) for offset in offsets), Int32(argument.rank)) + + def __c_pointers__(self) -> list[Any]: + pointers: list[Any] = [] + for field in self._fields: + pointers.extend(get_c_pointers(field)) + return pointers + + def __get_mlir_types__(self) -> list[Any]: + types: list[Any] = [] + for field in self._fields: + types.extend(get_mlir_types(field)) + return types + + def __extract_mlir_values__(self) -> list[ir.Value]: + values: list[ir.Value] = [] + for field in self._fields: + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: list[ir.Value]) -> SymmetricBufferHost: + value_index = 0 + rebuilt = [] + for field in self._fields: + field_value_count = len(get_mlir_types(field)) + rebuilt.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) + value_index += field_value_count + if value_index != len(values): + raise ValueError(f"Consumed {value_index} MLIR values, got {len(values)}.") + + result = object.__new__(SymmetricBufferHost) + object.__setattr__(result, "base_address", rebuilt[0]) + object.__setattr__(result, "offsets", tuple(rebuilt[1:-1])) + object.__setattr__(result, "rank", rebuilt[-1]) + object.__setattr__(result, "max_ranks", self._argument.max_ranks) + return result + + +__all__ = ["SymmetricBufferDevice", "SymmetricBufferHost"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm.py new file mode 100644 index 000000000..a24a0997f --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm.py @@ -0,0 +1,2138 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Metadata-push routing and fused non-deterministic token communication.""" + +import dataclasses +import os +from typing import Callable, ClassVar, Literal, Optional, Tuple, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 +from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF + +from ...api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.dsl_helpers import mark_alignment, smem_exclusive_prefix +from ...helpers.flag_batch import make_flag_batch_tracker +from ...helpers.iket_compat import iket +from ...helpers.software_sync import NvlinkBarrier +from ...helpers.ptx_helpers import ( + cp_async_bulk_s2g, + cp_reduce_async_bulk_add_bf16_s2g, + cp_reduce_async_bulk_add_u32_s2g, + nanosleep, + read_clock64, + red_add_relaxed_sys_s32, + stg_b64, + stg_f32, + tma_load_1d, +) +from ...helpers.smem_workspace import SmemWorkspace +from ...helpers.utils import ceil_div, round_up +from ...quant_def import CombineFormat, QuantKind +from ..token_protocol import TokenSrcMetadata +from .symmetric_buffer import SymmetricBufferDevice + + +TokenBackMode = Literal["epi_warps", "standalone_warps", "reuse_dispatch_warps"] +TokenBackScheduleMode = Literal["static", "atomic_counter"] + + +@dataclasses.dataclass(frozen=True) +class TokenCommArgs: + """Device views materialized inside the fused kernel region. + + ``pre_reduced_activation`` is the per-topk combine staging plane: TokenComm's own symmetric region, except + under in-kernel top-k reduction where it degenerates to a view of the caller's 2D output. That REDG + accumulates, so in that mode the incoming content is the caller's accumulation base -- zero, or a + shared-expert result. + """ + + activation: cute.Tensor + activation_sf: cute.Tensor + pre_reduced_activation: cute.Tensor + pre_reduced_activation_sf: Optional[cute.Tensor] + peer_rank_ptr_mapper: SymmetricBufferDevice + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "TokenCommArgs": + if values: + raise ValueError(f"TokenCommArgs expected no MLIR values, got {len(values)}.") + return self + + +@dataclasses.dataclass(frozen=True) +class _SortedElement: + flat_topk_index: Int32 + topk_score: Optional[cutlass.Float32] + + def pack(self) -> Union[Int64, Int32]: + if cutlass.const_expr(self.topk_score is None): + return self.flat_topk_index + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + scratch[0] = self.flat_topk_index + cute.recast_tensor(scratch, cutlass.Float32)[1] = self.topk_score + return cute.recast_tensor(scratch, cutlass.Int64)[0] + + @classmethod + def from_packed(cls, packed: Union[Int64, Int32]) -> "_SortedElement": + if cutlass.const_expr(type(packed).width == 32): + return cls(flat_topk_index=Int32(packed), topk_score=None) + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + cute.recast_tensor(scratch, cutlass.Int64)[0] = packed + return cls(flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1]) + + +@cute.jit +def _copy_atom(dtype, num_bits_per_copy: int): + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy) + + +class _MetadataPushRouter(KernelComponent): + """Sort and push routing metadata into each destination rank's final pool.""" + + router_smem_limit_bytes: ClassVar[int] = 227 * 1024 + router_warps_per_cta: ClassVar[int] = 16 + + sizes_by_rank_region = "nvlink.token_comm.sizes_by_rank" + sizes_region = "nvlink.token_comm.sizes" + sizes_ready_region = "nvlink.token_comm.sizes_ready" + metadata_ready_region = "nvlink.token_comm.metadata_ready" + sorted_metadata_region = "nvlink.token_comm.sorted_metadata" + sorted_scores_region = "nvlink.token_comm.sorted_scores" + pool_expert_base_region = "nvlink.token_comm.pool_expert_base" + token_src_metadata_region = "nvlink.token_comm.token_src_metadata" + fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" + source_expert_base_region = "nvlink.token_comm.source_expert_base" + push_destination_base_region = "nvlink.token_comm.push_destination_base" + sorted_metadata_ready_region = "nvlink.token_comm.sorted_metadata_ready" + push_table_ready_region = "nvlink.token_comm.push_table_ready" + router_size_counter_region = "nvlink.token_comm.router_size_counter" + router_histogram_done_region = "nvlink.token_comm.router_histogram_done" + source_base_ready_region = "nvlink.token_comm.source_base_ready" + + router_data_histogram_region = "nvlink.token_comm.router_smem.data_histogram" + router_data_prefix_region = "nvlink.token_comm.router_smem.data_prefix" + router_data_warp_totals_region = "nvlink.token_comm.router_smem.data_warp_totals" + router_data_sorted_region = "nvlink.token_comm.router_smem.data_sorted" + router_data_base_region = "nvlink.token_comm.router_smem.data_base" + router_helper_size_matrix_region = "nvlink.token_comm.router_smem.helper_size_matrix" + router_helper_totals_region = "nvlink.token_comm.router_smem.helper_totals" + router_helper_prefix_region = "nvlink.token_comm.router_smem.helper_prefix" + router_helper_warp_totals_region = "nvlink.token_comm.router_smem.helper_warp_totals" + router_helper_load_mbarrier_region = "nvlink.token_comm.router_smem.helper_load_mbarrier" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "token_padding_block": int, + "promised_launchable_sm_count": int, + "router_smem_limit_bytes": OptionalRequirement(int), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + + self.token_padding_block = impl_desc["token_padding_block"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.router_smem_limit_bytes = impl_desc.get("router_smem_limit_bytes", 227 * 1024) + + self._validate_router_configuration() + self.expert_count_padded = round_up(self.expert_count, 4) + self.expert_count_with_trash = self.expert_count_padded + 1 + self.router_elements_per_lane, self.router_data_cta_count = self._router_launch_configuration() + self.router_tokens_per_cta = self.router_elements_per_lane * self.router_warps_per_cta * 32 + self.router_push_cta_count = ceil_div(self.expert_count, self.router_warps_per_cta) + self.router_grid_cta_count = max(self.router_data_cta_count + 1, self.router_push_cta_count) + if self.router_grid_cta_count > self.promised_launchable_sm_count: + raise ValueError( + "Router grid exceeds promised_launchable_sm_count; all metadata-push CTAs must be concurrently resident." + ) + self.worst_case_token_count = self.worst_case_padded_tokens(self.token_padding_block) + self._router_smem_workspace = self._build_router_smem_workspace() + + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_grid_thread_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + + def _validate_router_configuration(self) -> None: + positive_fields = ( + "world_size", + "expert_count", + "topk", + "max_tokens_per_rank", + "token_padding_block", + "promised_launchable_sm_count", + "router_smem_limit_bytes", + ) + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.expert_count % self.world_size != 0: + raise ValueError( + f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." + ) + if self.expert_count > 16384: + raise NotImplementedError("TokenComm supports at most 16384 global experts.") + if self.topk > self.expert_count: + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + def worst_case_padded_tokens(self, block: int) -> int: + source_token_capacity = self.world_size * self.max_tokens_per_rank + routes_per_source_token = min(self.topk, self.experts_per_rank) + route_capacity = source_token_capacity * routes_per_source_token + active_expert_capacity = min(self.experts_per_rank, route_capacity) + route_budget_blocks = active_expert_capacity + (route_capacity - active_expert_capacity) // block + expert_bound_blocks = active_expert_capacity * int(ceil_div(source_token_capacity, block)) + return min(route_budget_blocks, expert_bound_blocks) * block + + def _router_launch_configuration(self) -> Tuple[int, int]: + def next_power_of_two(value: int) -> int: + return 1 << (max(value, 1) - 1).bit_length() + + routed_token_capacity = self.max_tokens_per_rank * self.topk + minimum_cta_capacity = 2048 + maximum_cta_capacity = 16384 + maximum_data_cta_count = 128 + maximum_supported_tokens = maximum_cta_capacity * maximum_data_cta_count + if routed_token_capacity > maximum_supported_tokens: + raise NotImplementedError(f"The router supports at most {maximum_supported_tokens} routed tokens per rank.") + cta_capacity = min(maximum_cta_capacity, next_power_of_two(max(routed_token_capacity, minimum_cta_capacity))) + elements_per_lane = cta_capacity // (self.router_warps_per_cta * 32) + data_cta_count = ceil_div(routed_token_capacity, cta_capacity) + return elements_per_lane, data_cta_count + + def _build_router_smem_workspace(self) -> SmemWorkspace: + workspace = SmemWorkspace() + workspace.register_mbarrier(self.router_helper_load_mbarrier_region, 1) + overlay = workspace.create_overlay("nvlink.token_comm.router_smem.role") + data_lifetime = overlay.add_lifetime("data_cta") + data_lifetime.register_tensor(self.router_data_histogram_region, cutlass.Int32, (self.expert_count_with_trash,)) + data_lifetime.register_tensor( + self.router_data_prefix_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 + ) + data_lifetime.register_tensor(self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,)) + data_lifetime.register_tensor( + self.router_data_sorted_region, + (cutlass.Int64 if self.apply_topk_at_fc1 else cutlass.Int32), + (self.router_tokens_per_cta,), + byte_alignment=16, + ) + if self.router_data_cta_count > 1: + data_lifetime.register_tensor( + self.router_data_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + + helper_lifetime = overlay.add_lifetime("helper_cta") + helper_lifetime.register_tensor( + self.router_helper_size_matrix_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + stride=(self.expert_count_padded, 1), + byte_alignment=16, + ) + helper_lifetime.register_tensor( + self.router_helper_totals_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_prefix_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 + ) + workspace.finalize(max_bytes=self.router_smem_limit_bytes) + return workspace + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router_smem_workspace + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + """Register router-private state and Router-to-Main outputs.""" + self._register_router_workspace(workspace) + + def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: + maximum_routed_tokens = self.max_tokens_per_rank * self.topk + workspace.register( + self.sizes_by_rank_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + buffer_space="shared", + stride=(self.expert_count_padded, 1), + ) + workspace.register( + self.sizes_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="shared", reset="tail_reset" + ) + workspace.register(self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.metadata_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.sorted_metadata_region, cutlass.Int64, (maximum_routed_tokens,), buffer_space="local") + if self.apply_topk_at_fc1: + workspace.register( + self.sorted_scores_region, cutlass.Float32, (maximum_routed_tokens,), buffer_space="local" + ) + workspace.register( + self.token_src_metadata_region, + cutlass.Int64, + (self.worst_case_token_count,), + buffer_space="shared", + byte_alignment=16, + ) + if self.apply_topk_at_fc1: + workspace.register( + self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" + ) + workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local") + workspace.register( + self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.push_destination_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.sorted_metadata_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register(self.push_table_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") + if self.router_data_cta_count > 1: + workspace.register( + self.router_size_counter_region, + cutlass.Int32, + (self.expert_count_with_trash,), + buffer_space="local", + reset="tail_reset", + ) + workspace.register( + self.router_histogram_done_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register( + self.source_base_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "_MetadataPushRouter": + if values: + raise ValueError("_MetadataPushRouter carries no MLIR values.") + return self + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + """Launch counting-sort DATA, size-exchange HELPER, and metadata PUSH roles.""" + if cutlass.const_expr(self.apply_topk_at_fc1 and topk_scores is None): + raise ValueError("apply_topk_at_fc1 requires router topk_scores.") + peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() + self._router_kernel( + topk_indices, + topk_scores, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper, + device_workspace, + ).launch( + grid=[self.router_grid_cta_count, 1, 1], + block=[self.router_warps_per_cta * 32, 1, 1], + min_blocks_per_mp=1, + stream=stream, + ) + + @cute.kernel + def _router_kernel( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper: SymmetricBufferDevice, + device_workspace: DeviceWorkspace, + ) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + linear_cta_idx, _, _ = cute.arch.block_idx() + cute.arch.griddepcontrol_launch_dependents() + block_thread_count = self.router_warps_per_cta * 32 + grid_thread_idx = thread_idx + linear_cta_idx * block_thread_count + warp_idx = cute.arch.make_warp_uniform(thread_idx // Int32(32)) + lane_idx = thread_idx % Int32(32) + + storage_type = self._router_smem_workspace.storage_class() + smem_allocator = cutlass.utils.SmemAllocator() + storage = smem_allocator.allocate(storage_type) + smem_base = storage.buffer.data_ptr() + + device_workspace.assign_device_members(local_workspace, shared_workspace) + self._device_workspace = device_workspace + self._peer_rank_ptr_mapper = peer_rank_ptr_mapper + self._router_local_rank = local_rank + self._router_thread_idx = thread_idx + self._router_linear_cta_idx = linear_cta_idx + self._router_grid_thread_idx = grid_thread_idx + self._router_warp_idx = warp_idx + self._router_lane_idx = lane_idx + + if cutlass.const_expr(self.router_data_cta_count == 1): + self._router_single_cta(topk_indices, topk_scores, smem_base) + else: + self._router_multiple_ctas(topk_indices, topk_scores, smem_base) + if linear_cta_idx < Int32(self.router_push_cta_count): + self._router_push_metadata() + + device_workspace.remove_device_members() + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_grid_thread_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + + @cute.jit + def _router_single_cta( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + + zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + expert = Int32(zero_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_with_trash): + histogram[expert] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + within_expert_indices = self._build_histogram(expert_registers, histogram) + iket.range_pop() + + iket.range_push("router.prefix_and_publish") + publish_sizes = self._broadcast_sizes_to_peers( + cute.make_tensor(histogram.iterator, cute.make_layout((self.expert_count_padded,))) + ) + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(histogram.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + publish_sizes() + iket.range_pop() + + iket.range_push("router.sort") + self._sort_router_elements( + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.write_out") + self._dump_contiguous_router_output(sorted_elements, total_valid_routes) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_single_cta(smem_base) + + @cute.jit + def _router_multiple_ctas( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + dump_base = self._router_smem_workspace.tensor(self.router_data_base_region, smem_base) + + zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + expert = Int32(zero_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_with_trash): + histogram[expert] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + within_expert_indices = self._build_histogram(expert_registers, histogram) + iket.range_pop() + + iket.range_push("router.reserve_and_prefix") + size_counter = self._device_workspace.ptr(self.router_size_counter_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + dump_base[expert] = Int32( + cute.arch.atomic_add(size_counter + expert, histogram[expert], sem="relaxed", scope="gpu") + ) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.router_histogram_done_region), Int32(1), sem="release", scope="gpu" + ) + + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(histogram.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.sort") + self._sort_router_elements( + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.wait_source_base") + source_base_ready = self._device_workspace.ptr(self.source_base_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.write_out") + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + dump_base[expert] = dump_base[expert] + source_expert_base[expert] + cute.arch.sync_threads() + self._dump_router_output_by_expert(histogram, prefix, dump_base, sorted_elements) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_multiple_ctas(smem_base) + + @cute.jit + def _router_helper_single_cta(self, smem_base: cute.Pointer) -> None: + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + size_counter = self._device_workspace.tensor(self.router_size_counter_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + + histogram_done = self._device_workspace.ptr(self.router_histogram_done_region) + iket.range_push("router.wait_histogram") + if self._router_thread_idx == Int32(0): + while cute.arch.load(histogram_done, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.broadcast_sizes") + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + totals[expert] = size_counter[expert] + cute.arch.sync_threads() + + publish_sizes = self._broadcast_sizes_to_peers(totals) + smem_exclusive_prefix( + totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.source_base_ready_region), Int32(1), sem="release", scope="gpu" + ) + publish_sizes() + iket.range_pop() + + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_push_metadata(self) -> None: + block_thread_count = self.router_warps_per_cta * 32 + sorted_metadata_ready = self._device_workspace.ptr(self.sorted_metadata_ready_region) + push_table_ready = self._device_workspace.ptr(self.push_table_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(sorted_metadata_ready, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + while cute.arch.load(push_table_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + + global_expert = self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx + if global_expert < Int32(self.expert_count): + sizes_by_rank = self._device_workspace.tensor(self.sizes_by_rank_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + route_count = sizes_by_rank[self._router_local_rank, global_expert] + source_begin = source_expert_base[global_expert] + destination_begin = push_destination_base[global_expert] + destination_rank = global_expert // Int32(self.experts_per_rank) + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) + destination_metadata_address = ( + self._device_workspace.ptr(self.token_src_metadata_region).toint() + peer_offset + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + source_scores = self._device_workspace.ptr(self.sorted_scores_region) + destination_scores_address = ( + self._device_workspace.ptr(self.fc1_topk_scores_region).toint() + peer_offset + ) + route_round_count = (route_count + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + if route < route_count: + source_position = source_begin + route + destination_position = destination_begin + route + metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) + stg_b64( + destination_metadata_address + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), + metadata, + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + score = cute.arch.load(source_scores + source_position, cutlass.Float32) + stg_f32(destination_scores_address + Int64(destination_position) * Int64(4), score) + + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + # Keep notifier threads behind the leader's system fence. + cute.arch.sync_threads() + metadata_ready_address = self._device_workspace.ptr(self.metadata_ready_region).toint() + rank_round_count = ceil_div(self.world_size, block_thread_count) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = Int32(rank_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(metadata_ready_address, destination_rank, Int64(0)), Int32(1) + ) + + @cute.jit + def _broadcast_sizes_to_peers(self, smem_expert_counts: cute.Tensor) -> Callable[[], None]: + block_thread_count = self.router_warps_per_cta * 32 + row_bytes = Int32(self.expert_count_padded * 4) + matrix_address = self._device_workspace.ptr(self.sizes_by_rank_region).toint() + total_address = self._device_workspace.ptr(self.sizes_region).toint() + rank_round_count = ceil_div(self.world_size, self.router_warps_per_cta) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = self._router_warp_idx + Int32(rank_round * self.router_warps_per_cta) + if destination_rank < Int32(self.world_size): + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + destination_row_address = ( + matrix_address + + peer_offset + + Int64(Int32(self._router_local_rank) * Int32(self.expert_count_padded)) * Int64(4) + ) + destination_row = cute.make_ptr( + cutlass.Int32, destination_row_address, AddressSpace.gmem, assumed_align=16 + ) + destination_total = cute.make_ptr( + cutlass.Int32, total_address + peer_offset, AddressSpace.gmem, assumed_align=16 + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_row, smem_expert_counts.iterator, row_bytes) + cp_reduce_async_bulk_add_u32_s2g(destination_total, smem_expert_counts.iterator, row_bytes) + cute.arch.cp_async_bulk_commit_group() + + def finalize() -> None: + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + cute.arch.sync_threads() + ready_address = self._device_workspace.ptr(self.sizes_ready_region).toint() + ready_round_count = ceil_div(self.world_size, block_thread_count) + for ready_round in cutlass.range_constexpr(ready_round_count): + destination_rank = Int32(ready_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(ready_address, destination_rank, Int64(0)), Int32(1) + ) + + return finalize + + @cute.jit + def _compute_push_tables(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + owner_expert_begin = Int32(self._router_local_rank) * Int32(self.experts_per_rank) + matrix_bytes = self.world_size * self.expert_count_padded * 4 + + size_matrix = self._router_smem_workspace.tensor(self.router_helper_size_matrix_region, smem_base) + padded_totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + load_mbarrier = self._router_smem_workspace.ptr(self.router_helper_load_mbarrier_region, smem_base) + sizes = self._device_workspace.tensor(self.sizes_region) + + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_init(load_mbarrier, 1) + + sizes_ready = self._device_workspace.ptr(self.sizes_ready_region) + iket.range_push("router.wait_sizes_ready") + if self._router_thread_idx == Int32(0): + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.world_size): + nanosleep(150) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.load_sizes_and_prefix") + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(load_mbarrier, Int32(matrix_bytes)) + tma_load_1d( + size_matrix.iterator, + self._device_workspace.ptr(self.sizes_by_rank_region), + load_mbarrier, + Int32(matrix_bytes), + ) + + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + expert_size = sizes[expert] + padded_totals[expert] = ( + (expert_size + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) + ) * Int32(self.token_padding_block) + cute.arch.sync_threads() + smem_exclusive_prefix( + padded_totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + pool_expert_base = self._device_workspace.tensor(self.pool_expert_base_region) + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + pool_expert_base[local_expert] = prefix[owner_expert_begin + local_expert] - prefix[owner_expert_begin] + + cute.arch.mbarrier_wait(load_mbarrier, 0) + iket.range_pop() + + iket.range_push("router.build_push_destinations") + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + global_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if global_expert < Int32(self.expert_count): + destination_rank = global_expert // Int32(self.experts_per_rank) + destination_expert_begin = destination_rank * Int32(self.experts_per_rank) + destination_pool_base = prefix[global_expert] - prefix[destination_expert_begin] + local_ring_position = ( + Int32(self._router_local_rank) - destination_rank + Int32(self.world_size) + ) % Int32(self.world_size) + source_ring_offset = Int32(0) + for ring_position in cutlass.range_constexpr(self.world_size): + source_rank = (destination_rank + Int32(ring_position)) % Int32(self.world_size) + if Int32(ring_position) < local_ring_position: + source_ring_offset = source_ring_offset + size_matrix[source_rank, global_expert] + push_destination_base[global_expert] = destination_pool_base + source_ring_offset + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.push_table_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + + @cute.jit + def _load_router_inputs( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor] + ) -> Tuple[cute.Tensor, Optional[cute.Tensor]]: + elements_per_vector = 128 // topk_indices.dtype.width + grid_thread_count = self.router_data_cta_count * self.router_warps_per_cta * 32 + tile_span = elements_per_vector * grid_thread_count + maximum_elements = self.max_tokens_per_rank * self.topk + actual_token_count = Int32(self.max_tokens_per_rank) + actual_elements = Int32(maximum_elements) + load_round_count = ceil_div(maximum_elements, tile_span) + elements_per_thread = load_round_count * elements_per_vector + + topk_flat = cute.make_tensor(topk_indices.iterator, cute.make_layout((maximum_elements,))) + topk_vectors = cute.logical_divide(cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None)) + load_atom = _copy_atom(topk_indices.dtype, 128) + expert_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Int32) + if cutlass.const_expr(topk_indices.dtype.width == 64): + raw_indices = cute.make_rmem_tensor((elements_per_thread,), topk_indices.dtype) + raw_vectors = cute.zipped_divide(raw_indices, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + load_atom, + mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), + raw_vectors[None, load_round], + ) + else: + expert_vectors = cute.zipped_divide(expert_registers, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + load_atom, + mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), + expert_vectors[None, load_round], + ) + + score_registers = None + if cutlass.const_expr(self.apply_topk_at_fc1): + score_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Float32) + scores_flat = cute.make_tensor(topk_scores.iterator, cute.make_layout((maximum_elements,))) + score_vectors = cute.logical_divide( + cute.zipped_divide(scores_flat, (tile_span,)), (elements_per_vector, None) + ) + score_atom = _copy_atom(cutlass.Float32, elements_per_vector * 32) + score_register_vectors = cute.zipped_divide(score_registers, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + score_atom, + mark_alignment(score_vectors[(None, self._router_grid_thread_idx), load_round], 16), + score_register_vectors[None, load_round], + ) + + expert_registers_u32 = cute.recast_tensor(expert_registers, cutlass.Uint32) + if cutlass.const_expr(topk_indices.dtype.width == 64): + raw_indices_i32 = cute.recast_tensor(raw_indices, cutlass.Int32) + for register_idx in cutlass.range_constexpr(elements_per_thread): + if cutlass.const_expr(topk_indices.dtype.width == 64): + expert_registers[register_idx] = raw_indices_i32[2 * register_idx] + token_idx, _ = self._router_value_coordinate(register_idx, topk_indices.dtype) + is_invalid = (expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count)) | ( + token_idx >= actual_token_count + ) + if is_invalid: + expert_registers[register_idx] = Int32(self.expert_count_padded) + return expert_registers, score_registers + + @cute.jit + def _build_histogram(self, expert_registers: cute.Tensor, histogram: cute.Tensor) -> cute.Tensor: + register_count = cute.size(expert_registers) + within_expert_indices = cute.make_rmem_tensor((register_count,), cutlass.Int32) + for register_idx in cutlass.range_constexpr(register_count): + within_expert_indices[register_idx] = Int32( + cute.arch.atomic_add( + histogram.iterator + expert_registers[register_idx], Int32(1), sem="relaxed", scope="cta" + ) + ) + cute.arch.sync_threads() + return within_expert_indices + + @cute.jit + def _sort_router_elements( + self, + expert_registers: cute.Tensor, + within_expert_indices: cute.Tensor, + score_registers: Optional[cute.Tensor], + sorted_elements: cute.Tensor, + expert_run_starts: cute.Tensor, + topk_index_type: type, + ) -> None: + register_count = cute.size(expert_registers) + for register_idx in cutlass.range_constexpr(register_count): + token_idx, topk_slot = self._router_value_coordinate(register_idx, topk_index_type) + flat_topk_index = token_idx * Int32(self.topk) + topk_slot + destination = expert_run_starts[expert_registers[register_idx]] + within_expert_indices[register_idx] + if cutlass.const_expr(self.apply_topk_at_fc1): + sorted_elements[destination] = _SortedElement(flat_topk_index, score_registers[register_idx]).pack() + else: + sorted_elements[destination] = _SortedElement(flat_topk_index, None).pack() + + @cute.jit + def _router_value_coordinate(self, register_idx: int, topk_index_type: type) -> Tuple[Int32, Int32]: + elements_per_vector = 128 // topk_index_type.width + tile_span = elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 + flat_index = Int32( + register_idx // elements_per_vector * tile_span + register_idx % elements_per_vector + ) + self._router_grid_thread_idx * Int32(elements_per_vector) + return (flat_index // Int32(self.topk), flat_index % Int32(self.topk)) + + @cute.jit + def _dump_contiguous_router_output(self, sorted_elements: cute.Tensor, total_valid_routes: Int32) -> None: + block_thread_count = self.router_warps_per_cta * 32 + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32(block_thread_count) + for dump_round in cutlass.range(dump_round_count, unroll=4): + position = Int32(dump_round * block_thread_count) + self._router_thread_idx + predicate = Int32(position < total_valid_routes) + element = _SortedElement.from_packed(sorted_elements[position]) + metadata = TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=(element.flat_topk_index // Int32(self.topk)), + src_topk=(element.flat_topk_index % Int32(self.topk)), + ) + stg_b64(metadata_address + Int64(position) * Int64(TokenSrcMetadata.nbytes), metadata.pack(), predicate) + if cutlass.const_expr(self.apply_topk_at_fc1): + stg_f32(score_address + Int64(position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def _dump_router_output_by_expert( + self, + histogram: cute.Tensor, + expert_run_starts: cute.Tensor, + expert_dump_bases: cute.Tensor, + sorted_elements: cute.Tensor, + ) -> None: + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + expert_round_count = ceil_div(self.expert_count_padded, self.router_warps_per_cta) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = self._router_warp_idx + Int32(expert_round * self.router_warps_per_cta) + if expert < Int32(self.expert_count_padded): + run_begin = expert_run_starts[expert] + run_length = histogram[expert] + dump_begin = expert_dump_bases[expert] + route_round_count = (run_length + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + predicate = Int32(route < run_length) + element = _SortedElement.from_packed(sorted_elements[predicate * (run_begin + route)]) + output_position = dump_begin + route + metadata = TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=(element.flat_topk_index // Int32(self.topk)), + src_topk=(element.flat_topk_index % Int32(self.topk)), + ) + stg_b64( + metadata_address + Int64(output_position) * Int64(TokenSrcMetadata.nbytes), + metadata.pack(), + predicate, + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + stg_f32(score_address + Int64(output_position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + """Return this rank's contiguous expert-size view.""" + sizes = device_workspace.tensor(self.sizes_region) + expert_begin = local_rank * Int32(self.experts_per_rank) + return cute.make_tensor(sizes.iterator + expert_begin, cute.make_layout((self.experts_per_rank,))) + + @property + def metadata_ready_target(self) -> int: + return self.router_push_cta_count * self.world_size + + @cute.jit + def sizes_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.sizes_region) + + @cute.jit + def pool_expert_base_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.pool_expert_base_region) + + @cute.jit + def token_src_metadata_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.token_src_metadata_region) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + sizes_ready = device_workspace.ptr(self.sizes_ready_region) + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.world_size): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def wait_for_metadata_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + metadata_ready = device_workspace.ptr(self.metadata_ready_region) + while cute.arch.load(metadata_ready, Int32, sem="acquire", scope="sys") != Int32( + self.metadata_ready_target + ): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.token_src_metadata_region) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.apply_topk_at_fc1): + return None + return device_workspace.tensor(self.fc1_topk_scores_region) + + +class TokenCommNonDeterministic(KernelComponent): + """Public fused-kernel token communication component.""" + + transfer_warp_count: ClassVar[int] = 4 + transfer_thread_count: ClassVar[int] = transfer_warp_count * 32 + standalone_chunk_bytes: ClassVar[int] = 2048 + minimum_pacing_window_cycles: ClassVar[int] = 512 + standalone_max_backoff_cycles: ClassVar[int] = 500 + adaptive_minimum_sleep_cycles: ClassVar[int] = 50 + transfer_lifetime_barrier_id: ClassVar[int] = 9 + grid_sync_barrier_id: ClassVar[int] = 10 + standalone_size_barrier_id: ClassVar[int] = 11 + token_in_size_barrier_id: ClassVar[int] = 12 + + fc1_ready_region = "nvlink.token_comm.fc1_ready" + fc1_activation_region = "nvlink.token_comm.fc1_activation" + fc1_activation_sf_region = "nvlink.token_comm.fc1_activation_sf" + fc2_done_region = "nvlink.token_comm.fc2_done" + fc2_activation_region = "nvlink.token_comm.fc2_activation" + fc2_activation_sf_region = "nvlink.token_comm.fc2_activation_sf" + pre_reduced_activation_region = "nvlink.token_comm.pre_reduced_activation" + pre_reduced_activation_sf_region = "nvlink.token_comm.pre_reduced_activation_sf" + token_back_schedule_region = "nvlink.token_comm.token_back_schedule" + + token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" + token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" + expert_sizes_smem_region = "nvlink.token_comm.main_smem.expert_sizes" + token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" + token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" + token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" + token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "hidden_size": int, + "quant_kind": str, + "combine_format": CombineFormat, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "token_padding_block": int, + "sf_padding_block": int, + "tokens_per_fc1_ready_slot": int, + "fc2_done_signals_per_token_tile": int, + "promised_launchable_sm_count": int, + "token_in_flag_batch": int, + "token_back_mode": str, + "token_back_schedule_mode": str, + "reduce_topk_in_kernel": bool, + "router_smem_limit_bytes": OptionalRequirement(int), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.hidden_size = problem_desc["hidden_size"] + self.quant_kind = QuantKind(problem_desc["quant_kind"]) + self.combine_format = problem_desc["combine_format"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.tokens_per_fc1_ready_slot = impl_desc["tokens_per_fc1_ready_slot"] + self.fc2_done_signals_per_token_tile = impl_desc["fc2_done_signals_per_token_tile"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.token_in_flag_batch = impl_desc["token_in_flag_batch"] + self.token_back_mode: TokenBackMode = impl_desc["token_back_mode"] + self.token_back_schedule_mode: TokenBackScheduleMode = impl_desc["token_back_schedule_mode"] + self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] + self.router_smem_limit_bytes = impl_desc.get("router_smem_limit_bytes", 227 * 1024) + + self._validate_configuration() + self._router = _MetadataPushRouter(problem_desc, impl_desc) + self._nvlink_barrier = NvlinkBarrier(world_size=self.world_size, barrier_id=self.grid_sync_barrier_id) + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + + def _validate_configuration(self) -> None: + positive_fields = ( + "world_size", + "expert_count", + "topk", + "max_tokens_per_rank", + "hidden_size", + "token_padding_block", + "sf_padding_block", + "tokens_per_fc1_ready_slot", + "promised_launchable_sm_count", + "router_smem_limit_bytes", + ) + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.expert_count % self.world_size != 0: + raise ValueError( + f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." + ) + if self.expert_count > 16384: + raise NotImplementedError("TokenComm supports at most 16384 global experts.") + if self.topk > self.expert_count: + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") + if self.token_back_mode not in ("epi_warps", "standalone_warps", "reuse_dispatch_warps"): + raise ValueError(f"Unsupported token_back_mode {self.token_back_mode!r}.") + if self.token_back_schedule_mode not in ("static", "atomic_counter"): + raise ValueError( + f"token_back_schedule_mode must be static or atomic_counter, got {self.token_back_schedule_mode!r}." + ) + if not 1 <= self.token_in_flag_batch <= 32: + raise ValueError(f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}.") + if self.tokens_per_fc1_ready_slot % self.token_padding_block != 0: + raise ValueError("tokens_per_fc1_ready_slot must be divisible by token_padding_block.") + if self.token_back_enabled and self.fc2_done_signals_per_token_tile <= 0: + raise ValueError("fc2_done_signals_per_token_tile must be positive when token-back is enabled.") + if self.reduce_topk_in_kernel and self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("In-kernel top-k reduction requires BF16 combine data.") + element_block = self.activation_sf_vector_size * 4 + if self.hidden_size % element_block != 0: + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + if self.sf_padding_block % 128 != 0: + raise ValueError("sf_padding_block must be a multiple of 128.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + @property + def activation_dtype(self) -> type: + return self.quant_kind.activation_dtype + + @property + def activation_sf_dtype(self) -> type: + return self.quant_kind.sf_dtype + + @property + def activation_sf_vector_size(self) -> int: + return self.quant_kind.sf_vec_size + + @property + def bytes_per_token(self) -> int: + return self.hidden_size * int(self.activation_dtype.width) // 8 + + @property + def activation_sf_hidden_padded(self) -> int: + valid_hidden = self.hidden_size // self.activation_sf_vector_size + elements_per_16_bytes = 128 // int(self.activation_sf_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def combine_sf_hidden_padded(self) -> int: + if not self.combine_format.is_quantized: + return 0 + valid_hidden = self.hidden_size // self.combine_format.scale_block + elements_per_16_bytes = 128 // int(self.combine_format.scale_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def token_back_push_data(self) -> bool: + return self.token_back_mode != "epi_warps" + + @property + def token_back_push_sf(self) -> bool: + return self.combine_format.is_quantized + + @property + def token_back_enabled(self) -> bool: + return self.token_back_push_data or self.token_back_push_sf + + @property + def worst_case_token_count(self) -> int: + return self._router.worst_case_token_count + + @property + def worst_case_sf_token_count(self) -> int: + return self._router.worst_case_padded_tokens(self.sf_padding_block) + + @property + def max_fc1_ready_slot_count(self) -> int: + return self._router.worst_case_padded_tokens(self.tokens_per_fc1_ready_slot) // self.tokens_per_fc1_ready_slot + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router.router_smem_workspace + + @property + def expert_count_padded(self) -> int: + return self._router.expert_count_padded + + @property + def expert_count_with_trash(self) -> int: + return self._router.expert_count_with_trash + + @property + def router_elements_per_lane(self) -> int: + return self._router.router_elements_per_lane + + @property + def router_warps_per_cta(self) -> int: + return self._router.router_warps_per_cta + + @property + def router_data_cta_count(self) -> int: + return self._router.router_data_cta_count + + @property + def router_tokens_per_cta(self) -> int: + return self._router.router_tokens_per_cta + + @property + def router_push_cta_count(self) -> int: + return self._router.router_push_cta_count + + @property + def router_grid_cta_count(self) -> int: + return self._router.router_grid_cta_count + + @property + def metadata_ready_target(self) -> int: + return self._router.metadata_ready_target + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + self._router.register_device_workspace(workspace) + self._register_main_workspace(workspace) + self._nvlink_barrier.register_device_workspace(workspace) + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + self._router.launch_router( + topk_indices, + topk_scores, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper_host, + device_workspace, + stream, + ) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + return self._router.local_expert_sizes(device_workspace, local_rank) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + self._router.wait_for_sizes_ready(device_workspace, sleep_cycles) + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return self._router.token_src_metadata_tensor(device_workspace) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return self._router.fc1_topk_scores_tensor(device_workspace) + + @cute.jit + def assign_device_members( + self, + *, + device_workspace: DeviceWorkspace, + token_comm_args: TokenCommArgs, + local_rank: Int32, + linear_cta_idx: Int32, + ) -> None: + self._device_workspace = device_workspace + self._token_comm_args = token_comm_args + self._local_rank = local_rank + self._linear_cta_idx = linear_cta_idx + thread_idx, _, _ = cute.arch.thread_idx() + transfer_thread_idx = thread_idx % Int32(self.transfer_thread_count) + self._transfer_warp_idx = cute.arch.make_warp_uniform(transfer_thread_idx // Int32(32)) + self._lane_idx = transfer_thread_idx % Int32(32) + self._nvlink_barrier.assign_device_members(device_workspace, token_comm_args.peer_rank_ptr_mapper) + + def remove_device_members(self) -> None: + self._nvlink_barrier.remove_device_members() + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "TokenCommNonDeterministic": + if values: + raise ValueError("TokenCommNonDeterministic carries no MLIR values.") + return self + + def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: + """Register the fused kernel's GMEM regions. Three groups, each stating its own existence condition. + + FC1 pools, always present: the dispatched payload ``token_in`` pulls in, addressed in POOL index space + (per-expert padded runs of tokens routed into this rank's experts, from every rank). + + Token-back machinery, iff ``token_back_enabled``: the transfer warps' counters plus the rank-local FC2 + staging they read, also in pool index space. ``epi_warps`` needs none of it, because there the FC2 + epilogue reaches the peers itself. The scale plane is registered separately from the data plane and is + staged locally in EVERY mode: pushing scales per token would scatter one warp's 32 lanes across up to 32 + ranks and explode the NVLink request count. + + Combine plane, iff ``not reduce_topk_in_kernel``: the symmetric per-topk landing zone peers deliver into, + addressed in (source token, source topk) space. It is the DESTINATION of the round trip whose source is + the staging above, so its condition is deliberately independent of who performs the transfer -- + ``epi_warps`` with a separate reduce registers this plane while registering no token-back machinery at + all. Peers reach it through the symmetric heap, so it cannot be a caller tensor: only this component + knows the wire dtype and the padded scale stride. Its ``data`` reset keeps it out of both the host zero + prefix and the per-launch tail reset, since every cell it exposes is rewritten each launch and the plane + is far too large to be worth clearing. + """ + workspace.register( + self.fc1_ready_region, + cutlass.Int32, + (self.max_fc1_ready_slot_count,), + buffer_space="local", + reset="tail_reset", + ) + activation_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc1_activation_region, + self.activation_dtype, + (activation_element_count,), + buffer_space="local", + byte_alignment=128, + ) + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded + workspace.register( + self.fc1_activation_sf_region, + self.activation_sf_dtype, + (activation_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + + if self.token_back_enabled: + workspace.register( + self.fc2_done_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" + ) + if self.token_back_push_data: + fc2_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc2_activation_region, + self.combine_format.act_dtype, + (fc2_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_push_sf: + fc2_sf_element_count = self.worst_case_token_count * self.combine_sf_hidden_padded + workspace.register( + self.fc2_activation_sf_region, + self.combine_format.scale_dtype, + (fc2_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": + workspace.register( + self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + + if not self.reduce_topk_in_kernel: + workspace.register( + self.pre_reduced_activation_region, + self.combine_format.act_dtype, + (self.max_tokens_per_rank, self.topk, self.hidden_size), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + if self.combine_format.is_quantized: + workspace.register( + self.pre_reduced_activation_sf_region, + self.combine_format.scale_dtype, + (self.max_tokens_per_rank, self.topk, self.combine_sf_hidden_padded), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + + def register_smem_regions(self, workspace: SmemWorkspace) -> None: + workspace.register_mbarrier(self.token_in_mbarrier_region, self.transfer_warp_count) + if self.token_back_enabled: + workspace.register_mbarrier(self.token_back_mbarrier_region, self.transfer_warp_count) + workspace.register_tensor( + self.expert_sizes_smem_region, cutlass.Int32, (self.experts_per_rank,), byte_alignment=16 + ) + transfer_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.transfer") + token_in_lifetime = transfer_overlay.add_lifetime("token_in") + token_in_lifetime.register_tensor( + self.token_in_activation_smem_region, + self.activation_dtype, + (self.transfer_warp_count, self.hidden_size), + byte_alignment=16, + ) + token_in_lifetime.register_tensor( + self.token_in_sf_smem_region, + self.activation_sf_dtype, + (self.transfer_warp_count, (self.activation_sf_vector_size, self.activation_sf_hidden_padded)), + stride=(self.activation_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + if not self.token_back_enabled: + return + + if self.token_back_mode == "standalone_warps": + token_back_lifetime = workspace.create_overlay( + "nvlink.token_comm.main_smem.standalone_token_back" + ).add_lifetime("token_back") + else: + token_back_lifetime = transfer_overlay.add_lifetime("token_back") + + if self.token_back_mode == "standalone_warps": + available_bytes_per_warp = self.standalone_chunk_bytes + else: + activation_bytes = self.bytes_per_token + sf_bytes = self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 + available_bytes_per_warp = activation_bytes + sf_bytes + + if self.token_back_push_data: + bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + if self.token_back_mode == "standalone_warps": + chunk_bytes = self.standalone_chunk_bytes + elif available_bytes_per_warp < bytes_per_output_token: + chunk_bytes = self.bytes_per_token + else: + chunk_bytes = bytes_per_output_token + if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: + raise ValueError("Token-back data chunk bytes must divide one row.") + chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) + token_back_lifetime.register_tensor( + self.token_back_activation_smem_region, + self.combine_format.act_dtype, + (self.transfer_warp_count, chunk_elements), + byte_alignment=16, + ) + if self.token_back_push_sf: + sf_row_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + if sf_row_bytes > available_bytes_per_warp: + raise ValueError("Token-back scale row exceeds its per-warp stage.") + token_back_lifetime.register_tensor( + self.token_back_sf_smem_region, + self.combine_format.scale_dtype, + (self.transfer_warp_count, (self.combine_format.scale_block, self.combine_sf_hidden_padded)), + stride=(self.combine_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + + @cute.jit + def fc1_ready_counter_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.fc1_ready_region) + + @cute.jit + def fc1_activation_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return cute.make_tensor( + device_workspace.ptr(self.fc1_activation_region), + cute.make_layout((self.worst_case_token_count, self.hidden_size), stride=(self.hidden_size, 1)), + ) + + @cute.jit + def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + layout = tile_atom_to_shape_SF( + (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size + ) + return cute.make_tensor(device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1])) + + @cute.jit + def fc2_done_counter_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_enabled): + return None + return device_workspace.tensor(self.fc2_done_region) + + @cute.jit + def fc2_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_push_data): + return None + return cute.make_tensor( + device_workspace.ptr(self.fc2_activation_region), + cute.make_layout( + (self.worst_case_token_count, 1, self.hidden_size), stride=(self.hidden_size, self.hidden_size, 1) + ), + ) + + @cute.jit + def fc2_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_push_sf): + return None + return cute.make_tensor( + device_workspace.ptr(self.fc2_activation_sf_region), + cute.make_layout( + ( + self.worst_case_token_count, + 1, + (self.combine_format.scale_block, self.hidden_size // self.combine_format.scale_block), + ), + stride=(self.combine_sf_hidden_padded, self.combine_sf_hidden_padded, (0, 1)), + ), + ) + + @cute.jit + def pre_reduced_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """The (tokens, topk, hidden) combine staging plane, or None under in-kernel top-k reduction.""" + if cutlass.const_expr(self.reduce_topk_in_kernel): + return None + return device_workspace.tensor(self.pre_reduced_activation_region) + + @cute.jit + def pre_reduced_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """The scale plane parallel to ``pre_reduced_activation_tensor``; only quantized wire formats carry one.""" + if cutlass.const_expr(self.reduce_topk_in_kernel or not self.combine_format.is_quantized): + return None + return device_workspace.tensor(self.pre_reduced_activation_sf_region) + + @cute.jit + def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Wait for pushed metadata, then pull activation payloads into local pools.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + global_warp_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_warp_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + + sizes = self._router.sizes_tensor(self._device_workspace) + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + + iket.range_push("token_in.wait_sizes_ready") + self._router.wait_for_sizes_ready(self._device_workspace) + iket.range_pop() + iket.range_push("token_in.stage_sizes") + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + owner_expert_begin = self._local_rank * Int32(self.experts_per_rank) + source_sizes = cute.make_tensor(sizes.iterator + owner_expert_begin, cute.make_layout((self.experts_per_rank,))) + copy_elements = 4 if self.experts_per_rank % 4 == 0 else 1 + source_size_vectors = cute.zipped_divide( + (mark_alignment(source_sizes, 16) if cutlass.const_expr(copy_elements == 4) else source_sizes), + (copy_elements,), + ) + destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) + size_vector_count = cute.size(destination_size_vectors, mode=[1]) + size_copy_atom = _copy_atom(cutlass.Int32, copy_elements * 32) + size_copy_rounds = ceil_div(size_vector_count, self.transfer_thread_count) + size_copy_registers = cute.make_rmem_tensor((copy_elements, size_copy_rounds), cutlass.Int32) + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + if vector_idx < Int32(size_vector_count): + cute.copy( + size_copy_atom, + source_size_vectors[None, vector_idx], + size_copy_registers[None, size_copy_round], + ) + iket.range_pop() + + iket.range_push("token_in.wait_metadata_ready") + self._router.wait_for_metadata_ready(self._device_workspace) + iket.range_pop() + + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + if vector_idx < Int32(size_vector_count): + cute.copy( + size_copy_atom, + size_copy_registers[None, size_copy_round], + destination_size_vectors[None, vector_idx], + ) + iket.range_push("token_in.size_barrier") + token_in_size_barrier = pipeline.NamedBarrier( + barrier_id=self.token_in_size_barrier_id, num_threads=self.transfer_thread_count + ) + token_in_size_barrier.arrive_and_wait() + iket.range_pop() + if cutlass.const_expr(self.token_back_mode == "standalone_warps"): + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive() + + iket.range_push("token_in.pull_payload") + token_in_mbarriers = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_mbarrier = token_in_mbarriers + transfer_warp_idx + warp_activation_stage = token_in_activation[transfer_warp_idx, None] + warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] + if lane_idx == Int32(0): + cute.arch.mbarrier_init(warp_mbarrier, 1) + cute.arch.sync_warp() + + fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) + fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) + fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) + activation_bytes = cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 + activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 + sf_copy_elements = 4 + source_sf_values = cute.slice_(warp_sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) + + next_dense_token = global_warp_idx + expert_valid_begin = Int32(0) + expert_sf_begin = Int32(0) + expert_ready_slot_begin = Int32(0) + pull_phase = Int32(0) + flag_tracker = make_flag_batch_tracker( + use_async=self.token_in_flag_batch == 1, + flag_address=Int64(0), + accumulated_flags=Int32(0), + phase=Int32(0), + thread_idx=lane_idx, + ) + + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + pull_count = Int32(0) + if next_dense_token < expert_valid_end: + pull_count = (expert_valid_end - next_dense_token + global_warp_count - Int32(1)) // global_warp_count + + for pull_round in cutlass.range(pull_count, unroll=1): + dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count + token_in_expert = dense_token_idx - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + sf_token_idx = expert_sf_begin + token_in_expert + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + remote_activation_address = ( + self._token_comm_args.activation.iterator.toint() + + peer_offset + + Int64(metadata.src_token) * Int64(self.bytes_per_token) + ) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, + ) + remote_sf_row = remote_sf[Int64(metadata.src_token), None] + + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) + ) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address), + warp_mbarrier, + Int32(activation_bytes), + ) + tma_load_1d( + warp_sf_stage.iterator, remote_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_activation, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + + destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] + destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) + destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) + sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) + for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): + sf_vector_idx = Int32(sf_round * 32) + lane_idx + if sf_vector_idx < Int32(sf_vector_count): + cute.copy( + sf_copy_atom, + source_sf_vectors[None, sf_vector_idx], + destination_sf_vectors[None, sf_vector_idx], + ) + + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + ready_address = (fc1_ready_counter + ready_slot_idx).toint() + flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) + cute.arch.sync_warp() + pull_phase = pull_phase ^ Int32(1) + + next_dense_token = next_dense_token + pull_count * global_warp_count + expert_valid_begin = expert_valid_end + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) + expert_ready_slot_begin = expert_ready_slot_begin + ( + (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot) + ) + local_expert = local_expert + Int32(1) + + flag_tracker.fire() + cute.arch.sync_warp() + iket.range_pop() + if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): + iket.range_push("token_in.transfer_barrier") + transfer_lifetime_barrier = pipeline.NamedBarrier( + barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count + ) + transfer_lifetime_barrier.arrive_and_wait() + iket.range_pop() + + @cute.jit + def _stateless_pace(self, reference_window: Int32, current_window: Int32) -> None: + sleep_cycles = Int32(0) + if current_window < reference_window: + sleep_cycles = reference_window - current_window + elif current_window > reference_window: + sleep_cycles = cutlass.min(current_window - reference_window, Int32(self.standalone_max_backoff_cycles)) + if sleep_cycles > Int32(0): + nanosleep(sleep_cycles) + + @cute.jit + def _adaptive_pace(self, average_window: Int32, current_window: Int32, low_window: int, high_window: int) -> Int32: + sleep_cycles = Int32(0) + if current_window > average_window: + sleep_cycles = current_window - average_window + average_window = average_window + ((current_window - average_window + Int32(3)) // Int32(4)) + if sleep_cycles > Int32(high_window): + sleep_cycles = Int32(high_window) + else: + sleep_cycles = average_window - current_window + average_window = average_window - ((average_window - current_window + Int32(3)) // Int32(4)) + # with cute.arch.elect_one(): + # cute.printf("avg_window: {}, current_window: {}, sleep_cycles: {}", average_window, current_window, sleep_cycles) + if sleep_cycles > Int32(self.adaptive_minimum_sleep_cycles): + nanosleep(sleep_cycles) + if average_window > Int32(high_window): + average_window = Int32(high_window) + if average_window < Int32(low_window): + average_window = Int32(low_window) + return average_window + + @cute.jit + def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Push completed FC2 data and scale rows to source ranks.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + if cutlass.const_expr(not self.token_back_enabled): + return + if cutlass.const_expr( + self.combine_format.is_quantized and self._token_comm_args.pre_reduced_activation_sf is None + ): + raise ValueError("Quantized token-back requires a scale destination.") + + global_worker_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_worker_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + token_back_mbarriers = smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) + worker_mbarrier = token_back_mbarriers + transfer_warp_idx + if cutlass.const_expr(self.token_back_push_data): + token_back_activation = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) + worker_activation_stage = token_back_activation[transfer_warp_idx, None] + activation_chunk_bytes = ( + cute.cosize(worker_activation_stage) * int(self.combine_format.act_dtype.width) // 8 + ) + if cutlass.const_expr(self.token_back_push_sf): + token_back_sf = smem_workspace.tensor(self.token_back_sf_smem_region, smem_base) + worker_sf_stage = token_back_sf[transfer_warp_idx, (None, None)] + sf_chunk_bytes = cute.cosize(worker_sf_stage) * int(self.combine_format.scale_dtype.width) // 8 + if lane_idx == Int32(0): + cute.arch.mbarrier_init(worker_mbarrier, 1) + cute.arch.sync_warp() + + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + if cutlass.const_expr(self.token_back_mode == "standalone_warps"): + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive_and_wait() + + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + fc2_done = self._device_workspace.ptr(self.fc2_done_region) + if cutlass.const_expr(self.token_back_push_data): + fc2_activation_pointer = self._device_workspace.ptr(self.fc2_activation_region) + output_token_bytes = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) + data_window_unit = activation_chunk_bytes * 2 + reuse_data_pacing_enabled = ( + self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles + ) + # Preserve the empirical low:initial:high ratio of 1:2.5:5. + data_average_window = Int32(data_window_unit) + data_low_window = data_window_unit * 2 // 5 + data_high_window = data_window_unit * 2 + if cutlass.const_expr(self.token_back_push_sf): + fc2_sf_pointer = self._device_workspace.ptr(self.fc2_activation_sf_region) + output_sf_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) + sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) + reuse_sf_pacing_enabled = ( + self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles + ) + sf_average_window = Int32(sf_window_unit) + sf_low_window = sf_window_unit * 2 // 5 + sf_high_window = sf_window_unit * 2 + + next_dense_token = global_worker_idx - global_worker_count + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): + next_dense_token = Int32(0) + next_dense_token = self.next_token(next_dense_token) + expert_valid_begin = Int32(0) + transfer_phase = Int32(0) + + iket.range_push("token_back.work") + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + if next_dense_token < expert_valid_end: + token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( + self.tokens_per_fc1_ready_slot + ) + completion_target = token_tile_count * Int32(self.fc2_done_signals_per_token_tile) + iket.range_push("token_back.wait_fc2") + while cute.arch.load(fc2_done + local_expert, Int32, sem="acquire", scope="gpu") < completion_target: + nanosleep(500) + iket.range_pop() + + while next_dense_token < expert_valid_end: + token_in_expert = next_dense_token - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + destination_topk = metadata.src_topk + if cutlass.const_expr(self.reduce_topk_in_kernel): + destination_topk = Int32(0) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + is_remote_token = metadata.src_rank != self._local_rank + + if cutlass.const_expr(self.token_back_push_data): + iket.range_push("token_back.push_data") + local_activation_address = fc2_activation_pointer.toint() + Int64(pool_token_idx) * Int64( + output_token_bytes + ) + remote_activation = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation.dtype, + self._token_comm_args.pre_reduced_activation.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation.layout, + ) + destination_row = remote_activation[Int64(metadata.src_token), destination_topk, None] + for chunk_idx in cutlass.range_constexpr(activation_chunk_count): + chunk_byte_offset = Int64(chunk_idx * activation_chunk_bytes) + chunk_bytes_this_round = min( + activation_chunk_bytes, output_token_bytes - chunk_idx * activation_chunk_bytes + ) + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_data_pacing_enabled = ( + self.token_back_mode != "reuse_dispatch_warps" + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) + cute.arch.sync_warp() + if cutlass.const_expr(reuse_data_pacing_enabled or stateless_data_pacing_enabled): + if is_remote_token: + round_start_clock = read_clock64() + else: + round_start_clock = round_start_clock + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) + tma_load_1d( + worker_activation_stage.iterator, + local_activation_address + chunk_byte_offset, + worker_mbarrier, + current_chunk_bytes, + ) + cute.arch.mbarrier_wait(worker_mbarrier, transfer_phase) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_row.iterator.toint() + chunk_byte_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + if cutlass.const_expr(self.reduce_topk_in_kernel): + cp_reduce_async_bulk_add_bf16_s2g( + destination_chunk, worker_activation_stage.iterator, current_chunk_bytes + ) + else: + cp_async_bulk_s2g( + destination_chunk, worker_activation_stage.iterator, current_chunk_bytes + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0) + transfer_phase = transfer_phase ^ Int32(1) + cute.arch.sync_warp() + if cutlass.const_expr(reuse_data_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + data_average_window = self._adaptive_pace( + data_average_window, current_window, data_low_window, data_high_window + ) + elif cutlass.const_expr(stateless_data_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + self._stateless_pace(Int32(current_window_unit), current_window) + iket.range_pop() + + if cutlass.const_expr(self.token_back_push_sf): + iket.range_push("token_back.push_sf") + local_sf_address = fc2_sf_pointer.toint() + Int64(pool_token_idx) * Int64(output_sf_bytes) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation_sf.dtype, + self._token_comm_args.pre_reduced_activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation_sf.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation_sf.layout, + ) + destination_sf_row = remote_sf[Int64(metadata.src_token), destination_topk, None] + for chunk_idx in cutlass.range_constexpr(sf_chunk_count): + chunk_byte_offset = Int64(chunk_idx * sf_chunk_bytes) + chunk_bytes_this_round = min(sf_chunk_bytes, output_sf_bytes - chunk_idx * sf_chunk_bytes) + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_sf_pacing_enabled = ( + self.token_back_mode != "reuse_dispatch_warps" + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) + cute.arch.sync_warp() + if cutlass.const_expr(reuse_sf_pacing_enabled or stateless_sf_pacing_enabled): + if is_remote_token: + round_start_clock = read_clock64() + else: + round_start_clock = round_start_clock + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) + tma_load_1d( + worker_sf_stage.iterator, + local_sf_address + chunk_byte_offset, + worker_mbarrier, + current_chunk_bytes, + ) + cute.arch.mbarrier_wait(worker_mbarrier, transfer_phase) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_sf_row.iterator.toint() + chunk_byte_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_chunk, worker_sf_stage.iterator, current_chunk_bytes) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0) + transfer_phase = transfer_phase ^ Int32(1) + cute.arch.sync_warp() + if cutlass.const_expr(reuse_sf_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + sf_average_window = self._adaptive_pace( + sf_average_window, current_window, sf_low_window, sf_high_window + ) + elif cutlass.const_expr(stateless_sf_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + self._stateless_pace(Int32(current_window_unit), current_window) + iket.range_pop() + + cute.arch.sync_warp() + next_dense_token = self.next_token(next_dense_token) + + expert_valid_begin = expert_valid_end + local_expert = local_expert + Int32(1) + iket.range_pop() + # with cute.arch.elect_one(): + # cute.printf(" final data_average_window: {} ", data_average_window) + + @cute.jit + def next_token(self, current_token: Int32) -> Int32: + global_worker_count = self.promised_launchable_sm_count * self.transfer_warp_count + schedule_counter = None + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): + schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): + claimed_token = Int32(0) + if self._lane_idx == Int32(0): + claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") + return Int32(cute.arch.shuffle_sync(claimed_token, Int32(0))) + return current_token + global_worker_count + + @cute.jit + def reset_tail(self) -> None: + """Reset communication state with the four token-in transfer warps.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + iket.range_push("tail.nvlink_drain") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + total_reset_threads = self.promised_launchable_sm_count * self.transfer_thread_count + global_reset_thread = self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx + iket.range_push("tail.reset_workspace") + self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) + self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) + iket.range_pop() + iket.range_push("tail.nvlink_publish") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + if cutlass.const_expr(os.environ.get("MEGA_USE_NCU", "0") == "1"): + iket.range_push("tail.ncu_finalize") + self._nvlink_barrier.finalize( + 2, + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + ) + iket.range_pop() + + +__all__ = ["TokenBackMode", "TokenBackScheduleMode", "TokenCommArgs", "TokenCommNonDeterministic"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm_deterministic.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm_deterministic.py new file mode 100644 index 000000000..3238fe64e --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm_deterministic.py @@ -0,0 +1,2258 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Metadata-push routing with a fixed token sequence and fused communication.""" + +import dataclasses +import os +from typing import Callable, ClassVar, Optional, Tuple, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils +from cutlass._mlir.dialects import llvm +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 +from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF + +from ...api import ImplDesc, KernelComponent, ProblemDesc +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.dsl_helpers import smem_exclusive_prefix +from ...helpers.flag_batch import GpuReleaseFlagBatchTracker +from ...helpers.iket_compat import iket +from ...helpers.software_sync import NvlinkBarrier +from ...helpers.ptx_helpers import ( + cp_async_bulk_s2g, + cp_reduce_async_bulk_add_bf16_s2g, + cp_reduce_async_bulk_add_u32_s2g, + nanosleep, + read_clock64, + red_add_relaxed_sys_s32, + stg_b64, + stg_f32, + tma_load_1d, +) +from ...helpers.smem_workspace import SmemWorkspace +from ...helpers.utils import ceil_div, round_up +from ...quant_def import CombineFormat, QuantKind +from ..token_protocol import TokenSrcMetadata +from .symmetric_buffer import SymmetricBufferDevice +from .token_comm import TokenBackMode, TokenBackScheduleMode, TokenCommArgs + + +_quant_spec = { + "nvfp4": (cutlass.Float4E2M1FN, cutlass.Float8E4M3FN, 16), + "mxfp4": (cutlass.Float4E2M1FN, cutlass.Float8E8M0FNU, 32), + "mxfp8_e4m3": (cutlass.Float8E4M3FN, cutlass.Float8E8M0FNU, 32), + "mxfp8_e5m2": (cutlass.Float8E5M2, cutlass.Float8E8M0FNU, 32), +} + + +@dataclasses.dataclass(frozen=True) +class _ReceiveCapacity: + raw_route_count: int + padded_route_count: int + + +def _compute_receive_capacity( + *, + world_size: int, + max_tokens_per_rank: int, + topk: int, + max_recv_size_per_rank: int, +) -> _ReceiveCapacity: + return _ReceiveCapacity( + raw_route_count=world_size * max_tokens_per_rank * topk, + # max_recv_size_per_rank is the caller-provided physical pool size. Per-expert + # alignment consumes rows inside this budget rather than extending the pool. + padded_route_count=max_recv_size_per_rank, + ) + + +@dataclasses.dataclass(frozen=True) +class _SortedElement: + flat_topk_index: Int32 + topk_score: Optional[cutlass.Float32] + + def pack(self) -> Union[Int64, Int32]: + if cutlass.const_expr(self.topk_score is None): + return self.flat_topk_index + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + scratch[0] = self.flat_topk_index + cute.recast_tensor(scratch, cutlass.Float32)[1] = self.topk_score + return cute.recast_tensor(scratch, cutlass.Int64)[0] + + @classmethod + def from_packed(cls, packed: Union[Int64, Int32]) -> "_SortedElement": + if cutlass.const_expr(type(packed).width == 32): + return cls(flat_topk_index=Int32(packed), topk_score=None) + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + cute.recast_tensor(scratch, cutlass.Int64)[0] = packed + return cls(flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1]) + + +@cute.jit +def _copy_atom(dtype, num_bits_per_copy: int): + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy) + + +@cute.jit +def _mark_alignment(tensor: cute.Tensor, byte_alignment: int) -> cute.Tensor: + pointer = tensor.iterator + return cute.make_tensor( + cute.make_ptr(pointer.dtype, pointer.toint(), pointer.memspace, assumed_align=byte_alignment), tensor.layout + ) + + +@cute.jit +def _device_trap() -> None: + """Terminate the current kernel after its failure state is published.""" + llvm.inline_asm( + res=None, + operands_=[], + asm_string="trap;", + constraints="", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +class _MetadataPushRouter(KernelComponent): + """Sort and push routing metadata into each destination rank's final pool.""" + + router_smem_limit_bytes: ClassVar[int] = 227 * 1024 + router_warps_per_cta: ClassVar[int] = 16 + + sizes_by_rank_region = "nvlink.token_comm.sizes_by_rank" + sizes_region = "nvlink.token_comm.sizes" + sizes_ready_region = "nvlink.token_comm.sizes_ready" + metadata_ready_region = "nvlink.token_comm.metadata_ready" + sorted_metadata_region = "nvlink.token_comm.sorted_metadata" + sorted_scores_region = "nvlink.token_comm.sorted_scores" + pool_expert_base_region = "nvlink.token_comm.pool_expert_base" + token_src_metadata_region = "nvlink.token_comm.token_src_metadata" + fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" + source_expert_base_region = "nvlink.token_comm.source_expert_base" + push_destination_base_region = "nvlink.token_comm.push_destination_base" + sorted_metadata_ready_region = "nvlink.token_comm.sorted_metadata_ready" + push_table_ready_region = "nvlink.token_comm.push_table_ready" + router_histogram_done_region = "nvlink.token_comm.router_histogram_done" + router_cta_histograms_region = "nvlink.token_comm.router_cta_histograms" + source_base_ready_region = "nvlink.token_comm.source_base_ready" + + router_data_histogram_region = "nvlink.token_comm.router_smem.data_histogram" + router_data_totals_region = "nvlink.token_comm.router_smem.data_totals" + router_data_prefix_region = "nvlink.token_comm.router_smem.data_prefix" + router_data_warp_totals_region = "nvlink.token_comm.router_smem.data_warp_totals" + router_data_sorted_region = "nvlink.token_comm.router_smem.data_sorted" + router_data_base_region = "nvlink.token_comm.router_smem.data_base" + router_helper_size_matrix_region = "nvlink.token_comm.router_smem.helper_size_matrix" + router_helper_totals_region = "nvlink.token_comm.router_smem.helper_totals" + router_helper_prefix_region = "nvlink.token_comm.router_smem.helper_prefix" + router_helper_warp_totals_region = "nvlink.token_comm.router_smem.helper_warp_totals" + router_helper_load_mbarrier_region = "nvlink.token_comm.router_smem.helper_load_mbarrier" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "max_recv_size_per_rank": int, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return {"token_padding_block": int, "promised_launchable_sm_count": int, "drop_on_overflow": bool} + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.max_recv_size_per_rank = problem_desc["max_recv_size_per_rank"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + + self.token_padding_block = impl_desc["token_padding_block"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.drop_on_overflow = impl_desc["drop_on_overflow"] + + self._validate_router_configuration() + token_capacity = self.receive_capacity() + self.raw_route_count = token_capacity.raw_route_count + self.worst_case_token_count = token_capacity.padded_route_count + self.expert_count_padded = round_up(self.expert_count, 4) + self.expert_count_with_trash = self.expert_count_padded + 1 + self.router_elements_per_lane, self.router_data_cta_count = self._router_launch_configuration() + self.router_tokens_per_cta = self.router_elements_per_lane * self.router_warps_per_cta * 32 + self.router_push_cta_count = ceil_div(self.expert_count, self.router_warps_per_cta) + self.router_grid_cta_count = max(self.router_data_cta_count + 1, self.router_push_cta_count) + if self.router_grid_cta_count > self.promised_launchable_sm_count: + raise ValueError( + "Router grid exceeds promised_launchable_sm_count; all metadata-push CTAs must be concurrently resident." + ) + self._router_smem_workspace = self._build_router_smem_workspace() + + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_grid_thread_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + self._overflow_flag = None + + def _validate_router_configuration(self) -> None: + if type(self.max_recv_size_per_rank) is not int: + raise TypeError(f"max_recv_size_per_rank must be an int, got {type(self.max_recv_size_per_rank).__name__}.") + if type(self.drop_on_overflow) is not bool: + raise TypeError(f"drop_on_overflow must be a bool, got {type(self.drop_on_overflow).__name__}.") + positive_fields = ( + "world_size", + "expert_count", + "topk", + "max_tokens_per_rank", + "max_recv_size_per_rank", + "token_padding_block", + "promised_launchable_sm_count", + ) + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.expert_count % self.world_size != 0: + raise ValueError( + f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." + ) + if self.expert_count > 16384: + raise NotImplementedError("TokenComm supports at most 16384 global experts.") + if self.topk > self.expert_count: + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + def receive_capacity(self) -> _ReceiveCapacity: + return _compute_receive_capacity( + world_size=self.world_size, + max_tokens_per_rank=self.max_tokens_per_rank, + topk=self.topk, + max_recv_size_per_rank=self.max_recv_size_per_rank, + ) + + def _router_launch_configuration(self) -> Tuple[int, int]: + def next_power_of_two(value: int) -> int: + return 1 << (max(value, 1) - 1).bit_length() + + routed_token_capacity = self.max_tokens_per_rank * self.topk + minimum_cta_capacity = 2048 + maximum_cta_capacity = 16384 + maximum_data_cta_count = 128 + maximum_supported_tokens = maximum_cta_capacity * maximum_data_cta_count + if routed_token_capacity > maximum_supported_tokens: + raise NotImplementedError(f"The router supports at most {maximum_supported_tokens} routed tokens per rank.") + cta_capacity = min(maximum_cta_capacity, next_power_of_two(max(routed_token_capacity, minimum_cta_capacity))) + elements_per_lane = cta_capacity // (self.router_warps_per_cta * 32) + data_cta_count = ceil_div(routed_token_capacity, cta_capacity) + return elements_per_lane, data_cta_count + + def _build_router_smem_workspace(self) -> SmemWorkspace: + workspace = SmemWorkspace() + workspace.register_mbarrier(self.router_helper_load_mbarrier_region, 1) + overlay = workspace.create_overlay("nvlink.token_comm.router_smem.role") + data_lifetime = overlay.add_lifetime("data_cta") + data_lifetime.register_tensor( + self.router_data_histogram_region, + cutlass.Int32, + (self.expert_count_with_trash, self.router_warps_per_cta + 1), + stride=(self.router_warps_per_cta + 1, 1), + ) + data_lifetime.register_tensor( + self.router_data_totals_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 + ) + data_lifetime.register_tensor( + self.router_data_prefix_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 + ) + data_lifetime.register_tensor(self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,)) + data_lifetime.register_tensor( + self.router_data_sorted_region, + (cutlass.Int64 if self.apply_topk_at_fc1 else cutlass.Int32), + (self.router_tokens_per_cta,), + byte_alignment=16, + ) + if self.router_data_cta_count > 1: + data_lifetime.register_tensor( + self.router_data_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + + helper_lifetime = overlay.add_lifetime("helper_cta") + helper_lifetime.register_tensor( + self.router_helper_size_matrix_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + stride=(self.expert_count_padded, 1), + byte_alignment=16, + ) + helper_lifetime.register_tensor( + self.router_helper_totals_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_prefix_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 + ) + workspace.finalize(max_bytes=self.router_smem_limit_bytes) + return workspace + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router_smem_workspace + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + """Register router-private state and Router-to-Main outputs.""" + self._register_router_workspace(workspace) + + def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: + maximum_routed_tokens = self.max_tokens_per_rank * self.topk + workspace.register( + self.sizes_by_rank_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + buffer_space="shared", + stride=(self.expert_count_padded, 1), + ) + workspace.register( + self.sizes_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="shared", reset="tail_reset" + ) + workspace.register(self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.metadata_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.sorted_metadata_region, cutlass.Int64, (maximum_routed_tokens,), buffer_space="local") + if self.apply_topk_at_fc1: + workspace.register( + self.sorted_scores_region, cutlass.Float32, (maximum_routed_tokens,), buffer_space="local" + ) + workspace.register( + self.token_src_metadata_region, + cutlass.Int64, + (self.worst_case_token_count,), + buffer_space="shared", + byte_alignment=16, + ) + if self.apply_topk_at_fc1: + workspace.register( + self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" + ) + workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local") + workspace.register( + self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.push_destination_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.sorted_metadata_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register(self.push_table_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") + if self.router_data_cta_count > 1: + workspace.register( + self.router_cta_histograms_region, + cutlass.Int32, + (self.router_data_cta_count, self.expert_count_padded), + buffer_space="local", + stride=(self.expert_count_padded, 1), + ) + workspace.register( + self.router_histogram_done_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register( + self.source_base_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "_MetadataPushRouter": + if values: + raise ValueError("_MetadataPushRouter carries no MLIR values.") + return self + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + overflow_flag: cute.Tensor, + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + """Launch counting-sort DATA, size-exchange HELPER, and metadata PUSH roles.""" + if cutlass.const_expr(self.apply_topk_at_fc1 and topk_scores is None): + raise ValueError("apply_topk_at_fc1 requires router topk_scores.") + if cutlass.const_expr(overflow_flag.iterator.dtype is not cutlass.Int32): + raise TypeError("overflow_flag must use Int32 elements.") + if cutlass.const_expr(cute.size(overflow_flag) != 1): + raise ValueError("overflow_flag must contain exactly one element.") + peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() + self._router_kernel( + topk_indices, + topk_scores, + overflow_flag, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper, + device_workspace, + ).launch( + grid=[self.router_grid_cta_count, 1, 1], + block=[self.router_warps_per_cta * 32, 1, 1], + min_blocks_per_mp=1, + stream=stream, + ) + + @cute.kernel + def _router_kernel( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + overflow_flag: cute.Tensor, + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper: SymmetricBufferDevice, + device_workspace: DeviceWorkspace, + ) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + linear_cta_idx, _, _ = cute.arch.block_idx() + cute.arch.griddepcontrol_launch_dependents() + block_thread_count = self.router_warps_per_cta * 32 + grid_thread_idx = thread_idx + linear_cta_idx * block_thread_count + warp_idx = cute.arch.make_warp_uniform(thread_idx // Int32(32)) + lane_idx = thread_idx % Int32(32) + + storage_type = self._router_smem_workspace.storage_class() + smem_allocator = cutlass.utils.SmemAllocator() + storage = smem_allocator.allocate(storage_type) + smem_base = storage.buffer.data_ptr() + + device_workspace.assign_device_members(local_workspace, shared_workspace) + self._device_workspace = device_workspace + self._peer_rank_ptr_mapper = peer_rank_ptr_mapper + self._router_local_rank = local_rank + self._router_thread_idx = thread_idx + self._router_linear_cta_idx = linear_cta_idx + self._router_grid_thread_idx = grid_thread_idx + self._router_warp_idx = warp_idx + self._router_lane_idx = lane_idx + self._overflow_flag = overflow_flag + + if cutlass.const_expr(self.router_data_cta_count == 1): + self._router_single_cta(topk_indices, topk_scores, smem_base) + else: + self._router_multiple_ctas(topk_indices, topk_scores, smem_base) + if linear_cta_idx < Int32(self.router_push_cta_count): + self._router_push_metadata() + + device_workspace.remove_device_members() + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_grid_thread_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + self._overflow_flag = None + + @cute.jit + def _router_single_cta( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + totals = self._router_smem_workspace.tensor(self.router_data_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + + histogram_element_count = self.expert_count_with_trash * (self.router_warps_per_cta + 1) + histogram_flat = cute.make_tensor(histogram.iterator, cute.make_layout((histogram_element_count,))) + zero_round_count = ceil_div(histogram_element_count, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + slot = Int32(zero_round * block_thread_count) + self._router_thread_idx + if slot < Int32(histogram_element_count): + histogram_flat[slot] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + match_masks = self._build_histogram(expert_registers, histogram) + self._prefix_warp_histogram(histogram, totals) + iket.range_pop() + + iket.range_push("router.prefix_and_publish") + publish_sizes = self._broadcast_sizes_to_peers( + cute.make_tensor(totals.iterator, cute.make_layout((self.expert_count_padded,))) + ) + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(totals.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + publish_sizes() + iket.range_pop() + + iket.range_push("router.sort") + self._sort_router_elements( + expert_registers, match_masks, score_registers, sorted_elements, prefix, histogram, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.write_out") + self._dump_contiguous_router_output(sorted_elements, total_valid_routes) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_single_cta(smem_base) + + @cute.jit + def _router_multiple_ctas( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + totals = self._router_smem_workspace.tensor(self.router_data_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + dump_base = self._router_smem_workspace.tensor(self.router_data_base_region, smem_base) + + histogram_element_count = self.expert_count_with_trash * (self.router_warps_per_cta + 1) + histogram_flat = cute.make_tensor(histogram.iterator, cute.make_layout((histogram_element_count,))) + zero_round_count = ceil_div(histogram_element_count, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + slot = Int32(zero_round * block_thread_count) + self._router_thread_idx + if slot < Int32(histogram_element_count): + histogram_flat[slot] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + match_masks = self._build_histogram(expert_registers, histogram) + self._prefix_warp_histogram(histogram, totals) + iket.range_pop() + + iket.range_push("router.reserve_and_prefix") + cta_histograms = self._device_workspace.tensor(self.router_cta_histograms_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + cta_histograms[self._router_linear_cta_idx, expert] = totals[expert] + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.router_histogram_done_region), Int32(1), sem="release", scope="gpu" + ) + + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(totals.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.sort") + self._sort_router_elements( + expert_registers, match_masks, score_registers, sorted_elements, prefix, histogram, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.wait_source_base") + source_base_ready = self._device_workspace.ptr(self.source_base_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.write_out") + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + cta_histograms = self._device_workspace.tensor(self.router_cta_histograms_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + dump_base[expert] = source_expert_base[expert] + cta_histograms[self._router_linear_cta_idx, expert] + cute.arch.sync_threads() + self._dump_router_output_by_expert(totals, prefix, dump_base, sorted_elements) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_multiple_ctas(smem_base) + + @cute.jit + def _router_helper_single_cta(self, smem_base: cute.Pointer) -> None: + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + cta_histograms = self._device_workspace.tensor(self.router_cta_histograms_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + + histogram_done = self._device_workspace.ptr(self.router_histogram_done_region) + iket.range_push("router.wait_histogram") + if self._router_thread_idx == Int32(0): + while cute.arch.load(histogram_done, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.broadcast_sizes") + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + running = Int32(0) + for cta_idx in cutlass.range_constexpr(self.router_data_cta_count): + count = cta_histograms[cta_idx, expert] + cta_histograms[cta_idx, expert] = running + running = running + count + totals[expert] = running + cute.arch.sync_threads() + + publish_sizes = self._broadcast_sizes_to_peers(totals) + smem_exclusive_prefix( + totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.source_base_ready_region), Int32(1), sem="release", scope="gpu" + ) + publish_sizes() + iket.range_pop() + + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_push_metadata(self) -> None: + block_thread_count = self.router_warps_per_cta * 32 + sorted_metadata_ready = self._device_workspace.ptr(self.sorted_metadata_ready_region) + push_table_ready = self._device_workspace.ptr(self.push_table_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(sorted_metadata_ready, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + while cute.arch.load(push_table_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + + global_expert = self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx + if global_expert < Int32(self.expert_count): + sizes_by_rank = self._device_workspace.tensor(self.sizes_by_rank_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + route_count = sizes_by_rank[self._router_local_rank, global_expert] + source_begin = source_expert_base[global_expert] + destination_begin = push_destination_base[global_expert] + destination_rank = global_expert // Int32(self.experts_per_rank) + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) + destination_metadata_address = ( + self._device_workspace.ptr(self.token_src_metadata_region).toint() + peer_offset + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + source_scores = self._device_workspace.ptr(self.sorted_scores_region) + destination_scores_address = ( + self._device_workspace.ptr(self.fc1_topk_scores_region).toint() + peer_offset + ) + route_round_count = (route_count + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + if route < route_count: + source_position = source_begin + route + destination_position = destination_begin + route + physical_store = Int32(destination_position < Int32(self.worst_case_token_count)) + metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) + stg_b64( + destination_metadata_address + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), + metadata, + physical_store, + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + score = cute.arch.load(source_scores + source_position, cutlass.Float32) + stg_f32( + destination_scores_address + Int64(destination_position) * Int64(4), score, physical_store + ) + + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + # Keep notifier threads behind the leader's system fence. + cute.arch.sync_threads() + metadata_ready_address = self._device_workspace.ptr(self.metadata_ready_region).toint() + rank_round_count = ceil_div(self.world_size, block_thread_count) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = Int32(rank_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(metadata_ready_address, destination_rank, Int64(0)), Int32(1) + ) + + @cute.jit + def _broadcast_sizes_to_peers(self, smem_expert_counts: cute.Tensor) -> Callable[[], None]: + block_thread_count = self.router_warps_per_cta * 32 + row_bytes = Int32(self.expert_count_padded * 4) + matrix_address = self._device_workspace.ptr(self.sizes_by_rank_region).toint() + total_address = self._device_workspace.ptr(self.sizes_region).toint() + rank_round_count = ceil_div(self.world_size, self.router_warps_per_cta) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = self._router_warp_idx + Int32(rank_round * self.router_warps_per_cta) + if destination_rank < Int32(self.world_size): + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + destination_row_address = ( + matrix_address + + peer_offset + + Int64(Int32(self._router_local_rank) * Int32(self.expert_count_padded)) * Int64(4) + ) + destination_row = cute.make_ptr( + cutlass.Int32, destination_row_address, AddressSpace.gmem, assumed_align=16 + ) + destination_total = cute.make_ptr( + cutlass.Int32, total_address + peer_offset, AddressSpace.gmem, assumed_align=16 + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_row, smem_expert_counts.iterator, row_bytes) + cp_reduce_async_bulk_add_u32_s2g(destination_total, smem_expert_counts.iterator, row_bytes) + cute.arch.cp_async_bulk_commit_group() + + def finalize() -> None: + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + cute.arch.sync_threads() + ready_address = self._device_workspace.ptr(self.sizes_ready_region).toint() + ready_round_count = ceil_div(self.world_size, block_thread_count) + for ready_round in cutlass.range_constexpr(ready_round_count): + destination_rank = Int32(ready_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(ready_address, destination_rank, Int64(0)), Int32(1) + ) + + return finalize + + @cute.jit + def _compute_push_tables(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + owner_expert_begin = Int32(self._router_local_rank) * Int32(self.experts_per_rank) + matrix_bytes = self.world_size * self.expert_count_padded * 4 + + size_matrix = self._router_smem_workspace.tensor(self.router_helper_size_matrix_region, smem_base) + padded_totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + load_mbarrier = self._router_smem_workspace.ptr(self.router_helper_load_mbarrier_region, smem_base) + sizes = self._device_workspace.tensor(self.sizes_region) + + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_init(load_mbarrier, 1) + + sizes_ready = self._device_workspace.ptr(self.sizes_ready_region) + iket.range_push("router.wait_sizes_ready") + if self._router_thread_idx == Int32(0): + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.raw_sizes_ready_target): + nanosleep(150) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.load_sizes_and_prefix") + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(load_mbarrier, Int32(matrix_bytes)) + tma_load_1d( + size_matrix.iterator, + self._device_workspace.ptr(self.sizes_by_rank_region), + load_mbarrier, + Int32(matrix_bytes), + ) + + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + expert_size = sizes[expert] + padded_totals[expert] = ( + (expert_size + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) + ) * Int32(self.token_padding_block) + cute.arch.sync_threads() + smem_exclusive_prefix( + padded_totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + pool_expert_base = self._device_workspace.tensor(self.pool_expert_base_region) + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + pool_expert_base[local_expert] = prefix[owner_expert_begin + local_expert] - prefix[owner_expert_begin] + + cute.arch.mbarrier_wait(load_mbarrier, 0) + iket.range_pop() + + iket.range_push("router.build_push_destinations") + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + global_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if global_expert < Int32(self.expert_count): + destination_rank = global_expert // Int32(self.experts_per_rank) + destination_expert_begin = destination_rank * Int32(self.experts_per_rank) + destination_pool_base = prefix[global_expert] - prefix[destination_expert_begin] + local_ring_position = ( + Int32(self._router_local_rank) - destination_rank + Int32(self.world_size) + ) % Int32(self.world_size) + source_ring_offset = Int32(0) + for ring_position in cutlass.range_constexpr(self.world_size): + source_rank = (destination_rank + Int32(ring_position)) % Int32(self.world_size) + if Int32(ring_position) < local_ring_position: + source_ring_offset = source_ring_offset + size_matrix[source_rank, global_expert] + push_destination_base[global_expert] = destination_pool_base + source_ring_offset + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.apply_receive_limit") + if self._router_thread_idx == Int32(0): + warp_totals[0] = Int32(0) + cute.arch.sync_threads() + thread_local_total = Int32(0) + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + thread_local_total = ( + thread_local_total + padded_totals[owner_expert_begin + local_expert] + ) + if thread_local_total > Int32(0): + cute.arch.atomic_add(warp_totals.iterator, thread_local_total, scope="cta") + cute.arch.sync_threads() + padded_local_total = warp_totals[0] + did_overflow = Int32(padded_local_total > Int32(self.max_recv_size_per_rank)) + + if cutlass.const_expr(self.drop_on_overflow): + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + # Keep every raw row that still lands inside the statically allocated padded pool. + raw_size = sizes[owner_expert_begin + local_expert] + retained_size = Int32(0) + remaining_capacity = Int32(self.worst_case_token_count) - pool_expert_base[local_expert] + if remaining_capacity > Int32(0): + retained_size = cutlass.min(raw_size, remaining_capacity) + if did_overflow: + sizes[owner_expert_begin + local_expert] = retained_size + + if self._router_thread_idx == Int32(0): + cute.arch.store(self._overflow_flag.iterator, did_overflow, sem="relaxed", scope="sys") + cute.arch.sync_threads() + + if cutlass.const_expr(self.drop_on_overflow): + self._publish_push_table_and_sizes_ready(sizes_ready) + else: + if did_overflow: + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + _device_trap() + else: + self._publish_push_table_and_sizes_ready(sizes_ready) + iket.range_pop() + + @cute.jit + def _publish_push_table_and_sizes_ready(self, sizes_ready: cute.Pointer) -> None: + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_gpu() + cute.arch.atomic_add( + self._device_workspace.ptr(self.push_table_ready_region), Int32(1), sem="relaxed", scope="gpu" + ) + cute.arch.atomic_add(sizes_ready, Int32(1), sem="relaxed", scope="gpu") + + @cute.jit + def _load_router_inputs( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor] + ) -> Tuple[cute.Tensor, Optional[cute.Tensor]]: + elements_per_vector = 128 // topk_indices.dtype.width + grid_thread_count = self.router_data_cta_count * self.router_warps_per_cta * 32 + tile_span = elements_per_vector * grid_thread_count + maximum_elements = self.max_tokens_per_rank * self.topk + actual_token_count = Int32(self.max_tokens_per_rank) + actual_elements = Int32(maximum_elements) + load_round_count = ceil_div(maximum_elements, tile_span) + elements_per_thread = load_round_count * elements_per_vector + + topk_flat = cute.make_tensor(topk_indices.iterator, cute.make_layout((maximum_elements,))) + topk_vectors = cute.logical_divide(cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None)) + load_atom = _copy_atom(topk_indices.dtype, 128) + expert_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Int32) + if cutlass.const_expr(topk_indices.dtype.width == 64): + raw_indices = cute.make_rmem_tensor((elements_per_thread,), topk_indices.dtype) + raw_vectors = cute.zipped_divide(raw_indices, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + load_atom, + _mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), + raw_vectors[None, load_round], + ) + else: + expert_vectors = cute.zipped_divide(expert_registers, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + load_atom, + _mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), + expert_vectors[None, load_round], + ) + + score_registers = None + if cutlass.const_expr(self.apply_topk_at_fc1): + score_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Float32) + scores_flat = cute.make_tensor(topk_scores.iterator, cute.make_layout((maximum_elements,))) + score_vectors = cute.logical_divide( + cute.zipped_divide(scores_flat, (tile_span,)), (elements_per_vector, None) + ) + score_atom = _copy_atom(cutlass.Float32, elements_per_vector * 32) + score_register_vectors = cute.zipped_divide(score_registers, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + score_atom, + _mark_alignment(score_vectors[(None, self._router_grid_thread_idx), load_round], 16), + score_register_vectors[None, load_round], + ) + + expert_registers_u32 = cute.recast_tensor(expert_registers, cutlass.Uint32) + if cutlass.const_expr(topk_indices.dtype.width == 64): + raw_indices_i32 = cute.recast_tensor(raw_indices, cutlass.Int32) + for register_idx in cutlass.range_constexpr(elements_per_thread): + if cutlass.const_expr(topk_indices.dtype.width == 64): + expert_registers[register_idx] = raw_indices_i32[2 * register_idx] + token_idx, _ = self._router_value_coordinate(register_idx, topk_indices.dtype) + is_invalid = (expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count)) | ( + token_idx >= actual_token_count + ) + if is_invalid: + expert_registers[register_idx] = Int32(self.expert_count_padded) + return expert_registers, score_registers + + @cute.jit + def _build_histogram(self, expert_registers: cute.Tensor, histogram: cute.Tensor) -> cute.Tensor: + register_count = cute.size(expert_registers) + match_masks = cute.make_rmem_tensor((register_count,), cutlass.Int32) + lane_mask_less_than = (Int32(1) << self._router_lane_idx) - Int32(1) + for register_idx in cutlass.range_constexpr(register_count): + expert = expert_registers[register_idx] + match_mask = Int32(cute.arch.match_sync(0xFFFFFFFF, expert, kind="any")) + match_masks[register_idx] = match_mask + histogram_slot = expert * Int32(self.router_warps_per_cta + 1) + self._router_warp_idx + rank_in_group = Int32(cute.arch.popc(match_mask & lane_mask_less_than)) + if rank_in_group == Int32(0): + cute.arch.atomic_add( + histogram.iterator + histogram_slot, Int32(cute.arch.popc(match_mask)), sem="relaxed", scope="cta" + ) + cute.arch.sync_threads() + return match_masks + + @cute.jit + def _prefix_warp_histogram(self, histogram: cute.Tensor, totals: cute.Tensor) -> None: + """Turn per-warp counts into fixed warp bases and expert totals.""" + block_thread_count = self.router_warps_per_cta * 32 + expert_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_with_trash): + running = Int32(0) + for warp in cutlass.range_constexpr(self.router_warps_per_cta): + count = histogram[expert, warp] + histogram[expert, warp] = running + running = running + count + totals[expert] = running + cute.arch.sync_threads() + + @cute.jit + def _sort_router_elements( + self, + expert_registers: cute.Tensor, + match_masks: cute.Tensor, + score_registers: Optional[cute.Tensor], + sorted_elements: cute.Tensor, + expert_run_starts: cute.Tensor, + warp_histogram: cute.Tensor, + topk_index_type: type, + ) -> None: + """Stable CTA-local counting-sort scatter. + + The cursor for an expert is private to one warp. ``match.any`` gives + every equal-expert lane a fixed rank, and only that group's first lane + advances the cursor. Therefore the output order is a pure function of + warp id, register round, and lane id rather than atomic arrival order. + """ + register_count = cute.size(expert_registers) + lane_mask_less_than = (Int32(1) << self._router_lane_idx) - Int32(1) + for register_idx in cutlass.range_constexpr(register_count): + expert = expert_registers[register_idx] + match_mask = match_masks[register_idx] + group_size = Int32(cute.arch.popc(match_mask)) + rank_in_group = Int32(cute.arch.popc(match_mask & lane_mask_less_than)) + warp_base = warp_histogram[expert, self._router_warp_idx] + token_idx, topk_slot = self._router_value_coordinate(register_idx, topk_index_type) + flat_topk_index = token_idx * Int32(self.topk) + topk_slot + destination = expert_run_starts[expert] + warp_base + rank_in_group + + # Every lane in an equal-expert group reads the old cursor before + # its leader advances it for the next register round. + cute.arch.sync_warp() + if rank_in_group == Int32(0): + warp_histogram[expert, self._router_warp_idx] = warp_base + group_size + if cutlass.const_expr(self.apply_topk_at_fc1): + sorted_elements[destination] = _SortedElement(flat_topk_index, score_registers[register_idx]).pack() + else: + sorted_elements[destination] = _SortedElement(flat_topk_index, None).pack() + cute.arch.sync_warp() + + @cute.jit + def _router_value_coordinate(self, register_idx: int, topk_index_type: type) -> Tuple[Int32, Int32]: + elements_per_vector = 128 // topk_index_type.width + tile_span = elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 + flat_index = Int32( + register_idx // elements_per_vector * tile_span + register_idx % elements_per_vector + ) + self._router_grid_thread_idx * Int32(elements_per_vector) + return (flat_index // Int32(self.topk), flat_index % Int32(self.topk)) + + @cute.jit + def _dump_contiguous_router_output(self, sorted_elements: cute.Tensor, total_valid_routes: Int32) -> None: + block_thread_count = self.router_warps_per_cta * 32 + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32(block_thread_count) + for dump_round in cutlass.range(dump_round_count, unroll=4): + position = Int32(dump_round * block_thread_count) + self._router_thread_idx + predicate = Int32(position < total_valid_routes) + element = _SortedElement.from_packed(sorted_elements[position]) + metadata = TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=(element.flat_topk_index // Int32(self.topk)), + src_topk=(element.flat_topk_index % Int32(self.topk)), + ) + stg_b64(metadata_address + Int64(position) * Int64(TokenSrcMetadata.nbytes), metadata.pack(), predicate) + if cutlass.const_expr(self.apply_topk_at_fc1): + stg_f32(score_address + Int64(position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def _dump_router_output_by_expert( + self, + expert_totals: cute.Tensor, + expert_run_starts: cute.Tensor, + expert_dump_bases: cute.Tensor, + sorted_elements: cute.Tensor, + ) -> None: + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + expert_round_count = ceil_div(self.expert_count_padded, self.router_warps_per_cta) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = self._router_warp_idx + Int32(expert_round * self.router_warps_per_cta) + if expert < Int32(self.expert_count_padded): + run_begin = expert_run_starts[expert] + run_length = expert_totals[expert] + dump_begin = expert_dump_bases[expert] + route_round_count = (run_length + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + predicate = Int32(route < run_length) + element = _SortedElement.from_packed(sorted_elements[predicate * (run_begin + route)]) + output_position = dump_begin + route + metadata = TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=(element.flat_topk_index // Int32(self.topk)), + src_topk=(element.flat_topk_index % Int32(self.topk)), + ) + stg_b64( + metadata_address + Int64(output_position) * Int64(TokenSrcMetadata.nbytes), + metadata.pack(), + predicate, + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + stg_f32(score_address + Int64(output_position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + """Return this rank's contiguous expert-size view.""" + sizes = device_workspace.tensor(self.sizes_region) + expert_begin = local_rank * Int32(self.experts_per_rank) + return cute.make_tensor(sizes.iterator + expert_begin, cute.make_layout((self.experts_per_rank,))) + + @property + def metadata_ready_target(self) -> int: + return self.router_push_cta_count * self.world_size + + @property + def raw_sizes_ready_target(self) -> int: + return self.world_size + + @property + def published_sizes_ready_target(self) -> int: + return self.world_size + 1 + + @cute.jit + def sizes_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.sizes_region) + + @cute.jit + def pool_expert_base_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.pool_expert_base_region) + + @cute.jit + def token_src_metadata_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.token_src_metadata_region) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + sizes_ready = device_workspace.ptr(self.sizes_ready_region) + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="gpu") != Int32( + self.published_sizes_ready_target + ): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def wait_for_metadata_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + metadata_ready = device_workspace.ptr(self.metadata_ready_region) + while cute.arch.load(metadata_ready, Int32, sem="acquire", scope="sys") != Int32( + self.metadata_ready_target + ): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.token_src_metadata_region) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.apply_topk_at_fc1): + return None + return device_workspace.tensor(self.fc1_topk_scores_region) + + +class TokenCommDeterministic(KernelComponent): + """Fused communication with a fixed logical token-position sequence.""" + + transfer_warp_count: ClassVar[int] = 4 + transfer_thread_count: ClassVar[int] = transfer_warp_count * 32 + standalone_chunk_bytes: ClassVar[int] = 2048 + minimum_pacing_window_cycles: ClassVar[int] = 512 + standalone_max_backoff_cycles: ClassVar[int] = 500 + adaptive_minimum_sleep_cycles: ClassVar[int] = 50 + transfer_lifetime_barrier_id: ClassVar[int] = 9 + grid_sync_barrier_id: ClassVar[int] = 10 + standalone_size_barrier_id: ClassVar[int] = 11 + token_in_size_barrier_id: ClassVar[int] = 12 + + fc1_ready_region = "nvlink.token_comm.fc1_ready" + fc1_activation_region = "nvlink.token_comm.fc1_activation" + fc1_activation_sf_region = "nvlink.token_comm.fc1_activation_sf" + fc2_done_region = "nvlink.token_comm.fc2_done" + fc2_activation_region = "nvlink.token_comm.fc2_activation" + fc2_activation_sf_region = "nvlink.token_comm.fc2_activation_sf" + pre_reduced_activation_region = "nvlink.token_comm.pre_reduced_activation" + pre_reduced_activation_sf_region = "nvlink.token_comm.pre_reduced_activation_sf" + token_back_schedule_region = "nvlink.token_comm.token_back_schedule" + + token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" + token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" + expert_sizes_smem_region = "nvlink.token_comm.main_smem.expert_sizes" + token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" + token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" + token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" + token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "max_recv_size_per_rank": int, + "hidden_size": int, + "quant_kind": str, + "combine_format": CombineFormat, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "token_padding_block": int, + "sf_padding_block": int, + "tokens_per_fc1_ready_slot": int, + "fc2_done_signals_per_token_tile": int, + "promised_launchable_sm_count": int, + "token_in_flag_batch": int, + "token_back_mode": str, + "token_back_schedule_mode": str, + "reduce_topk_in_kernel": bool, + "drop_on_overflow": bool, + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.max_recv_size_per_rank = problem_desc["max_recv_size_per_rank"] + self.hidden_size = problem_desc["hidden_size"] + self.quant_kind = problem_desc["quant_kind"] + self.combine_format = problem_desc["combine_format"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.tokens_per_fc1_ready_slot = impl_desc["tokens_per_fc1_ready_slot"] + self.fc2_done_signals_per_token_tile = impl_desc["fc2_done_signals_per_token_tile"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.token_in_flag_batch = impl_desc["token_in_flag_batch"] + self.token_back_mode: TokenBackMode = impl_desc["token_back_mode"] + self.token_back_schedule_mode: TokenBackScheduleMode = impl_desc["token_back_schedule_mode"] + self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] + self.drop_on_overflow = impl_desc["drop_on_overflow"] + + self._validate_configuration() + self._router = _MetadataPushRouter(problem_desc, impl_desc) + self._nvlink_barrier = NvlinkBarrier(world_size=self.world_size, barrier_id=self.grid_sync_barrier_id) + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + + def _validate_configuration(self) -> None: + positive_fields = ("hidden_size", "token_padding_block", "sf_padding_block", "tokens_per_fc1_ready_slot") + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.quant_kind not in _quant_spec: + raise ValueError(f"Unsupported quant_kind {self.quant_kind!r}.") + if self.token_back_mode not in ("epi_warps", "standalone_warps", "reuse_dispatch_warps"): + raise ValueError(f"Unsupported token_back_mode {self.token_back_mode!r}.") + if self.token_back_schedule_mode not in ("static", "atomic_counter"): + raise ValueError( + f"token_back_schedule_mode must be static or atomic_counter, got {self.token_back_schedule_mode!r}." + ) + if not 1 <= self.token_in_flag_batch <= 32: + raise ValueError(f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}.") + if self.tokens_per_fc1_ready_slot % self.token_padding_block != 0: + raise ValueError("tokens_per_fc1_ready_slot must be divisible by token_padding_block.") + _pad_lo = min(self.sf_padding_block, self.token_padding_block) + _pad_hi = max(self.sf_padding_block, self.token_padding_block) + if _pad_hi % _pad_lo != 0: + raise ValueError("sf_padding_block and token_padding_block must be power-of-two multiples of each other.") + if self.token_back_enabled and self.fc2_done_signals_per_token_tile <= 0: + raise ValueError("fc2_done_signals_per_token_tile must be positive when token-back is enabled.") + if self.reduce_topk_in_kernel and self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("In-kernel top-k reduction requires BF16 combine data.") + element_block = self.activation_sf_vector_size * 4 + if self.hidden_size % element_block != 0: + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + if self.sf_padding_block % 128 != 0: + raise ValueError("sf_padding_block must be a multiple of 128.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + @property + def activation_dtype(self) -> type: + return _quant_spec[self.quant_kind][0] + + @property + def activation_sf_dtype(self) -> type: + return _quant_spec[self.quant_kind][1] + + @property + def activation_sf_vector_size(self) -> int: + return _quant_spec[self.quant_kind][2] + + @property + def bytes_per_token(self) -> int: + return self.hidden_size * int(self.activation_dtype.width) // 8 + + @property + def activation_sf_hidden_padded(self) -> int: + valid_hidden = self.hidden_size // self.activation_sf_vector_size + elements_per_16_bytes = 128 // int(self.activation_sf_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def combine_sf_hidden_padded(self) -> int: + if not self.combine_format.is_quantized: + return 0 + valid_hidden = self.hidden_size // self.combine_format.scale_block + elements_per_16_bytes = 128 // int(self.combine_format.scale_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def token_back_push_data(self) -> bool: + return self.token_back_mode != "epi_warps" + + @property + def token_back_push_sf(self) -> bool: + return self.combine_format.is_quantized + + @property + def token_back_enabled(self) -> bool: + return self.token_back_push_data or self.token_back_push_sf + + @property + def worst_case_token_count(self) -> int: + return self._router.worst_case_token_count + + @property + def worst_case_sf_token_count(self) -> int: + return self.worst_case_token_count + + @property + def max_fc1_ready_slot_count(self) -> int: + return ceil_div(self.worst_case_token_count, self.tokens_per_fc1_ready_slot) + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router.router_smem_workspace + + @property + def expert_count_padded(self) -> int: + return self._router.expert_count_padded + + @property + def expert_count_with_trash(self) -> int: + return self._router.expert_count_with_trash + + @property + def router_elements_per_lane(self) -> int: + return self._router.router_elements_per_lane + + @property + def router_warps_per_cta(self) -> int: + return self._router.router_warps_per_cta + + @property + def router_data_cta_count(self) -> int: + return self._router.router_data_cta_count + + @property + def router_tokens_per_cta(self) -> int: + return self._router.router_tokens_per_cta + + @property + def router_push_cta_count(self) -> int: + return self._router.router_push_cta_count + + @property + def router_grid_cta_count(self) -> int: + return self._router.router_grid_cta_count + + @property + def metadata_ready_target(self) -> int: + return self._router.metadata_ready_target + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + self._router.register_device_workspace(workspace) + self._register_main_workspace(workspace) + self._nvlink_barrier.register_device_workspace(workspace) + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + overflow_flag: cute.Tensor, + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + self._router.launch_router( + topk_indices, + topk_scores, + overflow_flag, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper_host, + device_workspace, + stream, + ) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + return self._router.local_expert_sizes(device_workspace, local_rank) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + self._router.wait_for_sizes_ready(device_workspace, sleep_cycles) + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return self._router.token_src_metadata_tensor(device_workspace) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return self._router.fc1_topk_scores_tensor(device_workspace) + + @cute.jit + def assign_device_members( + self, + *, + device_workspace: DeviceWorkspace, + token_comm_args: TokenCommArgs, + local_rank: Int32, + linear_cta_idx: Int32, + ) -> None: + self._device_workspace = device_workspace + self._token_comm_args = token_comm_args + self._local_rank = local_rank + self._linear_cta_idx = linear_cta_idx + thread_idx, _, _ = cute.arch.thread_idx() + transfer_thread_idx = thread_idx % Int32(self.transfer_thread_count) + self._transfer_warp_idx = cute.arch.make_warp_uniform(transfer_thread_idx // Int32(32)) + self._lane_idx = transfer_thread_idx % Int32(32) + self._nvlink_barrier.assign_device_members(device_workspace, token_comm_args.peer_rank_ptr_mapper) + + def remove_device_members(self) -> None: + self._nvlink_barrier.remove_device_members() + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "TokenCommDeterministic": + if values: + raise ValueError("TokenCommDeterministic carries no MLIR values.") + return self + + def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: + workspace.register( + self.fc1_ready_region, + cutlass.Int32, + (self.max_fc1_ready_slot_count,), + buffer_space="local", + reset="tail_reset", + ) + activation_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc1_activation_region, + self.activation_dtype, + (activation_element_count,), + buffer_space="local", + byte_alignment=128, + ) + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded + workspace.register( + self.fc1_activation_sf_region, + self.activation_sf_dtype, + (activation_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_enabled: + workspace.register( + self.fc2_done_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" + ) + if self.token_back_push_data: + fc2_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc2_activation_region, + self.combine_format.act_dtype, + (fc2_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_push_sf: + fc2_sf_element_count = self.worst_case_token_count * self.combine_sf_hidden_padded + workspace.register( + self.fc2_activation_sf_region, + self.combine_format.scale_dtype, + (fc2_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": + workspace.register( + self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + if not self.reduce_topk_in_kernel: + workspace.register( + self.pre_reduced_activation_region, + self.combine_format.act_dtype, + (self.max_tokens_per_rank, self.topk, self.hidden_size), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + if self.combine_format.is_quantized: + workspace.register( + self.pre_reduced_activation_sf_region, + self.combine_format.scale_dtype, + (self.max_tokens_per_rank, self.topk, self.combine_sf_hidden_padded), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + + def register_smem_regions(self, workspace: SmemWorkspace) -> None: + workspace.register_mbarrier(self.token_in_mbarrier_region, self.transfer_warp_count) + if self.token_back_enabled: + workspace.register_mbarrier(self.token_back_mbarrier_region, self.transfer_warp_count) + workspace.register_tensor( + self.expert_sizes_smem_region, cutlass.Int32, (self.experts_per_rank,), byte_alignment=16 + ) + transfer_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.transfer") + token_in_lifetime = transfer_overlay.add_lifetime("token_in") + token_in_lifetime.register_tensor( + self.token_in_activation_smem_region, + self.activation_dtype, + (self.transfer_warp_count, self.hidden_size), + byte_alignment=16, + ) + token_in_lifetime.register_tensor( + self.token_in_sf_smem_region, + self.activation_sf_dtype, + (self.transfer_warp_count, (self.activation_sf_vector_size, self.activation_sf_hidden_padded)), + stride=(self.activation_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + if not self.token_back_enabled: + return + + if self.token_back_mode == "standalone_warps": + token_back_lifetime = workspace.create_overlay( + "nvlink.token_comm.main_smem.standalone_token_back" + ).add_lifetime("token_back") + else: + token_back_lifetime = transfer_overlay.add_lifetime("token_back") + + if self.token_back_mode == "standalone_warps": + available_bytes_per_warp = self.standalone_chunk_bytes + else: + activation_bytes = self.bytes_per_token + sf_bytes = self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 + available_bytes_per_warp = activation_bytes + sf_bytes + + if self.token_back_push_data: + bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + if self.token_back_mode == "standalone_warps": + chunk_bytes = self.standalone_chunk_bytes + elif available_bytes_per_warp < bytes_per_output_token: + chunk_bytes = self.bytes_per_token + else: + chunk_bytes = bytes_per_output_token + if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: + raise ValueError("Token-back data chunk bytes must divide one row.") + chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) + token_back_lifetime.register_tensor( + self.token_back_activation_smem_region, + self.combine_format.act_dtype, + (self.transfer_warp_count, chunk_elements), + byte_alignment=16, + ) + if self.token_back_push_sf: + sf_row_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + if sf_row_bytes > available_bytes_per_warp: + raise ValueError("Token-back scale row exceeds its per-warp stage.") + token_back_lifetime.register_tensor( + self.token_back_sf_smem_region, + self.combine_format.scale_dtype, + (self.transfer_warp_count, (self.combine_format.scale_block, self.combine_sf_hidden_padded)), + stride=(self.combine_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + + @cute.jit + def fc1_ready_counter_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.fc1_ready_region) + + @cute.jit + def fc1_activation_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return cute.make_tensor( + device_workspace.ptr(self.fc1_activation_region), + cute.make_layout((self.worst_case_token_count, self.hidden_size), stride=(self.hidden_size, 1)), + ) + + @cute.jit + def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + layout = tile_atom_to_shape_SF( + (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size + ) + return cute.make_tensor(device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1])) + + @cute.jit + def fc2_done_counter_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_enabled): + return None + return device_workspace.tensor(self.fc2_done_region) + + @cute.jit + def fc2_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_push_data): + return None + return cute.make_tensor( + device_workspace.ptr(self.fc2_activation_region), + cute.make_layout( + (self.worst_case_token_count, 1, self.hidden_size), stride=(self.hidden_size, self.hidden_size, 1) + ), + ) + + @cute.jit + def fc2_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_push_sf): + return None + return cute.make_tensor( + device_workspace.ptr(self.fc2_activation_sf_region), + cute.make_layout( + ( + self.worst_case_token_count, + 1, + (self.combine_format.scale_block, self.hidden_size // self.combine_format.scale_block), + ), + stride=(self.combine_sf_hidden_padded, self.combine_sf_hidden_padded, (0, 1)), + ), + ) + + @cute.jit + def pre_reduced_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """Return the source-domain combine plane when top-k reduction is separate.""" + if cutlass.const_expr(self.reduce_topk_in_kernel): + return None + return device_workspace.tensor(self.pre_reduced_activation_region) + + @cute.jit + def pre_reduced_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """Return the source-domain combine scale plane when one is required.""" + if cutlass.const_expr(self.reduce_topk_in_kernel or not self.combine_format.is_quantized): + return None + return device_workspace.tensor(self.pre_reduced_activation_sf_region) + + @cute.jit + def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Wait for pushed metadata, then pull activation payloads into local pools.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + global_warp_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_warp_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + + sizes = self._router.sizes_tensor(self._device_workspace) + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + + iket.range_push("token_in.wait_sizes_ready") + self._router.wait_for_sizes_ready(self._device_workspace) + iket.range_pop() + iket.range_push("token_in.stage_sizes") + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + owner_expert_begin = self._local_rank * Int32(self.experts_per_rank) + source_sizes = cute.make_tensor(sizes.iterator + owner_expert_begin, cute.make_layout((self.experts_per_rank,))) + copy_elements = 4 if self.experts_per_rank % 4 == 0 else 1 + source_size_vectors = cute.zipped_divide( + (_mark_alignment(source_sizes, 16) if cutlass.const_expr(copy_elements == 4) else source_sizes), + (copy_elements,), + ) + destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) + size_vector_count = cute.size(destination_size_vectors, mode=[1]) + size_copy_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp( + cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL if copy_elements == 4 else cute.nvgpu.LoadCacheMode.ALWAYS + ), + cutlass.Int32, + num_bits_per_copy=copy_elements * 32, + ) + size_copy_rounds = ceil_div(size_vector_count, self.transfer_thread_count) + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + if vector_idx < Int32(size_vector_count): + cute.copy( + size_copy_atom, source_size_vectors[None, vector_idx], destination_size_vectors[None, vector_idx] + ) + cute.arch.cp_async_commit_group() + iket.range_pop() + + iket.range_push("token_in.wait_metadata_ready") + self._router.wait_for_metadata_ready(self._device_workspace) + iket.range_pop() + + cute.arch.cp_async_wait_group(0) + iket.range_push("token_in.size_barrier") + token_in_size_barrier = pipeline.NamedBarrier( + barrier_id=self.token_in_size_barrier_id, num_threads=self.transfer_thread_count + ) + token_in_size_barrier.arrive_and_wait() + iket.range_pop() + if cutlass.const_expr(self.token_back_mode == "standalone_warps"): + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive() + + iket.range_push("token_in.pull_payload") + token_in_mbarriers = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_mbarrier = token_in_mbarriers + transfer_warp_idx + warp_activation_stage = token_in_activation[transfer_warp_idx, None] + warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] + if lane_idx == Int32(0): + cute.arch.mbarrier_init(warp_mbarrier, 1) + cute.arch.sync_warp() + + fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) + fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) + fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) + activation_bytes = cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 + activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 + sf_copy_elements = 4 + source_sf_values = cute.slice_(warp_sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) + + next_dense_token = global_warp_idx + expert_valid_begin = Int32(0) + expert_sf_begin = Int32(0) + expert_ready_slot_begin = Int32(0) + pull_phase = Int32(0) + flag_tracker = GpuReleaseFlagBatchTracker( + flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(0), thread_idx=lane_idx + ) + + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + pull_count = Int32(0) + if next_dense_token < expert_valid_end: + pull_count = (expert_valid_end - next_dense_token + global_warp_count - Int32(1)) // global_warp_count + + for pull_round in cutlass.range(pull_count, unroll=1): + dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count + token_in_expert = dense_token_idx - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + sf_token_idx = expert_sf_begin + token_in_expert + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + remote_activation_address = ( + self._token_comm_args.activation.iterator.toint() + + peer_offset + + Int64(metadata.src_token) * Int64(self.bytes_per_token) + ) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, + ) + remote_sf_row = remote_sf[Int64(metadata.src_token), None] + + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) + ) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address), + warp_mbarrier, + Int32(activation_bytes), + ) + tma_load_1d( + warp_sf_stage.iterator, remote_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_activation, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + + destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] + destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) + destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) + sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) + for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): + sf_vector_idx = Int32(sf_round * 32) + lane_idx + if sf_vector_idx < Int32(sf_vector_count): + cute.copy( + sf_copy_atom, + source_sf_vectors[None, sf_vector_idx], + destination_sf_vectors[None, sf_vector_idx], + ) + + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + ready_address = (fc1_ready_counter + ready_slot_idx).toint() + flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) + cute.arch.sync_warp() + pull_phase = pull_phase ^ Int32(1) + + next_dense_token = next_dense_token + pull_count * global_warp_count + expert_valid_begin = expert_valid_end + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) + expert_ready_slot_begin = expert_ready_slot_begin + ( + (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot) + ) + local_expert = local_expert + Int32(1) + + flag_tracker.fire() + cute.arch.sync_warp() + iket.range_pop() + if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): + iket.range_push("token_in.transfer_barrier") + transfer_lifetime_barrier = pipeline.NamedBarrier( + barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count + ) + transfer_lifetime_barrier.arrive_and_wait() + iket.range_pop() + + @cute.jit + def _stateless_pace(self, reference_window: Int32, current_window: Int32) -> None: + sleep_cycles = Int32(0) + if current_window < reference_window: + sleep_cycles = reference_window - current_window + elif current_window > reference_window: + sleep_cycles = cutlass.min(current_window - reference_window, Int32(self.standalone_max_backoff_cycles)) + if sleep_cycles > Int32(0): + nanosleep(sleep_cycles) + + @cute.jit + def _adaptive_pace(self, average_window: Int32, current_window: Int32, low_window: int, high_window: int) -> Int32: + sleep_cycles = Int32(0) + if current_window > average_window: + average_window = average_window + ((current_window - average_window + Int32(3)) // Int32(4)) + sleep_cycles = current_window - average_window + if sleep_cycles > Int32(high_window): + sleep_cycles = Int32(high_window) + else: + average_window = average_window - ((average_window - current_window + Int32(3)) // Int32(4)) + sleep_cycles = average_window - current_window + if sleep_cycles > Int32(self.adaptive_minimum_sleep_cycles): + nanosleep(sleep_cycles) + if average_window > Int32(high_window): + average_window = Int32(high_window) + if average_window < Int32(low_window): + average_window = Int32(low_window) + return average_window + + @cute.jit + def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Push completed FC2 data and scale rows to source ranks.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + if cutlass.const_expr(not self.token_back_enabled): + return + if cutlass.const_expr( + self.combine_format.is_quantized and self._token_comm_args.pre_reduced_activation_sf is None + ): + raise ValueError("Quantized token-back requires a scale destination.") + + global_worker_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_worker_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + token_back_mbarriers = smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) + worker_mbarrier = token_back_mbarriers + transfer_warp_idx + if cutlass.const_expr(self.token_back_push_data): + token_back_activation = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) + worker_activation_stage = token_back_activation[transfer_warp_idx, None] + activation_chunk_bytes = ( + cute.cosize(worker_activation_stage) * int(self.combine_format.act_dtype.width) // 8 + ) + if cutlass.const_expr(self.token_back_push_sf): + token_back_sf = smem_workspace.tensor(self.token_back_sf_smem_region, smem_base) + worker_sf_stage = token_back_sf[transfer_warp_idx, (None, None)] + sf_chunk_bytes = cute.cosize(worker_sf_stage) * int(self.combine_format.scale_dtype.width) // 8 + if lane_idx == Int32(0): + cute.arch.mbarrier_init(worker_mbarrier, 1) + cute.arch.sync_warp() + + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + if cutlass.const_expr(self.token_back_mode == "standalone_warps"): + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive_and_wait() + + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + fc2_done = self._device_workspace.ptr(self.fc2_done_region) + if cutlass.const_expr(self.token_back_push_data): + fc2_activation_pointer = self._device_workspace.ptr(self.fc2_activation_region) + output_token_bytes = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) + data_window_unit = ceil_div(activation_chunk_bytes * 2, 3) + reuse_data_pacing_enabled = ( + self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles + ) + # Preserve the empirical low:initial:high ratio of 1:2.5:5. + data_average_window = Int32(data_window_unit) + data_low_window = data_window_unit * 2 // 5 + data_high_window = data_window_unit * 2 + print(f"[{data_window_unit}, {data_low_window}, {data_high_window}]") + if cutlass.const_expr(self.token_back_push_sf): + fc2_sf_pointer = self._device_workspace.ptr(self.fc2_activation_sf_region) + output_sf_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) + sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) + reuse_sf_pacing_enabled = ( + self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles + ) + sf_average_window = Int32(sf_window_unit) + sf_low_window = sf_window_unit * 2 // 5 + sf_high_window = sf_window_unit * 2 + + next_dense_token = global_worker_idx - global_worker_count + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): + next_dense_token = Int32(0) + next_dense_token = self.next_token(next_dense_token) + expert_valid_begin = Int32(0) + transfer_phase = Int32(0) + + iket.range_push("token_back.work") + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + if next_dense_token < expert_valid_end: + token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( + self.tokens_per_fc1_ready_slot + ) + completion_target = token_tile_count * Int32(self.fc2_done_signals_per_token_tile) + iket.range_push("token_back.wait_fc2") + while cute.arch.load(fc2_done + local_expert, Int32, sem="acquire", scope="gpu") < completion_target: + nanosleep(500) + iket.range_pop() + + while next_dense_token < expert_valid_end: + token_in_expert = next_dense_token - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + destination_topk = metadata.src_topk + if cutlass.const_expr(self.reduce_topk_in_kernel): + destination_topk = Int32(0) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + is_remote_token = metadata.src_rank != self._local_rank + + if cutlass.const_expr(self.token_back_push_data): + iket.range_push("token_back.push_data") + local_activation_address = fc2_activation_pointer.toint() + Int64(pool_token_idx) * Int64( + output_token_bytes + ) + remote_activation = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation.dtype, + self._token_comm_args.pre_reduced_activation.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation.layout, + ) + destination_row = remote_activation[Int64(metadata.src_token), destination_topk, None] + for chunk_idx in cutlass.range_constexpr(activation_chunk_count): + chunk_byte_offset = Int64(chunk_idx * activation_chunk_bytes) + chunk_bytes_this_round = min( + activation_chunk_bytes, output_token_bytes - chunk_idx * activation_chunk_bytes + ) + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_data_pacing_enabled = ( + self.token_back_mode != "reuse_dispatch_warps" + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) + if cutlass.const_expr(reuse_data_pacing_enabled or stateless_data_pacing_enabled): + if is_remote_token: + round_start_clock = read_clock64() + else: + round_start_clock = round_start_clock + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) + tma_load_1d( + worker_activation_stage.iterator, + local_activation_address + chunk_byte_offset, + worker_mbarrier, + current_chunk_bytes, + ) + cute.arch.mbarrier_wait(worker_mbarrier, transfer_phase) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_row.iterator.toint() + chunk_byte_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + if cutlass.const_expr(self.reduce_topk_in_kernel): + cp_reduce_async_bulk_add_bf16_s2g( + destination_chunk, worker_activation_stage.iterator, current_chunk_bytes + ) + else: + cp_async_bulk_s2g( + destination_chunk, worker_activation_stage.iterator, current_chunk_bytes + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + transfer_phase = transfer_phase ^ Int32(1) + if cutlass.const_expr(reuse_data_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + data_average_window = self._adaptive_pace( + data_average_window, current_window, data_low_window, data_high_window + ) + elif cutlass.const_expr(stateless_data_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + self._stateless_pace(Int32(current_window_unit), current_window) + iket.range_pop() + + if cutlass.const_expr(self.token_back_push_sf): + iket.range_push("token_back.push_sf") + local_sf_address = fc2_sf_pointer.toint() + Int64(pool_token_idx) * Int64(output_sf_bytes) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation_sf.dtype, + self._token_comm_args.pre_reduced_activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation_sf.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation_sf.layout, + ) + destination_sf_row = remote_sf[Int64(metadata.src_token), destination_topk, None] + for chunk_idx in cutlass.range_constexpr(sf_chunk_count): + chunk_byte_offset = Int64(chunk_idx * sf_chunk_bytes) + chunk_bytes_this_round = min(sf_chunk_bytes, output_sf_bytes - chunk_idx * sf_chunk_bytes) + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_sf_pacing_enabled = ( + self.token_back_mode != "reuse_dispatch_warps" + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) + if cutlass.const_expr(reuse_sf_pacing_enabled or stateless_sf_pacing_enabled): + if is_remote_token: + round_start_clock = read_clock64() + else: + round_start_clock = round_start_clock + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) + tma_load_1d( + worker_sf_stage.iterator, + local_sf_address + chunk_byte_offset, + worker_mbarrier, + current_chunk_bytes, + ) + cute.arch.mbarrier_wait(worker_mbarrier, transfer_phase) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_sf_row.iterator.toint() + chunk_byte_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_chunk, worker_sf_stage.iterator, current_chunk_bytes) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + transfer_phase = transfer_phase ^ Int32(1) + if cutlass.const_expr(reuse_sf_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + sf_average_window = self._adaptive_pace( + sf_average_window, current_window, sf_low_window, sf_high_window + ) + elif cutlass.const_expr(stateless_sf_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + self._stateless_pace(Int32(current_window_unit), current_window) + iket.range_pop() + + cute.arch.sync_warp() + next_dense_token = self.next_token(next_dense_token) + + expert_valid_begin = expert_valid_end + local_expert = local_expert + Int32(1) + iket.range_pop() + + @cute.jit + def next_token(self, current_token: Int32) -> Int32: + global_worker_count = self.promised_launchable_sm_count * self.transfer_warp_count + schedule_counter = None + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): + schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): + claimed_token = Int32(0) + if self._lane_idx == Int32(0): + claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") + return Int32(cute.arch.shuffle_sync(claimed_token, Int32(0))) + return current_token + global_worker_count + + @cute.jit + def reset_tail(self) -> None: + """Reset communication state with the four token-in transfer warps.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + iket.range_push("tail.nvlink_drain") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + total_reset_threads = self.promised_launchable_sm_count * self.transfer_thread_count + global_reset_thread = self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx + iket.range_push("tail.reset_workspace") + self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) + self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) + iket.range_pop() + iket.range_push("tail.nvlink_publish") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=False, + epilogue_grid_sync=True, + ) + iket.range_pop() + if cutlass.const_expr(os.environ.get("MEGA_USE_NCU", "0") == "1"): + iket.range_push("tail.ncu_finalize") + self._nvlink_barrier.finalize( + 2, + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + ) + iket.range_pop() + + +__all__ = ["TokenCommDeterministic"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/token_protocol.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/token_protocol.py new file mode 100644 index 000000000..289c5fe47 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/token_protocol.py @@ -0,0 +1,59 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Cross-rank token metadata protocol.""" + +import dataclasses +from typing import ClassVar, Union + +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 + +@dataclasses.dataclass(frozen=True) +class TokenSrcMetadata: + """One i64 routing record: rank:u16, topk:u16, token:u32.""" + + src_rank: Int32 + src_token: Int32 + src_topk: Int32 + + nbytes: ClassVar[int] = 8 + + def pack(self) -> Int64: + high = (Int64(self.src_rank) << Int64(16)) | Int64(self.src_topk) + return (high << Int64(32)) | ( + Int64(self.src_token) & Int64(0xFFFFFFFF) + ) + + @staticmethod + def _pointer(address: Union[cute.Pointer, Int64]) -> cute.Pointer: + raw_address = ( + address if isinstance(address, Int64) else address.toint() + ) + return cute.make_ptr( + Int64, + raw_address, + AddressSpace.gmem, + assumed_align=8, + ) + + def store(self, address: Union[cute.Pointer, Int64]) -> None: + cute.arch.store(self._pointer(address), self.pack(), scope="gpu") + + @classmethod + def load( + cls, + address: Union[cute.Pointer, Int64], + ) -> "TokenSrcMetadata": + packed = Int64( + cute.arch.load(cls._pointer(address), Int64, scope="gpu") + ) + high = packed >> Int64(32) + return cls( + src_rank=Int32((high >> Int64(16)) & Int64(0xFFFF)), + src_token=Int32(packed & Int64(0xFFFFFFFF)), + src_topk=Int32(high & Int64(0xFFFF)), + ) + + +__all__ = ["TokenSrcMetadata"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/__init__.py new file mode 100644 index 000000000..780c03c78 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/__init__.py @@ -0,0 +1,50 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Low-level workspace, synchronization, and PTX helpers.""" + +from .device_workspace import DeviceWorkspace +from .dsl_helpers import spin_peek, spin_wait +from .flag_batch import GpuAsyncReleaseFlagBatchTracker, GpuReleaseFlagBatchTracker, make_flag_batch_tracker +from .iket_compat import iket +from .ptx_helpers import ( + cvt_f32_to_fp8_to_f32, + cvt_f32x4_to_f8x4_pack_i32, + stg_e8m0_from_f32, + stg_e8m0x8_from_f32, +) +from .smem_workspace import SmemWorkspace +from .software_sync import NvlinkBarrier, SoftwareGridSync +from .utils import ( + IntegerType, + ceil_div, + cosize_from_shape_stride_tuples, + product, + row_major_stride, + round_up, + validate_static_integer_tuple, +) + +__all__ = [ + "DeviceWorkspace", + "GpuAsyncReleaseFlagBatchTracker", + "GpuReleaseFlagBatchTracker", + "IntegerType", + "NvlinkBarrier", + "SmemWorkspace", + "SoftwareGridSync", + "ceil_div", + "cosize_from_shape_stride_tuples", + "cvt_f32_to_fp8_to_f32", + "cvt_f32x4_to_f8x4_pack_i32", + "make_flag_batch_tracker", + "product", + "row_major_stride", + "round_up", + "iket", + "spin_peek", + "spin_wait", + "stg_e8m0_from_f32", + "stg_e8m0x8_from_f32", + "validate_static_integer_tuple", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/constants.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/constants.py new file mode 100644 index 000000000..839ba88d6 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/constants.py @@ -0,0 +1,28 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Specification-defined numeric constants shared across kernel components.""" + + +Log2E = 1.4426950408889634 +Fp32Max = 3.40282346638528859812e38 + +Nvfp4E2M1Max = 6.0 +Fp8E4M3FNMax = 448.0 +Fp8E5M2Max = 57344.0 + +Nvfp4E2M1RcpLimit = 1.0 / Nvfp4E2M1Max +Fp8E4M3RcpLimit = 1.0 / Fp8E4M3FNMax +Fp8E5M2RcpLimit = 1.0 / Fp8E5M2Max + + +__all__ = [ + "Fp32Max", + "Fp8E4M3FNMax", + "Fp8E4M3RcpLimit", + "Fp8E5M2Max", + "Fp8E5M2RcpLimit", + "Log2E", + "Nvfp4E2M1Max", + "Nvfp4E2M1RcpLimit", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/cute_py_helpers.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/cute_py_helpers.py new file mode 100644 index 000000000..13736388e --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/cute_py_helpers.py @@ -0,0 +1,520 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import dataclass +from math import gcd +from typing import Optional, Tuple, Union + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import OperandMajorMode + +from .smem_workspace import SmemRegion + + +# Every block-scaled tcgen05 MMA kind accumulates in F32 and no other accumulator is legal, so +# this belongs to the instruction family rather than being a caller choice. +tcgen05_block_scaled_acc_dtype = cutlass.Float32 + + +@dataclass(frozen=True) +class Tcgen05MmaInstruction: + a_type: type[cutlass.Numeric] + b_type: type[cutlass.Numeric] + instruction_mnk: Tuple[int, int, int] + participates: int + acc_type: type[cutlass.Numeric] = cutlass.Float32 + sfa_type: Optional[type[cutlass.Numeric]] = None + sfb_type: Optional[type[cutlass.Numeric]] = None + sf_vec_size: Optional[int] = None + + +@dataclass(frozen=True) +class Tcgen05TmemPlan: + """Column-level TMEM placement and accumulator staging contract.""" + + allocation_columns: int + accumulator_columns: int + accumulator_stage_columns: int + accumulator_stage_count: int + accumulator_stage_stride_columns: int + accumulator_pipeline_stages: int + sfa_columns: int + sfb_columns: int + + +def make_tcgen05_tmem_plan( + mma_instruction: Tcgen05MmaInstruction, arch: str, mma_tiler_mnk: Tuple[int, int, int] +) -> Tcgen05TmemPlan: + """Plan TMEM for canonical dense or block-scaled TCGen05 MMA. + + This planner requires A and B to reside in SMEM, per-CTA M to equal 128, + and tile M to equal instruction M. A-from-TMEM, sparse MMA, B-reuse, and + custom atom layouts or permutations are outside its contract. + SM100/SM103 use the canonical two-stage overlap for 256-column block-scaled + accumulators; other supported cases maximize disjoint accumulator stages. + """ + if mma_instruction.participates not in (1, 2): + raise ValueError(f"TCGen05 MMA participates must be one or two, got {mma_instruction.participates}.") + if len(mma_instruction.instruction_mnk) != 3 or len(mma_tiler_mnk) != 3: + raise ValueError("instruction_mnk and mma_tiler_mnk must each contain three dimensions.") + if any(dimension <= 0 for dimension in (*mma_instruction.instruction_mnk, *mma_tiler_mnk)): + raise ValueError("MMA instruction and tiler dimensions must be positive.") + + instruction_m, instruction_n, instruction_k = mma_instruction.instruction_mnk + tile_m, tile_n, tile_k = mma_tiler_mnk + if instruction_m % mma_instruction.participates != 0 or instruction_n % mma_instruction.participates != 0: + raise ValueError("MMA instruction M and N must be divisible by participates.") + if tile_m % instruction_m != 0 or tile_n % instruction_n != 0 or tile_k % instruction_k != 0: + raise ValueError("MMA instruction dimensions must divide mma_tiler_mnk.") + if tile_m != instruction_m: + raise ValueError("TCGen05 TMEM planning does not support M repetition or B-reuse.") + if tile_m // mma_instruction.participates != 128: + raise ValueError("TCGen05 TMEM planning requires per-CTA M to equal 128.") + + sf_fields = (mma_instruction.sfa_type, mma_instruction.sfb_type, mma_instruction.sf_vec_size) + has_scale_factors = any(value is not None for value in sf_fields) + if has_scale_factors and any(value is None for value in sf_fields): + raise ValueError("sfa_type, sfb_type, and sf_vec_size must be provided together.") + + sfa_columns = 0 + sfb_columns = 0 + if has_scale_factors: + sf_vec_size = mma_instruction.sf_vec_size + if not isinstance(sf_vec_size, int) or isinstance(sf_vec_size, bool) or sf_vec_size <= 0: + raise ValueError("sf_vec_size must be a positive Python int.") + if instruction_k % sf_vec_size != 0: + raise ValueError("Instruction K must be divisible by sf_vec_size.") + if tile_k % (sf_vec_size * 4) != 0: + raise ValueError("Tile K must contain complete block-scaled basic chunks.") + sfa_columns = tile_k // sf_vec_size + sfb_columns = max(tile_n // 128, 1) * tile_k // sf_vec_size + + tmem_column_capacity = cute.arch.get_max_tmem_alloc_cols(arch) + accumulator_stage_columns = tile_n + scale_factor_columns = sfa_columns + sfb_columns + accumulator_stage_count = (tmem_column_capacity - scale_factor_columns) // accumulator_stage_columns + if accumulator_stage_count < 1: + raise ValueError("Scale-factor TMEM leaves no accumulator stage.") + + accumulator_stage_stride_columns = accumulator_stage_columns + accumulator_pipeline_stages = accumulator_stage_count + accumulator_columns = accumulator_stage_columns * accumulator_stage_count + + arch_number = _parse_arch_number(arch) + use_sm100_overlap = ( + arch_number in (100, 103) + and has_scale_factors + and accumulator_stage_columns == 256 + and accumulator_stage_count == 1 + and scale_factor_columns <= 64 + ) + if use_sm100_overlap: + accumulator_stage_count = 2 + accumulator_stage_stride_columns = accumulator_stage_columns - scale_factor_columns + accumulator_pipeline_stages = 1 + accumulator_columns = accumulator_stage_columns + accumulator_stage_stride_columns + + used_columns = accumulator_columns + scale_factor_columns + allocation_columns = _round_tmem_allocation_columns(used_columns, tmem_column_capacity) + return Tcgen05TmemPlan( + allocation_columns=allocation_columns, + accumulator_columns=accumulator_columns, + accumulator_stage_columns=accumulator_stage_columns, + accumulator_stage_count=accumulator_stage_count, + accumulator_stage_stride_columns=accumulator_stage_stride_columns, + accumulator_pipeline_stages=accumulator_pipeline_stages, + sfa_columns=sfa_columns, + sfb_columns=sfb_columns, + ) + + +def _parse_arch_number(arch: str) -> int: + if not isinstance(arch, str): + raise TypeError(f"arch must be a string, got {type(arch)}.") + normalized = arch.lower() + if normalized.startswith("sm_"): + normalized = normalized[3:] + elif normalized.startswith("sm"): + normalized = normalized[2:] + digits = [] + for character in normalized: + if not character.isdigit(): + break + digits.append(character) + if not digits: + raise ValueError(f"Cannot parse architecture {arch!r}.") + return int("".join(digits)) + + +def _round_tmem_allocation_columns(used_columns: int, capacity_columns: int) -> int: + if used_columns <= 0: + raise ValueError("TMEM usage must be positive.") + if used_columns <= 512: + allocation_columns = max(32, 1 << (used_columns - 1).bit_length()) + else: + allocation_columns = _round_up(used_columns, 32) + if allocation_columns > capacity_columns: + raise ValueError(f"TMEM plan needs {allocation_columns} columns, exceeding {capacity_columns}.") + return allocation_columns + + +def tcgen05_smem_alloc_type( + dtype: type[cutlass.Numeric], peer_dtype: type[cutlass.Numeric], arch: str +) -> type[cutlass.Numeric]: + """SMEM container type for one block-scaled TCGen05 operand. + + Blackwell mixed-width MMA consumes a uniform byte-per-element SMEM image, so its narrow operand + arrives through U4_UNPACK_U8. Rubin consumes mixed FP4 directly from packed SMEM. A 6-bit + operand always uses U6_UNPACK_U8 because no packed U6 TMA format exists. + """ + arch_number = _parse_arch_number(arch) + if arch_number not in (100, 103, 107): + raise ValueError(f"Unsupported TCGen05 architecture {arch!r}.") + needs_mixed_width_unpack = arch_number in (100, 103) and dtype.width != peer_dtype.width + needs_unpack = needs_mixed_width_unpack or 6 in (dtype.width, peer_dtype.width) + return cutlass.Int8 if (needs_unpack and dtype.width < 8) else dtype + + +def make_smem_layouts( + mma_inst: Tcgen05MmaInstruction, + mma_tiler_mnk: Tuple[int, int, int], + stages: Union[int, Tuple[int, ...]], + ab_gmem_major_modes: Tuple[OperandMajorMode, OperandMajorMode], + arch: str, +) -> Union[Tuple[SmemRegion, SmemRegion], Tuple[SmemRegion, SmemRegion, SmemRegion, SmemRegion]]: + """Derive workspace-ready TCGen05 operand regions without an MLIR context.""" + has_scale_factors = any(value is not None for value in (mma_inst.sfa_type, mma_inst.sfb_type, mma_inst.sf_vec_size)) + if has_scale_factors and any( + value is None for value in (mma_inst.sfa_type, mma_inst.sfb_type, mma_inst.sf_vec_size) + ): + raise ValueError("sfa_type, sfb_type, and sf_vec_size must be provided together.") + + operand_count = 4 if has_scale_factors else 2 + if isinstance(stages, int): + operand_stages = (stages,) * operand_count + else: + operand_stages = stages + if len(operand_stages) != operand_count: + raise ValueError(f"Expected {operand_count} stage counts, got {len(operand_stages)}.") + if any(stage <= 0 for stage in operand_stages): + raise ValueError(f"All stage counts must be positive, got {operand_stages}.") + + if mma_inst.participates not in (1, 2): + raise ValueError(f"TCGen05 MMA participates must be one or two, got {mma_inst.participates}.") + if len(mma_inst.instruction_mnk) != 3 or len(mma_tiler_mnk) != 3: + raise ValueError("instruction_mnk and mma_tiler_mnk must each contain three dimensions.") + if any(dimension <= 0 for dimension in (*mma_inst.instruction_mnk, *mma_tiler_mnk)): + raise ValueError("MMA instruction and tiler dimensions must be positive.") + + inst_m, inst_n, inst_k = mma_inst.instruction_mnk + tile_m, tile_n, tile_k = mma_tiler_mnk + if inst_m % mma_inst.participates != 0 or inst_n % mma_inst.participates != 0: + raise ValueError("MMA instruction M and N must be divisible by participates.") + if tile_m % inst_m != 0 or tile_n % inst_n != 0 or tile_k % inst_k != 0: + raise ValueError("MMA instruction dimensions must divide mma_tiler_mnk.") + + a_region = _make_ab_region( + mma_inst.a_type, + tcgen05_smem_alloc_type(mma_inst.a_type, mma_inst.b_type, arch), + tile_m // mma_inst.participates, + tile_k, + (inst_m // mma_inst.participates, inst_k), + operand_stages[0], + ab_gmem_major_modes[0], + ) + b_region = _make_ab_region( + mma_inst.b_type, + tcgen05_smem_alloc_type(mma_inst.b_type, mma_inst.a_type, arch), + tile_n // mma_inst.participates, + tile_k, + (inst_n // mma_inst.participates, inst_k), + operand_stages[1], + ab_gmem_major_modes[1], + ) + if not has_scale_factors: + return a_region, b_region + + sfa_region = _make_sf_region( + mma_inst.sfa_type, + inst_m // mma_inst.participates // 128, + tile_m // inst_m, + tile_k // inst_k, + inst_k, + mma_inst.sf_vec_size, + operand_stages[2], + ) + sfb_region = _make_sf_region( + mma_inst.sfb_type, + _round_up(inst_n, 128) // 128, + tile_n // inst_n, + tile_k // inst_k, + inst_k, + mma_inst.sf_vec_size, + operand_stages[3], + ) + return a_region, b_region, sfa_region, sfb_region + + +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +def _canonical_mode(shape, stride): + if not isinstance(shape, tuple): + return (1, 0) if shape == 1 else (shape, stride) + + result_shape = [] + result_stride = [] + for current_shape, current_stride in zip(shape, stride): + if current_shape == 1: + continue + if result_shape and result_shape[-1] * result_stride[-1] == current_stride: + result_shape[-1] *= current_shape + else: + result_shape.append(current_shape) + result_stride.append(current_stride) + if not result_shape: + return 1, 0 + if len(result_shape) == 1: + return result_shape[0], result_stride[0] + return tuple(result_shape), tuple(result_stride) + + +def _shape_size(shape) -> int: + if isinstance(shape, tuple): + result = 1 + for child_shape in shape: + result *= _shape_size(child_shape) + return result + return shape + + +def _prefix_tile_profile(shape, tile_extent: int): + """Represent a scalar prefix tile using the source mode boundaries.""" + if tile_extent <= 0 or _shape_size(shape) % tile_extent != 0: + raise ValueError(f"Tile extent {tile_extent} must divide mode {shape}.") + if not isinstance(shape, tuple): + return tile_extent + + tile_modes = [] + remaining_extent = tile_extent + for child_shape in shape: + if remaining_extent == 1: + break + + child_size = _shape_size(child_shape) + if remaining_extent >= child_size and remaining_extent % child_size == 0: + tile_modes.append(child_shape) + remaining_extent //= child_size + elif child_size % remaining_extent == 0: + tile_modes.append(_prefix_tile_profile(child_shape, remaining_extent)) + remaining_extent = 1 + else: + raise ValueError(f"Tile extent {tile_extent} does not divide a prefix of mode {shape}.") + + if remaining_extent != 1: + raise ValueError(f"Tile extent {tile_extent} exceeds the prefix of mode {shape}.") + if not tile_modes: + return 1 + if len(tile_modes) == 1: + return tile_modes[0] + return tuple(tile_modes) + + +def _flatten_mode(shape, stride): + if isinstance(shape, tuple): + result = [] + for child_shape, child_stride in zip(shape, stride): + result.extend(_flatten_mode(child_shape, child_stride)) + return result + return [(shape, stride)] + + +def _flatten_shape(shape): + if isinstance(shape, tuple): + result = [] + for child_shape in shape: + result.extend(_flatten_shape(child_shape)) + return result + return [shape] + + +def _rebuild_stride_like(shape, flat_strides): + stride_iterator = iter(flat_strides) + + def rebuild(current_shape): + if isinstance(current_shape, tuple): + return tuple(rebuild(child_shape) for child_shape in current_shape) + return next(stride_iterator) + + return rebuild(shape) + + +def _divide_mode_by_tile(shape, stride, tile_shape): + """Apply a static tile profile and return its stride and canonical rest.""" + remaining_modes = _flatten_mode(shape, stride) + tile_strides = [] + for tile_extent in _flatten_shape(tile_shape): + if not remaining_modes: + raise ValueError(f"Tile {tile_shape} exceeds mode {shape}.") + + mode_extent, mode_stride = remaining_modes[0] + if tile_extent <= 0 or mode_extent % tile_extent != 0: + raise ValueError(f"Tile mode {tile_extent} must divide source mode {mode_extent}.") + + tile_strides.append(mode_stride if tile_extent > 1 else 0) + rest_extent = mode_extent // tile_extent + if rest_extent > 1: + remaining_modes[0] = (rest_extent, mode_stride * tile_extent) + else: + remaining_modes.pop(0) + + rest_shape, rest_stride = _canonical_mode( + tuple(mode_extent for mode_extent, _ in remaining_modes), + tuple(mode_stride for _, mode_stride in remaining_modes), + ) + return (_rebuild_stride_like(tile_shape, tile_strides), rest_shape, rest_stride) + + +def _tiled_divide_2d(shape, stride, tile_shape): + """Tile a static rank-2 layout and group value modes before rest modes.""" + if ( + not isinstance(shape, tuple) + or not isinstance(stride, tuple) + or not isinstance(tile_shape, tuple) + or len(shape) != 2 + or len(stride) != 2 + or len(tile_shape) != 2 + ): + raise ValueError("A 2-D tiled divide requires rank-2 shape, stride, and tile.") + + mn_tile_stride, mn_rest_shape, mn_rest_stride = _divide_mode_by_tile(shape[0], stride[0], tile_shape[0]) + k_tile_stride, k_rest_shape, k_rest_stride = _divide_mode_by_tile(shape[1], stride[1], tile_shape[1]) + return ( + ((tile_shape[0], tile_shape[1]), mn_rest_shape, k_rest_shape), + ((mn_tile_stride, k_tile_stride), mn_rest_stride, k_rest_stride), + ) + + +def _make_ab_region( + dtype: type[cutlass.Numeric], + alloc_dtype: type[cutlass.Numeric], + mn_extent: int, + k_extent: int, + value_shape: Tuple[int, int], + stages: int, + major_mode: OperandMajorMode, +) -> SmemRegion: + """Plan one operand's SMEM region. + + ``dtype`` is the logical element type the MMA sees; ``alloc_dtype`` is the container it + occupies in SMEM. The two differ only for a sub-byte operand loaded through the unpacking TMA + (see ``tcgen05_smem_alloc_type``), where the layout must be sized and swizzled for 1-byte + containers while the major-mode rule still follows the logical type. + """ + if dtype.width in (4, 6) and major_mode != OperandMajorMode.K: + raise ValueError(f"{dtype} TCGen05 operands require K-major SMEM.") + + leading_extent = k_extent if major_mode == OperandMajorMode.K else mn_extent + leading_bits = leading_extent * alloc_dtype.width + if leading_bits % 8 != 0: + raise ValueError("The leading dimension must occupy a whole number of bytes.") + swizzle_bytes = gcd(leading_bits // 8, 128) + swizzle_by_bytes = {16: (0, 4, 3), 32: (1, 4, 3), 64: (2, 4, 3), 128: (3, 4, 3)} + if swizzle_bytes not in swizzle_by_bytes: + raise ValueError(f"Unsupported leading dimension size {leading_bits // 8} bytes.") + swizzle = swizzle_by_bytes[swizzle_bytes] + if major_mode == OperandMajorMode.MN and alloc_dtype.width == 32 and swizzle_bytes == 128: + swizzle = (2, 5, 2) + + chunk_elements = swizzle_bytes * 8 // alloc_dtype.width + if leading_extent % chunk_elements != 0: + raise ValueError(f"Leading extent {leading_extent} must be divisible by {chunk_elements}.") + repeats = leading_extent // chunk_elements + stage_stride = mn_extent * k_extent + + if major_mode == OperandMajorMode.K: + mn_shape = mn_extent + mn_stride = chunk_elements if repeats > 1 else k_extent + if repeats > 1: + k_shape = (chunk_elements, repeats) + k_stride = (1, mn_extent * chunk_elements) + else: + k_shape = k_extent + k_stride = 1 + else: + k_shape = k_extent + k_stride = chunk_elements if repeats > 1 else mn_extent + if repeats > 1: + mn_shape = (chunk_elements, repeats) + mn_stride = (1, k_extent * chunk_elements) + else: + mn_shape = mn_extent + mn_stride = 1 + + mma_mn_shape = _prefix_tile_profile(mn_shape, value_shape[0]) + mma_k_shape = _prefix_tile_profile(k_shape, value_shape[1]) + mma_shape, mma_stride = _tiled_divide_2d((mn_shape, k_shape), (mn_stride, k_stride), (mma_mn_shape, mma_k_shape)) + return SmemRegion( + name="", + kind="tensor", + dtype=alloc_dtype, + shape=(*mma_shape, stages), + stride=(*mma_stride, stage_stride if stages > 1 else 0), + swizzle=swizzle, + byte_alignment=128, + ) + + +def _make_sf_region( + dtype: type[cutlass.Numeric], + instruction_mn_blocks: int, + mn_iterations: int, + k_iterations: int, + instruction_k: int, + sf_vec_size: int, + stages: int, +) -> SmemRegion: + if instruction_mn_blocks <= 0: + raise ValueError("A scale-factor instruction must cover at least one 128-element MN block.") + if instruction_k % sf_vec_size != 0: + raise ValueError("instruction K must be divisible by sf_vec_size.") + + cta_mn_blocks = instruction_mn_blocks * mn_iterations + cta_k = instruction_k * k_iterations + basic_chunk_k = sf_vec_size * 4 + if cta_k % basic_chunk_k != 0: + raise ValueError("The CTA K extent must contain complete block-scaled basic chunks.") + basic_chunk_repetitions = cta_k // basic_chunk_k + + full_mn_shape = ((32, 4), cta_mn_blocks) + full_mn_stride = ((16, 4), basic_chunk_repetitions * 512 if cta_mn_blocks > 1 else 0) + full_k_shape = ((sf_vec_size, 4), basic_chunk_repetitions) + full_k_stride = ((0, 1), 512 if basic_chunk_repetitions > 1 else 0) + + mma_mn_shape = ((32, 4), instruction_mn_blocks) + mma_k_shape = (sf_vec_size, _prefix_tile_profile((4, basic_chunk_repetitions), instruction_k // sf_vec_size)) + mma_shape, mma_stride = _tiled_divide_2d( + (full_mn_shape, full_k_shape), (full_mn_stride, full_k_stride), (mma_mn_shape, mma_k_shape) + ) + stage_stride = cta_mn_blocks * basic_chunk_repetitions * 512 + return SmemRegion( + name="", + kind="tensor", + dtype=dtype, + shape=(*mma_shape, stages), + stride=(*mma_stride, stage_stride if stages > 1 else 0), + swizzle=(0, 4, 3), + byte_alignment=128, + ) + + +__all__ = [ + "Tcgen05MmaInstruction", + "Tcgen05TmemPlan", + "make_smem_layouts", + "make_tcgen05_tmem_plan", + "tcgen05_block_scaled_acc_dtype", + "tcgen05_smem_alloc_type", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py new file mode 100644 index 000000000..0e75bc695 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py @@ -0,0 +1,277 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Declarative GMEM workspace shared by host layout and device access.""" + +import dataclasses +from typing import Any, Dict, List, Literal, Optional, Tuple, Type + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 + +from .utils import ( + ceil_div, + cosize_from_shape_stride_tuples, + flatten_shape_stride, + is_nested_shape, + is_power_of_two, + ordered_stride, + row_major_stride, + round_up, + validate_static_integer_tuple, +) + + +BufferResetAttr = Literal["data", "zero_on_first_allocate", "tail_reset"] +BufferSpace = Literal["local", "shared"] + +_reset_order = {"tail_reset": 0, "zero_on_first_allocate": 1, "data": 2} + + +def _footprint_from_shape_stride(shape: Tuple, stride: Tuple) -> int: + """Elements a region must own, as opposed to the ones its layout can address.""" + leaf_pairs = flatten_shape_stride(shape, stride) if shape else [] + claimed = max((size * step for size, step in leaf_pairs), default=1) + # Layouts whose leaves overlap can address past what any single leaf tiles, + # so the cosize stays a floor. + return int(max(claimed, cosize_from_shape_stride_tuples(shape, stride))) + + +@dataclasses.dataclass(frozen=True) +class DeviceRegion: + """One typed region in a local or symmetric GMEM workspace.""" + + name: str + dtype: Type[cutlass.Numeric] + shape: Tuple + buffer_space: BufferSpace + stride: Optional[Tuple] = None + mem_order: Optional[Tuple[int, ...]] = None + byte_alignment: int = 16 + reset: BufferResetAttr = "data" + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("A workspace region needs a non-empty name.") + if self.buffer_space not in ("local", "shared"): + raise ValueError(f"Invalid buffer space {self.buffer_space!r}.") + if self.reset not in _reset_order: + raise ValueError(f"Invalid reset policy {self.reset!r}.") + if not is_power_of_two(self.byte_alignment): + raise ValueError(f"Region {self.name!r} alignment must be a positive power of two.") + validate_static_integer_tuple(self.shape, field_name=f"{self.name}.shape") + if self.stride is not None and self.mem_order is not None: + raise ValueError(f"Region {self.name!r} accepts either stride or mem_order, not both.") + if self.stride is None and self.mem_order is None: + if is_nested_shape(self.shape) or len(self.shape) != 1: + raise ValueError(f"Region {self.name!r} needs an explicit layout.") + object.__setattr__(self, "stride", row_major_stride(self.shape)) + if is_nested_shape(self.shape) and self.mem_order is not None: + raise ValueError(f"Nested region {self.name!r} needs an explicit stride.") + if self.stride is not None: + validate_static_integer_tuple(self.stride, field_name=f"{self.name}.stride") + if self.mem_order is not None: + validate_static_integer_tuple(self.mem_order, field_name=f"{self.name}.mem_order") + expected = tuple(range(len(self.shape))) + if tuple(sorted(self.mem_order)) != expected: + raise ValueError( + f"Region {self.name!r} mem_order must be a permutation of {expected}, got {self.mem_order}." + ) + resolved_stride, _ = ordered_stride(self.shape, self.mem_order) + object.__setattr__(self, "stride", resolved_stride) + + +class DeviceWorkspace: + """Single source of truth for GMEM region layout and device pointer derivation.""" + + def __init__(self) -> None: + self._registered: Dict[BufferSpace, List[DeviceRegion]] = {"local": [], "shared": []} + self._region_by_name: Dict[str, DeviceRegion] = {} + self._offset: Dict[str, int] = {} + self._stride: Dict[str, Tuple] = {} + self._cosize: Dict[str, int] = {} + self._nbytes: Dict[str, int] = {} + self._byte_alignment: Dict[str, int] = {} + self._total: Dict[BufferSpace, int] = {"local": 0, "shared": 0} + self._zero_leading: Dict[BufferSpace, int] = {"local": 0, "shared": 0} + self._tail_leading: Dict[BufferSpace, int] = {"local": 0, "shared": 0} + self._base: Dict[BufferSpace, Any] = {"local": None, "shared": None} + self._finalized = False + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "DeviceWorkspace": + return self + + @property + def finalized(self) -> bool: + return self._finalized + + def register( + self, + name: str, + dtype: Type[cutlass.Numeric], + shape: Tuple, + *, + buffer_space: BufferSpace, + stride: Optional[Tuple] = None, + mem_order: Optional[Tuple[int, ...]] = None, + byte_alignment: int = 16, + reset: BufferResetAttr = "data", + ) -> None: + if self._finalized: + raise RuntimeError("Cannot register a region after finalize().") + if name in self._region_by_name or any( + region.name == name for regions in self._registered.values() for region in regions + ): + raise ValueError(f"Duplicate workspace region {name!r}.") + self._registered[buffer_space].append( + DeviceRegion( + name=name, + dtype=dtype, + shape=shape, + buffer_space=buffer_space, + stride=stride, + mem_order=mem_order, + byte_alignment=byte_alignment, + reset=reset, + ) + ) + + def finalize(self) -> None: + if self._finalized: + raise RuntimeError("DeviceWorkspace.finalize() may only be called once.") + for buffer_space in ("local", "shared"): + ordered = sorted(self._registered[buffer_space], key=lambda region: _reset_order[region.reset]) + cursor = 0 + for region_index, region in enumerate(ordered): + stride = region.stride + if stride is None: + raise RuntimeError(f"Region {region.name!r} stride was not resolved.") + cosize = cosize_from_shape_stride_tuples(region.shape, stride) + nbytes = (_footprint_from_shape_stride(region.shape, stride) * int(region.dtype.width) + 7) // 8 + if region.reset == "tail_reset" and ( + region_index == 0 or ordered[region_index - 1].reset != "tail_reset" + ): + cursor = round_up(cursor, 16) + byte_alignment = region.byte_alignment + cursor = round_up(cursor, byte_alignment) + self._region_by_name[region.name] = region + self._offset[region.name] = cursor + self._stride[region.name] = stride + self._cosize[region.name] = cosize + self._nbytes[region.name] = nbytes + self._byte_alignment[region.name] = byte_alignment + cursor += nbytes + is_last_tail_reset_region = region.reset == "tail_reset" and ( + region_index + 1 == len(ordered) or ordered[region_index + 1].reset != "tail_reset" + ) + if is_last_tail_reset_region: + cursor = round_up(cursor, 16) + if region.reset != "data": + self._zero_leading[buffer_space] = cursor + if is_last_tail_reset_region: + self._tail_leading[buffer_space] = cursor + self._total[buffer_space] = round_up(cursor, 16) + self._finalized = True + + def regions(self, buffer_space: BufferSpace) -> Tuple[DeviceRegion, ...]: + return tuple(self._registered[buffer_space]) + + def region(self, name: str) -> DeviceRegion: + self._require_finalized() + return self._region_by_name[name] + + def offset(self, name: str) -> int: + self._require_finalized() + return self._offset[name] + + def stride(self, name: str) -> Tuple: + self._require_finalized() + return self._stride[name] + + def cosize(self, name: str) -> int: + self._require_finalized() + return self._cosize[name] + + def nbytes(self, name: str) -> int: + self._require_finalized() + return self._nbytes[name] + + def byte_alignment(self, name: str) -> int: + self._require_finalized() + return self._byte_alignment[name] + + def total_bytes(self, buffer_space: BufferSpace) -> int: + self._require_finalized() + return self._total[buffer_space] + + @property + def local_and_shared_bytes(self) -> Tuple[int, int]: + self._require_finalized() + return self._total["local"], self._total["shared"] + + def zero_on_allocate_bytes(self, buffer_space: BufferSpace) -> int: + self._require_finalized() + return self._zero_leading[buffer_space] + + def tail_reset_bytes(self, buffer_space: BufferSpace) -> int: + self._require_finalized() + return self._tail_leading[buffer_space] + + @property + def require_zero_workspace_leading_bytes(self) -> Tuple[int, int]: + self._require_finalized() + return self._zero_leading["local"], self._zero_leading["shared"] + + @cute.jit + def assign_device_members( + self, local_workspace: cute.Pointer, shared_workspace: Optional[cute.Pointer] = None + ) -> None: + self._base["local"] = local_workspace + self._base["shared"] = shared_workspace + + def remove_device_members(self) -> None: + self._base = {"local": None, "shared": None} + + @cute.jit + def ptr(self, name: str) -> cute.Pointer: + region = self._region_by_name[name] + base = self._base[region.buffer_space] + address = base.toint() + Int64(self._offset[name]) + return cute.make_ptr(region.dtype, address, AddressSpace.gmem, assumed_align=self._byte_alignment[name]) + + @cute.jit + def tensor(self, name: str) -> cute.Tensor: + region = self._region_by_name[name] + return cute.make_tensor(self.ptr(name), cute.make_layout(region.shape, stride=self._stride[name])) + + @cute.jit + def reset_tail(self, tid: Int32, total_threads: int) -> None: + self.reset_tail_space("local", tid, total_threads) + self.reset_tail_space("shared", tid, total_threads) + + @cute.jit + def reset_tail_space(self, buffer_space: BufferSpace, tid: Int32, total_threads: int) -> None: + num_vectors = self._tail_leading[buffer_space] // 16 + if cutlass.const_expr(num_vectors > 0): + vectors = cute.make_tensor( + cute.make_ptr(Int32, self._base[buffer_space].toint(), AddressSpace.gmem, assumed_align=16), + cute.make_layout((num_vectors, 4), stride=(4, 1)), + ) + zero = cute.make_rmem_tensor((4,), Int32) + for element in cutlass.range_constexpr(4): + zero[element] = Int32(0) + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Int32, num_bits_per_copy=128) + reset_round_count = ceil_div(num_vectors, total_threads) + for reset_round in cutlass.range_constexpr(reset_round_count): + vector_index = Int32(reset_round * total_threads) + tid + if vector_index < Int32(num_vectors): + cute.copy(store_atom, zero, vectors[vector_index, None]) + + def _require_finalized(self) -> None: + if not self._finalized: + raise RuntimeError("DeviceWorkspace must be finalized first.") diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/dsl_helpers.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/dsl_helpers.py new file mode 100644 index 000000000..7221b2954 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/dsl_helpers.py @@ -0,0 +1,238 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""General-purpose CuTe DSL helpers.""" + +from typing import Callable, Literal, Optional + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import Pointer +from cutlass.cutlass_dsl import Boolean, Int16, Int32 + +from .ptx_helpers import nanosleep +from .utils import ceil_div + + +@cute.jit +def smem_exclusive_prefix( + input_tensor: cute.Tensor, + output_tensor: cute.Tensor, + warp_totals: cute.Tensor, + block_thread_count: int, + thread_idx: Int32, + lane_idx: Int32, + warp_idx: Int32, +) -> Int32: + """Compute a CTA-wide exclusive prefix over an Int32 SMEM tensor.""" + num_elements = cute.size(input_tensor) + scan_rows = num_elements // 4 + num_warps = block_thread_count // 32 + input_vectors = cute.make_tensor(input_tensor.iterator, cute.make_layout((scan_rows, 4), stride=(4, 1))) + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128) + + values = cute.make_rmem_tensor((4,), cutlass.Int32) + carry = Int32(0) + for segment in cutlass.range_constexpr(ceil_div(scan_rows, block_thread_count)): + row = Int32(segment * block_thread_count) + thread_idx + if row < Int32(scan_rows): + row_slice = input_vectors[row, None] + pointer = row_slice.iterator + aligned_row_slice = cute.make_tensor( + cute.make_ptr(pointer.dtype, pointer.toint(), pointer.memspace, assumed_align=16), + row_slice.layout, + ) + cute.copy(load_atom, aligned_row_slice, values) + else: + for element in cutlass.range_constexpr(4): + values[element] = Int32(0) + + local_prefix = (Int32(0), values[0], values[0] + values[1], values[0] + values[1] + values[2]) + lane_total = local_prefix[3] + values[3] + inclusive = lane_total + for step_log in cutlass.range_constexpr(5): + step = Int32(1 << step_log) + previous = Int32(cute.arch.shuffle_sync(inclusive, lane_idx - step)) + if lane_idx >= step: + inclusive = inclusive + previous + lane_base = inclusive - lane_total + warp_total = Int32(cute.arch.shuffle_sync(inclusive, Int32(31))) + if lane_idx == Int32(0): + warp_totals[warp_idx] = warp_total + cute.arch.sync_threads() + + region_total = Int32(0) + if lane_idx < Int32(num_warps): + region_total = warp_totals[lane_idx] + inclusive_region = region_total + for step_log in cutlass.range_constexpr(5): + step = Int32(1 << step_log) + previous = Int32(cute.arch.shuffle_sync(inclusive_region, lane_idx - step)) + if lane_idx >= step: + inclusive_region = inclusive_region + previous + warp_base = Int32(cute.arch.shuffle_sync(inclusive_region - region_total, warp_idx)) + segment_total = Int32(cute.arch.shuffle_sync(inclusive_region, Int32(31))) + base = carry + warp_base + lane_base + if row < Int32(scan_rows): + first_element = row * Int32(4) + for element in cutlass.range_constexpr(4): + output_tensor[first_element + Int32(element)] = base + local_prefix[element] + carry = carry + segment_total + cute.arch.sync_threads() + return carry + + +@cute.jit +def mark_alignment(tensor: cute.Tensor, byte_alignment: int) -> cute.Tensor: + pointer = tensor.iterator + return cute.make_tensor( + cute.make_ptr(pointer.dtype, pointer.toint(), pointer.memspace, assumed_align=byte_alignment), tensor.layout + ) + + +@cute.jit +def spin_peek(pointer: Pointer, condition: Callable[[Int32], Boolean], scope: str = "gpu") -> Boolean: + """Perform one acquire load and test its value.""" + value = cute.arch.load(pointer, pointer.dtype, sem="acquire", scope=scope) + return Boolean(condition(value)) + + +@cute.jit +def spin_wait( + pointer: Pointer, + condition: Callable[[Int32], Boolean], + scope: str = "gpu", + sleep_cycles: int = 150, + peek_status: Optional[Boolean] = None, +) -> None: + """Wait until an acquire-loaded value satisfies the condition.""" + wait_required = Boolean(True) + if cutlass.const_expr(peek_status is not None): + wait_required = not peek_status + + if wait_required: + value = cute.arch.load(pointer, pointer.dtype, sem="acquire", scope=scope) + while not condition(value): + nanosleep(sleep_cycles) + value = cute.arch.load(pointer, pointer.dtype, sem="acquire", scope=scope) + + +def _tma_multicast_pattern( + cluster_mn: tuple[int, int], mma_cta_count: int, tensor_role: Literal["a", "b", "sfa", "sfb"] +) -> int: + cluster_m, cluster_n = cluster_mn + if tensor_role in ("a", "sfa"): + return sum(1 << (cluster_n_index * cluster_m) for cluster_n_index in range(cluster_n)) + if tensor_role == "b": + return sum(1 << (cluster_m_index * mma_cta_count) for cluster_m_index in range(cluster_m // mma_cta_count)) + if tensor_role == "sfb": + return (1 << cluster_m) - 1 + raise ValueError(f"Unsupported TMA tensor role {tensor_role!r}.") + + +@cute.jit +def tma_multicast_mask( + preferred_cluster_mn: tuple[int, int], + fallback_cluster_mn: Optional[tuple[int, int]], + cta_coord_in_cluster: cute.Coord, + is_preferred: Optional[Boolean], + is_2cta: bool, + tensor_role: Literal["a", "b", "sfa", "sfb"], +) -> Int16: + """Build a preferred/fallback TMA multicast mask.""" + preferred_m, preferred_n = preferred_cluster_mn + if cutlass.const_expr(preferred_m <= 0 or preferred_n <= 0 or preferred_m * preferred_n > 16): + raise ValueError(f"Invalid preferred cluster shape {preferred_cluster_mn}.") + + mma_cta_count = 2 if cutlass.const_expr(is_2cta) else 1 + if cutlass.const_expr(preferred_m % mma_cta_count != 0): + raise ValueError("Preferred cluster M must be divisible by the MMA CTA count.") + + preferred_pattern = _tma_multicast_pattern(preferred_cluster_mn, mma_cta_count, tensor_role) + cta_m = Int32(cta_coord_in_cluster[0]) + if cutlass.const_expr(fallback_cluster_mn is None): + if cutlass.const_expr(tensor_role in ("a", "sfa")): + result = cute.arch.inline_ptx( + f"shl.b16 {{$w0}}, 0x{preferred_pattern:04x}, {{$r0}};", + write_only_types=[Int16], + read_only_args=[cta_m], + ) + else: + cta_n = Int32(cta_coord_in_cluster[1]) + mma_cta_index = Int32(0) if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) else cta_m % 2 + result = cute.arch.inline_ptx( + "{\n\t" + ".reg .u32 offset;\n\t" + f"mad.lo.u32 offset, {{$r0}}, {preferred_m}, {{$r1}};\n\t" + f"shl.b16 {{$w0}}, 0x{preferred_pattern:04x}, offset;\n\t" + "}", + write_only_types=[Int16], + read_only_args=[cta_n, mma_cta_index], + ) + return Int16(result) + + fallback_m, fallback_n = fallback_cluster_mn + if cutlass.const_expr(fallback_m <= 0 or fallback_n <= 0 or fallback_m * fallback_n > 16): + raise ValueError(f"Invalid fallback cluster shape {fallback_cluster_mn}.") + if cutlass.const_expr(preferred_m % fallback_m != 0 or preferred_n % fallback_n != 0): + raise ValueError("Preferred cluster dimensions must be divisible by fallback dimensions.") + if cutlass.const_expr(fallback_m % mma_cta_count != 0): + raise ValueError("Fallback cluster M must be divisible by the MMA CTA count.") + + fallback_pattern = _tma_multicast_pattern(fallback_cluster_mn, mma_cta_count, tensor_role) + if cutlass.const_expr(preferred_pattern == fallback_pattern): + if cutlass.const_expr(tensor_role in ("a", "sfa")): + result = cute.arch.inline_ptx( + f"shl.b16 {{$w0}}, 0x{preferred_pattern:04x}, {{$r0}};", + write_only_types=[Int16], + read_only_args=[cta_m], + ) + else: + cta_n = Int32(cta_coord_in_cluster[1]) + mma_cta_index = Int32(0) if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) else cta_m % 2 + result = cute.arch.inline_ptx( + "{\n\t" + ".reg .u32 offset;\n\t" + f"mad.lo.u32 offset, {{$r0}}, {preferred_m}, {{$r1}};\n\t" + f"shl.b16 {{$w0}}, 0x{preferred_pattern:04x}, offset;\n\t" + "}", + write_only_types=[Int16], + read_only_args=[cta_n, mma_cta_index], + ) + return Int16(result) + + if cutlass.const_expr(is_preferred is None): + raise ValueError("is_preferred is required when the preferred and fallback multicast patterns differ.") + + if cutlass.const_expr(tensor_role in ("a", "sfa")): + result = cute.arch.inline_ptx( + "{\n\t" + f"mov.b16 {{$w0}}, 0x{preferred_pattern:04x};\n\t" + f"@!{{$r0}} mov.b16 {{$w0}}, 0x{fallback_pattern:04x};\n\t" + "shl.b16 {$w0}, {$w0}, {$r1};\n\t" + "}", + write_only_types=[Int16], + read_only_args=[is_preferred, cta_m], + ) + else: + if cutlass.const_expr(fallback_pattern & preferred_pattern != fallback_pattern): + raise ValueError("Fallback B/SFB multicast pattern must be a subset of the preferred pattern.") + cta_n = Int32(cta_coord_in_cluster[1]) + mma_cta_index = Int32(0) if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) else cta_m % 2 + result = cute.arch.inline_ptx( + "{\n\t" + ".reg .u32 cluster_m, offset;\n\t" + f"mov.u32 cluster_m, {preferred_m};\n\t" + f"@!{{$r0}} mov.u32 cluster_m, {fallback_m};\n\t" + "mad.lo.u32 offset, {$r1}, cluster_m, {$r2};\n\t" + f"mov.b16 {{$w0}}, 0x{preferred_pattern:04x};\n\t" + f"@!{{$r0}} and.b16 {{$w0}}, {{$w0}}, 0x{fallback_pattern:04x};\n\t" + "shl.b16 {$w0}, {$w0}, offset;\n\t" + "}", + write_only_types=[Int16], + read_only_args=[is_preferred, cta_n, mma_cta_index], + ) + return Int16(result) + + +__all__ = ["mark_alignment", "smem_exclusive_prefix", "spin_peek", "spin_wait", "tma_multicast_mask"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/flag_batch.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/flag_batch.py new file mode 100644 index 000000000..e69d60d29 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/flag_batch.py @@ -0,0 +1,136 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Lane- and warp-distributed GPU release-counter batching.""" + +import dataclasses +from typing import Any, Union + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 + +from .ptx_helpers import red_add_release_gpu_s32, red_async_add_release_gpu_s32 + + +@dataclasses.dataclass(frozen=True) +class GpuReleaseFlagBatchTracker: + """Lane-distributed delayed publication state for synchronous GPU release counters.""" + + flag_address: Int64 + accumulated_flags: Int32 + phase: Int32 + thread_idx: Int32 + + @cute.jit + def _make(self, flag_address: Int64, accumulated_flags: Int32, phase: Int32) -> "GpuReleaseFlagBatchTracker": + return type(self)( + flag_address=flag_address, accumulated_flags=accumulated_flags, phase=phase, thread_idx=self.thread_idx + ) + + @cute.jit + def fire(self) -> None: + if self.flag_address != Int64(0): + pointer = cute.make_ptr(cutlass.Int32, self.flag_address, AddressSpace.gmem, assumed_align=4) + red_add_release_gpu_s32(pointer, Int32(1)) + + @cute.jit + def accumulate( + self, next_phase: Any, flush_threshold: int, flag_address: Int64, no_fire: bool = False + ) -> "GpuReleaseFlagBatchTracker": + if cutlass.const_expr(flush_threshold == 1): + if cutlass.const_expr(not no_fire): + lane_address = Int64(0) + if self.thread_idx == Int32(0): + lane_address = flag_address + self._make(flag_address=lane_address, accumulated_flags=Int32(1), phase=self.phase).fire() + return self._make(flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(next_phase)) + + current_address = self.flag_address + accumulated_flags = self.accumulated_flags + if self.thread_idx == accumulated_flags: + current_address = flag_address + accumulated_flags = accumulated_flags + Int32(1) + + if accumulated_flags == Int32(flush_threshold) or next_phase != self.phase: + if cutlass.const_expr(not no_fire): + self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=self.phase).fire() + accumulated_flags = Int32(0) + current_address = Int64(0) + + return self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=Int32(next_phase)) + + +@dataclasses.dataclass(frozen=True) +class GpuAsyncReleaseFlagBatchTracker: + """Loop-carried warp-uniform state for asynchronous GPU release counters.""" + + flag_address: Int64 + accumulated_flags: Int32 + phase: Int32 + warp_idx: Int32 + + @cute.jit + def _make(self, flag_address: Int64, accumulated_flags: Int32, phase: Int32) -> "GpuAsyncReleaseFlagBatchTracker": + return type(self)( + flag_address=flag_address, accumulated_flags=accumulated_flags, phase=phase, warp_idx=self.warp_idx + ) + + @cute.jit + def fire(self) -> None: + if self.flag_address != Int64(0): + pointer = cute.make_ptr(cutlass.Int32, self.flag_address, AddressSpace.gmem, assumed_align=4) + with cute.arch.elect_one(): + red_async_add_release_gpu_s32(pointer, Int32(1)) + + @cute.jit + def accumulate( + self, next_phase: Any, flush_threshold: int, flag_address: Int64, no_fire: bool = False + ) -> "GpuAsyncReleaseFlagBatchTracker": + if cutlass.const_expr(flush_threshold == 1): + if cutlass.const_expr(not no_fire): + warp_address = Int64(0) + if self.warp_idx == Int32(0): + warp_address = flag_address + self._make(flag_address=warp_address, accumulated_flags=Int32(1), phase=self.phase).fire() + return self._make(flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(next_phase)) + + current_address = self.flag_address + accumulated_flags = self.accumulated_flags + if self.warp_idx == accumulated_flags: + current_address = flag_address + accumulated_flags = accumulated_flags + Int32(1) + + if accumulated_flags == Int32(flush_threshold) or next_phase != self.phase: + if cutlass.const_expr(not no_fire): + self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=self.phase).fire() + accumulated_flags = Int32(0) + current_address = Int64(0) + + return self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=Int32(next_phase)) + + +@cute.jit +def make_flag_batch_tracker( + use_async: bool, *, flag_address: Int64, accumulated_flags: Int32, phase: Int32, thread_idx: Int32 +) -> Union[GpuReleaseFlagBatchTracker, GpuAsyncReleaseFlagBatchTracker]: + """Construct a lane-batched synchronous or warp-batched asynchronous tracker. + + ``thread_idx`` must be zero-based within the caller's cooperating publisher + group. The async tracker maps each contiguous group of 32 thread indices to + one warp-uniform publisher. + """ + if cutlass.const_expr(use_async): + return GpuAsyncReleaseFlagBatchTracker( + flag_address=flag_address, + accumulated_flags=accumulated_flags, + phase=phase, + warp_idx=cute.arch.make_warp_uniform(thread_idx // Int32(32)), + ) + return GpuReleaseFlagBatchTracker( + flag_address=flag_address, accumulated_flags=accumulated_flags, phase=phase, thread_idx=thread_idx + ) + + +__all__ = ["GpuAsyncReleaseFlagBatchTracker", "GpuReleaseFlagBatchTracker", "make_flag_batch_tracker"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/iket_compat.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/iket_compat.py new file mode 100644 index 000000000..a076067ae --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/iket_compat.py @@ -0,0 +1,38 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Compatibility wrapper for optional in-kernel event tracing support.""" + +try: + from cutlass.cute.experimental import iket +except (ImportError, NotImplementedError): + try: + from cutlass.cute import iket # type: ignore + except (ImportError, NotImplementedError): + class _IketShim: + """No-op IKET interface for toolchains without the dialect.""" + + @staticmethod + def range_push(_name, *_args, **_kwargs): + return None + + @staticmethod + def range_pop(*_args, **_kwargs): + return None + + @staticmethod + def range_start(_name, *_args, **_kwargs): + return None + + @staticmethod + def range_end(_token=None, *_args, **_kwargs): + return None + + @staticmethod + def mark(_name, *_args, **_kwargs): + return None + + iket = _IketShim() # type: ignore + + +__all__ = ["iket"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py new file mode 100644 index 000000000..ebf782ca0 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py @@ -0,0 +1,621 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Minimal inline-PTX primitives required by the greenfield NVFP4 kernels.""" + +from typing import Optional + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith, llvm, vector +from cutlass.cutlass_dsl import Float32, Int32, Int64, T, dsl_user_op + + +TmaCacheHintEvictFirst = 0x12F0000000000000 + + +def _address_value(pointer_or_address, *, loc=None, ip=None): + if isinstance(pointer_or_address, Int64): + return pointer_or_address.ir_value() + return pointer_or_address.toint(loc=loc, ip=ip).ir_value() + + +@dsl_user_op +def nanosleep(sleep_cycles: int, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> None: + """Suspend the calling thread for up to the requested clock cycles.""" + if cutlass.const_expr(hasattr(cute.arch, "nanosleep")): + cute.arch.nanosleep(sleep_time=sleep_cycles, loc=loc, ip=ip) + return + + llvm.inline_asm( + res=None, + operands_=[Int32(sleep_cycles).ir_value(loc=loc, ip=ip)], + asm_string="nanosleep.u32 $0;", + constraints="r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def exit(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> None: + llvm.inline_asm( + res=None, + operands_=[], + asm_string="exit;", + constraints="", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def read_clock64(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> Int64: + """Read the per-SM 64-bit cycle counter.""" + return Int64( + llvm.inline_asm( + T.i64(), [], "mov.u64 $0, %clock64;", "=l", has_side_effects=True, asm_dialect=0, loc=loc, ip=ip + ) + ) + + +@cute.jit +def movmatrix_b16(input_regs: cute.Tensor) -> cute.Tensor: + """Transpose every packed m8n8 b16 register fragment across the warp.""" + if cutlass.const_expr(input_regs.element_type.width != 32): + raise TypeError(f"movmatrix_b16 expects packed 32-bit registers, got {input_regs.element_type}.") + + input_words = cute.coalesce(cute.flatten(cute.recast_tensor(input_regs, Int32))) + output_words = cute.make_rmem_tensor((cute.size(input_words),), Int32) + for word_idx in cutlass.range_constexpr(cute.size(input_words)): + output_words[word_idx] = Int32( + llvm.inline_asm( + T.i32(), + [Int32(input_words[word_idx]).ir_value()], + "movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;", + "=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + output_regs = cute.make_rmem_tensor(input_regs.layout, input_regs.element_type) + output_regs_words = cute.coalesce(cute.flatten(cute.recast_tensor(output_regs, Int32))) + for word_idx in cutlass.range_constexpr(cute.size(output_words)): + output_regs_words[word_idx] = output_words[word_idx] + return output_regs + + +@dsl_user_op +def cvt_f32_to_fp8_to_f32( + value, fp8_type, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> Float32: + """Round one f32 through the selected FP8 format and widen it back.""" + if cutlass.const_expr(fp8_type is cutlass.Float8E8M0FNU): + downcast_instruction = "cvt.rp.satfinite.ue8m0x2.f32" + upcast_instruction = "cvt.rn.bf16x2.ue8m0x2" + elif cutlass.const_expr(fp8_type is cutlass.Float8E4M3FN): + downcast_instruction = "cvt.rn.satfinite.e4m3x2.f32" + upcast_instruction = "cvt.rn.bf16x2.e4m3x2" + elif cutlass.const_expr(fp8_type is cutlass.Float8E5M2): + downcast_instruction = "cvt.rn.satfinite.e5m2x2.f32" + upcast_instruction = "cvt.rn.bf16x2.e5m2x2" + else: + raise ValueError(f"Unsupported FP8 type {fp8_type}.") + + packed_bf16 = llvm.inline_asm( + T.i32(), + [Float32(value).ir_value(loc=loc, ip=ip)], + "{\n" + " .reg .b16 converted;\n" + f" {downcast_instruction} converted, 0f00000000, $1;\n" + f" {upcast_instruction} $0, converted;\n" + "}", + "=r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + bf16_pair_type = ir.Type.parse("vector<2xbf16>") + bf16_pair = llvm.bitcast(bf16_pair_type, packed_bf16, loc=loc, ip=ip) + rounded_bf16 = vector.extract(bf16_pair, [], [0], loc=loc, ip=ip) + return Float32(arith.extf(Float32.mlir_type, rounded_bf16, loc=loc, ip=ip)) + + +@dsl_user_op +def tma_load_1d(destination_smem, source_gmem, mbarrier_smem, num_bytes, *, loc=None, ip=None) -> None: + """Issue a cache-hinted 1D GMEM-to-SMEM bulk copy.""" + llvm.inline_asm( + None, + [ + destination_smem.toint(loc=loc, ip=ip).ir_value(), + _address_value(source_gmem, loc=loc, ip=ip), + num_bytes.ir_value(), + mbarrier_smem.toint(loc=loc, ip=ip).ir_value(), + Int64(TmaCacheHintEvictFirst).ir_value(), + ], + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [$0], [$1], $2, [$3], $4;", + "r,l,r,r,l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_i32_to_peer_cluster_smem_async( + smem_pointer, + value: Int32, + mbarrier_pointer, + destination_cta_rank, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + """Store one Int32 to peer SMEM and complete its transaction barrier.""" + smem_address = llvm.ptrtoint(T.i32(), smem_pointer.llvm_ptr, loc=loc, ip=ip) + mbarrier_address = llvm.ptrtoint(T.i32(), mbarrier_pointer.llvm_ptr, loc=loc, ip=ip) + llvm.inline_asm( + res=None, + operands_=[ + smem_address, + value.ir_value(loc=loc, ip=ip), + mbarrier_address, + Int32(destination_cta_rank).ir_value(loc=loc, ip=ip), + ], + asm_string="""{{ + .reg .u32 remote_addr; + .reg .u32 remote_mbar; + mapa.shared::cluster.u32 remote_addr, $0, $3; + mapa.shared::cluster.u32 remote_mbar, $2, $3; + st.async.shared::cluster.mbarrier::complete_tx::bytes.u32 [remote_addr], $1, [remote_mbar]; + }}""", + constraints="r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def mbarrier_arrive_expect_tx_on_peer( + mbarrier_pointer, + transaction_bytes: Int32, + destination_cta_rank, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + """Declare an expected peer-CTA SMEM transaction.""" + mbarrier_address = llvm.ptrtoint(T.i32(), mbarrier_pointer.llvm_ptr, loc=loc, ip=ip) + llvm.inline_asm( + res=None, + operands_=[ + mbarrier_address, + Int32(destination_cta_rank).ir_value(loc=loc, ip=ip), + transaction_bytes.ir_value(loc=loc, ip=ip), + ], + asm_string="""{{ + .reg .u32 remote_mbar; + mapa.shared::cluster.u32 remote_mbar, $0, $1; + mbarrier.arrive.expect_tx.shared::cluster.b64 _, [remote_mbar], $2; + }}""", + constraints="r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cp_async_bulk_s2g( + destination_gmem, + source_smem, + num_bytes, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + llvm.inline_asm( + None, + [ + destination_gmem.toint(loc=loc, ip=ip).ir_value(), + source_smem.toint(loc=loc, ip=ip).ir_value(), + num_bytes.ir_value(), + ], + "cp.async.bulk.global.shared::cta.bulk_group [$0], [$1], $2;", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cp_reduce_async_bulk_add_bf16_s2g( + destination_gmem, + source_smem, + num_bytes, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + llvm.inline_asm( + None, + [ + destination_gmem.toint(loc=loc, ip=ip).ir_value(), + source_smem.toint(loc=loc, ip=ip).ir_value(), + num_bytes.ir_value(), + ], + "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.noftz.bf16 [$0], [$1], $2;", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cp_reduce_async_bulk_add_u32_s2g( + destination_gmem, + source_smem, + num_bytes, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + llvm.inline_asm( + None, + [ + destination_gmem.toint(loc=loc, ip=ip).ir_value(), + source_smem.toint(loc=loc, ip=ip).ir_value(), + num_bytes.ir_value(), + ], + "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.u32 [$0], [$1], $2;", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def lds128_v4_b32(smem_pointer, *, loc=None, ip=None): + result = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 4), + [smem_pointer.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return tuple(Int32(llvm.extractvalue(T.i32(), result, [index])) for index in range(4)) + + +@dsl_user_op +def stg_f32(address: Int64, value: Float32, predicate: Optional[Int32] = None, *, loc=None, ip=None) -> None: + if predicate is None: + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value()], + "st.global.f32 [$0], $1;", + "l,f", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value(), predicate.ir_value()], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.f32 [$0], $1;\n\t}", + "l,f,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stg_b64(address: Int64, value: Int64, predicate: Optional[Int32] = None, *, loc=None, ip=None) -> None: + if predicate is None: + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value()], + "st.global.u64 [$0], $1;", + "l,l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value(), predicate.ir_value()], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.u64 [$0], $1;\n\t}", + "l,l,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add_relaxed_sys_s32(address: Int64, value: Int32, *, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value()], + "red.relaxed.sys.global.add.s32 [$0], $1;", + "l,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add_relaxed_sys_f32( + address, value: Float32, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> None: + llvm.inline_asm( + None, + [address.toint(loc=loc, ip=ip).ir_value(), value.ir_value(loc=loc, ip=ip)], + "red.relaxed.sys.global.add.f32 [$0], $1;", + "l,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add_release_sys_s32(address: Int64, value: Int32, *, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value()], + "red.release.sys.global.add.s32 [$0], $1;", + "l,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add_release_gpu_s32( + counter_pointer, value: Int32, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> None: + llvm.inline_asm( + None, + [counter_pointer.toint(loc=loc, ip=ip).ir_value(), value.ir_value()], + "red.release.gpu.global.add.s32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_async_add_release_gpu_s32( + counter_pointer, value: Int32, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> None: + llvm.inline_asm( + None, + [counter_pointer.toint(loc=loc, ip=ip).ir_value(), value.ir_value()], + "red.async.release.gpu.global.add.s32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_async_add_release_sys_u32(address: Int64, value: Int32, *, loc=None, ip=None) -> None: + """Fire-and-forget cross-rank counter bump carrying its own release ordering. + + The issuing warp does not wait for the L2/HBM round trip, and the release is + enforced by the memory system, so this replaces an explicit ``membar.sys`` + followed by a relaxed reduction. ``u32`` rather than ``s32`` only because a + counter bump is the same two's-complement add either way and this is the + spelling already proven on this path. + + Operands go through uniform registers, so the address must be warp-uniform: + a fan-out where lanes target different peers cannot use this. + """ + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value()], + "red.async.release.sys.global.add.u32 [$0], $1;", + "l,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add_relaxed_sys_v2_bf16x2( + address, value0, value1, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> None: + llvm.inline_asm( + None, + [address.toint(loc=loc, ip=ip).ir_value(), value0.ir_value(), value1.ir_value()], + "red.relaxed.sys.global.add.noftz.v2.bf16x2 [$0], {$1, $2};", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def cvt_f32x4_to_f8x4_pack_i32(fp32x4: cute.Tensor, fp8_type, *, loc=None, ip=None) -> Int32: + """Round four f32 lanes to the selected FP8 format and pack them into one i32.""" + fp32x4 = fp32x4.load() + src_vec4 = fp32x4.ir_value(loc=loc, ip=ip) if hasattr(fp32x4, "ir_value") else fp32x4 + + src0 = Float32(vector.extract(src_vec4, [], [0])).ir_value(loc=loc, ip=ip) + src1 = Float32(vector.extract(src_vec4, [], [1])).ir_value(loc=loc, ip=ip) + src2 = Float32(vector.extract(src_vec4, [], [2])).ir_value(loc=loc, ip=ip) + src3 = Float32(vector.extract(src_vec4, [], [3])).ir_value(loc=loc, ip=ip) + + if cutlass.const_expr(fp8_type is cutlass.Float8E8M0FNU): + cvt_instruction = "cvt.rp.satfinite.ue8m0x2.f32" + elif cutlass.const_expr(fp8_type is cutlass.Float8E4M3FN): + cvt_instruction = "cvt.rn.satfinite.e4m3x2.f32" + elif cutlass.const_expr(fp8_type is cutlass.Float8E5M2): + cvt_instruction = "cvt.rn.satfinite.e5m2x2.f32" + else: + raise ValueError(f"Unsupported FP8 type {fp8_type}.") + + packed_i32 = llvm.inline_asm( + T.i32(), + [src0, src1, src2, src3], + "{\n" + " .reg .b16 lo;\n" + " .reg .b16 hi;\n" + f" {cvt_instruction} lo, $2, $1;\n" + f" {cvt_instruction} hi, $4, $3;\n" + " mov.b32 $0, {lo, hi};\n" + "}", + "=r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return Int32(packed_i32) + + +@dsl_user_op +def stg_e8m0_from_f32(addr: Int64, fp32_val: Float32, *, loc=None, ip=None) -> None: + """Convert ``fp32_val`` to E8M0 via PTX and store the 1-byte result to global memory. + + Uses ``cvt.rp.satfinite.ue8m0x2.f32`` -- the correct fp32 -> E8M0 path -- rather + than the DSL's generic ``.to(Float8E8M0FNU)``, which does not lower correctly for + the non-IEEE-754 E8M0 type. + """ + llvm.inline_asm( + None, + [addr.ir_value(), fp32_val.ir_value()], + "{\n" + " .reg .b16 bf_lo;\n" + " .reg .u32 tmp;\n" + " cvt.rp.satfinite.ue8m0x2.f32 bf_lo, 0f00000000, $1;\n" + " cvt.u32.u16 tmp, bf_lo;\n" + " st.global.b8 [$0], tmp;\n" + "}", + "l,f", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stg_e8m0x8_from_f32( + addr: Int64, + v0: Float32, v1: Float32, v2: Float32, v3: Float32, + v4: Float32, v5: Float32, v6: Float32, v7: Float32, + *, loc=None, ip=None, +) -> None: + """Convert 8 fp32 values to E8M0 and store them as 8 contiguous bytes in one shot. + + Batched form of ``stg_e8m0_from_f32``: four ``cvt.rp.satfinite.ue8m0x2.f32`` each + pack two E8M0 bytes, the four ``.b16`` results are assembled into two ``.b32`` words, + and a single ``st.global.v2.u32`` writes all 8 bytes. ``addr`` must be 8-byte aligned; + output byte ``k`` holds E8M0(``v{k}``). + """ + llvm.inline_asm( + None, + [ + addr.ir_value(), + v0.ir_value(), v1.ir_value(), v2.ir_value(), v3.ir_value(), + v4.ir_value(), v5.ir_value(), v6.ir_value(), v7.ir_value(), + ], + "{\n" + " .reg .b16 p0, p1, p2, p3;\n" + " .reg .b32 w0, w1;\n" + " cvt.rp.satfinite.ue8m0x2.f32 p0, $2, $1;\n" + " cvt.rp.satfinite.ue8m0x2.f32 p1, $4, $3;\n" + " cvt.rp.satfinite.ue8m0x2.f32 p2, $6, $5;\n" + " cvt.rp.satfinite.ue8m0x2.f32 p3, $8, $7;\n" + " mov.b32 w0, {p0, p1};\n" + " mov.b32 w1, {p2, p3};\n" + " st.global.v2.u32 [$0], {w0, w1};\n" + "}", + "l,f,f,f,f,f,f,f,f", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +__all__ = [ + "TmaCacheHintEvictFirst", + "cp_async_bulk_s2g", + "cp_reduce_async_bulk_add_bf16_s2g", + "cp_reduce_async_bulk_add_u32_s2g", + "cvt_f32_to_fp8_to_f32", + "cvt_f32x4_to_f8x4_pack_i32", + "exit", + "lds128_v4_b32", + "mbarrier_arrive_expect_tx_on_peer", + "movmatrix_b16", + "nanosleep", + "read_clock64", + "red_add_relaxed_sys_f32", + "red_add_relaxed_sys_s32", + "red_add_relaxed_sys_v2_bf16x2", + "red_async_add_release_gpu_s32", + "red_async_add_release_sys_u32", + "red_add_release_gpu_s32", + "red_add_release_sys_s32", + "store_i32_to_peer_cluster_smem_async", + "stg_b64", + "stg_e8m0_from_f32", + "stg_e8m0x8_from_f32", + "stg_f32", + "tma_load_1d", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py new file mode 100644 index 000000000..db5cdb729 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py @@ -0,0 +1,431 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Pure-Python SMEM declarations with lifetime-aware overlay placement.""" + +import dataclasses +from typing import Dict, List, Literal, Optional, Tuple, Type, Union + +import cutlass +import cutlass.cute as cute + +from .utils import ( + cosize_from_shape_stride_tuples, + is_power_of_two, + row_major_stride, + round_up, + validate_static_integer_tuple, +) + + +SmemRegionKind = Literal["mbarrier", "tensor"] +SwizzleSpec = Tuple[int, int, int] + + +def _swizzle_alignment(swizzle: Optional[SwizzleSpec]) -> Optional[int]: + if swizzle is None or swizzle[0] == 0: + return None + num_bits, num_base, _ = swizzle + return 1 << (num_base + num_bits) + + +@dataclasses.dataclass(frozen=True) +class SmemRegion: + """One logical SMEM tensor with no assigned byte offset.""" + + name: str + kind: SmemRegionKind + dtype: Type[cutlass.Numeric] + shape: Tuple + stride: Tuple + swizzle: Optional[SwizzleSpec] + byte_alignment: int + + @property + def cosize(self) -> int: + return int(cosize_from_shape_stride_tuples(self.shape, self.stride)) + + @property + def nbytes(self) -> int: + return (self.cosize * int(self.dtype.width) + 7) // 8 + + +class SmemLifetime: + """One mutually exclusive use of an overlay's physical storage.""" + + def __init__( + self, + workspace: "SmemWorkspace", + overlay: "SmemOverlay", + name: str, + ) -> None: + self._workspace = workspace + self._overlay = overlay + self.name = name + self._regions: List[SmemRegion] = [] + + @property + def regions(self) -> Tuple[SmemRegion, ...]: + return tuple(self._regions) + + def register_tensor( + self, + name: str, + dtype: Type[cutlass.Numeric], + shape: Tuple, + *, + stride: Optional[Tuple] = None, + swizzle: Optional[SwizzleSpec] = None, + byte_alignment: Optional[int] = None, + ) -> SmemRegion: + region = self._workspace._make_region( + name=name, + kind="tensor", + dtype=dtype, + shape=shape, + stride=stride, + swizzle=swizzle, + byte_alignment=byte_alignment, + ) + self._workspace._claim_region_name(region) + self._regions.append(region) + return region + + +class SmemOverlay: + """One physical allocation shared by mutually exclusive lifetimes.""" + + def __init__(self, workspace: "SmemWorkspace", name: str) -> None: + self._workspace = workspace + self.name = name + self._lifetimes: List[SmemLifetime] = [] + self._lifetime_names: set[str] = set() + + @property + def lifetimes(self) -> Tuple[SmemLifetime, ...]: + return tuple(self._lifetimes) + + def add_lifetime(self, name: str) -> SmemLifetime: + if self._workspace.finalized: + raise RuntimeError("Cannot add an SMEM lifetime after finalize().") + if not name: + raise ValueError("An SMEM lifetime needs a non-empty name.") + if name in self._lifetime_names: + raise ValueError( + f"Duplicate lifetime {name!r} in overlay {self.name!r}." + ) + lifetime = SmemLifetime(self._workspace, self, name) + self._lifetimes.append(lifetime) + self._lifetime_names.add(name) + return lifetime + + +_TopLevelDeclaration = Union[SmemRegion, SmemOverlay] + + +class SmemWorkspace: + """Collect static SMEM declarations and finalize one physical placement.""" + + def __init__( + self, + *, + base_alignment: int = 1024, + total_alignment: int = 16, + ) -> None: + if not is_power_of_two(base_alignment): + raise ValueError("base_alignment must be a positive power of two.") + if not is_power_of_two(total_alignment): + raise ValueError("total_alignment must be a positive power of two.") + self.base_alignment = base_alignment + self.total_alignment = total_alignment + self._mbarriers: List[SmemRegion] = [] + self._declarations: List[_TopLevelDeclaration] = [] + self._region_by_name: Dict[str, SmemRegion] = {} + self._overlay_names: set[str] = set() + self._offset: Dict[str, int] = {} + self._total_bytes = 0 + self._finalized = False + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "SmemWorkspace": + return self + + @property + def finalized(self) -> bool: + return self._finalized + + @property + def total_bytes(self) -> int: + self._require_finalized() + return self._total_bytes + + def register_mbarrier( + self, + name: str, + count: int, + *, + byte_alignment: Optional[int] = None, + ) -> SmemRegion: + if count <= 0: + raise ValueError(f"Mbarrier region {name!r} needs a positive count.") + region = self._make_region( + name=name, + kind="mbarrier", + dtype=cutlass.Int64, + shape=(count,), + stride=(1,), + swizzle=None, + byte_alignment=byte_alignment, + ) + self._claim_region_name(region) + self._mbarriers.append(region) + return region + + def register_tensor( + self, + name: str, + dtype: Type[cutlass.Numeric], + shape: Tuple, + *, + stride: Optional[Tuple] = None, + swizzle: Optional[SwizzleSpec] = None, + byte_alignment: Optional[int] = None, + ) -> SmemRegion: + region = self._make_region( + name=name, + kind="tensor", + dtype=dtype, + shape=shape, + stride=stride, + swizzle=swizzle, + byte_alignment=byte_alignment, + ) + self._claim_region_name(region) + self._declarations.append(region) + return region + + def create_overlay(self, name: str) -> SmemOverlay: + if self._finalized: + raise RuntimeError("Cannot create an SMEM overlay after finalize().") + if not name: + raise ValueError("An SMEM overlay needs a non-empty name.") + if name in self._overlay_names or name in self._region_by_name: + raise ValueError(f"Duplicate SMEM overlay {name!r}.") + overlay = SmemOverlay(self, name) + self._overlay_names.add(name) + self._declarations.append(overlay) + return overlay + + def _make_region( + self, + *, + name: str, + kind: SmemRegionKind, + dtype: Type[cutlass.Numeric], + shape: Tuple, + stride: Optional[Tuple], + swizzle: Optional[SwizzleSpec], + byte_alignment: Optional[int], + ) -> SmemRegion: + if self._finalized: + raise RuntimeError("Cannot register an SMEM region after finalize().") + if not name: + raise ValueError("An SMEM region needs a non-empty name.") + validate_static_integer_tuple(shape, field_name=f"{name}.shape") + if stride is None: + stride = row_major_stride(shape) + validate_static_integer_tuple(stride, field_name=f"{name}.stride") + if len(shape) != len(stride): + raise ValueError( + f"SMEM region {name!r} shape and stride ranks differ." + ) + if swizzle is not None: + if len(swizzle) != 3 or not all( + isinstance(parameter, int) for parameter in swizzle + ): + raise TypeError( + f"SMEM region {name!r} swizzle must be three Python ints." + ) + if swizzle[0] < 0 or swizzle[1] < 0: + raise ValueError( + f"SMEM region {name!r} swizzle bits/base must be non-negative." + ) + + natural_alignment = max(1, (int(dtype.width) + 7) // 8) + explicit_alignment = ( + natural_alignment if byte_alignment is None else byte_alignment + ) + if not is_power_of_two(explicit_alignment): + raise ValueError( + f"SMEM region {name!r} alignment must be a positive power of two." + ) + swizzle_alignment = _swizzle_alignment(swizzle) + effective_alignment = max( + explicit_alignment, + natural_alignment, + 1 if swizzle_alignment is None else swizzle_alignment, + ) + if effective_alignment > self.base_alignment: + raise ValueError( + f"SMEM region {name!r} needs {effective_alignment}B alignment, " + f"but the workspace base only promises {self.base_alignment}B." + ) + return SmemRegion( + name=name, + kind=kind, + dtype=dtype, + shape=shape, + stride=stride, + swizzle=swizzle, + byte_alignment=effective_alignment, + ) + + def _claim_region_name(self, region: SmemRegion) -> None: + if ( + region.name in self._region_by_name + or region.name in self._overlay_names + ): + raise ValueError(f"Duplicate SMEM region {region.name!r}.") + self._region_by_name[region.name] = region + + def estimate_total_bytes(self) -> int: + """Exact size of what is registered so far, priced by running the real placement. + + ``_build_placement`` is side-effect free -- ``finalize`` is what stores its result -- so the layout can be + measured without consuming the workspace. Summing region sizes instead would miss the alignment padding + that only exists once regions are placed, and a budget derived from that undercount overspends the + workspace: ``finalize`` then rejects a plan the host arithmetic had already accepted. + """ + return self._build_placement()[1] + + def finalize(self, *, max_bytes: Optional[int] = None) -> None: + if self._finalized: + raise RuntimeError("SmemWorkspace.finalize() may only be called once.") + offsets, total_bytes = self._build_placement() + if max_bytes is not None and total_bytes > max_bytes: + raise ValueError( + f"SMEM plan needs {total_bytes} bytes, exceeding {max_bytes} bytes." + ) + self._offset = offsets + self._total_bytes = total_bytes + self._finalized = True + + def _build_placement(self) -> Tuple[Dict[str, int], int]: + offsets: Dict[str, int] = {} + cursor = 0 + for mbarrier in self._mbarriers: + cursor = round_up(cursor, mbarrier.byte_alignment) + offsets[mbarrier.name] = cursor + cursor += mbarrier.nbytes + for declaration in self._declarations: + if isinstance(declaration, SmemRegion): + cursor = round_up(cursor, declaration.byte_alignment) + offsets[declaration.name] = cursor + cursor += declaration.nbytes + continue + relative_offsets, overlay_alignment, overlay_bytes = ( + self._layout_overlay(declaration) + ) + cursor = round_up(cursor, overlay_alignment) + for region_name, relative_offset in relative_offsets.items(): + offsets[region_name] = cursor + relative_offset + cursor += overlay_bytes + return offsets, int(round_up(cursor, self.total_alignment)) + + def _layout_overlay( + self, + overlay: SmemOverlay, + ) -> Tuple[Dict[str, int], int, int]: + if not overlay.lifetimes: + raise ValueError(f"SMEM overlay {overlay.name!r} has no lifetimes.") + relative_offsets: Dict[str, int] = {} + overlay_alignment = 1 + overlay_bytes = 0 + for lifetime in overlay.lifetimes: + if not lifetime.regions: + raise ValueError( + f"SMEM lifetime {overlay.name}.{lifetime.name} has no regions." + ) + lifetime_cursor = 0 + for region in lifetime.regions: + overlay_alignment = max( + overlay_alignment, + region.byte_alignment, + ) + lifetime_cursor = round_up( + lifetime_cursor, + region.byte_alignment, + ) + relative_offsets[region.name] = lifetime_cursor + lifetime_cursor += region.nbytes + overlay_bytes = max(overlay_bytes, lifetime_cursor) + return relative_offsets, overlay_alignment, overlay_bytes + + def regions(self) -> Tuple[SmemRegion, ...]: + return tuple(self._region_by_name.values()) + + def region(self, name: str) -> SmemRegion: + return self._region_by_name[name] + + def offset(self, name: str) -> int: + self._require_finalized() + return self._offset[name] + + def nbytes(self, name: str) -> int: + return self._region_by_name[name].nbytes + + def byte_alignment(self, name: str) -> int: + return self._region_by_name[name].byte_alignment + + def storage_class(self) -> type: + self._require_finalized() + storage_bytes = max(self._total_bytes, 1) + base_alignment = self.base_alignment + + @cute.struct + class SmemStorage: + buffer: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int8, storage_bytes], + base_alignment, + ] + + return SmemStorage + + @cute.jit + def ptr(self, name: str, smem_base: cute.Pointer) -> cute.Pointer: + region = self._region_by_name[name] + swizzle = ( + None + if region.swizzle is None + else cute.make_swizzle(*region.swizzle) + ) + return cute.make_ptr( + region.dtype, + smem_base.toint() + self._offset[name], + smem_base.memspace, + assumed_align=region.byte_alignment, + swizzle_=swizzle, + ) + + @cute.jit + def tensor(self, name: str, smem_base: cute.Pointer) -> cute.Tensor: + region = self._region_by_name[name] + layout = cute.make_layout(region.shape, stride=region.stride) + return cute.make_tensor(self.ptr(name, smem_base), layout) + + def _require_finalized(self) -> None: + if not self._finalized: + raise RuntimeError("SmemWorkspace must be finalized first.") + + +__all__ = [ + "SmemLifetime", + "SmemOverlay", + "SmemRegion", + "SmemRegionKind", + "SmemWorkspace", + "SwizzleSpec", +] + diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/software_sync.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/software_sync.py new file mode 100644 index 000000000..af282e016 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/software_sync.py @@ -0,0 +1,279 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Software grid and NVLink synchronization for persistent kernels.""" + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import Int32, Int64 + +from .device_workspace import DeviceWorkspace +from .ptx_helpers import red_add_relaxed_sys_s32 + + +class SoftwareGridSync: + """Reusable device-local barrier over a phase-flipping GMEM counter.""" + + finish_sum_tag = 0x80000000 + grid_counter_region = "software_grid_sync.counter" + + def __init__(self, *, barrier_id: int) -> None: + self.barrier_id = barrier_id + self._grid_counter = None + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "SoftwareGridSync": + if values: + raise ValueError(f"SoftwareGridSync expected no MLIR values, got {len(values)}.") + return self + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + workspace.register( + self.grid_counter_region, + cutlass.Int32, + (1,), + buffer_space="local", + byte_alignment=16, + reset="zero_on_first_allocate", + ) + + @cute.jit + def assign_device_members(self, workspace: DeviceWorkspace) -> None: + self._grid_counter = workspace.ptr(self.grid_counter_region) + + def remove_device_members(self) -> None: + self._grid_counter = None + + @cute.jit + def _cta_rendezvous(self, participating_threads: int) -> None: + cute.arch.barrier(barrier_id=self.barrier_id, number_of_threads=participating_threads) + + @cute.jit + def sync( + self, + participating_threads: int, + actual_cta_count: Int32, + linear_cta_idx: Int32, + thread_idx_in_group: Int32, + ) -> None: + self._cta_rendezvous(participating_threads) + leader_delta = Int32(-self.finish_sum_tag) - (actual_cta_count - Int32(1)) + _inline_grid_sync( + self._grid_counter, + linear_cta_idx, + leader_delta, + Int32(1), + thread_idx_in_group, + ) + self._cta_rendezvous(participating_threads) + + +class NvlinkBarrier(SoftwareGridSync): + """Sense-reversing all-rank barrier layered over software grid sync.""" + + period = 4 + grid_counter_region = "nvlink.token_comm.grid_sync_counter" + phase_counter_region = "nvlink.token_comm.nvlink_phase_counter" + signal_region = "nvlink.token_comm.nvlink_signal" + + def __init__( + self, + *, + world_size: int, + barrier_id: int, + ) -> None: + super().__init__(barrier_id=barrier_id) + self.world_size = world_size + self._phase_counter = None + self._signal = None + self._peer_rank_ptr_mapper = None + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + super().register_device_workspace(workspace) + workspace.register( + self.phase_counter_region, + cutlass.Int32, + (1,), + buffer_space="local", + byte_alignment=16, + reset="zero_on_first_allocate", + ) + workspace.register( + self.signal_region, + cutlass.Int32, + (2,), + buffer_space="shared", + byte_alignment=16, + reset="zero_on_first_allocate", + ) + + @cute.jit + def assign_device_members( + self, + workspace: DeviceWorkspace, + peer_rank_ptr_mapper, + ) -> None: + super().assign_device_members(workspace) + self._phase_counter = workspace.ptr(self.phase_counter_region) + self._signal = workspace.ptr(self.signal_region) + self._peer_rank_ptr_mapper = peer_rank_ptr_mapper + + def remove_device_members(self) -> None: + super().remove_device_members() + self._phase_counter = None + self._signal = None + self._peer_rank_ptr_mapper = None + + @cute.jit + def arrive_and_wait( + self, + participating_threads: int, + actual_cta_count: Int32, + linear_cta_idx: Int32, + thread_idx_in_group: Int32, + *, + prologue_grid_sync: bool, + epilogue_grid_sync: bool, + ) -> None: + if cutlass.const_expr(prologue_grid_sync): + self.sync( + participating_threads, + actual_cta_count, + linear_cta_idx, + thread_idx_in_group, + ) + + if linear_cta_idx == Int32(0): + status = cute.arch.load( + self._phase_counter, + Int32, + sem="relaxed", + scope="gpu", + ) & Int32(3) + signal_phase = status & Int32(1) + signal_direction = status >> Int32(1) + signal_delta = Int32(1) + signal_target = Int32(self.world_size) + if signal_direction != Int32(0): + signal_delta = Int32(-1) + signal_target = Int32(0) + + self._cta_rendezvous(participating_threads) + if thread_idx_in_group == Int32(0): + cute.arch.fence_acq_rel_sys() + self._cta_rendezvous(participating_threads) + + rank_round_count = ( + self.world_size + participating_threads - 1 + ) // participating_threads + signal_base_address = self._signal.toint() + signal_byte_offset = Int64(signal_phase * Int32(4)) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = ( + Int32(rank_round * participating_threads) + + thread_idx_in_group + ) + if destination_rank < Int32(self.world_size): + destination_address = self._peer_rank_ptr_mapper.map( + signal_base_address, + destination_rank, + signal_byte_offset, + ) + red_add_relaxed_sys_s32( + destination_address, + signal_delta, + ) + + self._cta_rendezvous(participating_threads) + if thread_idx_in_group == Int32(0): + cute.arch.atomic_add( + self._phase_counter, + Int32(1), + sem="relaxed", + scope="gpu", + ) + local_signal = self._signal + signal_phase + while cute.arch.load( + local_signal, + Int32, + sem="acquire", + scope="sys", + ) != signal_target: + pass + + if cutlass.const_expr(epilogue_grid_sync): + self.sync( + participating_threads, + actual_cta_count, + linear_cta_idx, + thread_idx_in_group, + ) + + @cute.jit + def finalize( + self, + completed_calls: int, + participating_threads: int, + actual_cta_count: Int32, + linear_cta_idx: Int32, + thread_idx_in_group: Int32, + ) -> None: + padding_calls = (-completed_calls) % self.period + for _ in cutlass.range_constexpr(padding_calls): + self.arrive_and_wait( + participating_threads, + actual_cta_count, + linear_cta_idx, + thread_idx_in_group, + prologue_grid_sync=True, + epilogue_grid_sync=True, + ) + + +@cute.jit +def _inline_grid_sync( + counter, + linear_cta_idx, + leader_delta, + other_delta, + thread_idx_in_group, +) -> None: + llvm.inline_asm( + None, + [ + counter.toint().ir_value(), + Int32(linear_cta_idx).ir_value(), + leader_delta.ir_value(), + other_delta.ir_value(), + Int32(thread_idx_in_group).ir_value(), + ], + ( + "{\n\t" + ".reg .b32 %delta; .reg .b32 %old; .reg .b32 %current;\n\t" + ".reg .pred %not_leader; .reg .pred %is_cta0; " + ".reg .pred %waiting;\n\t" + "setp.ne.u32 %not_leader, $4, 0;\n\t" + "@%not_leader bra DONE;\n\t" + "setp.eq.u32 %is_cta0, $1, 0;\n\t" + "selp.b32 %delta, $2, $3, %is_cta0;\n\t" + "atom.release.gpu.global.add.u32 %old, [$0], %delta;\n\t" + "SPIN:\n\t" + "ld.relaxed.gpu.global.b32 %current, [$0];\n\t" + "xor.b32 %current, %current, %old;\n\t" + "and.b32 %current, %current, 0x80000000;\n\t" + "setp.eq.u32 %waiting, %current, 0;\n\t" + "@%waiting bra SPIN;\n\t" + "fence.acq_rel.gpu;\n\t" + "DONE:\n\t" + "}" + ), + "l,r,r,r,r", + has_side_effects=True, + asm_dialect=0, + ) + + +__all__ = ["NvlinkBarrier", "SoftwareGridSync"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py new file mode 100644 index 000000000..d1349e99b --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py @@ -0,0 +1,115 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Small integer and layout helpers shared by workspace implementations.""" + +from typing import Iterable, List, Tuple, Union + +import cutlass + + +IntegerType = Union[int, cutlass.Int32, cutlass.Int64, cutlass.Uint32, cutlass.Uint64] + + +def round_up(value: IntegerType, alignment: IntegerType) -> IntegerType: + return ((value + alignment - 1) // alignment) * alignment + + +def ceil_div(value: IntegerType, divisor: IntegerType) -> IntegerType: + return (value + divisor - 1) // divisor + + +def padded_expert_rows(token_count: IntegerType, padding_block: IntegerType) -> IntegerType: + """Row span one expert occupies in a block-padded pool. + + The single definition of a pool's per-expert stride. Communication components + write bases derived from it while the FC12 scheduler rebuilds the same bases + from ``expert_sizes`` at runtime; if the two ever disagree the metadata a + kernel reads no longer describes the rows it loads. + """ + return round_up(token_count, padding_block) + + +def is_power_of_two(value: int) -> bool: + return value > 0 and (value & (value - 1)) == 0 + + +def is_nested_shape(shape: Tuple) -> bool: + return any(isinstance(dimension, tuple) for dimension in shape) + + +def validate_static_integer_tuple(value: Tuple, *, field_name: str) -> None: + for element in value: + if isinstance(element, tuple): + validate_static_integer_tuple(element, field_name=field_name) + elif not isinstance(element, int): + raise TypeError(f"{field_name} must contain Python ints, got {type(element)}.") + + +def flatten_shape_stride(shape: Tuple, stride: Tuple) -> List[Tuple[IntegerType, IntegerType]]: + pairs: List[Tuple[IntegerType, IntegerType]] = [] + for size, step in zip(shape, stride): + if isinstance(size, tuple): + pairs.extend(flatten_shape_stride(size, step)) + else: + pairs.append((size, step)) + return pairs + + +def strides_equal_ignoring_singletons(shape, lhs_stride, rhs_stride) -> bool: + """Compare strides while ignoring leaves whose logical extent is one.""" + if isinstance(shape, tuple): + if not isinstance(lhs_stride, tuple) or not isinstance(rhs_stride, tuple): + return False + if len(shape) != len(lhs_stride) or len(shape) != len(rhs_stride): + return False + return all( + strides_equal_ignoring_singletons(child_shape, lhs_step, rhs_step) + for child_shape, lhs_step, rhs_step in zip(shape, lhs_stride, rhs_stride) + ) + return shape == 1 or lhs_stride == rhs_stride + + +def ordered_stride(shape: Tuple[int, ...], mem_order: Tuple[int, ...]) -> Tuple[Tuple[int, ...], int]: + stride = [0] * len(shape) + cosize = 1 + for mode in sorted(range(len(shape)), key=lambda index: mem_order[index]): + stride[mode] = cosize + cosize *= shape[mode] + return tuple(stride), cosize + + +def row_major_stride(shape: Tuple[int, ...]) -> Tuple[int, ...]: + if is_nested_shape(shape): + raise ValueError("A nested shape needs an explicit stride.") + stride, _ = ordered_stride(shape, tuple(reversed(range(len(shape))))) + return stride + + +def cosize_from_shape_stride_tuples(shape: Tuple, stride: Tuple) -> IntegerType: + leaf_pairs = flatten_shape_stride(shape, stride) if shape else [] + return 1 + sum((size - 1) * step for size, step in leaf_pairs) + + +def product(values: Iterable[IntegerType]) -> IntegerType: + result: IntegerType = 1 + for value in values: + result = result * value + return result + + +__all__ = [ + "IntegerType", + "ceil_div", + "cosize_from_shape_stride_tuples", + "flatten_shape_stride", + "is_nested_shape", + "is_power_of_two", + "ordered_stride", + "padded_expert_rows", + "product", + "row_major_stride", + "round_up", + "strides_equal_ignoring_singletons", + "validate_static_integer_tuple", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/__init__.py new file mode 100644 index 000000000..1b80fdf34 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/__init__.py @@ -0,0 +1,29 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Kernel source modules independent of repository-only runners.""" + +from .schedulers import ( + BlackwellFusedFc12Scheduler, + BlockPhase, + Fc12WorkTileState, + NonSwapAbFc12WorkTileInfo, + SchedulerBase, + SchedulerConsumer, + SchedulerWorkTileBase, + SwapAbFc12WorkTileInfo, + WorkIdAcquisitionMode, +) + + +__all__ = [ + "BlackwellFusedFc12Scheduler", + "BlockPhase", + "Fc12WorkTileState", + "NonSwapAbFc12WorkTileInfo", + "SchedulerBase", + "SchedulerConsumer", + "SchedulerWorkTileBase", + "SwapAbFc12WorkTileInfo", + "WorkIdAcquisitionMode", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py new file mode 100644 index 000000000..cedc19ea4 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py @@ -0,0 +1,3072 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Autonomous epilogue for the fused FC1+FC2 swap-AB MegaMoE kernel. + +Per-thread RMEM tensors flow between the transpose / SwiGLU / quantize / fc2 +store steps as bare ``cute.Tensor`` fragments; their thread distribution is a +fixed physical property of the surrounding atom sequence and is documented in +local comments. FC2 store mappings are finite ``FunctionMapping`` objects +evaluated at runtime to drive metadata lookup and destination pointer math. +""" + +import dataclasses +import math +from typing import Any, Callable, ClassVar, List, Literal, Optional, Tuple, Type, Union + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Int64, T +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import AddressSpace +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils + +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, vector + +from .....communication.nvlink_domain.symmetric_buffer import SymmetricBufferDevice +from .....communication.token_protocol import TokenSrcMetadata +from .....quant_def import CombineFormat, QuantKind +from .....api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from .....helpers.constants import Nvfp4E2M1RcpLimit, Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Fp32Max +from .....helpers.cute_py_helpers import tcgen05_block_scaled_acc_dtype +from .....helpers.dsl_helpers import mark_alignment +from .....helpers.flag_batch import GpuReleaseFlagBatchTracker +from .....helpers.ptx_helpers import ( + cp_async_bulk_s2g, + cp_reduce_async_bulk_add_bf16_s2g, + cvt_f32_to_fp8_to_f32, + movmatrix_b16, + red_add_relaxed_sys_v2_bf16x2, +) +from .....helpers.smem_workspace import SmemRegion, SmemWorkspace +from ....function_mapping import CoordinateSpace, FunctionMapping +from ....schedulers import BlockPhase, SchedulerConsumer, SwapAbFc12WorkTileInfo +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension + + +@dataclasses.dataclass(frozen=True) +class QuantImpl: + """Register-level block quantizer shared by the fc1 / fc2 epilogues. + + Returns ``(data_regs, sf_regs)`` ONLY: the caller pre-multiplies the topk + weight / global scale into ``prequant_reg`` beforehand and owns the data / + sf plane stores afterwards. ``sf_vec_direction`` selects the per-block amax + reduction: + + * ``regs_in_thread`` -- a block is ``sf_vec`` contiguous regs of + one thread; amax is thread-local (packed bf16x2 abs-max for combine, + fp32 ``fmax`` for orthodox). + * ``threads_with_the_same_reg`` -- a block is one reg across ``sf_vec`` warp + lanes; amax is a warp CREDUX (full warp for vec=32, lane-predicated + halves for vec=16) and is fp32-only, so bf16 is upconverted first. + * ``regs_in_pair_threads`` -- a block is ``sf_vec / 2`` regs in each of + two paired warps; amax is thread-local over the half, then exchanged with + the partner warp through SMEM. Orthodox mx only: fc1's TMEM transpose + hands one warp exactly 16 intermediate values per token, so a 32-wide + scale block necessarily straddles warps ``w`` and ``w ^ 1``. + ``prequant_reg`` must be 1D with size divisible by the per-thread block share + (``sf_vec``, or ``sf_vec / 2`` for the paired direction). Combine inputs are + bf16 (fc2's bf16 reorder regs); orthodox input is fp32 (swiglu). + """ + + quant_kind: Union[QuantKind, CombineFormat] + sf_vec_direction: Literal["regs_in_thread", "threads_with_the_same_reg", "regs_in_pair_threads"] + lane_idx: Optional[Any] = None + warp_idx: Optional[Any] = None + pair_exchange_barrier: Optional[Any] = None + + _directions: ClassVar[Tuple[str, ...]] = ("regs_in_thread", "threads_with_the_same_reg", "regs_in_pair_threads") + + # -- config / validation -------------------------------------------------- + + def __post_init__(self): + if isinstance(self.quant_kind, CombineFormat): + if not self.quant_kind.is_quantized: + raise ValueError(f"QuantImpl combine path needs a quantized CombineFormat, got {self.quant_kind}.") + elif not isinstance(self.quant_kind, QuantKind): + raise ValueError(f"quant_kind must be a QuantKind or a CombineFormat, got {self.quant_kind!r}.") + if self.sf_vec_direction not in self._directions: + raise ValueError(f"sf_vec_direction must be one of {self._directions}, got {self.sf_vec_direction!r}.") + if self.sf_vec_direction == "threads_with_the_same_reg" and self.lane_idx is None: + raise ValueError("across-lane quant needs lane_idx for the CREDUX half-warp predicate.") + if self.sf_vec_direction == "regs_in_pair_threads": + if isinstance(self.quant_kind, CombineFormat): + raise ValueError("The paired direction is orthodox-only; combine blocks never straddle warps.") + if self.sf_vec_size % 2 != 0: + raise ValueError(f"The paired direction needs an even sf_vec, got {self.sf_vec_size}.") + if self.lane_idx is None or self.warp_idx is None or self.pair_exchange_barrier is None: + raise ValueError("The paired direction needs lane_idx, warp_idx and pair_exchange_barrier.") + elif not isinstance(self.quant_kind, CombineFormat) and self.sf_vec_size != 16: + # A wider orthodox block cannot be reduced inside one thread: see the paired direction. + raise ValueError(f"Orthodox sf_vec {self.sf_vec_size} needs the paired direction.") + + @property + def data_dtype(self): + if isinstance(self.quant_kind, CombineFormat): + return self.quant_kind.act_dtype + # Orthodox output feeds fc2 as its activation, so it is the kind's activation type. + return self.quant_kind.activation_dtype + + @property + def scale_dtype(self): + if isinstance(self.quant_kind, CombineFormat): + return self.quant_kind.scale_dtype + return self.quant_kind.sf_dtype + + @property + def sf_vec_size(self) -> int: + if isinstance(self.quant_kind, CombineFormat): + return self.quant_kind.scale_block + return self.quant_kind.sf_vec_size + + @property + def _is_combine(self) -> bool: + return isinstance(self.quant_kind, CombineFormat) + + @property + def _data_rcp_limit(self) -> float: + # 1 / max representable magnitude of the data element type. + dt = self.data_dtype + if dt is cutlass.Float4E2M1FN: + return Nvfp4E2M1RcpLimit # 1/6 + if dt is cutlass.Float8E4M3FN: + return Fp8E4M3RcpLimit # 1/448 + return Fp8E5M2RcpLimit # 1/57344 + + @property + def _block_share_per_thread(self) -> int: + """How many of a block's elements this thread holds.""" + return self.sf_vec_size // 2 if self.sf_vec_direction == "regs_in_pair_threads" else self.sf_vec_size + + # -- dispatch ------------------------------------------------------------- + + @cute.jit + def __call__(self, prequant_reg: cute.Tensor, *, norm_const=None, smem_intermediate=None): + if cutlass.const_expr(cute.size(prequant_reg) % self._block_share_per_thread != 0): + raise ValueError("prequant_reg size must be divisible by this thread's share of a block.") + # Combine quantizes fc2's bf16 reorder regs; orthodox the fp32 swiglu. + expected_in = cutlass.BFloat16 if self._is_combine else cutlass.Float32 + if cutlass.const_expr(prequant_reg.element_type is not expected_in): + raise TypeError( + f"QuantImpl({self.quant_kind}) expects {expected_in} prequant input, got {prequant_reg.element_type}." + ) + needs_smem = cutlass.const_expr(self.sf_vec_direction == "regs_in_pair_threads") + if cutlass.const_expr(needs_smem != (smem_intermediate is not None)): + raise ValueError("smem_intermediate must be supplied for the paired direction and only for it.") + if cutlass.const_expr(not self._is_combine): + if cutlass.const_expr(self.sf_vec_direction == "regs_in_pair_threads"): + return self.mx_quant_regs_in_pair_threads_impl(prequant_reg, smem_intermediate) + return self.nvfp4_quant_impl(prequant_reg, norm_const=norm_const) + if cutlass.const_expr(self.data_dtype is cutlass.Float4E2M1FN): + if cutlass.const_expr(self.sf_vec_direction == "regs_in_thread"): + return self.nvfp4_combine_quant_regs_in_thread_impl(prequant_reg) + return self.nvfp4_combine_quant_threads_with_the_same_reg_impl(prequant_reg) + if cutlass.const_expr(self.sf_vec_direction == "regs_in_thread"): + return self.mxfp8_combine_quant_regs_in_thread_impl(prequant_reg) + return self.mxfp8_combine_quant_threads_with_the_same_reg_impl(prequant_reg) + + # -- impls ---------------------------------------------------------------- + + # regs_in_thread only; fc1 promises the vec direction via its TMEM transpose. + @cute.jit + def nvfp4_quant_impl( + self, prequant_reg: cute.Tensor, *, norm_const: Optional[cutlass.Float32] = None + ) -> Tuple[cute.Tensor, cute.Tensor]: + # fp32 in -> e2m1 data + e4m3 sfc. Mirrors the prior nvfp4_quant scale + # math (sfc -> capped/masked acc_scale); topk pre-mult + sf store are the + # caller's job now. + vec = self.sf_vec_size + n_blocks = cute.size(prequant_reg) // vec + data = cute.make_rmem_tensor((cute.size(prequant_reg),), cutlass.Float4E2M1FN) + sf = cute.make_rmem_tensor((n_blocks,), cutlass.Float8E4M3FN) + in_blocks = cute.zipped_divide(prequant_reg, (vec,)) # ((vec,), (n_blocks,)) + data_blocks = [] + sf_values = [] + rcp_limit = cutlass.Float32(self._data_rcp_limit) + for vec_block_idx in cutlass.range_constexpr(n_blocks): + block = in_blocks[None, vec_block_idx] + amax = self._amax_thread_fp32(block) + if cutlass.const_expr(norm_const is not None): + sfc_fp32 = amax * rcp_limit * norm_const + else: + sfc_fp32 = amax * rcp_limit + sfc_e4m3 = sfc_fp32.to(cutlass.Float8E4M3FN) + sfc_rt = cutlass.Float32(sfc_e4m3) + if cutlass.const_expr(norm_const is not None): + acc_scale = norm_const * cute.arch.rcp_approx(sfc_rt) + else: + acc_scale = cute.arch.rcp_approx(sfc_rt) + acc_scale = cute.arch.fmin(acc_scale, Fp32Max) + mask = cute.arch.fmin(sfc_rt * cutlass.Float32(1e30), cutlass.Float32(1.0)) + acc_scale = acc_scale * mask + sf_values.append(sfc_e4m3) + data_blocks.append(self._scale_to_data_ssa(block, acc_scale)) + self._store_packed_blocks(data, data_blocks) + sf.store(self._values_to_ssa(sf_values, cutlass.Float8E4M3FN)) + return data, sf + + @cute.jit + def mx_quant_regs_in_pair_threads_impl( + self, + prequant_reg: cute.Tensor, # (token_2_intermeidate_x) + smem_intermediate: cute.Tensor, # (token_blocks, epi_threads) + ) -> Tuple[cute.Tensor, cute.Tensor]: + half = self._block_share_per_thread + n_blocks = cute.size(prequant_reg) // half + data = cute.make_rmem_tensor((cute.size(prequant_reg),), self.data_dtype) + sf = cute.make_rmem_tensor((n_blocks,), self.scale_dtype) + in_blocks = cute.zipped_divide(prequant_reg, (half,)) # ((half,), (n_blocks,)) + + if cutlass.const_expr(cute.size(smem_intermediate, mode=[0]) != n_blocks): + raise ValueError( + f"The paired amax exchange needs {n_blocks} rows, got {cute.size(smem_intermediate, mode=[0])}." + ) + if cutlass.const_expr(smem_intermediate.stride[0] != 1): + # Thread-major would still be correct but would silently split the exchange into + # per-block scalar accesses. + raise ValueError("The paired amax exchange must be block-major to stay one vector access.") + + exchange_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=n_blocks * cutlass.Float32.width + ) + # Which warps pair up is epilogue knowledge: warp w owns one contiguous run of + # intermediate outputs, so a block that is twice that run wide joins adjacent warps. + partner_warp = self.warp_idx ^ cutlass.Int32(1) + own_thread = self.warp_idx * cutlass.Int32(32) + self.lane_idx + partner_thread = partner_warp * cutlass.Int32(32) + self.lane_idx + + half_amax = cute.make_rmem_tensor((n_blocks,), cutlass.Float32) + for vec_block_idx in cutlass.range_constexpr(n_blocks): + half_amax[vec_block_idx] = self._amax_thread_fp32(in_blocks[None, vec_block_idx]) + cute.copy(exchange_atom, cute.coalesce(half_amax), cute.coalesce(smem_intermediate[None, own_thread])) + + self.pair_exchange_barrier.arrive_and_wait() + + partner_amax = cute.make_rmem_tensor((n_blocks,), cutlass.Float32) + cute.copy(exchange_atom, cute.coalesce(smem_intermediate[None, partner_thread]), cute.coalesce(partner_amax)) + + data_blocks = [] + sf_values = [] + for vec_block_idx in cutlass.range_constexpr(n_blocks): + block_amax = cute.arch.fmax(half_amax[vec_block_idx], partner_amax[vec_block_idx]) + scale_e8m0, scale_f32 = self._e8m0(block_amax) + sf_values.append(scale_e8m0) + data_blocks.append(self._scale_to_data_ssa(in_blocks[None, vec_block_idx], self._enc_mxfp8(scale_f32))) + self._store_packed_blocks(data, data_blocks) + sf.store(self._values_to_ssa(sf_values, self.scale_dtype)) + return data, sf + + @cute.jit + def nvfp4_combine_quant_regs_in_thread_impl(self, prequant_reg: cute.Tensor) -> Tuple[cute.Tensor, cute.Tensor]: + # bf16 in -> e2m1 data + per-16 bf16 amax. amax found on bf16 (packed). + vec = self.sf_vec_size + n_blocks = cute.size(prequant_reg) // vec + data = cute.make_rmem_tensor((cute.size(prequant_reg),), cutlass.Float4E2M1FN) + sf = cute.make_rmem_tensor((n_blocks,), cutlass.BFloat16) + in_blocks = cute.zipped_divide(prequant_reg, (vec,)) # ((vec,), (n_blocks,)) + data_blocks = [] + sf_values = [] + for vec_block_idx in cutlass.range_constexpr(n_blocks): + block = in_blocks[None, vec_block_idx] + amax = self._amax_thread_bf16(block) + sf_values.append(amax) + decode_scale = cutlass.Float32(amax) * cutlass.Float32(self._data_rcp_limit) + data_blocks.append(self._scale_to_data_ssa(block, self._enc_nvfp4(decode_scale))) + self._store_packed_blocks(data, data_blocks) + sf.store(self._values_to_ssa(sf_values, cutlass.BFloat16)) + return data, sf + + # Mapping: (lane_idx, selected_sf_idx) -> (token_64, hidden_32) + # token_idx = lane_idx % 16 + selected_sf_idx * 16 + # hidden_idx = lane_idx // 16 * 16 + @cute.jit + def nvfp4_combine_quant_threads_with_the_same_reg_impl( + self, prequant_reg: cute.Tensor + ) -> Tuple[cute.Tensor, cute.Tensor]: + # UBLK has lane == hidden, so the warp's 32 lanes are 32 consecutive + # hidden. A scale block = sf_vec hidden, so the lanes split along hidden + # into 32 // sf_vec blocks of sf_vec lanes each (warp = blocks_per_warp * + # lanes_per_block, the EP x TP split). Only the sf_vec lanes inside a + # block share its CREDUX scale, so they pool the subtile tokens: sf_vec + # == 32 pools the whole warp, sf_vec < 32 pools fewer (more per lane). + lanes_per_block = self.sf_vec_size + n_tokens = cute.size(prequant_reg) + lane_in_block = self.lane_idx % cutlass.Int32(lanes_per_block) + data = cute.make_rmem_tensor((n_tokens,), cutlass.Float4E2M1FN) + selected_sf = cute.make_rmem_tensor((n_tokens // lanes_per_block,), cutlass.BFloat16) + scaled_vec = cute.full((n_tokens,), cutlass.Float32(0.0), cutlass.Float32) + for token_idx in cutlass.range_constexpr(n_tokens): + value = cutlass.Float32(prequant_reg[token_idx]) + amax_bf16 = self._amax_lane(value).to(cutlass.BFloat16) + slot = token_idx // lanes_per_block + if (token_idx % lanes_per_block) == lane_in_block: + selected_sf[slot] = amax_bf16 + else: + selected_sf[slot] = selected_sf[slot] + decode_scale = cutlass.Float32(amax_bf16) * cutlass.Float32(self._data_rcp_limit) + scaled_value = value * self._enc_nvfp4(decode_scale) + scaled_vec = cute.TensorSSA( + vector.insert(scaled_value.ir_value(), scaled_vec.ir_value(), [], [token_idx]), + (n_tokens,), + cutlass.Float32, + ) + self._store_packed_data(data, scaled_vec.to(cutlass.Float4E2M1FN)) + return data, selected_sf + + @cute.jit + def mxfp8_combine_quant_regs_in_thread_impl(self, prequant_reg: cute.Tensor) -> Tuple[cute.Tensor, cute.Tensor]: + # bf16 in -> e4m3/e5m2 data + per-32 e8m0. amax found on bf16 (packed). + vec = self.sf_vec_size + n_blocks = cute.size(prequant_reg) // vec + data = cute.make_rmem_tensor((cute.size(prequant_reg),), self.data_dtype) + sf = cute.make_rmem_tensor((n_blocks,), cutlass.Float8E8M0FNU) + in_blocks = cute.zipped_divide(prequant_reg, (vec,)) # ((vec,), (n_blocks,)) + data_blocks = [] + sf_values = [] + for vec_block_idx in cutlass.range_constexpr(n_blocks): + block = in_blocks[None, vec_block_idx] + # widen the native-bf16 amax to fp32 for the e8m0 round-up math. + scale_e8m0, scale_f32 = self._e8m0(cutlass.Float32(self._amax_thread_bf16(block))) + sf_values.append(scale_e8m0) + data_blocks.append(self._scale_to_data_ssa(block, self._enc_mxfp8(scale_f32))) + self._store_packed_blocks(data, data_blocks) + sf.store(self._values_to_ssa(sf_values, cutlass.Float8E8M0FNU)) + return data, sf + + # Mapping: (lane_idx, selected_sf_idx) -> (token_64, hidden_32) + # token_idx = lane_idx + selected_sf_idx * 32 + # hidden_idx = 0 + @cute.jit + def mxfp8_combine_quant_threads_with_the_same_reg_impl( + self, prequant_reg: cute.Tensor + ) -> Tuple[cute.Tensor, cute.Tensor]: + # UBLK has lane == hidden, so the warp's 32 lanes are 32 consecutive + # hidden. A scale block = sf_vec hidden, so the lanes split along hidden + # into 32 // sf_vec blocks of sf_vec lanes each (warp = blocks_per_warp * + # lanes_per_block, the EP x TP split). Only the sf_vec lanes inside a + # block share its CREDUX scale, so they pool the subtile tokens. mxfp8 + # sf_vec == 32 -> the whole warp is one block, all 32 lanes pool. + lanes_per_block = self.sf_vec_size + n_tokens = cute.size(prequant_reg) + lane_in_block = self.lane_idx % cutlass.Int32(lanes_per_block) + data = cute.make_rmem_tensor((n_tokens,), self.data_dtype) + selected_sf = cute.make_rmem_tensor((n_tokens // lanes_per_block,), cutlass.Float8E8M0FNU) + scaled_vec = cute.full((n_tokens,), cutlass.Float32(0.0), cutlass.Float32) + for token_idx in cutlass.range_constexpr(n_tokens): + value = cutlass.Float32(prequant_reg[token_idx]) + scale_e8m0, scale_f32 = self._e8m0(self._amax_lane(value)) + slot = token_idx // lanes_per_block + if (token_idx % lanes_per_block) == lane_in_block: + selected_sf[slot] = scale_e8m0 + else: + selected_sf[slot] = selected_sf[slot] + scaled_value = value * self._enc_mxfp8(scale_f32) + scaled_vec = cute.TensorSSA( + vector.insert(scaled_value.ir_value(), scaled_vec.ir_value(), [], [token_idx]), + (n_tokens,), + cutlass.Float32, + ) + self._store_packed_data(data, scaled_vec.to(self.data_dtype)) + return data, selected_sf + + # -- shared sub-steps ----------------------------------------------------- + + @cute.jit + def _scale_to_data_ssa(self, block: cute.Tensor, enc: cutlass.Float32) -> cute.TensorSSA: + block_f32 = block.load().to(cutlass.Float32) + enc_vec = cute.full_like(block_f32, enc, cutlass.Float32) + return (block_f32 * enc_vec).to(self.data_dtype) + + @cute.jit + def _concat_blocks_ssa(self, blocks, dtype: Type[cutlass.Numeric]) -> cute.TensorSSA: + values = [] + for block_idx in cutlass.range_constexpr(len(blocks)): + block = blocks[block_idx] + for elem_idx in cutlass.range_constexpr(cute.size(block.shape)): + values.append(block[elem_idx].ir_value()) + vec = vector.from_elements(T.vector(len(values), dtype.mlir_type), values) + return cute.TensorSSA(vec, (len(values),), dtype) + + @cute.jit + def _values_to_ssa(self, values, dtype: Type[cutlass.Numeric]) -> cute.TensorSSA: + vec = vector.from_elements( + T.vector(len(values), dtype.mlir_type), [values[i].ir_value() for i in range(len(values))] + ) + return cute.TensorSSA(vec, (len(values),), dtype) + + @cute.jit + def _concat_i32_blocks_ssa(self, blocks) -> cute.TensorSSA: + values = [] + for block_idx in cutlass.range_constexpr(len(blocks)): + packed_block = blocks[block_idx].bitcast(cutlass.Int32) + for elem_idx in cutlass.range_constexpr(cute.size(packed_block.shape)): + values.append(packed_block[elem_idx].ir_value()) + vec = vector.from_elements(T.vector(len(values), cutlass.Int32.mlir_type), values) + return cute.TensorSSA(vec, (len(values),), cutlass.Int32) + + @cute.jit + def _store_packed_blocks(self, data: cute.Tensor, blocks) -> None: + packed_data = cute.recast_tensor(data, cutlass.Int32) + packed_data.store(self._concat_i32_blocks_ssa(blocks)) + + @cute.jit + def _store_packed_data(self, data: cute.Tensor, data_ssa: cute.TensorSSA) -> None: + packed_data = cute.recast_tensor(data, cutlass.Int32) + packed_data.store(data_ssa.bitcast(cutlass.Int32)) + + @cute.jit + def _amax_thread_fp32(self, block: cute.Tensor) -> cutlass.Float32: + # max.xorsign.abs reduces |.| in one op per element; the result sign is + # the xor of the inputs (junk for an amax), so clear it at the end. + def max_abs(lhs: cutlass.Float32, rhs: cutlass.Float32) -> cutlass.Float32: + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [cutlass.Float32(lhs).ir_value(), cutlass.Float32(rhs).ir_value()], + "max.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + acc = block[0] + for elem_idx in cutlass.range_constexpr(1, cute.size(block)): + acc = max_abs(acc, block[elem_idx]) + mag_bits = cutlass.Int32(llvm.bitcast(T.i32(), cutlass.Float32(acc).ir_value())) & cutlass.Int32(0x7FFFFFFF) + return cutlass.Float32(llvm.bitcast(T.f32(), mag_bits.ir_value())) + + @cute.jit + def _amax_thread_bf16(self, block: cute.Tensor) -> cutlass.BFloat16: + # Packed bf16x2 abs-max: tree-reduce the pairs, then fold the survivor's + # two halves (high shifted into low). max.xorsign.abs leaves a junk sign, + # so the low bf16 is masked before being read back. The amax is natively + # bf16 -- exactly what the wire format stores. + def max_abs(lhs: cutlass.Int32, rhs: cutlass.Int32) -> cutlass.Int32: + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [cutlass.Int32(lhs).ir_value(), cutlass.Int32(rhs).ir_value()], + "max.xorsign.abs.bf16x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + pairs = cute.recast_tensor(block, cutlass.Int32) # (vec/2,) bf16x2 + acc = cutlass.Int32(pairs[0]) + for pair_idx in cutlass.range_constexpr(1, cute.size(pairs)): + acc = max_abs(acc, pairs[pair_idx]) + acc = max_abs(acc, acc >> cutlass.Int32(16)) + amax_bits = cute.make_rmem_tensor((1,), cutlass.Int32) + amax_bits[0] = acc & cutlass.Int32(0x7FFF) + return cute.recast_tensor(amax_bits, cutlass.BFloat16)[0] + + @cute.jit + def _amax_lane(self, v: cutlass.Float32) -> cutlass.Float32: + if cutlass.const_expr(self.sf_vec_size == 32): + return cute.arch.warp_redux_sync(v, "fmax", abs=True) + first_half = (self.lane_idx % cutlass.Int32(32)) < cutlass.Int32(16) + vsel = cutlass.Float32(0.0) + if first_half: + vsel = v + amax = cute.arch.warp_redux_sync(vsel, "fmax", abs=True) + if not first_half: + amax = cute.arch.warp_redux_sync(v, "fmax", abs=True) + return amax + + @cute.jit + def _e8m0(self, amax: cutlass.Float32) -> Tuple[cutlass.Float8E8M0FNU, cutlass.Float32]: + candidate = amax * cutlass.Float32(self._data_rcp_limit) + scale_f32 = cutlass.Float32(cvt_f32_to_fp8_to_f32(candidate, cutlass.Float8E8M0FNU)) + return scale_f32.to(cutlass.Float8E8M0FNU), scale_f32 + + @cute.jit + def _enc_nvfp4(self, decode_scale: cutlass.Float32) -> cutlass.Float32: + # rcp.approx.ftz with the fc1 cap+mask idiom (amax==0 -> 0, no inf*0 NaN). + enc = cute.arch.fmin(cute.arch.rcp_approx(decode_scale), Fp32Max) + mask = cute.arch.fmin(decode_scale * cutlass.Float32(1e30), cutlass.Float32(1.0)) + return enc * mask + + @cute.jit + def _enc_mxfp8(self, scale_f32: cutlass.Float32) -> cutlass.Float32: + # Skip nan + enc = cute.arch.fmin(cute.arch.rcp_approx(scale_f32), Fp32Max) + mask = cute.arch.fmin(scale_f32 * cutlass.Float32(1e30), cutlass.Float32(1.0)) + return enc * mask + + +# ============================================================================= +# Region tag +# ============================================================================= + + +class Region: + """Codegen-time region tag for a 16x32 sub-region within a 32x32 tile.""" + + Top = 0 + Bottom = 1 + + +# ============================================================================= +# TmemTranspose16x32 +# ============================================================================= + + +class _TmemTranspose16x32Core: + """Physical implementation of the 16x32 -> 32x16 TMEM in-place transpose. + + The transpose is a fixed sequence of tcgen05 32-bit element atoms; each + 32-bit slot is an fp32 SwiGLU-fold value for FC1. The (thread, reg) -> + (tmem_dp, tmem_col) input / output mapping is documented on the + ``TmemTranspose16x32`` subclass, which is the public entry point. + + Per-thread RMEM coordinate convention: + + - ``lane_idx`` -- warp lane id (= thread index within warp), in [0, 32). + - ``elem_idx`` -- per-thread reg index, in [0, 16). + """ + + _PermR1 = (0, 8, 2, 10, 4, 12, 6, 14, 1, 9, 3, 11, 5, 13, 7, 15) + _PermR3 = (0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15) + _PermR4 = (0, 8, 2, 10, 4, 12, 6, 14, 1, 9, 3, 11, 5, 13, 7, 15) + + _TmemRowStride = 1 << 16 + _io_dtype = cutlass.Float32 + + @staticmethod + def _tmem_layout(num_lanes: int, num_cols: int) -> cute.Layout: + return cute.make_layout( + (((num_lanes, num_cols), 1),), stride=(((_TmemTranspose16x32Core._TmemRowStride, 1), 0),) + ) + + @staticmethod + def _rmem_copy_view(rmem: cute.Tensor, num_regs: int, offset: int = 0) -> cute.Tensor: + return cute.make_tensor(rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),))) + + @staticmethod + def load_subtile_raw_acc( + tmem_subtile_tensor: cute.Tensor, + ) -> Tuple[cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor]: + """LDTM the entire 32-lane x 64-col raw acc region of one epi + subtile into 4 independent (16,) fp32 RMEM tensors. + + Used by the FC1 overlap-acc unroll path to extract all raw acc data + of the first 2 subtiles up front, so that the acc TMEM can be released + after the first subtile's 4 LDTMs. + + ``tmem_subtile_tensor`` is the (32 lanes, 64 cols) view onto a + single epi subtile's acc TMEM region (already offset by + ``warp_lane_offset + acc_stage_col_offset + subtile_col_offset``; + see ``SwapABGatedActEpilogue._subtile_local_tmem_tensor``). + + Returns a 4-tuple of (16,) fp32 RMEM tensors carrying the FC1 raw + LDTM distribution: + + [0] gate_lo / first-half top -- subtile cols 0..31, lanes 0..15 + [1] up_lo / first-half bot -- subtile cols 0..31, lanes 16..31 + [2] raw_top / second-half top -- subtile cols 32..63, lanes 0..15 + [3] raw_bot / second-half bot -- subtile cols 32..63, lanes 16..31 + + 4 atom calls of ``Ld16x64bOp(Repetition.x16) Float32`` -- the same + atom used by the per-subtile entry LDTM. Each output is in the + raw-LDTM input distribution consumed by ``TmemTranspose16x32``. + """ + atom_ld16x64 = cute.make_copy_atom( + tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), _TmemTranspose16x32Core._io_dtype + ) + + ptr = tmem_subtile_tensor.iterator + half_lane_off = 16 * _TmemTranspose16x32Core._TmemRowStride + + # 4 source 16-lane x 32-col views over the (32, 64) subtile region: + # first half (cols 0..31): top lanes 0..15 / bot lanes 16..31 + # second half (cols 32..63): top lanes 0..15 / bot lanes 16..31 + # All offsets are Python ints (compile-time const) so cute can + # const-fold them and infer the correct (>= 8 B / 2 col) ptr + # alignment that the LDTM atom requires. Using ``cutlass.Int32`` + # offsets here would wrap them as SSA values that cute treats as + # alignment-unknown, tripping the atom's verifier. + first_top_view = cute.make_tensor(ptr, _TmemTranspose16x32Core._tmem_layout(16, 32)) + first_bot_view = cute.make_tensor(ptr + half_lane_off, _TmemTranspose16x32Core._tmem_layout(16, 32)) + second_top_view = cute.make_tensor(ptr + 32, _TmemTranspose16x32Core._tmem_layout(16, 32)) + second_bot_view = cute.make_tensor(ptr + 32 + half_lane_off, _TmemTranspose16x32Core._tmem_layout(16, 32)) + + first_top = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) + first_bot = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) + second_top = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) + second_bot = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) + + cute.copy(atom_ld16x64, first_top_view, _TmemTranspose16x32Core._rmem_copy_view(first_top, 16)) + cute.copy(atom_ld16x64, first_bot_view, _TmemTranspose16x32Core._rmem_copy_view(first_bot, 16)) + cute.copy(atom_ld16x64, second_top_view, _TmemTranspose16x32Core._rmem_copy_view(second_top, 16)) + cute.copy(atom_ld16x64, second_bot_view, _TmemTranspose16x32Core._rmem_copy_view(second_bot, 16)) + + return (first_top, first_bot, second_top, second_bot) + + def __init__(self, tmem_ptr, region: int, reg_tensor: Optional[cute.Tensor] = None) -> None: + # The whole transpose is built from 32-bit element atoms; _io_dtype + # drives _src_regs / output / every LDTM/STTM atom below, so guard the + # invariant once here (tautological today, defensive against future + # dtype edits). + if cutlass.const_expr(self._io_dtype.width != 32): + raise TypeError( + f"{type(self).__name__} requires a 32-bit _io_dtype (the " + f"transpose uses 32-bit element atoms), got {self._io_dtype} " + f"(width {self._io_dtype.width})." + ) + + half_lane_off = 16 * self._TmemRowStride + if region == Region.Top: + src_ptr = tmem_ptr + dst_ptr = tmem_ptr + elif region == Region.Bottom: + src_ptr = tmem_ptr + half_lane_off + dst_ptr = tmem_ptr + 16 + else: + raise ValueError("region must be Region.Top or Region.Bottom") + + self.region = region + + self._tmem_src_full = cute.make_tensor(src_ptr, self._tmem_layout(16, 32)) + self._tmem_dst_full = cute.make_tensor(dst_ptr, self._tmem_layout(32, 16)) + self._tmem_dst_top = cute.make_tensor(dst_ptr, self._tmem_layout(16, 16)) + self._tmem_dst_bot = cute.make_tensor(dst_ptr + half_lane_off, self._tmem_layout(16, 16)) + + self._atom_ld16x64 = cute.make_copy_atom(tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), self._io_dtype) + self._atom_st16x128 = cute.make_copy_atom(tcgen05.St16x128bOp(tcgen05.Repetition.x8), self._io_dtype) + self._atom_st32x32 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), self._io_dtype) + self._atom_ld16x256 = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), self._io_dtype) + self._atom_ld16x128 = cute.make_copy_atom(tcgen05.Ld16x128bOp(tcgen05.Repetition.x4), self._io_dtype) + + self._src_regs = cute.make_rmem_tensor((16,), self._io_dtype) + # ``output`` is a bare (16,) RMEM fragment; its (lane_idx, elem_idx) + # distribution after all four rounds is the transpose output mapping + # documented on ``TmemTranspose16x32``. + self.output = cute.make_rmem_tensor((16,), self._io_dtype) + + # skip-R1.Load mode: ``reg_tensor`` must already be in the transpose + # input distribution (see ``TmemTranspose16x32`` / produced by + # ``load_subtile_raw_acc``); we copy it in lieu of the R1 LDTM. + # Weak entry guard (replaces the removed input contract): the transpose + # atoms are 32-bit element atoms over exactly 16 regs/lane, so the fed + # tensor must be a 32-bit element type of size 16. + self._reg_tensor = reg_tensor + if reg_tensor is not None: + if cutlass.const_expr(reg_tensor.element_type.width != 32): + raise TypeError( + f"{type(self).__name__} reg_tensor must be a 32-bit element " + f"type, got element type " + f"{reg_tensor.element_type} (width {reg_tensor.element_type.width})." + ) + if cutlass.const_expr(cute.size(reg_tensor) != 16): + raise ValueError( + f"{type(self).__name__} reg_tensor must hold exactly 16 elements, got {cute.size(reg_tensor)}." + ) + for r in range(16): + self._src_regs[r] = reg_tensor[r] + + # -- R1 ------------------------------------------------------------------ + + def r1_load(self) -> None: + """LDTM src region -> ``_src_regs``. No-op in skip-R1.Load mode.""" + if self._reg_tensor is not None: + return + cute.copy(self._atom_ld16x64, self._tmem_src_full, self._rmem_copy_view(self._src_regs, 16)) + + def r1_perm(self) -> None: + for r in range(16): + self.output[r] = self._src_regs[self._PermR1[r]] + + def r1_store(self) -> None: + cute.copy(self._atom_st16x128, self._rmem_copy_view(self.output, 16), self._tmem_src_full) + + # -- R2 ------------------------------------------------------------------ + + def r2_load(self) -> None: + cute.copy(self._atom_ld16x64, self._tmem_src_full, self._rmem_copy_view(self._src_regs, 16)) + + def r2_store(self) -> None: + cute.copy(self._atom_st32x32, self._rmem_copy_view(self._src_regs, 16), self._tmem_dst_full) + + # -- R3 ------------------------------------------------------------------ + + def r3_load_top(self) -> None: + cute.copy(self._atom_ld16x256, self._tmem_dst_top, self._rmem_copy_view(self._src_regs, 8, offset=0)) + + def r3_load_bot(self) -> None: + cute.copy(self._atom_ld16x256, self._tmem_dst_bot, self._rmem_copy_view(self._src_regs, 8, offset=8)) + + def r3_perm(self) -> None: + for r in range(16): + self.output[r] = self._src_regs[self._PermR3[r]] + + def r3_store(self) -> None: + cute.copy(self._atom_st32x32, self._rmem_copy_view(self.output, 16), self._tmem_dst_full) + + # -- R4 ------------------------------------------------------------------ + + def r4_load_top(self) -> None: + cute.copy(self._atom_ld16x128, self._tmem_dst_top, self._rmem_copy_view(self._src_regs, 8, offset=0)) + + def r4_load_bot(self) -> None: + cute.copy(self._atom_ld16x128, self._tmem_dst_bot, self._rmem_copy_view(self._src_regs, 8, offset=8)) + + def r4_perm(self) -> None: + for r in range(16): + self.output[r] = self._src_regs[self._PermR4[r]] + + def r4_store(self) -> None: + cute.copy(self._atom_st32x32, self._rmem_copy_view(self.output, 16), self._tmem_dst_full) + + def from_r1_perm_until_last_store(self) -> cute.Tensor: + self.r1_perm() + self.r1_store() + self.r2_load() + self.r2_store() + self.r3_load_top() + self.r3_load_bot() + self.r3_perm() + self.r3_store() + self.r4_load_top() + self.r4_load_bot() + self.r4_perm() + return self.output + + +class TmemTranspose16x32(_TmemTranspose16x32Core): + """FC1 16x32 -> 32x16 TMEM in-place transpose. + + The per-thread RMEM ``(lane_idx, elem_idx) -> (tmem_dp, tmem_col)`` mapping + is fixed by the underlying atom sequence. Each slot is an fp32 SwiGLU-fold + value and ``tmem_col`` is the intermediate-output index. + + Input distribution -- what each (lane_idx, elem_idx) reg holds on entry + (i.e. straight after the 16-dp x 32-col source LDTM, or as fed in via + ``reg_tensor`` / ``load_subtile_raw_acc`` for skip-R1.Load mode): + + tmem_dp = elem_idx * 2 + (lane_idx // 2) % 2 # in [0, 32) + tmem_col = (lane_idx % 2) * 8 + lane_idx // 4 # in [0, 16) + + Output distribution -- after all four rounds, the 32-dp x 16-col result has + each lane owning one full dp-row of 16 cols: + + tmem_dp = lane_idx # in [0, 32) + tmem_col = elem_idx # in [0, 16) + """ + + +# ============================================================================= +# TmemTranspose32x32Inplace +# ============================================================================= + + +class TmemTranspose32x32Inplace: + """fc1 epi 32x32 in-place TMEM transpose: two ``TmemTranspose16x32`` + sub-instances (``top`` = lanes 0..15, ``bot`` = lanes 16..31). + + Optional ``reg_tensor_top`` / ``reg_tensor_bot`` enable skip-R1.Load mode + for both halves; they must be provided or omitted together. + """ + + def __init__( + self, tmem_ptr, reg_tensor_top: Optional[cute.Tensor] = None, reg_tensor_bot: Optional[cute.Tensor] = None + ) -> None: + if (reg_tensor_top is None) != (reg_tensor_bot is None): + raise ValueError( + "TmemTranspose32x32Inplace: reg_tensor_top and reg_tensor_bot " + "must be provided or omitted together (both halves either " + "skip-R1.Load or do R1.Load)." + ) + self.top = TmemTranspose16x32(tmem_ptr, Region.Top, reg_tensor=reg_tensor_top) + self.bot = TmemTranspose16x32(tmem_ptr, Region.Bottom, reg_tensor=reg_tensor_bot) + + def from_r1_perm_until_last_store(self) -> Tuple[cute.Tensor, cute.Tensor]: + self.bot.r1_perm() + self.top.r1_perm() + self.bot.r1_store() + self.top.r1_store() + + self.bot.r2_load() + self.top.r2_load() + self.top.r2_store() + self.bot.r2_store() + + self.top.r3_load_top() + self.top.r3_load_bot() + self.bot.r3_load_top() + self.bot.r3_load_bot() + self.top.r3_perm() + self.bot.r3_perm() + self.top.r3_store() + self.bot.r3_store() + + self.top.r4_load_top() + self.top.r4_load_bot() + self.bot.r4_load_top() + self.bot.r4_load_bot() + self.top.r4_perm() + self.bot.r4_perm() + return self.top.output, self.bot.output + + +class TmemTranspose32x64B16Movm: + """FC2 warp-local 32-hidden x 64-token BF16 transpose using MOVM. + + Input is the flat ``[top, bottom]`` distribution produced by two + 16dp256bit accumulator loads followed by ``fc2_f2fp``. Output + ``(lane_idx, elem_idx)`` coordinates are: + + token = lane_idx + 32 * (elem_idx // 32) + hidden = elem_idx % 32 + + The fixed register permutation between MOVM and STTM is an SSA rename. It + keeps each thread's two complete hidden-32 rows without any lane exchange. + """ + + _tmem_row_stride = 1 << 16 + _store_reg_source_indices = ( + 0, + 2, + 1, + 3, + 16, + 18, + 17, + 19, + 8, + 10, + 9, + 11, + 24, + 26, + 25, + 27, + 4, + 6, + 5, + 7, + 20, + 22, + 21, + 23, + 12, + 14, + 13, + 15, + 28, + 30, + 29, + 31, + ) + + @staticmethod + def _tmem_layout(num_lanes: int, num_cols: int) -> cute.Layout: + return cute.make_layout( + (((num_lanes, num_cols), 1),), stride=(((TmemTranspose32x64B16Movm._tmem_row_stride, 1), 0),) + ) + + @staticmethod + def _rmem_copy_view(rmem: cute.Tensor, num_regs: int, offset: int = 0) -> cute.Tensor: + return cute.make_tensor(rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),))) + + @cute.jit + def __init__(self, tmem_ptr, reg_tensor: cute.Tensor) -> None: + if cutlass.const_expr(reg_tensor.element_type is not cutlass.BFloat16): + raise TypeError(f"{type(self).__name__} expects BF16 input after f2fp, got {reg_tensor.element_type}.") + if cutlass.const_expr(cute.size(reg_tensor) != 64): + raise ValueError(f"{type(self).__name__} expects 64 BF16 elements, got {cute.size(reg_tensor)}.") + + movm_words = movmatrix_b16(cute.recast_tensor(reg_tensor, cutlass.Int32)) + self._store_words = cute.make_rmem_tensor(movm_words.layout, movm_words.element_type) + for store_reg in cutlass.range_constexpr(32): + self._store_words[store_reg] = movm_words[self._store_reg_source_indices[store_reg]] + + half_lane_offset = 16 * self._tmem_row_stride + self._tmem_top = cute.make_tensor(tmem_ptr, self._tmem_layout(16, 32)) + self._tmem_bottom = cute.make_tensor(tmem_ptr + half_lane_offset, self._tmem_layout(16, 32)) + self._tmem_full = cute.make_tensor(tmem_ptr, self._tmem_layout(32, 32)) + self._store_atom = cute.make_copy_atom(tcgen05.St16x128bOp(tcgen05.Repetition.x8), cutlass.Float32) + self._load_atom = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) + + @cute.jit + def __call__(self) -> cute.Tensor: + movm_words_f32 = cute.recast_tensor(self._store_words, cutlass.Float32) + cute.copy(self._store_atom, self._rmem_copy_view(movm_words_f32, 16), self._tmem_top) + cute.copy(self._store_atom, self._rmem_copy_view(movm_words_f32, 16, offset=16), self._tmem_bottom) + + output_words = cute.make_rmem_tensor((32,), cutlass.Float32) + cute.copy(self._load_atom, self._tmem_full, self._rmem_copy_view(output_words, 32)) + return cute.recast_tensor(output_words, cutlass.BFloat16) + + +@dataclasses.dataclass(frozen=True) +class GatedActEpilogueArgs: + """Optional runtime tensors used by the gated-activation epilogue.""" + + fc1_alpha: Optional[cute.Tensor] + fc2_alpha: Optional[cute.Tensor] + fc1_norm_const: Optional[cute.Tensor] + # ----------------------------------- + # MoE domain (token, topk), deepgemm graph only? for transformer graph, we want reduce kernel to perform the score mul. + topk_scores: Optional[cute.Tensor] + + +class SwapABGatedActEpilogue(KernelComponent): + """Autonomous epilogue for the swap-AB SwiGLU NVFP4 kernel. + + ``run()`` is the single entry point the kernel calls inside the epi + warp body. The kernel's responsibility is reduced to: + + - allocate / free TMEM and build ``acc_tensor`` + - construct the AB / acc pipelines + - obtain the scheduler consumer + + Everything else (acc consumer state, task-tile loop, overlap rotation, + early release, TMA store commit / drain, per-subtile dispatch) lives + inside this class. + """ + + _EpilogueSyncWaitBarId = 1 # Arrive and wait only + _EpilogueAsyncBarIdBase = 4 # Some arrive, the others arrive and wait + _EpilogueFc1GateUpInterleave = 16 + _EpilogueTokenTileSize = 64 # Fundamentally the epi_tile_n + _EpilogueFc1IntermediateGateUpTileSize = 128 # Fundamentally epi_tile_m + _EpilogueFc1IntermediateDownTileSize = 64 # Fundamentally epi_tile_m // 2 + _EpilogueFc2HiddenTileSize = 128 # Fundamentally epi_tile_m + _EpilogueWarpCnt = 4 + # One warp owns this many intermediate_down outputs per token: the TMEM transpose gives each + # warp 32 accumulator rows, which the gate/up interleave halves. + _EpilogueFc1IntermediateDownPerWarp = 16 + smem_scratch_overlay: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.scratch" + fc1_staging_region: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.fc1_staging" + fc2_staging_region: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.fc2_staging" + # A 144 B FP8 stride rotates each row by four SMEM banks. + _Fc2UblkFp8RowStrideBytes = 144 + _ScratchByteAlignment = 128 + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + # The accumulator follows from the instruction family; see __init__. + "quant_kind": str, + "hidden_size": StaticOrRuntimeIntegerType, + "intermediate_gateup_size": StaticOrRuntimeIntegerType, + "combine_format": CombineFormat, + "gate_up_clamp": Optional[float], + # SiTU (Kimi K3) selects a different gated-activation core; optional so + # existing SwiGLU descriptors stay valid unchanged. See _resolve_situ_betas. + "situ_beta": OptionalRequirement(Optional[float]), + "situ_linear_beta": OptionalRequirement(Optional[float]), + } + + @classmethod + def impl_desc_require(cls) -> dict[str, object]: + return { + "mma_tiler_mnk": tuple, + "cluster_shape_mn": tuple, + "use_2cta_instrs": bool, + "fc2_use_bulk": bool, + "communication_enabled": bool, + "fc1_epi_flag_batch": int, + "fc2_epi_flag_batch": int, + "fc2_tma_stages": OptionalRequirement(int), + "reduce_topk_in_kernel": OptionalRequirement(bool), + "token_back_push_data": OptionalRequirement(bool), + } + + @staticmethod + def _resolve_situ_betas(problem_desc: ProblemDesc) -> Tuple[Optional[float], Optional[float]]: + """Resolve the SiTU (Kimi K3) betas from the ProblemDesc; ``(None, None)`` means SwiGLU. + + Both betas must be given together, and SiTU excludes ``gate_up_clamp`` -- matching DeepGEMM's + ``DG_HOST_ASSERT(not use_situ or not activation_clamp_opt.has_value())``. They are baked in at + codegen time exactly like ``gate_up_clamp``, so the enclosing KernelClass must also fold them + into its ``name()`` cache key. + """ + beta = problem_desc.get("situ_beta") + linear_beta = problem_desc.get("situ_linear_beta") + if (beta is None) != (linear_beta is None): + raise ValueError(f"situ_beta and situ_linear_beta must be set together, got {beta} and {linear_beta}.") + if beta is None: + return None, None + if beta <= 0 or linear_beta <= 0: + raise ValueError(f"SiTU beta parameters must be positive, got {beta} and {linear_beta}.") + if problem_desc["gate_up_clamp"] is not None: + raise ValueError("SiTU does not support gate_up_clamp; the two activation variants are exclusive.") + return float(beta), float(linear_beta) + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.quant_kind = QuantKind(problem_desc["quant_kind"]) + self.acc_dtype = tcgen05_block_scaled_acc_dtype + self.hidden_size = problem_desc["hidden_size"] + self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] + self.combine_format = problem_desc["combine_format"] + self.gate_up_clamp = problem_desc["gate_up_clamp"] + self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) + self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + self.cluster_shape_mn = impl_desc["cluster_shape_mn"] + self.use_2cta_instrs = impl_desc["use_2cta_instrs"] + self.fc2_use_bulk = impl_desc["fc2_use_bulk"] + self.communication_enabled = impl_desc["communication_enabled"] + self.fc1_epi_flag_batch = impl_desc["fc1_epi_flag_batch"] + self.fc2_epi_flag_batch = impl_desc["fc2_epi_flag_batch"] + + if self.communication_enabled: + for field_name in ("reduce_topk_in_kernel", "token_back_push_data"): + if field_name not in impl_desc: + raise KeyError(f"Communication-enabled Epilogue requires ImplDesc field {field_name!r}.") + self.reduce_topk_in_kernel = impl_desc.get("reduce_topk_in_kernel", False) + self.token_back_push_data = impl_desc.get("token_back_push_data", False) + if not self.communication_enabled and (self.reduce_topk_in_kernel or self.token_back_push_data): + raise ValueError("A communication-disabled Epilogue cannot enable communication policies.") + + # FC1 emits what FC2 consumes as its activation, in the kind's own scale format. + self.fc1_output_dtype = self.quant_kind.activation_dtype + self.fc1_output_sf_dtype = self.quant_kind.sf_dtype + self.sf_vec_size = self.quant_kind.sf_vec_size + # A 32-wide scale block spans two epilogue warps; see QuantImpl's paired direction. + self.needs_pair_amax_exchange = self.sf_vec_size > self._EpilogueFc1IntermediateDownPerWarp + self.token_back_push_sf = self.communication_enabled and self.combine_format.is_quantized + self.token_back_enabled = self.token_back_push_data or self.token_back_push_sf + self.fc2_output_is_local = not self.communication_enabled or self.token_back_push_data + self.fc2_use_tma = self.fc2_use_bulk and self.fc2_output_is_local + self.fc2_use_ublk = self.fc2_use_bulk and not self.fc2_output_is_local + if self.fc2_use_ublk and self.combine_format.act_dtype.width == 4: + raise ValueError("FC2 UBLK does not support an FP4 combine payload.") + if self.reduce_topk_in_kernel and self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("In-kernel top-k reduction requires a BF16 combine format.") + self.reduce_topk_in_epilogue = self.reduce_topk_in_kernel and not self.token_back_push_data + if not 1 <= self.fc1_epi_flag_batch <= 32 or not 1 <= self.fc2_epi_flag_batch <= 32: + raise ValueError("Epilogue flag batch sizes must be in [1, 32].") + self.cluster_tile_intermediate_downproj = self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] + + atom_thr_size = 2 if self.use_2cta_instrs else 1 + self.cta_tile_m = self._EpilogueFc2HiddenTileSize + self.cta_tile_n = self.mma_tiler_mnk[1] + self.cta_tile_k = self.mma_tiler_mnk[2] + assert self.mma_tiler_mnk[0] // atom_thr_size == self.cta_tile_m + assert self.cta_tile_n % self._EpilogueTokenTileSize == 0 + tmem_plan = impl_desc["tmem_plan"] + self.num_sfa_tmem_cols = tmem_plan.sfa_columns + self.num_sfb_tmem_cols = tmem_plan.sfb_columns + self.num_sf_tmem_cols = tmem_plan.sfa_columns + tmem_plan.sfb_columns + self.num_tmem_alloc_cols = tmem_plan.allocation_columns + self.num_accumulator_stages = tmem_plan.accumulator_stage_count + self.num_accumulator_pipeline_stages = tmem_plan.accumulator_pipeline_stages + self.accumulator_overlap_columns = ( + tmem_plan.accumulator_stage_columns - tmem_plan.accumulator_stage_stride_columns + ) + self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns + self.overlapping_accum = self.accumulator_overlap_columns > 0 + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) + self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) + + if isinstance(self.hidden_size, int) and self.hidden_size % (self.cta_tile_m * self.cluster_shape_mn[0]) == 0: + self.fc2_hidden_needs_predicate: bool = False + else: + self.fc2_hidden_needs_predicate: bool = True + + if isinstance(self.intermediate_gateup_size, int): + self.intermediate_downproj: Optional[int] = self.intermediate_gateup_size // 2 + else: + self.intermediate_downproj: Optional[int] = None + + self.subtile_cnt = self.cta_tile_n // self._EpilogueTokenTileSize + + # One staging stage per token subtile, each an (epi_tile_n, epi_tile_m // 2) quantized tile. + self.fc1_staging_stage_bytes = ( + self._EpilogueTokenTileSize * self._EpilogueFc1IntermediateDownTileSize * self.fc1_output_dtype.width // 8 + ) + self.fc1_staging_bytes = self.subtile_cnt * self.fc1_staging_stage_bytes + # The amax exchange plane is (token chunk, epilogue thread), block-major so each thread's + # column is one vector access. A thread holds one token per 32-lane chunk of the subtile. + # These slots borrow the current subtile's staging stage rather than costing their own + # bytes -- see the lifetime argument in fc1_quant. + self.fc1_amax_token_chunks = self._EpilogueTokenTileSize // 32 + self.fc1_amax_slot_count = ( + self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 if self.needs_pair_amax_exchange else 0 + ) + if self.fc1_amax_slot_count * 4 > self.fc1_staging_stage_bytes: + raise ValueError("The paired amax exchange does not fit in one fc1 staging stage.") + + requested_fc2_tma_stages = impl_desc.get("fc2_tma_stages") + if requested_fc2_tma_stages is not None and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt: + raise ValueError(f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}.") + if self.fc2_use_bulk: + single_stage_region = self._make_fc2_single_stage_region() + if self.fc2_use_ublk and single_stage_region.nbytes % 16 != 0: + raise ValueError("Each FC2 UBLK staging stage must occupy a multiple of 16 bytes.") + # Additional stages trade mainloop SMEM for store overlap. + self.fc2_tma_stages = requested_fc2_tma_stages if requested_fc2_tma_stages is not None else 1 + self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region(self.fc2_tma_stages) + else: + self.fc2_tma_stages = 0 + self.fc2_staging_spec = None + + @classmethod + def epilogue_sync_barrier(cls) -> pipeline.NamedBarrier: + """The one barrier every epilogue rendezvous uses: all four warps, arrive-and-wait. + + Reused rather than split per purpose because the participant set is always the same 128 + threads and no two uses are ever in flight together -- the tile-boundary rendezvous sits + outside the subtile loop, FC1's amax exchange and FC2's bulk-store handshake belong to + different work tiles. + """ + return pipeline.NamedBarrier(barrier_id=cls._EpilogueSyncWaitBarId, num_threads=32 * cls._EpilogueWarpCnt) + + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: + """Declare the epilogue scratch as one allocation shared by two exclusive lifetimes. + + FC1's staging tile and FC2's store tile belong to different work tiles, separated by the + tile-boundary TMA drain and rendezvous in ``run()``, so they never coexist. + """ + overlay = smem_workspace.create_overlay(self.smem_scratch_overlay) + overlay.add_lifetime("fc1_staging").register_tensor( + self.fc1_staging_region, cutlass.Int8, (self.fc1_staging_bytes,), byte_alignment=128 + ) + if self.fc2_staging_spec is not None: + overlay.add_lifetime("fc2_staging").register_tensor( + self.fc2_staging_region, + self.fc2_staging_spec.dtype, + self.fc2_staging_spec.shape, + stride=self.fc2_staging_spec.stride, + swizzle=self.fc2_staging_spec.swizzle, + byte_alignment=self.fc2_staging_spec.byte_alignment, + ) + + def fc1_staged_smem_layout( + self, n_stages: int, without_stage_mode: bool = False + ) -> Union[cute.Layout, cute.ComposedLayout]: + layout = sm100_utils.make_smem_layout_epi( + self.fc1_output_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize), + n_stages, + ) + if without_stage_mode: + return cute.select(layout, mode=[0, 1]) + return layout + + def fc2_tma_staged_smem_spec(self, n_stages: int) -> Tuple[Tuple, Tuple, Tuple[int, int, int]]: + """Return the bank-conflict-free token-major FC2 TMA staging layout.""" + stage_stride = self._EpilogueTokenTileSize * self._EpilogueFc2HiddenTileSize + wire_dtype = self.combine_format.act_dtype + if wire_dtype is cutlass.BFloat16: + shape = ((32, 2), (64, 2), n_stages) + stride = ((64, 2048), (1, 4096), stage_stride if n_stages > 1 else 0) + swizzle = (3, 4, 3) + elif wire_dtype.width == 8: + shape = ((32, 2), 128, n_stages) + stride = ((128, 4096), 1, stage_stride if n_stages > 1 else 0) + swizzle = (3, 4, 3) + elif wire_dtype.width == 4: + shape = ((32, 2), 128, n_stages) + stride = ((128, 4096), 1, stage_stride if n_stages > 1 else 0) + swizzle = (2, 4, 3) + else: + raise ValueError(f"Unsupported FC2 TMA staging dtype {wire_dtype}.") + return shape, stride, swizzle + + def fc2_tma_staged_smem_layout(self, n_stages: int, without_stage_mode: bool = False) -> cute.ComposedLayout: + shape, stride, swizzle = self.fc2_tma_staged_smem_spec(n_stages) + layout = cute.make_composed_layout(cute.make_swizzle(*swizzle), 0, cute.make_layout(shape, stride=stride)) + if without_stage_mode: + return cute.select(layout, mode=[0, 1]) + return layout + + def _make_fc2_single_stage_region(self) -> SmemRegion: + """Describe one FC2 bulk-staging subtile for TMA or UBLK.""" + if self.fc2_use_tma: + shape, stride, swizzle = self.fc2_tma_staged_smem_spec(1) + return SmemRegion( + name="", + kind="tensor", + dtype=self.combine_format.act_dtype, + shape=shape[:-1], + stride=stride[:-1], + swizzle=swizzle, + byte_alignment=128, + ) + row_stride_elements = ( + self._Fc2UblkFp8RowStrideBytes + if self.combine_format.act_dtype.width == 8 + else self._EpilogueFc2HiddenTileSize + ) + return SmemRegion( + name="", + kind="tensor", + dtype=self.combine_format.act_dtype, + shape=(self._EpilogueTokenTileSize, self._EpilogueFc2HiddenTileSize), + stride=(row_stride_elements, 1), + swizzle=None, + byte_alignment=16, + ) + + def _make_fc2_staging_region(self, stage_count: int) -> SmemRegion: + single_stage_region = self._make_fc2_single_stage_region() + return SmemRegion( + name="", + kind="tensor", + dtype=single_stage_region.dtype, + shape=(*single_stage_region.shape, stage_count), + stride=(*single_stage_region.stride, single_stage_region.cosize if stage_count > 1 else 0), + swizzle=single_stage_region.swizzle, + byte_alignment=single_stage_region.byte_alignment, + ) + + def prepare_tma_store_params( + self, fc1_output_template: cute.Tensor, fc2_output_template: cute.Tensor + ) -> Tuple[cute.CopyAtom, cute.Tensor, Optional[cute.CopyAtom], Optional[cute.Tensor]]: + """Build the FC1 TMA store and the optional local FC2 TMA store.""" + fc1_operation = cpasync.CopyBulkTensorTileS2GOp() + fc1_smem_layout = self.fc1_staged_smem_layout(1, without_stage_mode=True) + fc1_tile = (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize) + fc1_atom, fc1_tensor = cpasync.make_tiled_tma_atom( + fc1_operation, fc1_output_template, fc1_smem_layout, fc1_tile + ) + + if cutlass.const_expr(self.fc2_use_tma): + # Keep tiled rest modes dynamic so a single static tile cannot collapse to stride zero. + runtime_token_extent = cutlass.Int32(fc2_output_template.shape[0]) + runtime_hidden_extent = cutlass.Int32(fc2_output_template.shape[2]) + fc2_token_major_template = cute.make_tensor( + fc2_output_template.iterator, + cute.make_layout( + (runtime_token_extent, runtime_hidden_extent, cutlass.Int32(1)), + stride=(fc2_output_template.stride[0], fc2_output_template.stride[2], 0), + ), + ) + fc2_operation = cpasync.CopyBulkTensorTileS2GOp() + fc2_smem_layout = self.fc2_tma_staged_smem_layout(1, without_stage_mode=True) + fc2_tile = (self._EpilogueTokenTileSize, self._EpilogueFc2HiddenTileSize) + fc2_atom, fc2_tensor = cpasync.make_tiled_tma_atom( + fc2_operation, fc2_token_major_template, fc2_smem_layout, fc2_tile + ) + else: + fc2_atom = None + fc2_tensor = None + return fc1_atom, fc1_tensor, fc2_atom, fc2_tensor + + @cute.jit + def run( + self, + smem_workspace: SmemWorkspace, + smem_base: cute.Pointer, + tmem_ptr: cute.Pointer, + acc_pipeline, + # ── Sched ──────────────────────────────────────────────────────── + sched_consumer: SchedulerConsumer, + kernel_extension: BlockScaledSwapAbFc12Extension, + # ── tensors ────────────────────────────────── + tma_atom_fc1_output: cute.CopyAtom, + fc1_output: cute.Tensor, # Domain of fake (m, n, l) + fc1_output_sf: cute.Tensor, # Domain of fake (m, n, l) + tma_atom_fc2_output: Optional[cute.CopyAtom], + fc2_tma_output: Optional[cute.Tensor], # Domain (physical_token, hidden, l=1) + fc2_output: cute.Tensor, # MoE domain (token, topk, hidden) + fc1_done_counter: cute.Tensor, # 1D tensor + tidx: cutlass.Int32, + token_src_metadata: Optional[cute.Tensor], + fc2_done_counter: Optional[cute.Tensor], + fc2_output_sf: Optional[cute.Tensor], + peer_rank_ptr_mapper: Optional[SymmetricBufferDevice], + optional_epi_args: Optional[GatedActEpilogueArgs] = None, + ): + if cutlass.const_expr(not smem_workspace.finalized): + raise RuntimeError("SwapABGatedActEpilogue.run requires a finalized SmemWorkspace.") + if cutlass.const_expr(optional_epi_args is None): + optional_epi_args = GatedActEpilogueArgs( + fc1_alpha=None, fc2_alpha=None, fc1_norm_const=None, topk_scores=None + ) + fc1_staging_pointer = smem_workspace.ptr(self.fc1_staging_region, smem_base) + if cutlass.const_expr(self.fc2_tma_stages > 0): + fc2_smem_tensor = smem_workspace.tensor(self.fc2_staging_region, smem_base) + else: + fc2_smem_tensor = None + if cutlass.const_expr(self.fc2_use_tma and (tma_atom_fc2_output is None or fc2_tma_output is None)): + raise ValueError("FC2 TMA store requires a TMA atom and token-major output tensor.") + if cutlass.const_expr( + self.communication_enabled and (token_src_metadata is None or peer_rank_ptr_mapper is None) + ): + raise ValueError("Communication-enabled Epilogue requires token metadata and a peer pointer mapper.") + if cutlass.const_expr(self.token_back_enabled and fc2_done_counter is None): + raise ValueError("Token-back requires an FC2 done counter.") + if cutlass.const_expr(self.token_back_push_sf and fc2_output_sf is None): + raise ValueError("Quantized token-back requires an FC2 output scale tensor.") + tmem_acc = cute.make_tensor( + cute.recast_ptr(tmem_ptr, dtype=cutlass.Float32), + cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride), + ) + + fc1_epi = SwapABFc1Epilogue( + self, + tidx, + fc1_staging_pointer, + kernel_extension, + tma_atom_fc1_output, + fc1_output, + fc1_output_sf, + fc1_done_counter, + optional_epi_args, + ) + fc2_epi = SwapABFc2Epilogue( + self, + tidx, + fc2_smem_tensor, + tma_atom_fc2_output, + fc2_tma_output, + fc2_output, + token_src_metadata, + fc2_done_counter, + fc2_output_sf, + peer_rank_ptr_mapper, + optional_epi_args, + ) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_accumulator_pipeline_stages + ) + wait_only_named_barrier = self.epilogue_sync_barrier() + is_odd_turn = cutlass.Int32(1) + work_tile_info = sched_consumer.consume_work() + + flag_tracker = GpuReleaseFlagBatchTracker( + flag_address=Int64(0), + accumulated_flags=cutlass.Int32(0), + phase=cutlass.Int32(work_tile_info.phase), + thread_idx=tidx % (self._EpilogueWarpCnt * 32), + ) + + while work_tile_info.is_valid_tile: + if cutlass.const_expr(self.overlapping_accum): + tmem_stage_idx = acc_consumer_state.phase + else: + tmem_stage_idx = acc_consumer_state.index + tmem_acc_current = tmem_acc[None, None, tmem_stage_idx] + if work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1): + # The __call__ args should only take the while loop args, leave all loop irrevalent args to the init. + fc1_epi( + work_tile_info=work_tile_info, + tmem_acc_tensor=tmem_acc_current, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + is_odd_turn=is_odd_turn, + ) + else: + # The __call__ args should only take the while loop args, leave all loop irrevalent args to the init. + fc2_epi( + work_tile_info=work_tile_info, + tmem_acc_tensor=tmem_acc_current, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + is_odd_turn=is_odd_turn, + ) + prev_work_tile_info = work_tile_info + cur_was_linear1 = prev_work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + + acc_consumer_state.advance() + if cutlass.const_expr(self.overlapping_accum): + is_odd_turn = cutlass.Int32(1) - is_odd_turn + + work_tile_info = sched_consumer.consume_work() + + # Every asynchronous store commits at issue; drain before completion or scratch reuse. + cute.arch.cp_async_bulk_wait_group(0) + # _fence_rel_gpu() + wait_only_named_barrier.arrive_and_wait() + + # Publish completion for the work tile snapshotted above. + if cur_was_linear1: + flag_tracker = fc1_epi.signal_fc1_done(prev_work_tile_info, work_tile_info, flag_tracker) + else: + flag_tracker = fc2_epi.signal_fc2_done(prev_work_tile_info, work_tile_info, flag_tracker) + # Tail flush + flag_tracker.fire() + + +class _ImmutableAfterInit: + """Froze at the point calling `_freeze()`""" + + def __setattr__(self, name, value): + if self.__dict__.get("_frozen_", False): + raise AttributeError(f"{type(self).__name__} is immutable after __init__ (cannot set {name!r}).") + object.__setattr__(self, name, value) + + def _freeze(self) -> None: + object.__setattr__(self, "_frozen_", True) + + +# Device only object +class SwapABFc1Epilogue(_ImmutableAfterInit): + def __init__( + self, + base: SwapABGatedActEpilogue, + tidx: cutlass.Int32, + staging_pointer: cute.Pointer, + kernel_extension: BlockScaledSwapAbFc12Extension, + tma_atom_fc1_output: cute.CopyAtom, + fc1_output: cute.Tensor, # fake (m,n,l) domain + fc1_output_sf: cute.Tensor, # fake (m,n,l) domain + fc1_done_counter: cute.Tensor, # 1D tensor + optional_epi_args: GatedActEpilogueArgs, + ): + self.base = base + self.tidx = tidx % (base._EpilogueWarpCnt * 32) + self.warp_idx = self.tidx // 32 + self.lane_idx = self.tidx % 32 + # (token64, intermediate, stage). The swizzle travels with the layout instead of being + # spelled out here, so an 8-bit fc1 output picks up its own atom without a code change. + staged_layout = base.fc1_staged_smem_layout(base.subtile_cnt) + self.smem_tensor = cute.make_tensor( + cute.recast_ptr(staging_pointer, staged_layout.inner, dtype=base.fc1_output_dtype), staged_layout.outer + ) + # Kept unswizzled and untyped so fc1_quant can carve an fp32 amax-exchange view out of one + # staging stage without going through the staging tensor's swizzle. + self.staging_pointer = staging_pointer + self.kernel_extension = kernel_extension + self.fc1_tma_atom = tma_atom_fc1_output + self.fc1_output = fc1_output + self.fc1_output_sf = fc1_output_sf + self.fc1_done_counter = fc1_done_counter + self.optional_epi_args = optional_epi_args + self._freeze() + + def __getattr__(self, name): + return getattr(object.__getattribute__(self, "base"), name) + + def __extract_mlir_values__(self) -> List[ir.Value]: + # This object is a loop-invariant Python context wrapper, not a + # dynamic value. Keep it out of scf.while iter_args and reconstruct by + # identity across region boundaries. Any field that becomes a + # loop-carried SSA value must be passed explicitly to __call__ instead + # of being stored here. + return [] + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "SwapABFc1Epilogue": + assert len(values) == 0 + return self + + @cute.jit + def signal_fc1_done(self, work_tile_info, next_work_tile_info, flag_tracker): + # Only in-bound intermediate_downproj tiles signal; OOB -> null slot. + needs_intermediate_guard = ( + self.intermediate_downproj is None + or self.intermediate_downproj % self.cluster_tile_intermediate_downproj != 0 + ) + if cutlass.const_expr(needs_intermediate_guard): + in_bound = work_tile_info.tile_m_idx * self._EpilogueFc1IntermediateDownTileSize < self.fc1_output.shape[1] + else: + in_bound = True + slot = work_tile_info.cumulative_token_block_count + work_tile_info.tile_n_idx + flag_address = Int64(0) + if in_bound: + flag_address = (self.fc1_done_counter.iterator + slot).toint() + return flag_tracker.accumulate(next_work_tile_info.phase, self.fc1_epi_flag_batch, flag_address) + + @cute.jit + def __call__( + self, + work_tile_info: SwapAbFc12WorkTileInfo, + tmem_acc_tensor: cute.Tensor, # (cta_tile_m, cta_tile_n) + acc_pipeline, + acc_consumer_state, + is_odd_turn: cutlass.Int32, + ): + # (tokens_this_expert, intermediate_down, 1) + real_fc1_output, _ = self.kernel_extension.get_gmem_tensor("c", self.fc1_output, work_tile_info) + # (tokens_this_expert, intermediate_down, 1) + real_fc1_output_sf, _ = self.kernel_extension.get_gmem_tensor("sfc", self.fc1_output_sf, work_tile_info) + # subtile-irrevalent hoist out here. + if cutlass.const_expr(self.optional_epi_args.fc1_alpha is not None): + alpha_val = self.optional_epi_args.fc1_alpha[work_tile_info.expert_idx] + else: + alpha_val = None + if cutlass.const_expr(self.optional_epi_args.fc1_norm_const is not None): + norm_const = self.optional_epi_args.fc1_norm_const[work_tile_info.expert_idx] + else: + norm_const = None + # (cta_tile_m, cta_tile_n) -> (epi_tile_m, epi_tile_n, iters) + tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( + tmem_acc_tensor, (self._EpilogueFc1IntermediateGateUpTileSize, self._EpilogueTokenTileSize) + )[None, None, 0, None] + + acc_pipeline.consumer_wait(acc_consumer_state) + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + + # Overlap path preloads two subtiles before releasing acc TMEM. + unroll_tile_cnt = 2 if cutlass.const_expr(self.overlapping_accum) else 0 + remain_subtile_cnt = self.subtile_cnt - unroll_tile_cnt + + if cutlass.const_expr(unroll_tile_cnt > 0): + subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + subtile_idx_second = (cutlass.Int32(self.subtile_cnt + 1) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + + # preload_subtile_first: subtile_idx_first's raw PRE-transpose acc, LDTM'd by + # all 128 epi threads into 4 reg tensors == the 4 quadrants of the subtile's + # (128 tmem_dp x 64 tmem_col) footprint. Only these raw-TMEM offsets are + # guaranteed: + # reg[0]/reg[1], reg[2]/reg[3] : top vs bot -> 16 apart in tmem_dp + # reg[0]/reg[2], reg[1]/reg[3] : 1st vs 2nd half -> 32 apart in tmem_col + # (so reg[0..1] = the first 128x32, reg[2..3] = the second 128x32 of the 128x64.) + # The per-lane (lane_idx, elem_idx) -> (tmem_dp, tmem_col) layout INSIDE each + # reg tensor is opaque -- do not assume it; it only becomes well-defined once + # the tmem transpose consumes them. + preload_subtile_first: Tuple[cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor] = ( + _TmemTranspose16x32Core.load_subtile_raw_acc( + tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_first] + ) + ) + + # Release acc to next MMA unconditionally. + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + # preload_subtile_second: same 128 tmem_dp x 64 tmem_col footprint, but for + # subtile_idx_second (the other token subtile, not the 2nd col-half). Same + # quadrant/offset invariants and opaque per-lane layout as above. + preload_subtile_second: Tuple[cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor] = ( + _TmemTranspose16x32Core.load_subtile_raw_acc( + tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second] + ) + ) + + # Both unrolled subtiles borrow tmem_subtile_second as workspace. + preload_pair = (preload_subtile_first, preload_subtile_second) + subtile_idx_pair = (subtile_idx_first, subtile_idx_second) + for i in cutlass.range_constexpr(unroll_tile_cnt): + if subtile_idx_pair[i] * cutlass.Int32(self._EpilogueTokenTileSize) < valid_tokens: + self.run_subtile( + work_tile_info=work_tile_info, + subtile_idx=subtile_idx_pair[i], + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second], + preload_acc=preload_pair[i], + fc1_output=real_fc1_output, + fc1_output_sf=real_fc1_output_sf, + alpha_val=alpha_val, + norm_const=norm_const, + ) + + for i in cutlass.range(remain_subtile_cnt, unroll=1): + real_i = i + unroll_tile_cnt + if cutlass.const_expr(self.overlapping_accum): + subtile_idx = (cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + else: + subtile_idx = cutlass.Int32(real_i) + + if subtile_idx * cutlass.Int32(self._EpilogueTokenTileSize) < valid_tokens: + self.run_subtile( + work_tile_info=work_tile_info, + subtile_idx=subtile_idx, + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx], + preload_acc=None, + fc1_output=real_fc1_output, + fc1_output_sf=real_fc1_output_sf, + alpha_val=alpha_val, + norm_const=norm_const, + ) + + # Non-overlap-path release: at the natural task-tile boundary. + if cutlass.const_expr(not self.overlapping_accum): + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + @cute.jit + def run_subtile( + self, + work_tile_info: SwapAbFc12WorkTileInfo, + subtile_idx: cutlass.Int32, + # (intermedaite_gateup_tile, token_subtile), fundamentally (epi_tile_m, epi_tile_n) + tmem_subtile_tensor: cute.Tensor, + # Rmems preloaded from tmem, contract with downstream tmem trans. Do not assume mapping here. + preload_acc: Tuple[cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor], + # (tokens_this_expert, intermediate_down, 1) + fc1_output: cute.Tensor, + fc1_output_sf: cute.Tensor, + alpha_val: Optional[cutlass.Float32], + norm_const: Optional[cutlass.Float32], + ): + if cutlass.const_expr(self.optional_epi_args.topk_scores is not None): + # This means we need to perform DeepGEMM computation graph, topk_score at fc1 pre-quant + topk_score_tensor, _ = self.kernel_extension.get_gmem_tensor( + "topk", self.optional_epi_args.topk_scores, work_tile_info + ) # (tokens_this_expert) + else: + topk_score_tensor = None + + # Mapping of the transposed accumulator for orthodox NVFP4 output: + # (epi_tid, val_id) -> (token_idx, intermediate_down_idx) + # token_idx = epi_tid % 32 + val_id // 16 * 32 + # intermediate_down_idx = val_id % 16 + epi_tid // 32 * 16 + # Each thread holds (intermediate_down_16, token_2):(1, 16) + + # Step -1: preload topk scores. + current_two_token_idices = ( + work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + self.lane_idx, + work_tile_info.tile_n_idx * self.cta_tile_n + + subtile_idx * self._EpilogueTokenTileSize + + self.lane_idx + + 32, + ) + if cutlass.const_expr(topk_score_tensor is not None): + topk_scores = ( + topk_score_tensor[current_two_token_idices[0]], + topk_score_tensor[current_two_token_idices[1]], + ) + else: + topk_scores = None + + # Step 0: load tmem + if cutlass.const_expr(preload_acc is not None): + gate_token_0_32, up_token_0_32, gate_token_32_64, up_token_32_64 = preload_acc + else: + gate_token_0_32 = cute.make_rmem_tensor((16,), cutlass.Float32) + up_token_0_32 = cute.make_rmem_tensor((16,), cutlass.Float32) + gate_token_32_64 = cute.make_rmem_tensor((16,), cutlass.Float32) + up_token_32_64 = cute.make_rmem_tensor((16,), cutlass.Float32) + # Although hardcode is not right, but since the whole tmem transpose is too tricky, I have to hardcode... + # (epi_tile_m, epi_tile_n) -> (warp_local_epi_tile_m, epi_tile_n) + # tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, self.warp_idx), None] + tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, 0), None] + # (warp_local_epi_tile_m, epi_tile_n) -> (((16, 32), 1), (2, 2)) + tmem_subtile_tensor_in_first_load_view = cute.logical_divide( + cute.zipped_divide(tmem_subtile_tensor_per_warp, (16, 32)), ((16, 32), 1) + ) + atom = cute.make_copy_atom(tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), cutlass.Float32) + cute.copy( + atom, + wrap_into_copy_standard_layout(tmem_subtile_tensor_in_first_load_view[None, 0]), + wrap_into_copy_standard_layout(gate_token_0_32), + ) + cute.copy( + atom, + wrap_into_copy_standard_layout(tmem_subtile_tensor_in_first_load_view[None, 1]), + wrap_into_copy_standard_layout(up_token_0_32), + ) + cute.copy( + atom, + wrap_into_copy_standard_layout(tmem_subtile_tensor_in_first_load_view[None, 2]), + wrap_into_copy_standard_layout(gate_token_32_64), + ) + cute.copy( + atom, + wrap_into_copy_standard_layout(tmem_subtile_tensor_in_first_load_view[None, 3]), + wrap_into_copy_standard_layout(up_token_32_64), + ) + + # Step 1: perform swiglu on the first part, interleave with the second's 32x32 tmem transpose. + token_0_32_pre_quant_pre_trans = self.alpha_swiglu_clamp(gate_token_0_32, up_token_0_32, alpha_val) + + # gate_token_32_64 / up_token_32_64 are already in the transpose input + # distribution (see TmemTranspose16x32 / load_subtile_raw_acc). + token_32_64_tmem_trans = TmemTranspose32x32Inplace( + tmem_subtile_tensor.iterator, reg_tensor_top=gate_token_32_64, reg_tensor_bot=up_token_32_64 + ) + + # Transpose output: each lane holds (token_1, intermediate_16); tmem_dp + # = lane_idx (token), tmem_col = elem_idx (intermediate output idx). + gate_token_32_64_trans_pre_act, up_token_32_64_trans_pre_act = ( + token_32_64_tmem_trans.from_r1_perm_until_last_store() + ) + + token_32_64_pre_quant = self.alpha_swiglu_clamp( + gate_token_32_64_trans_pre_act, up_token_32_64_trans_pre_act, alpha_val + ) + + token_0_32_tmem_trans = TmemTranspose16x32( + tmem_subtile_tensor.iterator, Region.Top, reg_tensor=token_0_32_pre_quant_pre_trans + ) + token_0_32_pre_quant = token_0_32_tmem_trans.from_r1_perm_until_last_store() + + # Step 2: Quant + self.fc1_quant( + work_tile_info=work_tile_info, + two_token=(token_0_32_pre_quant, token_32_64_pre_quant), + topk_scores=topk_scores, + norm_const=norm_const, + intermediate_output_size=cute.size(fc1_output, 1), + fc1_output_sf=fc1_output_sf, + subtile_idx=subtile_idx, + ) + + # Step 3: TMASTG + # (token_64, intermeidate_64) + fc1_smem = self.smem_tensor[None, None, subtile_idx] + # (token, intermediate_down, l=1) -> (cta_token, cta_intermediate_down) + fc1_gmem_cta_view = cute.flat_divide(fc1_output, (self.cta_tile_n, self.cta_tile_m // 2))[ + None, None, work_tile_info.tile_n_idx, work_tile_info.tile_m_idx, 0 + ] + # (cta_token, cta_intermediate_down) -> (token_64, intermediate_64) + fc1_gmem_subtile_view = cute.flat_divide( + fc1_gmem_cta_view, (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize) + )[None, None, subtile_idx, 0] + tma_smem_src, tma_gmem_dst = cpasync.tma_partition( + self.fc1_tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(fc1_smem, 0, 2), + cute.group_modes(fc1_gmem_subtile_view, 0, 2), + ) + + subtile_bar_id = subtile_idx + cutlass.Int32(SwapABGatedActEpilogue._EpilogueAsyncBarIdBase) + tma_ready_to_read_smem_named_barrier = pipeline.NamedBarrier( + barrier_id=subtile_bar_id, num_threads=self._EpilogueWarpCnt * 32 + ) + cute.arch.fence_proxy("async.shared", space="cta") + if self.warp_idx == subtile_idx: + tma_ready_to_read_smem_named_barrier.arrive_and_wait() + with cute.arch.elect_one(): + # if work_tile_info.tile_m_idx * (self.cta_tile_m // 2) < cute.size(fc1_output, 1): + cute.copy(self.fc1_tma_atom, tma_smem_src, tma_gmem_dst) + cute.arch.cp_async_bulk_commit_group() + else: + tma_ready_to_read_smem_named_barrier.arrive() + + @cute.jit + def alpha_swiglu_clamp( + self, + gate_rmem: cute.Tensor, # Raw fc1 acc (pre-dequant); even-size 1D fp32 rmem + up_rmem: cute.Tensor, # Raw fc1 acc (pre-dequant); even-size 1D fp32 rmem + alpha_val: Optional[cutlass.Float32], + ) -> cute.Tensor: + # ── Input contract checks (compile-time): fp32, 1D, even-count, rmem ── + # Wrapped in const_expr so the DSL evaluates them at trace time and the + # raise fires during compilation rather than emitting a runtime branch. + for _name, _t in (("gate_rmem", gate_rmem), ("up_rmem", up_rmem)): + if cutlass.const_expr(_t.element_type is not cutlass.Float32): + raise TypeError(f"alpha_swiglu_clamp: {_name} must be Float32, got {_t.element_type}") + if cutlass.const_expr(_t.memspace != AddressSpace.rmem): + raise ValueError( + f"alpha_swiglu_clamp: {_name} must be a register (rmem) tensor, got address space {_t.memspace}" + ) + if cutlass.const_expr(cute.rank(_t) != 1): + raise ValueError(f"alpha_swiglu_clamp: {_name} must be 1D, got rank {cute.rank(_t)}") + if cutlass.const_expr(cute.size(_t) % 2 != 0): + raise ValueError(f"alpha_swiglu_clamp: {_name} element count must be even, got {cute.size(_t)}") + if cutlass.const_expr(cute.size(gate_rmem) != cute.size(up_rmem)): + raise ValueError( + "alpha_swiglu_clamp: gate_rmem and up_rmem must have equal size, got " + f"{cute.size(gate_rmem)} vs {cute.size(up_rmem)}" + ) + + # gate_rmem / up_rmem are the RAW fc1 fp32 accumulator (pre-dequant). + # Order follows the NVFP4 -> fp32 -> SwiGLU contract and MUST be: + # + # 1. dequant: gate = alpha * gate_raw ; up = alpha * up_raw + # (alpha = expert-wise global scale on the acc; None => alpha == 1.) + # 2. clamp the DEQUANTED (real) values, gpt-oss ``_apply_gate`` style: + # gate = min(gate, +limit) (upper bound only) + # up = clamp(up, -limit, +limit) (symmetric) + # 3. gated activation, either SwiGLU or SiTU (mutually exclusive; SiTU has no clamp): + # SwiGLU: out = up * gate * sigmoid(gate) + # SiTU: out = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) + # sigmoid(x) = rcp(1 + exp2(-x * log2e)) + # + # The symmetric up-clamp is a single ``min.xorsign.abs.f32`` (magnitude + # min(|up|, limit), sign = sign(up)^sign(limit) = sign(up) since limit>=0); + # the gate-clamp is a plain ``min.f32``. ``.xorsign.abs`` has no f32x2 form, + # so dequant+clamp run scalar while the swiglu core stays packed f32x2. + n = cute.size(gate_rmem) + out = cute.make_rmem_tensor((n,), cutlass.Float32) + log2_e = 1.4426950408889634 + + neg_log2e_pair = (cutlass.Float32(-log2_e), cutlass.Float32(-log2_e)) + one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) + if cutlass.const_expr(self.gate_up_clamp is not None): + limit = cutlass.Float32(self.gate_up_clamp) + + for i in cutlass.range_constexpr(0, n, 2): + g0 = gate_rmem[i] + g1 = gate_rmem[i + 1] + u0 = up_rmem[i] + u1 = up_rmem[i + 1] + + # 1) dequant raw acc to real values (skip entirely when alpha is None). + if cutlass.const_expr(alpha_val is not None): + alpha_pair = (alpha_val, alpha_val) + g0, g1 = cute.arch.mul_packed_f32x2((g0, g1), alpha_pair) + u0, u1 = cute.arch.mul_packed_f32x2((u0, u1), alpha_pair) + + # 2) clamp the real values (skip when no clamp configured). + if cutlass.const_expr(self.gate_up_clamp is not None): + # gate upper-clamp: min(gate, +limit) + g0 = cutlass.Float32( + llvm.inline_asm( + cutlass.Float32.mlir_type, + [g0.ir_value(), limit.ir_value()], + "min.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + g1 = cutlass.Float32( + llvm.inline_asm( + cutlass.Float32.mlir_type, + [g1.ir_value(), limit.ir_value()], + "min.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + # up symmetric-clamp: clamp(up, -limit, +limit) in one instruction + u0 = cutlass.Float32( + llvm.inline_asm( + cutlass.Float32.mlir_type, + [u0.ir_value(), limit.ir_value()], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + u1 = cutlass.Float32( + llvm.inline_asm( + cutlass.Float32.mlir_type, + [u1.ir_value(), limit.ir_value()], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + # 3) gated activation on the dequanted (and clamped) real values. + # sigmoid(x) = rcp(1 + exp2(-x * log2e)) -- shared by both cores. + def _sigmoid(p0, p1): + neg = cute.arch.mul_packed_f32x2((p0, p1), neg_log2e_pair) + e = (cute.math.exp2(neg[0], fastmath=True), cute.math.exp2(neg[1], fastmath=True)) + d = cute.arch.add_packed_f32x2(e, one_pair) + return (cute.arch.rcp_approx(d[0]), cute.arch.rcp_approx(d[1])) + + sigmoid_pair = _sigmoid(g0, g1) + + if cutlass.const_expr(self.situ_beta is None): + # SwiGLU: out = up * gate * sigmoid(gate) + ug = cute.arch.mul_packed_f32x2((u0, u1), (g0, g1)) + out_pair = cute.arch.mul_packed_f32x2(ug, sigmoid_pair) + else: + # SiTU (Kimi K3), matching HF ``modeling_kimi.py``: + # situ_gate = beta * tanh(gate / beta) * sigmoid(gate) + # situ_up = linear_beta * tanh(up / linear_beta) + # out = situ_gate * situ_up + # + # ``tanh(z) = 2 * sigmoid(2z) - 1`` keeps the whole core on the packed f32x2 path -- + # there is no packed tanh, so calling one would force this loop back to scalar. + # + # beta * tanh(x/beta) = beta * (2*sigmoid(2x/beta) - 1) = 2*beta*sigmoid(2x/beta) - beta + # so the reciprocals and the 2*beta factors fold at trace time. + inv_2beta = cutlass.Float32(2.0 / self.situ_beta) + two_beta = cutlass.Float32(2.0 * self.situ_beta) + neg_beta = cutlass.Float32(-self.situ_beta) + inv_2lbeta = cutlass.Float32(2.0 / self.situ_linear_beta) + two_lbeta = cutlass.Float32(2.0 * self.situ_linear_beta) + neg_lbeta = cutlass.Float32(-self.situ_linear_beta) + + gs = _sigmoid(*cute.arch.mul_packed_f32x2((g0, g1), (inv_2beta, inv_2beta))) + tanh_g = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(gs, (two_beta, two_beta)), (neg_beta, neg_beta) + ) + + us = _sigmoid(*cute.arch.mul_packed_f32x2((u0, u1), (inv_2lbeta, inv_2lbeta))) + tanh_u = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(us, (two_lbeta, two_lbeta)), (neg_lbeta, neg_lbeta) + ) + + situ_gate = cute.arch.mul_packed_f32x2(tanh_g, sigmoid_pair) + out_pair = cute.arch.mul_packed_f32x2(situ_gate, tanh_u) + + out[i] = out_pair[0] + out[i + 1] = out_pair[1] + + return out + + @cute.jit + def fc1_quant( + self, + work_tile_info: SwapAbFc12WorkTileInfo, + two_token: Tuple[cute.Tensor, cute.Tensor], # two rmem tensor, each fp32 @ (token_1, intermediate_16) + topk_scores: Optional[Tuple[cutlass.Float32, cutlass.Float32]], + norm_const: Optional[cutlass.Float32], + intermediate_output_size: cutlass.Int32, + fc1_output_sf: cute.Tensor, # MoE domain (token_this_rank, intermediate_down, 1) + subtile_idx: cutlass.Int32, + ): + # ``two_token`` are the two post-swiglu, transposed token rmem tensors; each lane holds one + # token's 16 intermediate-output values. half 0 -> token (lane), half 1 -> (lane+32). + # + # Those 16 values are a whole scale block at sf_vec 16 (nvfp4) but only half of one at + # sf_vec 32 (mx), where the other half sits in the paired warp -- hence the two directions + # below. Both are handed the two tokens in one call so the paired path needs a single + # exchange barrier per subtile rather than one per token. + # + # Per token (ported from PostSwigluHalf._gen_sfc_quantize + stg_sfc + r2s): + # 1. (Path A) pre-multiply topk weight into the values, if present. + # 2. block quant -> data regs + one scale factor (QuantImpl's job). + # 3. write the scale factor to fc1_output_sf[token, intermediate_idx, 0] + # (plain scalar store; predicated unless statically in-bound). + # 4. STS the quantized values into this subtile's shared output stage. + # norm_const is treated like alpha_val: None => behaves as 1.0 (factors const-elided, not + # multiplied by 1.0). It only exists for nvfp4; an e8m0 scale absorbs the rescale itself. + values_per_token = cute.size(two_token[0]) + if cutlass.const_expr(self.needs_pair_amax_exchange): + # The exchange slots live in THIS subtile's staging stage. That stage is dead right + # now -- the previous tile's copy of it was drained at the tile boundary, and this + # tile writes it only after the retire barrier below -- whereas every other stage may + # still have a TMA store reading it. Borrowing any other stage would corrupt it. + amax_exchange = mark_alignment( + cute.make_tensor( + cute.recast_ptr(self.staging_pointer, dtype=cutlass.Float32) + + subtile_idx * cutlass.Int32(self.fc1_staging_stage_bytes // 4), + cute.make_layout( + (self.fc1_amax_token_chunks, self._EpilogueWarpCnt * 32), stride=(1, self.fc1_amax_token_chunks) + ), + ), + 16, + ) + quant = QuantImpl( + self.quant_kind, + "regs_in_pair_threads", + lane_idx=self.lane_idx, + warp_idx=self.warp_idx, + pair_exchange_barrier=SwapABGatedActEpilogue.epilogue_sync_barrier(), + ) + else: + amax_exchange = None + quant = QuantImpl(self.quant_kind, "regs_in_thread") + + # Both warps of a pair derive the same scale, so only the even one stores it. The scale + # plane folds any coordinate inside a block onto that block's slot, so the per-warp base + # addresses the right entry for either vec size. + intermediate_idx = ( + work_tile_info.tile_m_idx * (self.cta_tile_m // 2) + + self.warp_idx * self._EpilogueFc1IntermediateDownPerWarp + ) + if cutlass.const_expr(self.needs_pair_amax_exchange): + stores_scale_factor = self.warp_idx % 2 == 0 + else: + stores_scale_factor = True + subtile_token_start = work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + token_idx_pair = (subtile_token_start + self.lane_idx, subtile_token_start + self.lane_idx + 32) + + # This subtile's (token, intermediate) shared output stage, tiled into (1, 16) blocks so + # each thread's cells slice out directly (zipped_divide + slice; avoids the ambiguous + # local_tile surface). + smem_stage = self.smem_tensor[None, None, subtile_idx] + # (token_64, intermediate_down_64) -> ((1, 16), (token_tile_size, warp_cnt)) + smem_tiled = cute.zipped_divide(smem_stage, (1, values_per_token)) + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.fc1_output_dtype, + num_bits_per_copy=values_per_token * self.fc1_output_dtype.width, + ) + + # 1) topk-weight pre-multiply (Path A) into a weighted scratch, both tokens back to back. + weighted = cute.make_rmem_tensor((2 * values_per_token,), cutlass.Float32) + for half in cutlass.range_constexpr(2): + tok = two_token[half] + base = half * values_per_token + if cutlass.const_expr(topk_scores is not None): + topk_pair = (topk_scores[half], topk_scores[half]) + for i in cutlass.range_constexpr(0, values_per_token, 2): + w0, w1 = cute.arch.mul_packed_f32x2((tok[i], tok[i + 1]), topk_pair) + weighted[base + i] = w0 + weighted[base + i + 1] = w1 + else: + for i in cutlass.range_constexpr(0, values_per_token): + weighted[base + i] = tok[i] + + # 2) Core block quant. One scale factor per token either way; the paired direction spends + # its exchange barrier inside this call. + data_regs, sf_regs = quant(weighted, norm_const=norm_const, smem_intermediate=amax_exchange) + data_by_token = cute.zipped_divide(data_regs, (values_per_token,)) + + # Retire the exchange: its slots are this stage's own bytes, so no thread may start + # writing the stage until every thread has read its partner's amax. Sitting after the + # quantization rather than before it puts the sync latency behind work that is already in + # flight (the e8m0 path alone goes through MUFU). + if cutlass.const_expr(self.needs_pair_amax_exchange): + SwapABGatedActEpilogue.epilogue_sync_barrier().arrive_and_wait() + + for half in cutlass.range_constexpr(2): + # 3) scale-factor store (predicate const-elided when statically in-bound, mirroring + # signal_fc1_done's intermediate predicate). + if stores_scale_factor: + if cutlass.const_expr( + self.intermediate_downproj is None + or self.intermediate_downproj % self.cluster_tile_intermediate_downproj != 0 + ): + if intermediate_idx < intermediate_output_size: + fc1_output_sf[token_idx_pair[half], intermediate_idx, 0] = sf_regs[half] + else: + fc1_output_sf[token_idx_pair[half], intermediate_idx, 0] = sf_regs[half] + + # 4) STS the quantized values into this subtile's shared output stage. + # ((1, 16), (token_tile_size, warp_cnt)) -> (16) + smem_thread_row = smem_tiled[(0, None), (self.lane_idx + 32 * half, self.warp_idx)] + cute.copy(store_atom, cute.coalesce(data_by_token[None, half]), cute.coalesce(smem_thread_row)) + + +@dataclasses.dataclass(frozen=True) +class Fc2ProcessPipeline: + tmem_acc_load: Callable + f2fp: Callable + post_f2fp_reorder: Callable + store_function: Callable + # Kept as a finer-grained, elem-level reading aid for the store-out layout + # (never evaluated); ``store_out_mapping`` is the per-issue form that the + # router actually evaluates at runtime to drive metadata / pointer math. + fc2_cta_tile_mapping: FunctionMapping + store_out_mapping: FunctionMapping + require_tmem_trans: bool + # SF plane per-issue mapping; None for the bf16 (unquantized) paths. + sf_store_out_mapping: Optional[FunctionMapping] = None + + +# Device only object +class SwapABFc2Epilogue(_ImmutableAfterInit): + def __init__( + self, + base: SwapABGatedActEpilogue, + tidx: cutlass.Int32, + smem_tensor: Optional[cute.Tensor], + tma_atom_fc2_output: Optional[cute.CopyAtom], + fc2_tma_output: Optional[cute.Tensor], + fc2_output: cute.Tensor, # MoE domain (token, topk, hidden) + token_src_metadata: Optional[cute.Tensor], + fc2_done_counter: Optional[cute.Tensor], + fc2_output_sf: Optional[cute.Tensor], + peer_rank_ptr_mapper: Optional[SymmetricBufferDevice], + optional_epi_args: GatedActEpilogueArgs, + ): + self.base = base + self.tidx = tidx % (base._EpilogueWarpCnt * 32) + self.warp_idx = self.tidx // 32 + self.lane_idx = self.tidx % 32 + self.tma_atom_fc2_output = tma_atom_fc2_output + self.fc2_tma_output = fc2_tma_output + self.fc2_output = fc2_output + self.token_src_metadata = token_src_metadata + self.fc2_done_counter = fc2_done_counter + self.fc2_output_sf = fc2_output_sf + self.peer_rank_ptr_mapper = peer_rank_ptr_mapper + self.optional_epi_args = optional_epi_args + if cutlass.const_expr(base.fc2_use_tma): + self.smem_tensor = smem_tensor + self.process_pipeline = make_fc2_tma_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ) + elif cutlass.const_expr(base.fc2_use_ublk): + self.smem_tensor = smem_tensor + self.process_pipeline = make_fc2_ublk_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ) + else: + self.smem_tensor = None + if cutlass.const_expr(base.reduce_topk_in_epilogue): + self.process_pipeline = make_fc2_redg_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ) + else: + self.process_pipeline = make_fc2_stg_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ) + self._freeze() + + def __getattr__(self, name): + return getattr(object.__getattribute__(self, "base"), name) + + def __extract_mlir_values__(self) -> List[ir.Value]: + # See SwapABFc1Epilogue.__extract_mlir_values__: this helper carries + # only loop-invariant Python context. It intentionally serializes no + # MLIR values, so changing it to store loop-carried state would be a + # correctness bug. + return [] + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "SwapABFc2Epilogue": + assert len(values) == 0 + return self + + @cute.jit + def signal_fc2_done(self, work_tile_info, next_work_tile_info, flag_tracker): + publish: cutlass.Constexpr = self.token_back_enabled + flag_address = Int64(0) + if cutlass.const_expr(publish): + flag_address = (self.fc2_done_counter.iterator + work_tile_info.expert_idx).toint() + no_fire: cutlass.Constexpr = not publish + return flag_tracker.accumulate(next_work_tile_info.phase, self.fc2_epi_flag_batch, flag_address, no_fire) + + @cute.jit + def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2OutputRouter": + task_tile_data_row_start = ( + work_tile_info.cumulative_data_physical_row + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) + ) + hidden_base_this_cta_tile = work_tile_info.tile_m_idx * cutlass.Int32(self.cta_tile_m) + valid_hidden_this_cta_tile = cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile + if valid_hidden_this_cta_tile < 0: + valid_hidden_this_cta_tile = 0 + if valid_hidden_this_cta_tile > self._EpilogueFc2HiddenTileSize: + valid_hidden_this_cta_tile = self._EpilogueFc2HiddenTileSize + + metadata = None + peer_rank_ptr_mapper = None + data_token_base = task_tile_data_row_start + if cutlass.const_expr(self.token_src_metadata is not None and not self.token_back_push_data): + metadata = cute.domain_offset((task_tile_data_row_start,), self.token_src_metadata) + peer_rank_ptr_mapper = self.peer_rank_ptr_mapper + data_token_base = None + + if cutlass.const_expr(self.combine_format.is_quantized): + base_outputs = (self.fc2_output, self.fc2_output_sf) + token_bases = (data_token_base, task_tile_data_row_start) + output_mappings = (self.process_pipeline.store_out_mapping, self.process_pipeline.sf_store_out_mapping) + else: + base_outputs = self.fc2_output + token_bases = data_token_base + output_mappings = self.process_pipeline.store_out_mapping + + return Fc2OutputRouter( + metadata=metadata, + token_bases=token_bases, + base_outputs=base_outputs, + hidden_base_this_cta_tile=hidden_base_this_cta_tile, + peer_rank_ptr_mapper=peer_rank_ptr_mapper, + valid_tokens_this_cta_tile=work_tile_info.valid_tokens_in_cta_tile, + valid_hidden_this_cta_tile=valid_hidden_this_cta_tile, + reduce_topk_in_epilogue=self.reduce_topk_in_epilogue, + output_mappings=output_mappings, + epi_tid=self.tidx, + combine_format=self.combine_format, + ).prefetch() + + @cute.jit + def __call__( + self, + work_tile_info: SwapAbFc12WorkTileInfo, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + acc_consumer_state, + is_odd_turn: cutlass.Int32, + ): + # subtile-irrelevant hoist: fc2 alpha scales raw fc2 accumulators before f2fp. + if cutlass.const_expr(self.optional_epi_args.fc2_alpha is not None): + alpha_val = self.optional_epi_args.fc2_alpha[work_tile_info.expert_idx] + else: + alpha_val = None + acc_ready = False + if not work_tile_info.peek_ready: + acc_ready = True + acc_pipeline.consumer_wait(acc_consumer_state) + fc2_output_router = self._make_output_router(work_tile_info) + # (cta_tile_m, cta_tile_n) -> (epi_tile_m, epi_tile_n, iters) + tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( + tmem_acc_tensor, (self._EpilogueFc2HiddenTileSize, self._EpilogueTokenTileSize) + )[None, None, 0, None] + + acc_pipeline.consumer_wait(acc_consumer_state, acc_ready) + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + + # Overlap path preloads two subtiles before releasing acc TMEM. + unroll_tile_cnt = ( + 2 if cutlass.const_expr(self.overlapping_accum and self.process_pipeline.require_tmem_trans) else 0 + ) + remain_subtile_cnt = self.subtile_cnt - unroll_tile_cnt + + if cutlass.const_expr(unroll_tile_cnt > 0): + subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + subtile_idx_second = (cutlass.Int32(self.subtile_cnt + 1) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + + # Each warp preloads its local 32-hidden x 64-token accumulator as two + # (32,) FP32 tensors: + # + # preload[0]: hidden 0..15 x token 0..63 + # preload[1]: hidden 16..31 x token 0..63 + # + # One LDTM.16dp256bit.x8 produces each tensor. Within it, every + # consecutive register pair holds adjacent token columns at one + # hidden coordinate. After F2FP, that pair is one BF16x2 MOVM input. + preload_subtile_first: Tuple[cute.Tensor, ...] = self.process_pipeline.tmem_acc_load( + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_first], epi=self + ) + + # Release acc to next MMA unconditionally. + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + # Same two-half mapping for the other token subtile. Both preloaded + # subtiles can now use this second subtile's TMEM as transpose workspace. + preload_subtile_second: Tuple[cute.Tensor, ...] = self.process_pipeline.tmem_acc_load( + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second], epi=self + ) + + # Both unrolled subtiles borrow tmem_subtile_second as workspace. + preload_pair = (preload_subtile_first, preload_subtile_second) + subtile_idx_pair = (subtile_idx_first, subtile_idx_second) + for i in cutlass.range_constexpr(unroll_tile_cnt): + if subtile_idx_pair[i] * cutlass.Int32(self._EpilogueTokenTileSize) < valid_tokens: + self.run_subtile( + work_tile_info=work_tile_info, + epilogue_iter_idx=cutlass.Int32(i), + subtile_idx=subtile_idx_pair[i], + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second], + preload_acc=preload_pair[i], + fc2_output_router=fc2_output_router, + alpha_val=alpha_val, + release_after_ldtm=False, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + ) + + if cutlass.const_expr(self.overlapping_accum and unroll_tile_cnt == 0): + release_after_ldtm = True + else: + release_after_ldtm = False + for i in cutlass.range(remain_subtile_cnt, unroll=1): + # for i in cutlass.range_constexpr(remain_subtile_cnt): + real_i = i + unroll_tile_cnt + if cutlass.const_expr(self.overlapping_accum): + subtile_idx = (cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + else: + subtile_idx = cutlass.Int32(real_i) + + if subtile_idx * cutlass.Int32(self._EpilogueTokenTileSize) < valid_tokens: + self.run_subtile( + work_tile_info=work_tile_info, + epilogue_iter_idx=real_i, + subtile_idx=subtile_idx, + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx], + preload_acc=None, + fc2_output_router=fc2_output_router, + alpha_val=alpha_val, + release_after_ldtm=release_after_ldtm, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + ) + release_after_ldtm = False + + # Non-overlap-path release: at the natural task-tile boundary. + if cutlass.const_expr(not self.overlapping_accum): + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + @cute.jit + def run_subtile( + self, + work_tile_info: SwapAbFc12WorkTileInfo, + epilogue_iter_idx: cutlass.Int32, + subtile_idx: cutlass.Int32, + # (hidden_tile, token_subtile), fundamentally (epi_tile_m, epi_tile_n) + tmem_subtile_tensor: cute.Tensor, + preload_acc: Optional[Tuple[cute.Tensor, ...]], + fc2_output_router: "Fc2OutputRouter", + alpha_val: Optional[cutlass.Float32], + release_after_ldtm: Union[cutlass.Boolean, bool], + acc_pipeline, + acc_consumer_state, + ): + process_pipeline = self.process_pipeline + if cutlass.const_expr(preload_acc is None): + loaded = process_pipeline.tmem_acc_load(tmem_subtile_tensor=tmem_subtile_tensor, epi=self) + if release_after_ldtm: + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + else: + loaded = preload_acc + + casted = process_pipeline.f2fp(*loaded, alpha_val=alpha_val) + # reorder returns a bare RMEM fragment in the store's expected pre-store + # distribution; reorder + store are paired 1:1 inside the pipeline. + pre_store = process_pipeline.post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_tensor) + process_pipeline.store_function( + epi=self, + subtile=pre_store, + work_tile_info=work_tile_info, + epilogue_iter_idx=epilogue_iter_idx, + subtile_idx=subtile_idx, + fc2_output_router=fc2_output_router, + ) + + +@dataclasses.dataclass(frozen=True) +class Fc2OutputRouter: + # One packed i64 TokenSrcMetadata record per pool token. None means a local write. + metadata: Optional[cute.Tensor] + # token + possible sf + token_bases: Union[Tuple[Optional[cutlass.Int32], cutlass.Int32], Optional[cutlass.Int32]] + base_outputs: Union[Tuple[cute.Tensor, cute.Tensor], cute.Tensor] # (token, topk, hidden) + hidden_base_this_cta_tile: Union[cutlass.Int32, int] + peer_rank_ptr_mapper: Optional[SymmetricBufferDevice] + valid_tokens_this_cta_tile: cutlass.Int32 + valid_hidden_this_cta_tile: Union[cutlass.Int32, int] + reduce_topk_in_epilogue: bool + # Per-issue (epi_tid, iter_idx) -> (token_cta_tile, hidden_cta_tile). Data + # mapping, or (data mapping, sf mapping) when quantized. + output_mappings: Union[Tuple[FunctionMapping, FunctionMapping], FunctionMapping] + epi_tid: cutlass.Int32 + combine_format: CombineFormat + # After metadata prefetch + dst_ptrs: Optional[cute.Tensor] = None # i64 x (copy_iters_this_thread_cta_tile), fundamentally the pointers. + valid: Optional[cute.Tensor] = None # (copy_iters_this_thread_cta_tile) + + @property + def data_output(self) -> cute.Tensor: + return self.base_outputs[0] if isinstance(self.base_outputs, tuple) else self.base_outputs + + @property + def sf_output(self) -> Optional[cute.Tensor]: + # Present iff quantized; (pool_token, 1, hidden // sf_vec) rank-local. + return self.base_outputs[1] if isinstance(self.base_outputs, tuple) else None + + @property + def data_token_base(self) -> Optional[cutlass.Int32]: + return self.token_bases[0] if isinstance(self.token_bases, tuple) else self.token_bases + + @property + def sf_token_base(self) -> Optional[cutlass.Int32]: + return self.token_bases[1] if isinstance(self.token_bases, tuple) else None + + @property + def data_mapping(self) -> FunctionMapping: + return self.output_mappings[0] if isinstance(self.output_mappings, tuple) else self.output_mappings + + @property + def sf_mapping(self) -> Optional[FunctionMapping]: + return self.output_mappings[1] if isinstance(self.output_mappings, tuple) else None + + def __post_init__(self) -> None: + if (self.metadata is None) == (self.data_token_base is None): + raise ValueError("Fc2OutputRouter requires exactly one of metadata or a (data) token base.") + if (self.metadata is None) != (self.peer_rank_ptr_mapper is None): + raise ValueError("Fc2OutputRouter requires peer_rank_ptr_mapper iff metadata is set.") + if self.reduce_topk_in_epilogue and self.metadata is None: + raise ValueError("Fc2OutputRouter reduction requires metadata routing.") + + @cute.jit + def prefetch(self) -> "Fc2OutputRouter": + # Only the metadata (comm) path prefetches a pointer array: its + # metadata-derived address has long-latency LDGs worth issuing early. + # The local (no-comm) path computes its affine address on demand in + # get_dst() -- no array, hence no runtime-indexed local-memory spill. + if cutlass.const_expr(self.metadata is None): + return self + copy_iters: cutlass.Constexpr[int] = self.data_mapping.domain.axis_size("iter_idx") + + valid = cute.make_rmem_tensor((copy_iters,), cutlass.Int32) + dst_ptrs = cute.make_rmem_tensor((copy_iters,), cutlass.Int64) + + # Compiler should be able to optimize the same token_copy_group's offset add. (Fundamental cse + strength_reduce) + # We should check the SASS to ensure this happens. + for iter_idx in cutlass.range_constexpr(copy_iters): + coord = self.data_mapping.evaluate(epi_tid=self.epi_tid, iter_idx=iter_idx) + token_in_tile = cutlass.Int32(coord["token_in_cta_tile"]) + hidden_in_tile = cutlass.Int32(coord["hidden_in_cta_tile"]) + + valid[iter_idx] = cutlass.Int32(0) + dst_ptrs[iter_idx] = cutlass.Int64(0) + + token_valid = token_in_tile < self.valid_tokens_this_cta_tile + hidden_valid = hidden_in_tile < cutlass.Int32(self.valid_hidden_this_cta_tile) + if token_valid and hidden_valid: + valid[iter_idx] = cutlass.Int32(1) + if cutlass.const_expr(self.metadata is None): + dst_tokens = self.data_token_base + token_in_tile + dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile + # Int64 token coord: dst_tokens*K*H overflows int32 once + # T*K*H exceeds 2^31 (data_output is (token, topk, hidden)). + dst_ptrs[iter_idx] = self.data_output[Int64(dst_tokens), None, dst_hidden].iterator.toint() + + else: + md = TokenSrcMetadata.load( + self.metadata.iterator.toint() + Int64(token_in_tile) * Int64(TokenSrcMetadata.nbytes) + ) + dst_rank = md.src_rank + dst_token = md.src_token + dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile + if cutlass.const_expr(not self.reduce_topk_in_epilogue): + dst_topk = md.src_topk + else: + dst_topk = 0 + # Int64 token coord: domain_offset on (token, topk, hidden) + # computes dst_token*K*H, which overflows int32 once T*K*H > 2^31. + dst_ptrs[iter_idx] = self.peer_rank_ptr_mapper.map_pointer( + cute.domain_offset((Int64(dst_token), dst_topk, dst_hidden), self.data_output).iterator, + dst_rank, + byte_alignment=32, + ).toint() + + return dataclasses.replace(self, dst_ptrs=dst_ptrs, valid=valid) + + @cute.jit + def get_data_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, cutlass.Int32]: + """Per-issue DATA destination: gmem pointer + validity predicate. + + The router owns ``data_output`` so the caller never re-assembles a + pointer from a raw int; it just builds its own copy tensor (STG) or + feeds the pointer to inline asm (REDG/UBLK). + + Alignment is unified at 32 B: only STG feeds this pointer to a real + ``cute.copy`` (256 b vector store, genuinely 32 B aligned); REDG/UBLK + only ``ptrtoint`` it for inline-asm issue, where the hint is inert. + """ + if cutlass.const_expr(self.metadata is None): + # no-comm: on-demand affine address (no prefetched array). The + # invariant base hoists out of the caller's loop via CSE; a + # constexpr iter folds the per-issue offset into the store. + coord = self.data_mapping.evaluate(epi_tid=self.epi_tid, iter_idx=iter_idx) + token_in_tile = cutlass.Int32(coord["token_in_cta_tile"]) + hidden_in_tile = cutlass.Int32(coord["hidden_in_cta_tile"]) + pred = cutlass.Int32(0) + addr = cutlass.Int64(0) + if token_in_tile < self.valid_tokens_this_cta_tile and hidden_in_tile < cutlass.Int32( + self.valid_hidden_this_cta_tile + ): + pred = cutlass.Int32(1) + dst_tokens = self.data_token_base + token_in_tile + dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile + # Int64 token coord: dst_tokens*K*H overflows int32 once T*K*H > 2^31. + addr = self.data_output[Int64(dst_tokens), None, dst_hidden].iterator.toint() + else: + # comm: read the pointer / validity prefetched by prefetch(). + addr = self.dst_ptrs[iter_idx] + pred = self.valid[iter_idx] + ptr = cute.make_ptr(self.data_output.element_type, addr, AddressSpace.gmem, assumed_align=32) + return ptr, pred + + @cute.jit + def get_sf_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, cutlass.Int32]: + """Per-issue SF destination: rank-local gmem pointer + validity predicate. + + SF never goes to a peer (it is staged locally and pushed token-contiguously + by the dispatch / standalone warps), so this is always the affine local + address -- no metadata routing, no prefetch. ``sf_output`` is the broadcast + plane ``(pool_token, 1, (sf_vec, hidden//sf_vec)):(., ., (0, 1))``, so the + logical hidden coordinate folds to its scale block on indexing. + """ + coord = self.sf_mapping.evaluate(epi_tid=self.epi_tid, iter_idx=iter_idx) + token_in_tile = cutlass.Int32(coord["token_in_cta_tile"]) + hidden_in_tile = cutlass.Int32(coord["hidden_in_cta_tile"]) + pred = cutlass.Int32(0) + addr = cutlass.Int64(0) + if token_in_tile < self.valid_tokens_this_cta_tile and hidden_in_tile < cutlass.Int32( + self.valid_hidden_this_cta_tile + ): + pred = cutlass.Int32(1) + sf_row = self.sf_token_base + token_in_tile + sf_hidden = hidden_in_tile + self.hidden_base_this_cta_tile + addr = self.sf_output[Int64(sf_row), None, sf_hidden].iterator.toint() + # Per-block scale offsets are element-granular; claim the scale dtype's + # natural element alignment (e8m0 1 B / bf16 2 B). + sf_ptr = cute.make_ptr(self.sf_output.element_type, addr, AddressSpace.gmem, assumed_align=4) + return sf_ptr, pred + + +def make_fc2_stg_cta_store_out_mapping( + combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int +): + assert cta_hidden_tile_size == 128 + assert cta_token_tile_size % 64 == 0 + wire_dtype = combine_format.act_dtype + assert wire_dtype.width in (4, 8, 16), "fc2 STG wire dtype must be fp4/fp8/bf16." + elems_per_stg = min(256 // wire_dtype.width, 32) + stgs_per_hidden32 = 32 // elems_per_stg + fundamental_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "elem_idx"), (128, cta_token_tile_size)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) + ), + function=lambda epi_tid, elem_idx: { + "token_in_cta_tile": epi_tid % 32 + elem_idx // 32 * 32, + "hidden_in_cta_tile": elem_idx % 32 + epi_tid // 32 * 32, + }, + ) + store_out_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, stgs_per_hidden32 * cta_token_tile_size // 32)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) + ), + function=lambda epi_tid, iter_idx: { + "token_in_cta_tile": epi_tid % 32 + iter_idx // stgs_per_hidden32 * 32, + "hidden_in_cta_tile": (iter_idx % stgs_per_hidden32) * elems_per_stg + epi_tid // 32 * 32, + }, + ) + sf_store_out_mapping = None + if combine_format.is_quantized: + + def stg_sf_mapping(epi_tid, iter_idx): + lane = epi_tid % 32 + warp = epi_tid // 32 + return {"token_in_cta_tile": lane + iter_idx * 32, "hidden_in_cta_tile": warp * 32} + + sf_store_out_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, cta_token_tile_size // 32)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) + ), + function=stg_sf_mapping, + ) + return store_out_mapping, sf_store_out_mapping, fundamental_mapping + + +def make_fc2_redg_cta_store_out_mapping( + combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int +): + assert cta_hidden_tile_size == 128 + assert cta_token_tile_size % 64 == 0 + # In-kernel reduce is bf16-only and never quantized, so there is no SF plane. + assert combine_format.act_dtype.width == 16 + assert not combine_format.is_quantized + + fundamental_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "elem_idx"), (128, cta_token_tile_size)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) + ), + function=lambda epi_tid, elem_idx: { + "token_in_cta_tile": ( + ((elem_idx // 4) // 16) * 64 + + (((elem_idx // 4) % 16) // 8) * 32 + + (((elem_idx // 4) % 8) // 4) * 16 + + (((elem_idx // 4) % 4) % 2) * 8 + + (epi_tid % 32) // 4 + ), + "hidden_in_cta_tile": ( + (epi_tid // 32) * 32 + (epi_tid % 4) * 4 + (((elem_idx // 4) % 4) // 2) * 16 + elem_idx % 4 + ), + }, + ) + # SIMT REDG emits one 8B red.v2.bf16x2 per 4 hidden elements. Each + # 64-token subtile contributes two token rows per lane and 8 hidden + # segments per token row. + store_out_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, cta_token_tile_size // 64 * 16)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) + ), + function=lambda epi_tid, iter_idx: { + "token_in_cta_tile": ( + (iter_idx // 16) * 64 + + ((iter_idx % 16) // 8) * 32 + + ((iter_idx % 8) // 4) * 16 + + ((iter_idx % 4) % 2) * 8 + + (epi_tid % 32) // 4 + ), + "hidden_in_cta_tile": ((epi_tid // 32) * 32 + (epi_tid % 4) * 4 + ((iter_idx % 4) // 2) * 16), + }, + ) + return store_out_mapping, None, fundamental_mapping + + +def make_fc2_ublk_store_out_mapping(combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int): + assert cta_hidden_tile_size == 128 + assert cta_token_tile_size % 64 == 0 + # UBLK pushes whole hidden rows by byte count, so the token/hidden mapping + # is element-indexed and dtype-independent (wire dtype only sets copy bytes). + assert combine_format.act_dtype.width in (4, 8, 16), "fc2 UBLK wire dtype must be fp4/fp8/bf16." + assert cta_token_tile_size <= 256 + max_token_cta_tile = 256 + fundamental_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "elem_idx"), (128, cta_token_tile_size)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (max_token_cta_tile, cta_hidden_tile_size) + ), + function=lambda epi_tid, elem_idx: { + "token_in_cta_tile": elem_idx // cta_hidden_tile_size * 32 + + epi_tid % 8 + + epi_tid // 32 * 8 + + ((epi_tid % 32) // 8) * 64, + "hidden_in_cta_tile": elem_idx % cta_hidden_tile_size, + }, + ) + # One row per thread covers two 64-row subtiles per iteration. + copy_iters = (cta_token_tile_size + 127) // 128 + store_out_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, copy_iters)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (max_token_cta_tile, cta_hidden_tile_size) + ), + function=lambda epi_tid, iter_idx: { + "token_in_cta_tile": (iter_idx * 128 + ((epi_tid % 32) // 16) * 64 + (epi_tid // 32) * 16 + epi_tid % 16), + "hidden_in_cta_tile": 0, + }, + ) + if combine_format.is_quantized: + # Quantized UBLK reuses the hidden-contiguous STG mappings for 128-bit R2S. + _, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size + ) + else: + sf_store_out_mapping = None + return store_out_mapping, sf_store_out_mapping, fundamental_mapping + + +# (...) -> ((atom_v, 1)) +@cute.jit +def wrap_into_copy_standard_layout(tensor: cute.Tensor): + tensor = cute.coalesce(cute.flatten(tensor)) + tensor = cute.append_ones(tensor, cute.rank(tensor) + 1) + tensor = cute.group_modes(tensor, 0, cute.rank(tensor) - 1) + tensor = cute.group_modes(tensor, 0, cute.rank(tensor)) + return tensor + + +@cute.jit +def fc2_f2fp(*tensors, alpha_val: Optional[cutlass.Float32] = None, **_) -> cute.Tensor: + reorder_dtype = cutlass.BFloat16 + total_size = 0 + for t in tensors: + total_size += cute.size(t) + converted_acc = cute.make_rmem_tensor((total_size,), reorder_dtype) + elems_processed = 0 + for t in tensors: + current_tensor_size = cute.size(t) + dst = cute.make_tensor(converted_acc.iterator + elems_processed, cute.make_layout((current_tensor_size,))) + if cutlass.const_expr(alpha_val is None): + dst.store(t.load().to(reorder_dtype)) + else: + if cutlass.const_expr(current_tensor_size % 2 != 0): + raise ValueError("fc2_f2fp expects even elements for each input tensor.") + scaled = cute.make_rmem_tensor((current_tensor_size,), cutlass.Float32) + for i in cutlass.range_constexpr(0, current_tensor_size, 2): + # scaled[i] = t[i] * alpha_val + s0, s1 = cute.arch.mul_packed_f32x2((t[i], t[i + 1]), (alpha_val, alpha_val)) + scaled[i] = s0 + scaled[i + 1] = s1 + dst.store(scaled.load().to(reorder_dtype)) + elems_processed += current_tensor_size + return converted_acc + + +@cute.jit +def post_f2fp_reorder_identity(*, casted: cute.Tensor, **_): + # UBLK: the f2fp output is already in the pre-store distribution (each lane + # owns one hidden element across the 64 subtile tokens); no reorder needed. + return casted + + +@cute.jit +def fc2_stg_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, **_): + atom_ld16x256 = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x8), cutlass.Float32) + ptr = tmem_subtile_tensor.iterator + half_lane_offset = 16 * TmemTranspose32x64B16Movm._tmem_row_stride + top_view = cute.make_tensor(ptr, TmemTranspose32x64B16Movm._tmem_layout(16, 64)) + bottom_view = cute.make_tensor(ptr + half_lane_offset, TmemTranspose32x64B16Movm._tmem_layout(16, 64)) + top = cute.make_rmem_tensor((32,), cutlass.Float32) + bottom = cute.make_rmem_tensor((32,), cutlass.Float32) + cute.copy(atom_ld16x256, top_view, TmemTranspose32x64B16Movm._rmem_copy_view(top, 32)) + cute.copy(atom_ld16x256, bottom_view, TmemTranspose32x64B16Movm._rmem_copy_view(bottom, 32)) + return top, bottom + + +@cute.jit +def fc2_ublk_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, epi, **_): + # UBLK consumes a warp-local 32-hidden x 64-token slice. The caller passes + # the CTA-level 128-hidden x 64-token subtile view, so select this epi + # warp's hidden block before issuing LDTM.x64. + tmem_subtile_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, epi.warp_idx), None] + raw_regs = cute.make_rmem_tensor((64,), cutlass.Float32) + atom_ld32x32_x64 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x64), cutlass.Float32) + cute.copy( + atom_ld32x32_x64, + wrap_into_copy_standard_layout(tmem_subtile_per_warp), + wrap_into_copy_standard_layout(raw_regs), + ) + return (raw_regs,) + + +@cute.jit +def fc2_stg_post_f2fp_reorder( + *, + casted: cute.Tensor, # (subtile_cnt,) + tmem_subtile_view: cute.Tensor, # (epi_tile_m, epi_tile_n) + **_, +): + if cutlass.const_expr(cute.size(casted) != 64): + raise NotImplementedError("fc2 stg pass expects 64 BF16 registers before store reorder.") + return TmemTranspose32x64B16Movm(tmem_subtile_view.iterator, casted)() + + +@cute.jit +def fc2_redg_post_f2fp_reorder(*, casted: cute.Tensor, tmem_subtile_view: cute.Tensor, **_): + # (epi_tid, elem_idx) -> (token_64, hidden_128), each thread hold token_2 x hidden_32 + natural = fc2_stg_post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_view) + core_matrix_reorder_sttm_atom = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + core_matrix_reorder_ldtm_atom = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), cutlass.Float32) + # ((16, 2), token_32_group) + token_groups = cute.logical_divide(cute.zipped_divide(natural, (32,)), (16, None)) + out = cute.make_rmem_tensor(token_groups.shape, casted.dtype) + out_as_i32 = cute.recast_tensor(out, cutlass.Float32) + # (32, 64) + tmem_warp = cute.flat_divide(tmem_subtile_view, (32, cute.size(tmem_subtile_view, 1)))[None, None, 0, 0] + # (16, 16, 16dp_group, token_32_groups). Note, this tmem can provide 2x cols since the original is bf16. + tmem_groups = cute.flat_divide(tmem_warp, (16, 16)) + for group_idx in cutlass.range_constexpr(cute.size(token_groups, 1)): + sttm_source = cute.recast_tensor(token_groups[None, group_idx], cutlass.Float32) + sttm_destination = tmem_groups[None, None, None, group_idx] + cute.copy( + core_matrix_reorder_sttm_atom, + wrap_into_copy_standard_layout(sttm_source), + wrap_into_copy_standard_layout(sttm_destination), + ) + cute.copy( + core_matrix_reorder_ldtm_atom, + wrap_into_copy_standard_layout(tmem_groups[None, None, 0, group_idx]), + wrap_into_copy_standard_layout(out_as_i32[(None, 0), group_idx]), + ) + cute.copy( + core_matrix_reorder_ldtm_atom, + wrap_into_copy_standard_layout(tmem_groups[None, None, 1, group_idx]), + wrap_into_copy_standard_layout(out_as_i32[(None, 1), group_idx]), + ) + return cute.coalesce(out) + + +@cute.jit +def fc2_stg_store_function( + *, + epi, + subtile: cute.Tensor, # Always BF16 before quantization + subtile_idx: cutlass.Int32, + fc2_output_router: Fc2OutputRouter, + **_, +): + if cutlass.const_expr(epi.combine_format.is_quantized): + data_subtile, sf_regs = QuantImpl(epi.combine_format, "regs_in_thread")(subtile) + else: + data_subtile = subtile + sf_regs = None + stg_width_elems: cutlass.Constexpr[int] = min(32, 256 // data_subtile.element_type.width) + stg_bits: cutlass.Constexpr[int] = stg_width_elems * data_subtile.element_type.width + copy_atom_vec = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=stg_bits) + elems_per_thread: cutlass.Constexpr[int] = cute.size(data_subtile) + if cutlass.const_expr(elems_per_thread % stg_width_elems != 0): + raise ValueError( + "fc2 STG store requires pre-store elems per thread to be divisible " + f"by STG issue width, got {elems_per_thread} and {stg_width_elems}." + ) + + if cutlass.const_expr(sf_regs is not None): + sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block + token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 + for token_group in cutlass.range_constexpr(token_groups_per_subtile): + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group + sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) + if sf_pred != cutlass.Int32(0): + sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) + for scale_idx in cutlass.range_constexpr(sf_scales_per_stg): + sf_dst[scale_idx] = sf_regs[token_group * sf_scales_per_stg + scale_idx] + + iters_per_subtile: cutlass.Constexpr[int] = elems_per_thread // stg_width_elems + copy_src = cute.zipped_divide(data_subtile, (stg_width_elems,)) + single_copy_layout = cute.make_layout(((stg_width_elems, 1),), stride=((1, 0),)) + subtile_iter_base = cutlass.Int32(subtile_idx) * cutlass.Int32(iters_per_subtile) + for local_iter in cutlass.range_constexpr(iters_per_subtile): + global_iter = subtile_iter_base + cutlass.Int32(local_iter) + dst_ptr, pred = fc2_output_router.get_data_dst(global_iter) + if pred != cutlass.Int32(0): + src_i = cute.make_tensor(copy_src[None, local_iter].iterator, single_copy_layout) + dst_i = cute.make_tensor(dst_ptr, single_copy_layout) + cute.copy(copy_atom_vec, cute.recast_tensor(src_i, cutlass.Int32), cute.recast_tensor(dst_i, cutlass.Int32)) + + +@cute.jit +def fc2_tma_store_function( + *, + epi, + subtile: cute.Tensor, # Always BF16 before quantization + work_tile_info: SwapAbFc12WorkTileInfo, + epilogue_iter_idx: cutlass.Int32, + subtile_idx: cutlass.Int32, + fc2_output_router: Fc2OutputRouter, + **_, +): + if cutlass.const_expr(epi.smem_tensor is None or epi.tma_atom_fc2_output is None or epi.fc2_tma_output is None): + raise ValueError("FC2 TMA store requires staged SMEM, a TMA atom, and a token-major output tensor.") + + if cutlass.const_expr(epi.combine_format.is_quantized): + data_subtile, sf_regs = QuantImpl(epi.combine_format, "regs_in_thread")(subtile) + else: + data_subtile = subtile + sf_regs = None + + if cutlass.const_expr(sf_regs is not None): + sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block + token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 + for token_group in cutlass.range_constexpr(token_groups_per_subtile): + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group + sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) + if sf_pred != cutlass.Int32(0): + sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) + for scale_idx in cutlass.range_constexpr(sf_scales_per_stg): + sf_dst[scale_idx] = sf_regs[token_group * sf_scales_per_stg + scale_idx] + + vector_elements: cutlass.Constexpr[int] = 128 // data_subtile.element_type.width + if cutlass.const_expr( + cute.rank(data_subtile) != 1 + or cute.size(data_subtile) != 2 * 32 + or data_subtile.stride[0] != 1 + or 32 % vector_elements != 0 + ): + raise ValueError("FC2 TMA store requires a contiguous two-token by 32-hidden register fragment.") + + stage_idx = epilogue_iter_idx % cutlass.Int32(epi.fc2_tma_stages) + smem_stage = epi.smem_tensor[None, None, stage_idx] + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128) + tiler_mn = (64, 128) + layout_copy_tv = cute.make_layout(((32, 4), (32, 2)), stride=((1, 2048), (64, 32))) + tiled_store = cute.make_tiled_copy(store_atom, layout_copy_tv, tiler_mn) + thread_store = tiled_store.get_slice(epi.tidx) + smem_partition = thread_store.partition_D(smem_stage) + register_partition = cute.composition(data_subtile, cute.make_layout(smem_partition.shape)) + cute.copy(tiled_store, register_partition, smem_partition) + + token_tile_idx = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_n_idx * cutlass.Int32(epi.cta_tile_n) + + subtile_idx * cutlass.Int32(epi._EpilogueTokenTileSize) + ) // cutlass.Int32(epi._EpilogueTokenTileSize) + tiled_output = cute.flat_divide(epi.fc2_tma_output, (epi._EpilogueTokenTileSize, epi._EpilogueFc2HiddenTileSize)) + gmem_subtile = tiled_output[None, None, token_tile_idx, work_tile_info.tile_m_idx, 0] + tma_smem_src, tma_gmem_dst = cpasync.tma_partition( + epi.tma_atom_fc2_output, + 0, + cute.make_layout(1), + cute.group_modes(smem_stage, 0, 2), + cute.group_modes(gmem_subtile, 0, 2), + ) + + epilogue_barrier = SwapABGatedActEpilogue.epilogue_sync_barrier() + cute.arch.fence_proxy("async.shared", space="cta") + epilogue_barrier.arrive_and_wait() + if epi.warp_idx == cutlass.Int32(0): + cute.copy(epi.tma_atom_fc2_output, tma_smem_src, tma_gmem_dst) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(epi.fc2_tma_stages - 1, read=True) + epilogue_barrier.arrive_and_wait() + + +@cute.jit +def fc2_ublk_store_function_impl( + *, + epi, + subtile: cute.Tensor, # Always bf16 pre-quant tensor + epilogue_iter_idx: cutlass.Int32, + subtile_idx: cutlass.Int32, + fc2_output_router: Fc2OutputRouter, + **_, +): + smem_tensor = epi.smem_tensor + if cutlass.const_expr(smem_tensor is None): + raise ValueError("fc2 UBLK store requires epi.smem_tensor.") + quantized: cutlass.Constexpr[bool] = epi.combine_format.is_quantized + if cutlass.const_expr(quantized): + # The transposed fragment gives each thread complete hidden-axis scale blocks. + data_subtile, sf_regs = QuantImpl(epi.combine_format, "regs_in_thread")(subtile) + sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block + token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 + for token_group in cutlass.range_constexpr(token_groups_per_subtile): + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group + sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) + if sf_pred != cutlass.Int32(0): + sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) + for scale_idx in cutlass.range_constexpr(sf_scales_per_stg): + sf_dst[scale_idx] = sf_regs[token_group * sf_scales_per_stg + scale_idx] + else: + data_subtile = subtile + + smem_read_write_bar = SwapABGatedActEpilogue.epilogue_sync_barrier() + warp_idx = epi.warp_idx + lane_idx = epi.lane_idx + stage_cnt: cutlass.Constexpr[int] = epi.fc2_tma_stages + stage_idx = epilogue_iter_idx % cutlass.Int32(stage_cnt) + smem_stage = smem_tensor[None, None, stage_idx] + # Reuse waits here so the previous copy overlaps the next subtile; run() drains the final commit. + cute.arch.cp_async_bulk_wait_group(stage_cnt - 1, read=True) + smem_read_write_bar.arrive_and_wait() + + if cutlass.const_expr(quantized): + vector_elements: cutlass.Constexpr[int] = 128 // data_subtile.element_type.width + if cutlass.const_expr( + cute.rank(data_subtile) != 1 + or cute.size(data_subtile) != 2 * 32 + or data_subtile.stride[0] != 1 + or 32 % vector_elements != 0 + ): + raise ValueError("FC2 UBLK store requires a contiguous two-token by 32-hidden register fragment.") + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128) + tiler_mn = (64, 128) + layout_copy_tv = cute.make_layout(((32, 4), (32, 2)), stride=((1, 2048), (64, 32))) + tiled_store = cute.make_tiled_copy(store_atom, layout_copy_tv, tiler_mn) + thread_store = tiled_store.get_slice(epi.tidx) + smem_partition = thread_store.partition_D(smem_stage) + register_partition = cute.composition(data_subtile, cute.make_layout(smem_partition.shape)) + cute.copy(tiled_store, register_partition, smem_partition) + else: + if cutlass.const_expr(cute.size(data_subtile) != epi._EpilogueTokenTileSize): + raise ValueError("BF16 UBLK staging requires one register per token in the 64-token subtile.") + warp_hidden_base = cutlass.Int32(warp_idx * 32) + for token_idx in cutlass.range_constexpr(epi._EpilogueTokenTileSize): + smem_stage[token_idx, warp_hidden_base + lane_idx] = data_subtile[token_idx] + + cute.arch.fence_proxy("async.shared", space="cta") + smem_read_write_bar.arrive_and_wait() + + copy_elems = cutlass.Int32(epi._EpilogueFc2HiddenTileSize) + if cutlass.const_expr(epi.fc2_hidden_needs_predicate): + copy_elems = cutlass.Int32(fc2_output_router.valid_hidden_this_cta_tile) + copy_bytes = copy_elems * epi.combine_format.act_dtype.width // 8 + scratch_row = warp_idx * cutlass.Int32(16) + lane_idx % cutlass.Int32(16) + copy_iters: cutlass.Constexpr[int] = fc2_output_router.data_mapping.domain.axis_size("iter_idx") + for ublk_iter_idx in cutlass.range_constexpr(copy_iters): + owned_subtile = cutlass.Int32(ublk_iter_idx * 2) + lane_idx // cutlass.Int32(16) + if owned_subtile == subtile_idx: + dst_ptr, pred = fc2_output_router.get_data_dst(ublk_iter_idx) + if pred != cutlass.Int32(0): + src_row = cute.slice_(smem_stage, (scratch_row, None)) + if cutlass.const_expr(epi.reduce_topk_in_epilogue): + cp_reduce_async_bulk_add_bf16_s2g(dst_ptr, src_row.iterator, copy_bytes) + else: + cp_async_bulk_s2g(dst_ptr, src_row.iterator, copy_bytes) + + cute.arch.cp_async_bulk_commit_group() + + +@cute.jit +def fc2_redg_store_function( + *, + epi, + subtile: cute.Tensor, # Always bf16; in-kernel reduce never quantizes + subtile_idx: cutlass.Int32, + fc2_output_router: Fc2OutputRouter, + **_, +): + redg_width_elems: cutlass.Constexpr[int] = 4 + elems_per_thread: cutlass.Constexpr[int] = cute.size(subtile) + if cutlass.const_expr(elems_per_thread % redg_width_elems != 0): + raise ValueError( + "fc2 REDG store requires pre-store elems per thread to be divisible " + f"by REDG issue width, got {elems_per_thread} and {redg_width_elems}." + ) + iters_per_subtile: cutlass.Constexpr[int] = elems_per_thread // redg_width_elems + subtile_iter_base = cutlass.Int32(subtile_idx) * cutlass.Int32(iters_per_subtile) + subtile_by_redg_issue = cute.zipped_divide(subtile, (redg_width_elems,)) + + for local_iter in cutlass.range_constexpr(iters_per_subtile): + global_iter = subtile_iter_base + cutlass.Int32(local_iter) + dst_ptr, pred = fc2_output_router.get_data_dst(global_iter) + if pred != cutlass.Int32(0): + bf16x4 = subtile_by_redg_issue[None, local_iter] + packed_bf16x2 = cute.recast_tensor(bf16x4, cutlass.Float32) + red_add_relaxed_sys_v2_bf16x2(dst_ptr, cutlass.Float32(packed_bf16x2[0]), cutlass.Float32(packed_bf16x2[1])) + + +def make_fc2_stg_process_pipeline( + *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int +) -> Fc2ProcessPipeline: + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size + ) + return Fc2ProcessPipeline( + tmem_acc_load=fc2_stg_tmem_acc_load, + f2fp=fc2_f2fp, + post_f2fp_reorder=fc2_stg_post_f2fp_reorder, + store_function=fc2_stg_store_function, + fc2_cta_tile_mapping=fundamental_mapping, + store_out_mapping=store_out_mapping, + sf_store_out_mapping=sf_store_out_mapping, + require_tmem_trans=True, + ) + + +def make_fc2_tma_process_pipeline( + *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int +) -> Fc2ProcessPipeline: + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size + ) + return Fc2ProcessPipeline( + tmem_acc_load=fc2_stg_tmem_acc_load, + f2fp=fc2_f2fp, + post_f2fp_reorder=fc2_stg_post_f2fp_reorder, + store_function=fc2_tma_store_function, + fc2_cta_tile_mapping=fundamental_mapping, + store_out_mapping=store_out_mapping, + sf_store_out_mapping=sf_store_out_mapping, + require_tmem_trans=True, + ) + + +def make_fc2_redg_process_pipeline( + *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int +) -> Fc2ProcessPipeline: + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_redg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size + ) + return Fc2ProcessPipeline( + tmem_acc_load=fc2_stg_tmem_acc_load, + f2fp=fc2_f2fp, + post_f2fp_reorder=fc2_redg_post_f2fp_reorder, + store_function=fc2_redg_store_function, + fc2_cta_tile_mapping=fundamental_mapping, + store_out_mapping=store_out_mapping, + sf_store_out_mapping=sf_store_out_mapping, + require_tmem_trans=True, + ) + + +def make_fc2_ublk_process_pipeline( + *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int +) -> Fc2ProcessPipeline: + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_ublk_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size + ) + if combine_format.is_quantized: + if combine_format.act_dtype.width != 8: + raise ValueError("FC2 UBLK quantization requires an FP8 combine payload.") + tmem_acc_load = fc2_stg_tmem_acc_load + post_f2fp_reorder = fc2_stg_post_f2fp_reorder + else: + tmem_acc_load = fc2_ublk_tmem_acc_load + post_f2fp_reorder = post_f2fp_reorder_identity + return Fc2ProcessPipeline( + tmem_acc_load=tmem_acc_load, + f2fp=fc2_f2fp, + post_f2fp_reorder=post_f2fp_reorder, + store_function=fc2_ublk_store_function_impl, + fc2_cta_tile_mapping=fundamental_mapping, + store_out_mapping=store_out_mapping, + sf_store_out_mapping=sf_store_out_mapping, + require_tmem_trans=combine_format.is_quantized, + ) + + +__all__ = [ + "Fc2OutputRouter", + "Fc2ProcessPipeline", + "GatedActEpilogueArgs", + "QuantImpl", + "SwapABGatedActEpilogue", + "make_fc2_redg_process_pipeline", + "make_fc2_stg_process_pipeline", + "make_fc2_tma_process_pipeline", + "make_fc2_ublk_process_pipeline", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py new file mode 100644 index 000000000..c45dc66cf --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Block-scaled SwapAb adapter between FC12 scheduling and kernel tensor views.""" + +import dataclasses +from typing import ClassVar, List, Literal, Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass.cute.typing import Pointer +from cutlass.cutlass_dsl import Int32, extract_mlir_values, new_from_mlir_values +from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF + +from .....helpers.dsl_helpers import spin_peek, spin_wait +from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo, peek_ready_bit + + +TensorRole = Literal["a", "b", "sfa", "sfb", "c", "sfc", "topk"] + + +@cute.jit +def _rewrite_tensor_shape(tensor: cute.Tensor, new_shape: Tuple) -> cute.Tensor: + return cute.make_tensor(tensor.iterator, cute.make_layout(new_shape, stride=tensor.stride)) + + +@dataclasses.dataclass(frozen=True) +class BlockScaledSwapAbFc12Extension: + """Kernel-owned work-tile preparation and GMEM view adapter.""" + + work_tile_type: ClassVar[type] = SwapAbFc12WorkTileInfo + + sf_vec_size: int + fc1_done_counter_pointer: Pointer + fc2_spin_threshold: Int32 + fc1_ready_counter_pointer: Optional[Pointer] = None + + def __post_init__(self) -> None: + if self.sf_vec_size <= 0: + raise ValueError(f"sf_vec_size must be positive, got {self.sf_vec_size}.") + object.__setattr__(self, "fc2_spin_threshold", Int32(self.fc2_spin_threshold)) + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + values.extend(extract_mlir_values(self.fc1_done_counter_pointer)) + values.extend(extract_mlir_values(self.fc2_spin_threshold)) + if self.fc1_ready_counter_pointer is not None: + values.extend(extract_mlir_values(self.fc1_ready_counter_pointer)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "BlockScaledSwapAbFc12Extension": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + fc1_done_counter_pointer = rebuild(self.fc1_done_counter_pointer) + fc2_spin_threshold = rebuild(self.fc2_spin_threshold) + fc1_ready_counter_pointer = ( + rebuild(self.fc1_ready_counter_pointer) if self.fc1_ready_counter_pointer is not None else None + ) + if value_index != len(values): + raise ValueError( + f"BlockScaledSwapAbFc12Extension MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return type(self)( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter_pointer, + fc2_spin_threshold=fc2_spin_threshold, + fc1_ready_counter_pointer=fc1_ready_counter_pointer, + ) + + @cute.jit + def prepare_work_tile(self, work_tile: SwapAbFc12WorkTileInfo) -> SwapAbFc12WorkTileInfo: + """Pack kernel readiness observations into the published tile flags.""" + phase_and_flags = work_tile.phase_and_flags + if work_tile.is_valid_tile: + counter_slot = work_tile.cumulative_token_block_count + work_tile.tile_n_idx + is_fc1 = work_tile.phase == Int32(BlockPhase.Linear1) + is_fc2 = work_tile.phase == Int32(BlockPhase.Linear2) + + if cutlass.const_expr(self.fc1_ready_counter_pointer is not None): + if is_fc1: + counter_pointer = self.fc1_ready_counter_pointer + counter_slot + peek_flag = Int32(0) + if spin_peek(counter_pointer, lambda value: value >= work_tile.valid_tokens_in_cta_tile): + peek_flag = Int32(peek_ready_bit) + phase_and_flags = work_tile.phase_and_flags | peek_flag + + if is_fc2: + counter_pointer = self.fc1_done_counter_pointer + counter_slot + peek_flag = Int32(0) + if spin_peek(counter_pointer, lambda value: value >= self.fc2_spin_threshold): + peek_flag = Int32(peek_ready_bit) + phase_and_flags = work_tile.phase_and_flags | peek_flag + + return SwapAbFc12WorkTileInfo( + expert_idx=work_tile.expert_idx, + tile_m_idx=work_tile.tile_m_idx, + tile_n_idx=work_tile.tile_n_idx, + cumulative_data_physical_row=(work_tile.cumulative_data_physical_row), + cumulative_sf_physical_row=(work_tile.cumulative_sf_physical_row), + cumulative_token_block_count=(work_tile.cumulative_token_block_count), + valid_tokens_in_cta_tile=(work_tile.valid_tokens_in_cta_tile), + phase_and_flags=phase_and_flags, + ) + + @cute.jit + def wait_for_input(self, work_tile: SwapAbFc12WorkTileInfo) -> None: + """Wait until the current FC1 input tile is ready.""" + if cutlass.const_expr(self.fc1_ready_counter_pointer is not None): + counter_slot = work_tile.cumulative_token_block_count + work_tile.tile_n_idx + counter_pointer = self.fc1_ready_counter_pointer + counter_slot + spin_wait( + counter_pointer, + lambda value: value >= work_tile.valid_tokens_in_cta_tile, + peek_status=work_tile.peek_ready, + ) + + @cute.jit + def get_gmem_tensor( + self, tensor_name: TensorRole, tensor: cute.Tensor, work_tile: SwapAbFc12WorkTileInfo + ) -> Tuple[cute.Tensor, Optional[Pointer]]: + """Resolve one kernel tensor to its current expert/task-tile view.""" + expert_idx = work_tile.expert_idx + data_token_offset = work_tile.cumulative_data_physical_row + sf_token_offset = work_tile.cumulative_sf_physical_row + shape = tensor.shape + stride = tensor.stride + singleton = Int32(1) + + if cutlass.const_expr(tensor_name == "a"): + result = cute.domain_offset((0, 0, expert_idx), tensor) + return (_rewrite_tensor_shape(result, (shape[0], shape[1], singleton)), None) + + if cutlass.const_expr(tensor_name == "b"): + result = cute.domain_offset((data_token_offset, 0, 0), tensor) + return (_rewrite_tensor_shape(result, (shape[0], shape[1], singleton)), None) + + if cutlass.const_expr(tensor_name == "sfa"): + result = cute.domain_offset((0, 0, expert_idx), tensor) + per_expert_shape = (shape[0], shape[1], singleton) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, self.sf_vec_size) + return (cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), None) + + if cutlass.const_expr(tensor_name in ("sfb", "sfc")): + result = cute.domain_offset((sf_token_offset, 0, 0), tensor) + per_expert_shape = (shape[0], shape[1], singleton) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, self.sf_vec_size) + return (cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), None) + + if cutlass.const_expr(tensor_name == "c"): + result = cute.domain_offset((data_token_offset, 0, 0), tensor) + return (_rewrite_tensor_shape(result, (shape[0], shape[1], singleton)), None) + + if cutlass.const_expr(tensor_name == "topk"): + return (cute.domain_offset((data_token_offset,), tensor), None) + + raise ValueError(f"Unknown tensor_name: {tensor_name!r}.") + + +__all__ = ["BlockScaledSwapAbFc12Extension", "TensorRole"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/topk_reduce.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/topk_reduce.py new file mode 100644 index 000000000..582c4cdc4 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/topk_reduce.py @@ -0,0 +1,484 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Device combine reduce: collapse per-(token, topk) fc2 cells into one row. + +The combine step writes one fc2 output per ``(token, topk)`` cell; this reduces +over the topk axis into the token-centric ``(token, hidden)`` output. The wire +format is described by ``CombineFormat``: + + bf16 -- no staging: bf16 terms reduced directly. + 32e4m3xe8m0 -- MXFP8: fp8 e4m3 data + per-32 e8m0 (power-of-2) scale. + 16e2m1xbf16 -- fp4 e2m1 data + per-16 bf16 amax (one level, no global); + dequant per element x = fp4 * (amax * (1 / 6)). + +Task partition: each worker owns one ``(token, hidden_tile)`` and loops topk; the +flat worker index decodes into ``(token_idx, hidden_tile_idx)`` via a constant +divide by ``hidden_tiles``. The per-block scale is broadcast to a logical +per-hidden view (stride 0) so it tiles by the same worker index as the data. The +activation load stays in the topk loop (too large to hoist); the small scale and +score loads are hoisted ahead of the loop when topk is small. +""" + +import os +from typing import ClassVar, Dict, Optional + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Float32, Int32, T +from cutlass._mlir.dialects import llvm + +from .....quant_def import CombineFormat +from .....helpers.constants import Nvfp4E2M1RcpLimit +from .....helpers.dsl_helpers import mark_alignment + + +# --------------------------------------------------------------------------- +# fp4 (e2m1) -> fp32 register decode. +# +# Blackwell has no e2m1->f32 upconvert: the framework's ``term.load().to(f32)`` +# lowers to an ALU subnormal-normalization path (~60% DRAM SOL). Both helpers +# below force a table-driven decode instead; ``e2m1_reg`` (N e2m1 codes, N % 8 +# == 0) is read as packed b32 words and N fp32 values are written into +# ``fp32_reg`` in code order. The 16 e2m1 values are exact in fp32, so the two +# decoders are bit-for-bit identical (cross-check the optimal one against the +# cvt one over all 16 codes). +# --------------------------------------------------------------------------- + + +@cute.jit +def cvt_e2m1_to_fp32_cvt_ptx(e2m1_reg: cute.Tensor, fp32_reg: cute.Tensor) -> None: + """Decode via the e2m1->f16 HW cvt (``cvt.rn.f16x2.e2m1x2``) then widen f16->f32. + + Safe baseline: the per-pair ``cvt`` instruction is itself a HW PRMT+F2FP, so + the e2m1->f16 step already avoids ALU normalization; f16->f32 is one cheap + ``cvt.f32.f16`` per element. + """ + src_words = cute.recast_tensor(e2m1_reg, Int32) # (N,) e2m1 -> (N/8,) b32 + for w in cutlass.range_constexpr(cute.size(src_words)): + res = llvm.inline_asm( + llvm.StructType.get_literal([T.f32()] * 8), + [src_words[w].ir_value()], + "{\n" + " .reg .b8 b0, b1, b2, b3;\n" + " .reg .b32 p0, p1, p2, p3;\n" + " .reg .b16 c0, d0, c1, d1, c2, d2, c3, d3;\n" + " mov.b32 {b0, b1, b2, b3}, $8;\n" + " cvt.rn.f16x2.e2m1x2 p0, b0;\n" + " cvt.rn.f16x2.e2m1x2 p1, b1;\n" + " cvt.rn.f16x2.e2m1x2 p2, b2;\n" + " cvt.rn.f16x2.e2m1x2 p3, b3;\n" + " mov.b32 {c0, d0}, p0;\n" + " mov.b32 {c1, d1}, p1;\n" + " mov.b32 {c2, d2}, p2;\n" + " mov.b32 {c3, d3}, p3;\n" + " cvt.f32.f16 $0, c0;\n" + " cvt.f32.f16 $1, d0;\n" + " cvt.f32.f16 $2, c1;\n" + " cvt.f32.f16 $3, d1;\n" + " cvt.f32.f16 $4, c2;\n" + " cvt.f32.f16 $5, d2;\n" + " cvt.f32.f16 $6, c3;\n" + " cvt.f32.f16 $7, d3;\n" + "}", + "=f,=f,=f,=f,=f,=f,=f,=f,r", + has_side_effects=False, + ) + for i in cutlass.range_constexpr(8): + fp32_reg[w * 8 + i] = Float32(llvm.extractvalue(T.f32(), res, [i])) + + +@cute.jit +def cvt_e2m1_to_fp32_optimal_ptx(e2m1_reg: cute.Tensor, fp32_reg: cute.Tensor) -> None: + """Decode via a register-resident bf16 LUT + PRMT, landing fp32 directly. + + The 8 e2m1 magnitudes ``{0,.5,1,1.5,2,3,4,6}`` are exact bf16, so a PRMT + byte-gather of the per-magnitude hi/lo bytes builds the bf16; ``fp32 = + bf16 << 16`` makes the widen free. Sign (nibble bit3) is spread to the 4 + output-byte MSBs with one ``prmt`` of ``{word<<4, word}`` (selector 0x5140), + avoiding the non-uniform shift the 4-bit-vs-8-bit stride would otherwise need. + """ + src_words = cute.recast_tensor(e2m1_reg, Int32) # (N,) e2m1 -> (N/8,) b32 + dst_words = cute.recast_tensor(fp32_reg, Int32) # write fp32 bit patterns + for w in cutlass.range_constexpr(cute.size(src_words)): + res = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 8), + [src_words[w].ir_value()], + "{\n" + " .reg .b32 ha, hb, la, lb, inh, wl, ih, il, hl, ll, hh, lh, sl, sh, p0, p1, p2, p3;\n" + " mov.b32 ha, 0x3F3F3F00;\n" # hi byte LUT, magnitudes 0..3 + " mov.b32 hb, 0x40404040;\n" # hi byte LUT, magnitudes 4..7 + " mov.b32 la, 0xC0800000;\n" # lo byte LUT, magnitudes 0..3 + " mov.b32 lb, 0xC0804000;\n" # lo byte LUT, magnitudes 4..7 + " shr.b32 inh, $8, 16;\n" # high 4 elements -> low 16 bits + " and.b32 il, $8, 0x00007777;\n" # low 4 magnitude indices (clear sign) + " and.b32 ih, inh, 0x00007777;\n" # high 4 magnitude indices + " prmt.b32 hl, ha, hb, il;\n" # hi bytes for e0..e3 + " prmt.b32 ll, la, lb, il;\n" # lo bytes for e0..e3 + " prmt.b32 hh, ha, hb, ih;\n" # hi bytes for e4..e7 + " prmt.b32 lh, la, lb, ih;\n" # lo bytes for e4..e7 + " shl.b32 wl, $8, 4;\n" + " prmt.b32 sl, wl, $8, 0x5140;\n" # gather s0..s3 to byte MSBs + " and.b32 sl, sl, 0x80808080;\n" + " or.b32 hl, hl, sl;\n" + " shl.b32 wl, inh, 4;\n" + " prmt.b32 sh, wl, inh, 0x5140;\n" # gather s4..s7 to byte MSBs + " and.b32 sh, sh, 0x80808080;\n" + " or.b32 hh, hh, sh;\n" + " prmt.b32 p0, ll, hl, 0x5140;\n" # {bf16(e0), bf16(e1)} + " prmt.b32 p1, ll, hl, 0x7362;\n" # {bf16(e2), bf16(e3)} + " prmt.b32 p2, lh, hh, 0x5140;\n" # {bf16(e4), bf16(e5)} + " prmt.b32 p3, lh, hh, 0x7362;\n" # {bf16(e6), bf16(e7)} + " shl.b32 $0, p0, 16;\n" + " and.b32 $1, p0, 0xFFFF0000;\n" + " shl.b32 $2, p1, 16;\n" + " and.b32 $3, p1, 0xFFFF0000;\n" + " shl.b32 $4, p2, 16;\n" + " and.b32 $5, p2, 0xFFFF0000;\n" + " shl.b32 $6, p3, 16;\n" + " and.b32 $7, p3, 0xFFFF0000;\n" + "}", + "=r,=r,=r,=r,=r,=r,=r,=r,r", + has_side_effects=False, + ) + for i in cutlass.range_constexpr(8): + dst_words[w * 8 + i] = Int32(llvm.extractvalue(T.i32(), res, [i])) + + +class TopkReduce: + """Combine reduce for a fixed ``(hidden, num_topk, combine_format)``. + + ``__init__`` pins the static shape and format (and the derived launch + geometry); ``__call__`` (a ``@cute.jit`` launcher) sizes a 1D grid from the + runtime token count and dispatches the format's kernel. The caller owns the + torch->cute conversion and the ``cute.compile`` / ``aot_compile``. + """ + + _threads: ClassVar[int] = 128 + # combine_format.name -> hidden elements per worker (one LDG of data: + # bf16 8*2B=16B, e4m3 16*1B=16B, e2m1 16*0.5B=8B). For quantized formats this + # stays <= the scale block, so each worker reads exactly one scale entry. + _hidden_per_thread: ClassVar[Dict[str, int]] = {"bf16": 8, "32e4m3xe8m0": 16, "16e2m1xbf16": 16} + # topk count at/below which the scale + score loads are hoisted ahead of the + # topk loop (small enough to not bloat registers; a CTA-broadcast read). + _prefetch_limit: ClassVar[int] = 16 + + def __init__(self, hidden: int, num_topk: int, combine_format: CombineFormat) -> None: + self.hidden = int(hidden) + self.num_topk = int(num_topk) + self.combine_format = combine_format + self.hidden_per_thread = self._hidden_per_thread[combine_format.name] + # hidden must tile cleanly both into worker slices and into scale blocks. + align = max(combine_format.scale_block or self.hidden_per_thread, self.hidden_per_thread) + if self.hidden % align != 0: + raise ValueError( + f"hidden ({self.hidden}) must be divisible by max(scale_block, " + f"hidden_per_thread) = {align} for combine_format {combine_format}." + ) + self.hidden_tiles = self.hidden // self.hidden_per_thread + # tail guard only needed when the worker count per token is not a whole + # number of CTAs; prefetch only when topk is small enough to hoist. + self.require_predicate = self.hidden_tiles % self._threads != 0 + self.prefetch = self.num_topk <= self._prefetch_limit + + # -- launcher ------------------------------------------------------------- + + @cute.jit + def __call__( + self, + combine_quant: cute.Tensor, # (token, topk, hidden) + combine_sf: Optional[cute.Tensor], # (token, topk, hidden) + reduced_output: cute.Tensor, # (token, hidden) + topk_score: Optional[cute.Tensor], # (token, topk) + stream: cuda.CUstream, + ): + threads = self._threads + total_workers = reduced_output.shape[0] * self.hidden_tiles + grid = [(total_workers + threads - 1) // threads, 1, 1] + block = [threads, 1, 1] + + combine_quant = cute.make_tensor( + combine_quant.iterator, + cute.make_layout((combine_quant.shape[0], self.num_topk, self.hidden), stride=combine_quant.stride), + ) + reduced_output = cute.make_tensor( + reduced_output.iterator, + cute.make_layout((reduced_output.shape[0], self.hidden), stride=reduced_output.stride), + ) + if cutlass.const_expr(topk_score is not None): + topk_score = cute.make_tensor( + topk_score.iterator, cute.make_layout((topk_score.shape[0], self.num_topk), stride=topk_score.stride) + ) + + if cutlass.const_expr(not self.combine_format.is_quantized): + self._reduce_bf16(combine_quant, topk_score, reduced_output).launch(grid=grid, block=block, stream=stream) + return + + # The mega kernel hands sf in already as the depth-2 broadcast layout; a + # plain (torch) sf is depth-1 and gets its hidden mode split into + # (sf_vec, hidden/sf_vec):(0, s_h) so logical hidden h reads block h//sf_vec. + sf_vec = self.combine_format.scale_block + if cutlass.const_expr(cute.depth(combine_sf.layout) >= 2): + sf = cute.make_tensor( + combine_sf.iterator, + cute.make_layout( + (combine_sf.shape[0], self.num_topk, (sf_vec, self.hidden // sf_vec)), stride=combine_sf.stride + ), + ) + else: + sf = cute.make_tensor( + combine_sf.iterator, + cute.make_layout( + (combine_sf.shape[0], self.num_topk, (sf_vec, self.hidden // sf_vec)), + stride=(combine_sf.stride[0], combine_sf.stride[1], (0, combine_sf.stride[2])), + ), + ) + + if cutlass.const_expr(self.combine_format.act_dtype is cutlass.Float8E4M3FN): + self._reduce_mxfp8(combine_quant, sf, topk_score, reduced_output).launch( + grid=grid, block=block, stream=stream + ) + else: + self._reduce_fp4(combine_quant, sf, topk_score, reduced_output).launch( + grid=grid, block=block, stream=stream + ) + + # -- kernels -------------------------------------------------------------- + + @cute.kernel + def _reduce_bf16(self, combine_output: cute.Tensor, topk_score: Optional[cute.Tensor], reduced_output: cute.Tensor): + threads = self._threads + hidden_per_thread = self.hidden_per_thread + hidden_tiles = self.hidden_tiles + num_topk: cutlass.Constexpr[int] = self.num_topk + needs_guard = self.require_predicate + prefetch = self.prefetch + out_dtype = reduced_output.element_type + + worker_idx = cute.arch.block_idx()[0] * Int32(threads) + cute.arch.thread_idx()[0] + token_idx = worker_idx // hidden_tiles + hidden_tile_idx = worker_idx % hidden_tiles + + score_dtype = topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 + score_reg = cute.make_rmem_tensor((num_topk,), score_dtype) + + if (not needs_guard) or token_idx < reduced_output.shape[0]: + # (token, topk, hidden) -> (topk, hidden_per_thread) + terms = cute.zipped_divide(combine_output[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] + + if cutlass.const_expr(topk_score is not None): + if cutlass.const_expr(prefetch): + cute.autovec_copy(topk_score[token_idx, None], score_reg) + else: + for k in cutlass.range_constexpr(num_topk): + score_reg[k] = score_dtype(1) + + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=128) + acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) + + for k in cutlass.range_constexpr(0, num_topk, 1): + term = cute.make_rmem_tensor((hidden_per_thread,), cutlass.BFloat16) + cute.copy( + load_atom, mark_alignment(terms[k, None], hidden_per_thread * cutlass.BFloat16.width // 8), term + ) + if cutlass.const_expr(topk_score is not None and not prefetch): + score_reg[k] = topk_score[token_idx, Int32(k)] + score_pair = (Float32(score_reg[k]), Float32(score_reg[k])) + + for i in cutlass.range_constexpr(0, hidden_per_thread, 2): + value_pair = (Float32(term[i]), Float32(term[i + 1])) + if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(value_pair, score_pair, (acc[i], acc[i + 1])) + else: + if cutlass.const_expr(topk_score is not None): + acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2(value_pair, score_pair) + else: + acc[i] = value_pair[0] + acc[i + 1] = value_pair[1] + + out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) + out.store(acc.load().to(out_dtype)) + cute.copy( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), out_dtype, num_bits_per_copy=128), + out, + mark_alignment(dst, hidden_per_thread * out_dtype.width // 8), + ) + + @cute.kernel + def _reduce_mxfp8( + self, + combine_quant: cute.Tensor, + combine_sf: cute.Tensor, # depth-2 broadcast view: logical (token, topk, hidden) e8m0 + topk_score: Optional[cute.Tensor], + reduced_output: cute.Tensor, + ): + threads = self._threads + hidden_per_thread = self.hidden_per_thread + hidden_tiles = self.hidden_tiles + num_topk: cutlass.Constexpr[int] = self.num_topk + needs_guard = self.require_predicate + prefetch = self.prefetch + out_dtype = reduced_output.element_type + + worker_idx = cute.arch.block_idx()[0] * Int32(threads) + cute.arch.thread_idx()[0] + token_idx = worker_idx // hidden_tiles + hidden_tile_idx = worker_idx % hidden_tiles + + score_dtype = topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 + score_reg = cute.make_rmem_tensor((num_topk,), score_dtype) + scale_reg = cute.make_rmem_tensor((num_topk,), cutlass.Float8E8M0FNU) + + if (not needs_guard) or token_idx < reduced_output.shape[0]: + # (token, topk, hidden) -> (topk, hidden_per_thread) + codes = cute.zipped_divide(combine_quant[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] + # (token, topk, hidden) -> (topk, hidden_per_thread) + sf = cute.zipped_divide(combine_sf[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] + + if cutlass.const_expr(topk_score is not None): + if cutlass.const_expr(prefetch): + cute.autovec_copy(topk_score[token_idx, None], score_reg) + else: + for k in cutlass.range_constexpr(num_topk): + score_reg[k] = score_dtype(1) + if cutlass.const_expr(prefetch): + cute.autovec_copy(sf[None, 0], scale_reg) # one scale per topk slot (stride-0 broadcast) + + fp8_dtype = self.combine_format.act_dtype + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=128) + acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) + + for k in cutlass.range_constexpr(0, num_topk, 1): + term = cute.make_rmem_tensor((hidden_per_thread,), fp8_dtype) + cute.copy(load_atom, mark_alignment(codes[k, None], hidden_per_thread * fp8_dtype.width // 8), term) + value = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) + value.store(term.load().to(cutlass.Float32)) + + if cutlass.const_expr(not prefetch): + scale_reg[k] = sf[k, 0] + if cutlass.const_expr(topk_score is not None): + score_reg[k] = topk_score[token_idx, Int32(k)] + + scale = Float32(scale_reg[k]) # e8m0 -> f32 + scale_pair = (scale, scale) + score_pair = (Float32(score_reg[k]), Float32(score_reg[k])) + + for i in cutlass.range_constexpr(0, hidden_per_thread, 2): + dequant_pair = cute.arch.mul_packed_f32x2((value[i], value[i + 1]), scale_pair) + if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(dequant_pair, score_pair, (acc[i], acc[i + 1])) + else: + if cutlass.const_expr(topk_score is not None): + acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2(dequant_pair, score_pair) + else: + acc[i] = dequant_pair[0] + acc[i + 1] = dequant_pair[1] + + out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) + out.store(acc.load().to(out_dtype)) + cute.copy( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), out_dtype, num_bits_per_copy=256), + out, + mark_alignment(dst, hidden_per_thread * out_dtype.width // 8), + ) + + @cute.kernel + def _reduce_fp4( + self, + combine_quant: cute.Tensor, # (token, topk, hidden) e2m1 (logical) + combine_sf: cute.Tensor, # depth-2 broadcast view: logical (token, topk, hidden) bf16 amax + topk_score: Optional[cute.Tensor], + reduced_output: cute.Tensor, + ): + threads = self._threads + hidden_per_thread = self.hidden_per_thread + hidden_tiles = self.hidden_tiles + num_topk: cutlass.Constexpr[int] = self.num_topk + needs_guard = self.require_predicate + prefetch = self.prefetch + out_dtype = reduced_output.element_type + + worker_idx = cute.arch.block_idx()[0] * Int32(threads) + cute.arch.thread_idx()[0] + token_idx = worker_idx // hidden_tiles + hidden_tile_idx = worker_idx % hidden_tiles + + score_dtype = topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 + score_reg = cute.make_rmem_tensor((num_topk,), score_dtype) + scale_reg = cute.make_rmem_tensor((num_topk,), cutlass.BFloat16) + + if (not needs_guard) or token_idx < reduced_output.shape[0]: + # (token, topk, hidden) -> (topk, hidden_per_thread) + codes = cute.zipped_divide(combine_quant[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] + # (token, topk, hidden) -> (topk, hidden_per_thread) + sf = cute.zipped_divide(combine_sf[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] + + if cutlass.const_expr(topk_score is not None): + if cutlass.const_expr(prefetch): + cute.autovec_copy(topk_score[token_idx, None], score_reg) + else: + for k in cutlass.range_constexpr(num_topk): + score_reg[k] = score_dtype(1) + if cutlass.const_expr(prefetch): + cute.autovec_copy(sf[None, 0], scale_reg) + + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Float4E2M1FN, num_bits_per_copy=64) + acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) + + for k in cutlass.range_constexpr(0, num_topk, 1): + term = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float4E2M1FN) + cute.copy( + load_atom, mark_alignment(codes[k, None], hidden_per_thread * cutlass.Float4E2M1FN.width // 8), term + ) + value = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) + # Dev-only knob (MEGA_F4CVT_USE_MANUAL): manual LUT/PRMT decode vs + # the HW cvt path, so both SASS forms can be compared on device; + # one is kept once chosen. Read inline on purpose -- never a + # customer-facing option. Both decoders are bit-exact. + if cutlass.const_expr(os.environ.get("MEGA_F4CVT_USE_MANUAL", "0") == "1"): + cvt_e2m1_to_fp32_optimal_ptx(term, value) + else: + cvt_e2m1_to_fp32_cvt_ptx(term, value) + + if cutlass.const_expr(not prefetch): + scale_reg[k] = sf[k, 0] + if cutlass.const_expr(topk_score is not None): + score_reg[k] = topk_score[token_idx, Int32(k)] + + # amax (bf16) -> per-element scale; (1/6) folds the fp4 grid max. + scale = Float32(scale_reg[k]) * Float32(Nvfp4E2M1RcpLimit) + scale_pair = (scale, scale) + score_pair = (Float32(score_reg[k]), Float32(score_reg[k])) + + for i in cutlass.range_constexpr(0, hidden_per_thread, 2): + dequant_pair = cute.arch.mul_packed_f32x2((value[i], value[i + 1]), scale_pair) + if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(dequant_pair, score_pair, (acc[i], acc[i + 1])) + elif cutlass.const_expr(topk_score is not None): + acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2(dequant_pair, score_pair) + else: + acc[i] = dequant_pair[0] + acc[i + 1] = dequant_pair[1] + + out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) + out.store(acc.load().to(out_dtype)) + cute.copy( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), out_dtype, num_bits_per_copy=256), + out, + mark_alignment(dst, hidden_per_thread * out_dtype.width // 8), + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/function_mapping.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/function_mapping.py new file mode 100644 index 000000000..83cacc303 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/function_mapping.py @@ -0,0 +1,171 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Pure-Python finite coordinate mappings used by kernel code generation.""" + +import inspect +from dataclasses import dataclass +from math import prod +from typing import Callable, Mapping, Sequence + + +class FunctionMappingError(ValueError): + """Raised when a finite function mapping is malformed.""" + + +@dataclass(frozen=True) +class CoordinateSpace: + """A finite named coordinate space with axis 0 linearized fastest.""" + + names: tuple[str, ...] + sizes: tuple[int, ...] + + def __post_init__(self) -> None: + names = tuple(self.names) + sizes = tuple(self.sizes) + if not names or len(names) != len(sizes): + raise FunctionMappingError( + f"Coordinate-space rank mismatch: {names!r}, {sizes!r}." + ) + if any(not isinstance(name, str) or not name for name in names): + raise FunctionMappingError("Axis names must be non-empty strings.") + if len(set(names)) != len(names): + raise FunctionMappingError(f"Axis names must be unique: {names!r}.") + if any(not isinstance(size, int) or size <= 0 for size in sizes): + raise FunctionMappingError( + f"Axis sizes must be positive Python ints: {sizes!r}." + ) + object.__setattr__(self, "names", names) + object.__setattr__(self, "sizes", sizes) + + @property + def size(self) -> int: + return prod(self.sizes) + + def axis_size(self, name: str) -> int: + try: + return self.sizes[self.names.index(name)] + except ValueError as error: + raise KeyError(f"Unknown coordinate axis {name!r}.") from error + + def delinearize(self, linear_index: int) -> tuple[int, ...]: + if linear_index < 0 or linear_index >= self.size: + raise FunctionMappingError( + f"Linear index {linear_index} is outside [0, {self.size})." + ) + remaining = linear_index + coordinate = [] + for size in self.sizes: + coordinate.append(remaining % size) + remaining //= size + return tuple(coordinate) + + def coordinates(self) -> tuple[tuple[int, ...], ...]: + return tuple(self.delinearize(index) for index in range(self.size)) + + +MappingResult = int | Sequence[int] | Mapping[str, int] + + +@dataclass(frozen=True) +class FunctionMapping: + """A finite coordinate mapping generated by a deterministic Python function.""" + + domain: CoordinateSpace + codomain: CoordinateSpace + function: Callable[..., MappingResult] + + def __post_init__(self) -> None: + if not callable(self.function): + raise FunctionMappingError("FunctionMapping.function must be callable.") + self._validate_signature() + for domain_coordinate in self.domain.coordinates(): + arguments = dict(zip(self.domain.names, domain_coordinate)) + result = self.function(**arguments) + self._normalize_result(result, validate_static=True) + + def _validate_signature(self) -> None: + signature = inspect.signature(self.function) + parameters = signature.parameters + unsupported = [ + name + for name, parameter in parameters.items() + if parameter.kind + in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ) + or parameter.default is not inspect.Parameter.empty + ] + if unsupported: + raise FunctionMappingError( + f"Unsupported mapping parameters {unsupported!r}." + ) + if set(parameters) != set(self.domain.names): + raise FunctionMappingError( + f"Mapping parameters {tuple(parameters)!r} must match " + f"domain axes {self.domain.names!r}." + ) + + def _normalize_result( + self, + result, + *, + validate_static: bool, + ) -> dict[str, object]: + if isinstance(result, Mapping): + if set(result) != set(self.codomain.names): + raise FunctionMappingError( + f"Mapping result keys {tuple(result)!r} must match " + f"codomain axes {self.codomain.names!r}." + ) + coordinate = { + name: result[name] for name in self.codomain.names + } + elif isinstance(result, Sequence) and not isinstance( + result, + (str, bytes), + ): + if len(result) != len(self.codomain.names): + raise FunctionMappingError( + f"Mapping result rank {len(result)} must equal " + f"{len(self.codomain.names)}." + ) + coordinate = dict(zip(self.codomain.names, result)) + elif len(self.codomain.names) == 1: + coordinate = {self.codomain.names[0]: result} + else: + raise FunctionMappingError( + "A multi-axis mapping must return a sequence or named mapping." + ) + + if validate_static: + for name, value in coordinate.items(): + if not isinstance(value, int): + raise FunctionMappingError( + f"Static result {name!r} must be int, got {type(value)}." + ) + size = self.codomain.axis_size(name) + if value < 0 or value >= size: + raise FunctionMappingError( + f"Static result {name!r}={value} is outside [0, {size})." + ) + return coordinate + + def evaluate(self, **domain_coordinate) -> dict[str, object]: + if set(domain_coordinate) != set(self.domain.names): + raise FunctionMappingError( + f"Mapping arguments {tuple(domain_coordinate)!r} must match " + f"domain axes {self.domain.names!r}." + ) + return self._normalize_result( + self.function(**domain_coordinate), + validate_static=False, + ) + + +__all__ = [ + "CoordinateSpace", + "FunctionMapping", + "FunctionMappingError", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/__init__.py new file mode 100644 index 000000000..f763de2fc --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Rubin kernel implementations.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/__init__.py new file mode 100644 index 000000000..6152b881d --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Rubin SM107 training kernel package.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/__init__.py new file mode 100644 index 000000000..14b0bb10d --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Rubin training MegaMoE (mxfp8) kernels. + +Organised into two subpackages: + +* ``fwd_glu`` -- the forward fused FC1+SwiGLU+FC2 MoE kernel. +* ``bwd_dglu`` -- the backward fused dfc2+dswiglu+dfc1 MoE kernel. + +The forward symbols are re-exported here so existing ``rubin.training.mega`` +importers keep resolving after the fwd_glu/bwd_dglu reorg. +""" + +from .fwd_glu import ( + Fc2OutputDest, + GluMxFp8Fc12SchedExtension, + GluMxfp8Epilogue, + Sm107MegaMoEMxfp8GluKernel, + Sm107Mxfp8GluFc12Kernel, + TensorRole, +) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/__init__.py new file mode 100644 index 000000000..81bd8e3bd --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Rubin training MegaMoE (mxfp8 dGLU backward) kernel components. + +The backward fused dfc2 + dswiglu + dfc1 MoE kernel used by the training path. Computes +grad_x (activation gradient -> output_activation) and dprob (routing-weight gradient, +pool region). Built on top of the forward GLU package (subclasses the shared +``Fc2OutputDest`` peer-store) and mirrors the forward mega structure. +""" + +from .dglu_mxfp8_fc12_epilogue import DgluMxfp8Epilogue +from .dglu_mxfp8_fc12_extension import DgluMxFp8Fc12SchedExtension +from .dglu_mxfp8_fc12_kernel import Sm107Mxfp8DgluDfc21Kernel +from .dglu_mxfp8_mega_moe_kernel import Sm107MegaMoEMxfp8DgluKernel + + +__all__ = [ + "DgluMxfp8Epilogue", + "DgluMxFp8Fc12SchedExtension", + "Sm107Mxfp8DgluDfc21Kernel", + "Sm107MegaMoEMxfp8DgluKernel", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py new file mode 100644 index 000000000..55d78e996 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py @@ -0,0 +1,1542 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +from typing import Optional, Tuple, Type, Union + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import AddressSpace +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils + +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith as _arith +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import dsl_user_op, Int32 as _epi_Int32, Int64 +from cutlass.cute.typing import Float32 + +from ......helpers.iket_compat import iket +from ......helpers.flag_batch import GpuReleaseFlagBatchTracker +from ......helpers.ptx_helpers import ( + red_add_relaxed_sys_f32 as _red_add_relaxed_sys_f32, + red_add_relaxed_sys_v2_bf16x2 as _red_add_relaxed_sys_v2_bf16x2, + stg_e8m0_from_f32, + stg_e8m0x8_from_f32, +) +from ..helpers.utils import swiglu_act, dswiglu_act, quant_sfd_row, quant_sfd_col +from ......quant_def import CombineFormat +from .....schedulers import BlockPhase +from ..tmem_transpose import _TmemTranspose16x32Core +from ..fwd_glu.glu_mxfp8_fc12_epilogue import Fc2OutputDest + +Fc1GateUpInterleave = 32 +EpilogueTileN = 32 +Fc1EpilogueOutputTileM = 128 +Fc1EpilogueOutputTileN = 128 +WarpThreadCount = 32 +EpiWarpCount = 4 + + +@cute.jit +def dprob_reduce_gmem( + real_dprob: cute.Tensor, + dprob_val: cutlass.Float32, + is_valid: bool, + expert_local_token_idx, + system_scope: bool = False, +) -> None: + """Atomically reduce per-tile dprob accumulator into GMEM.""" + if is_valid: + if cutlass.const_expr(system_scope): + _red_add_relaxed_sys_f32( + real_dprob.iterator + expert_local_token_idx, + cutlass.Float32(dprob_val), + ) + else: + cute.arch.atomic_add( + real_dprob.iterator + expert_local_token_idx, + cutlass.Float32(dprob_val), + sem="relaxed", + scope="gpu", + ) + +# ============================================================================= +# DgluMxfp8Epilogue +# ============================================================================= + +class DgluMxfp8Epilogue: + + _SubtileBarIdBase = 4 + + def __init__( + self, + *, + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + use_2cta_instrs: bool, + sf_vec_size: int, + fc1_output_dtype: Type[cutlass.Numeric], + fc1_output_layout: utils.LayoutEnum, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + sf_dtype: Type[cutlass.Numeric] = cutlass.Float8E8M0FNU, + epilog_sync_bar_id: int = 1, + epilogue_warp_ids: Tuple[int, ...] = (0, 1, 2, 3), + static_expert_shape: Optional[Tuple[int, int, int]] = None, + token_back_by_dispatch: bool = False, + epi_flag_batch: Optional[Tuple[int, int]] = (1, 1), + dfc2_recompute: bool = False, + dfc2_col_output: bool = False, + fc2_in_kernel_topk_reduce: bool = False, + combine_format: Optional[CombineFormat] = None, + combine_hidden: Optional[int] = None, + act_func: str = "swiglu", + gate_up_clamp: Optional[float] = None, + ) -> None: + self._act_func = act_func + self._gate_up_clamp = ( + cutlass.Float32(gate_up_clamp) if gate_up_clamp is not None else None + ) + self.fc1_output_dtype = fc1_output_dtype + self.fc1_output_layout = fc1_output_layout + self.acc_dtype = acc_dtype + self.sf_dtype = sf_dtype + self._sf_vec_size = sf_vec_size + self._epilog_sync_bar_id = epilog_sync_bar_id + self._epilogue_warp_ids = epilogue_warp_ids + self._use_2cta_instrs = use_2cta_instrs + + self._atom_thr_size = 2 if use_2cta_instrs else 1 + self._cta_tile_m = mma_tiler_mnk[0] // self._atom_thr_size + self._cta_tile_n = mma_tiler_mnk[1] + self._mma_tiler_k = mma_tiler_mnk[2] + self._mma_tiler = tuple(mma_tiler_mnk) # for partition_C in the C-load + self._cta_tile_n_sfb = ((mma_tiler_mnk[1] + 127) // 128) * 128 + self._static_expert_shape = static_expert_shape + if ( + static_expert_shape is not None + and static_expert_shape[2] % (self._cta_tile_m * cluster_shape_mn[0]) == 0 + ): + self._fc2_stg_needs_predicate: bool = False + else: + self._fc2_stg_needs_predicate: bool = True + + self._epi_tile = (EpilogueTileN, Fc1EpilogueOutputTileM) + self._subtile_cnt = self._cta_tile_n // 2 // EpilogueTileN + + self._num_acc_stage = 2 + self._num_acc_pipeline_stages = self._num_acc_stage + + k = self._mma_tiler_k + self._num_sfa_tmem_cols = self._cta_tile_m * k // sf_vec_size * 4 // 4 // 128 + self._num_sfb_tmem_cols = ( + self._cta_tile_n_sfb * k // sf_vec_size * 4 // 4 // 128 + ) + self._num_sf_tmem_cols = 32 # self._num_sfa_tmem_cols + self._num_sfb_tmem_cols + + self._num_accumulator_tmem_cols = self._cta_tile_n * self._num_acc_stage + + self._token_back_by_dispatch = token_back_by_dispatch + # In-kernel top-k reduce + self._reduce_topk_in_epilogue = ( + fc2_in_kernel_topk_reduce and not token_back_by_dispatch + ) + _fc1_batch, _fc2_batch = (1, 1) if epi_flag_batch is None else epi_flag_batch + self._epi_fc1_batch = max(1, min(32, int(_fc1_batch))) + self._epi_fc2_batch = max(1, min(32, int(_fc2_batch))) + + self._dfc2_recompute = dfc2_recompute + self._dfc2_col_output = dfc2_col_output + # One PipelineTmaStore stage holds every data plane produced by a dFC2 subtile + self._d_output_slots = ( + 2 + + (2 if dfc2_col_output else 0) + + (1 if dfc2_recompute else 0) + ) + + # combine_format determines the dfc1 (final grad_x) combine encoding. + if combine_format is None: + combine_format = CombineFormat.parse("bf16") + self._combine_format = combine_format + self._combine_mxfp8 = combine_format.is_quantized + # sf_block_pad for the dfc1 MXFP8 combine + if self._combine_mxfp8 and combine_hidden is not None: + _hidden_dfc1 = combine_hidden + _sf_blocks_dfc1 = _hidden_dfc1 // EpilogueTileN + self._dfc1_sf_block_pad = ((_sf_blocks_dfc1 + 15) // 16) * 16 + self._hidden_dfc1 = _hidden_dfc1 + else: + self._dfc1_sf_block_pad = 0 + self._hidden_dfc1 = 0 + # batching stg.64 SF + self._dfc1_sf_batch8 = ( + self._combine_mxfp8 + and self._hidden_dfc1 > 0 + and (self._hidden_dfc1 % self._cta_tile_n == 0) + and (self._cta_tile_n // EpilogueTileN == 8) + ) + + pass + + # -- Codegen-time queries -- + + @property + def epi_tile(self) -> Tuple[int, int]: + return self._epi_tile + + @property + def num_acc_pipeline_stages(self) -> int: + return self._num_acc_pipeline_stages + + @property + def num_acc_stage(self) -> int: + return self._num_acc_stage + + @property + def d_output_slots(self) -> int: + return self._d_output_slots + + @property + def subtile_cnt(self) -> int: + return self._subtile_cnt + + @property + def cta_tile_n(self) -> int: + return self._cta_tile_n + + @property + def num_sf_tmem_cols(self) -> int: + return self._num_sf_tmem_cols + + @property + def num_sfa_tmem_cols(self) -> int: + return self._num_sfa_tmem_cols + + @property + def num_sfb_tmem_cols(self) -> int: + return self._num_sfb_tmem_cols + + @property + def num_accumulator_tmem_cols(self) -> int: + return self._num_accumulator_tmem_cols + + def staged_smem_layout( + self, + n_stages: int, + ) -> Union[cute.Layout, cute.ComposedLayout]: + return sm100_utils.make_smem_layout_epi( + self.fc1_output_dtype, + self.fc1_output_layout, + self._epi_tile, + n_stages, + ) + + @property + def smem_layout_one_stage(self) -> Union[cute.Layout, cute.ComposedLayout]: + staged = self.staged_smem_layout(1) + return cute.select(staged, mode=[0, 1]) + + @property + def bytes_per_stage(self) -> int: + return cute.size_in_bytes(self.fc1_output_dtype, self.smem_layout_one_stage) + + # -- grad_y1 (dfc2 output) sD staging: reference-style shared (128×32) box -- + @property + def d_epi_tile(self) -> Tuple[int, int]: + return (self._cta_tile_m, EpilogueTileN) + + def d_staged_smem_layout( + self, + n_stages: int, + ) -> Union[cute.Layout, cute.ComposedLayout]: + return sm100_utils.make_smem_layout_epi( + self.fc1_output_dtype, + self.fc1_output_layout, + self.d_epi_tile, + n_stages, + ) + + @property + def d_smem_layout_one_stage(self) -> Union[cute.Layout, cute.ComposedLayout]: + staged = self.d_staged_smem_layout(1) + return cute.select(staged, mode=[0, 1]) + + @property + def d_bytes_per_stage(self) -> int: + return cute.size_in_bytes(self.fc1_output_dtype, self.d_smem_layout_one_stage) + + # forward pre-activation (dswiglu C input) staging + @property + def preact_epi_tile(self) -> Tuple[int, int]: + return (self._cta_tile_m, EpilogueTileN) + + def preact_staged_smem_layout( + self, n_stages: int + ) -> Union[cute.Layout, cute.ComposedLayout]: + return sm100_utils.make_smem_layout_epi( + cutlass.BFloat16, + self.fc1_output_layout, + self.preact_epi_tile, + n_stages, + ) + + @property + def preact_smem_layout_one_stage(self) -> Union[cute.Layout, cute.ComposedLayout]: + staged = self.preact_staged_smem_layout(1) + return cute.select(staged, mode=[0, 1]) + + @property + def preact_bytes_per_stage(self) -> int: + return cute.size_in_bytes(cutlass.BFloat16, self.preact_smem_layout_one_stage) + + @cute.jit + def _store_aux_row_smem(self, r_data: cute.Tensor, s_data: cute.Tensor) -> None: + """Store one 32-byte token row as two swizzle-safe 128-bit segments.""" + segment_elements = 16 + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.fc1_output_dtype, + num_bits_per_copy=128, + ) + r_segments = cute.zipped_divide(r_data, (segment_elements,)) + s_segments = cute.zipped_divide(s_data, (segment_elements,)) + for segment in cutlass.range_constexpr(EpilogueTileN // segment_elements): + cute.copy( + store_atom, + cute.coalesce(r_segments[None, segment]), + cute.coalesce(s_segments[None, segment]), + ) + + @cute.jit + def _run_dfc2_task_tile( + self, + work_tile_info, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + acc_consumer_state, + sched_ext, + gmem_fc1_output: cute.Tensor, + gmem_fc1_output_sf: cute.Tensor, + tma_atom_fc1_recompute: cute.CopyAtom, + gmem_fc1_recompute: cute.Tensor, + gmem_fc1_recompute_sf: cute.Tensor, + tma_atom_fc1_col_output: cute.CopyAtom, + gmem_fc1_col_output: cute.Tensor, + gmem_fc1_col_output_sf: cute.Tensor, + c_pipeline, + smem_preact_buffer: cute.Tensor, + c_consumer_state, + smem_d_buffer: cute.Tensor, + tma_atom_grad_y1: cute.CopyAtom, + warp_idx: int, + tidx, + norm_const, + gmem_topk_scores: cute.Tensor, + gmem_beta: cute.Tensor, + gmem_dprob: cute.Tensor, + d_pipeline, + d_num_stage, + token_comm_args=None, + ): + """dfc2 task-tile — c_pipeline CONSUMER. """ + real_fc1_output, _ = sched_ext.get_gmem_tensor("d", gmem_fc1_output, work_tile_info) + real_fc1_output_sf, _ = sched_ext.get_gmem_tensor("sfd", gmem_fc1_output_sf, work_tile_info) + if cutlass.const_expr(token_comm_args is None): + real_dprob, _ = sched_ext.get_gmem_tensor("topk", gmem_dprob, work_tile_info) + else: + real_dprob = None + if cutlass.const_expr(self._dfc2_recompute): + real_fc1_recompute, _ = sched_ext.get_gmem_tensor("recompute", gmem_fc1_recompute, work_tile_info) + real_fc1_recompute_sf, _ = sched_ext.get_gmem_tensor("sfrecompute", gmem_fc1_recompute_sf, work_tile_info) + else: + real_fc1_recompute = None + real_fc1_recompute_sf = None + if cutlass.const_expr(self._dfc2_col_output): + real_fc1_col_output, _ = sched_ext.get_gmem_tensor("col_output", gmem_fc1_col_output, work_tile_info) + real_fc1_col_output_sf, _ = sched_ext.get_gmem_tensor("sfcol_output", gmem_fc1_col_output_sf, work_tile_info) + else: + real_fc1_col_output = None + real_fc1_col_output_sf = None + + acc_pipeline.consumer_wait(acc_consumer_state) + iket.range_push("mxfp8_dfc2_epi_tile") + + subtile_cnt = self._cta_tile_n // EpilogueTileN # 8 (256 / 32) + start_subtile = 0 + tmem_t = self._subtile_dfc12_tmem_tensor( + tmem_acc_tensor, cutlass.Int32(start_subtile), warp_idx, + ) + tmem_forward_cols = EpilogueTileN + + rmem_sf = cute.make_rmem_tensor( + cute.make_layout(2 * (self._cta_tile_n // EpilogueTileN)).shape, self.acc_dtype, + ) + # fc1_recompute SF accumulator: ONE SF per subtile (recompute N = half of dfc2). + rmem_sf_recompute = cute.make_rmem_tensor( + cute.make_layout(self._cta_tile_n // EpilogueTileN).shape, self.acc_dtype, + ) + # fc1_col_output SF accumulator + rmem_sf_col_output = cute.make_rmem_tensor( + cute.make_layout(2 * (self._cta_tile_n // EpilogueTileN)).shape, self.acc_dtype, + ) + thread_in_warp = tidx % WarpThreadCount + token_row_in_cta = cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + expert_local_token_idx = ( + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + token_row_in_cta + ) + + # beta / prob / dprob setup + beta_val = gmem_beta[work_tile_info.expert_idx] + # mProb: load from topk_scores for valid tokens; default 1.0 (unused) for invalid. + rmem_prob = cute.make_rmem_tensor(cute.make_layout(1).shape, self.acc_dtype) + rmem_prob[0] = cutlass.Float32(1.0) + if token_row_in_cta < valid_tokens: + real_topk, _ = sched_ext.get_gmem_tensor("topk", gmem_topk_scores, work_tile_info) + rmem_prob[0] = real_topk[expert_local_token_idx] + + # Per-tile dprob accumulator (single scalar). + dprob = cutlass.Float32(0.0) + + # Output col-strips per N-tile. + n_col_strips_per_tile = (self._cta_tile_n * 2) // EpilogueTileN + base_token_tile = work_tile_info.tile_m_idx # 128-row tile index + + _epilog_sync = pipeline.NamedBarrier( + barrier_id=self._epilog_sync_bar_id, + num_threads=WarpThreadCount * len(self._epilogue_warp_ids), + ) + + # Build tiled copies for SMEM↔REG (reference pattern). + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self._mma_tiler, + self.fc1_output_layout, + self.fc1_output_dtype, + self.acc_dtype, + self.d_epi_tile, + self._use_2cta_instrs, + ) + tAcc_epi = cute.flat_divide( + tmem_acc_tensor[((None, None), 0, 0)], + self.d_epi_tile, + ) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_rAcc_full = thr_copy_t2r.partition_D(tAcc_epi) + copy_atom_s2r = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16) + tiled_copy_s2r = cute.make_tiled_copy_D(copy_atom_s2r, tiled_copy_t2r) + thr_copy_s2r = tiled_copy_s2r.get_slice(tidx) + tRS_sPre = thr_copy_s2r.partition_D(smem_preact_buffer) + + r_layout = cute.make_layout(((1, EpilogueTileN,), 1, 1,), stride=((0, 1,), 0, 0,)) + r_gate_bf = cute.make_rmem_tensor(r_layout, cutlass.BFloat16) + r_up_bf = cute.make_rmem_tensor(r_layout, cutlass.BFloat16) + copy_atom_r2s = sm100_utils.get_smem_store_op( + self.fc1_output_layout, self.fc1_output_dtype, self.acc_dtype, tiled_copy_t2r + ) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + + for i in cutlass.range(0, subtile_cnt, 1): + subtile_idx = cutlass.Int32(i) + c_consumer_state, subtile_dprob = self._run_dfc2_subtile( + subtile_idx=subtile_idx, + subtile_i=i, + t_subtile=tmem_t, + smem_d=smem_d_buffer, + tiled_copy_r2s=tiled_copy_r2s, + tiled_copy_s2r=tiled_copy_s2r, + tRS_sPre=tRS_sPre, + c_pipeline=c_pipeline, + c_consumer_state=c_consumer_state, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + r_gate_bf=r_gate_bf, + r_up_bf=r_up_bf, + work_tile_info=work_tile_info, + warp_idx=warp_idx, + tidx=tidx, + norm_const=norm_const, + rmem_sf=rmem_sf, + rmem_sf_recompute=rmem_sf_recompute, + real_fc1_recompute=real_fc1_recompute, + rmem_sf_col_output=rmem_sf_col_output, + real_fc1_col_output=real_fc1_col_output, + beta=beta_val, + prob=rmem_prob[0], + epilog_sync=_epilog_sync, + d_pipeline=d_pipeline, + d_num_stage=d_num_stage, + ) + dprob = dprob + subtile_dprob + + tmem_t = self._advance_fc2_tmem_tensor(tmem_t, tmem_forward_cols) + + # BARRIER: fence_proxy makes R2S (stmatrix) writes visible to TMA + cute.arch.fence_proxy("async.shared", space="cta") + _epilog_sync.arrive_and_wait() + + # Compute GMEM tile pointers for this subtile. + gate_col_idx = ( + work_tile_info.tile_n_idx * cutlass.Int32(n_col_strips_per_tile) + + subtile_idx * cutlass.Int32(2) + ) + g_gate = cute.local_tile( + real_fc1_output, + (self._cta_tile_m, EpilogueTileN, 1), + (base_token_tile, gate_col_idx, cutlass.Int32(0)), + )[(None, None, 0)] + g_up = cute.local_tile( + real_fc1_output, + (self._cta_tile_m, EpilogueTileN, 1), + (base_token_tile, gate_col_idx + cutlass.Int32(1), cutlass.Int32(0)), + )[(None, None, 0)] + g_col_gate = None + g_col_up = None + if cutlass.const_expr(self._dfc2_col_output): + g_col_gate = cute.local_tile( + real_fc1_col_output, + (self._cta_tile_m, EpilogueTileN, 1), + (base_token_tile, gate_col_idx, cutlass.Int32(0)), + )[(None, None, 0)] + g_col_up = cute.local_tile( + real_fc1_col_output, + (self._cta_tile_m, EpilogueTileN, 1), + (base_token_tile, gate_col_idx + cutlass.Int32(1), cutlass.Int32(0)), + )[(None, None, 0)] + g_recompute = None + if cutlass.const_expr(self._dfc2_recompute): + recompute_col_idx = ( + work_tile_info.tile_n_idx + * cutlass.Int32(self._cta_tile_n // EpilogueTileN) + + subtile_idx + ) + g_recompute = cute.local_tile( + real_fc1_recompute, + (self._cta_tile_m, EpilogueTileN, 1), + (base_token_tile, recompute_col_idx, cutlass.Int32(0)), + )[(None, None, 0)] + # TMA issue (warp 0 only). + d_outputs_per_stage = cutlass.const_expr(self._d_output_slots) + d_n_stages = cutlass.const_expr(d_num_stage // d_outputs_per_stage) + d_slot = cutlass.Int32(d_outputs_per_stage) * ( + cutlass.Int32(i) % cutlass.Int32(d_n_stages) + ) + if warp_idx == self._epilogue_warp_ids[0]: + self.tma_store_dfc2_outputs( + smem_d_buffer, + tma_atom_grad_y1, + g_gate, + g_up, + tma_atom_fc1_col_output, + g_col_gate, + g_col_up, + tma_atom_fc1_recompute, + g_recompute, + valid_tokens, + d_pipeline, + d_slot, + ) + + self._acc_pipeline_consumer_release(acc_pipeline, acc_consumer_state, True) + + valid_inter = real_fc1_output.shape[1] + self._stg_sf_dfc2(rmem_sf, real_fc1_output_sf, work_tile_info, tidx, valid_inter) + + # fc1_recompute SFs + if cutlass.const_expr(self._dfc2_recompute): + valid_inter_recompute = real_fc1_recompute.shape[1] + self._stg_sf_recompute( + rmem_sf_recompute, real_fc1_recompute_sf, + work_tile_info, tidx, valid_inter_recompute, valid_tokens, + ) + # fc1_col_output SFs + if cutlass.const_expr(self._dfc2_col_output): + valid_inter_col_output = real_fc1_col_output.shape[1] + self._stg_sf_col_output( + rmem_sf_col_output, real_fc1_col_output_sf, + work_tile_info, tidx, valid_inter_col_output, valid_tokens, + ) + # MegaMoE maps the receiver-pool row back to the source rank's combine slot. + if cutlass.const_expr(token_comm_args is not None): + if token_row_in_cta < valid_tokens: + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + metadata_u32 = cute.recast_tensor( + token_comm_args.token_src_metadata, cutlass.Uint32, + ) + dprob_output_dest = Fc2OutputDest( + tensor=token_comm_args.dprob_output, + metadata=metadata_u32, + peer_rank_ptr_mapper=token_comm_args.peer_rank_ptr_mapper, + ) + dest_row = dprob_output_dest.resolve_token_row(pool_token_global) + dprob_reduce_gmem( + dest_row, + dprob, + True, + cutlass.Int32(0), + system_scope=True, + ) + else: + dprob_reduce_gmem( + real_dprob, + dprob, + token_row_in_cta < valid_tokens, + expert_local_token_idx, + ) + iket.range_pop() + return c_consumer_state + + @cute.jit + def _dfc2_load_c( + self, + c_pipeline, + c_consumer_state, + tiled_copy_s2r, + tRS_sPre: cute.Tensor, + r_gate_bf, + r_up_bf, + ): + """Consume gate then up c_pipeline stages into register tiles.""" + c_pipeline.consumer_wait(c_consumer_state) + c_slot = 2 * c_consumer_state.index + cute.copy( + tiled_copy_s2r, + tRS_sPre[(None, None, None, c_slot)], + r_gate_bf, + ) + cute.copy( + tiled_copy_s2r, + tRS_sPre[(None, None, None, c_slot + 1)], + r_up_bf, + ) + cute.arch.fence_proxy("async.shared", space="cta") + c_pipeline.consumer_release(c_consumer_state) + c_consumer_state.advance() + + return c_consumer_state + + @cute.jit + def _run_dfc2_subtile( + self, + subtile_idx, + subtile_i, + t_subtile: cute.Tensor, + smem_d: cute.Tensor, + tiled_copy_r2s, + tiled_copy_s2r, + tRS_sPre: cute.Tensor, + c_pipeline, + c_consumer_state, + acc_pipeline, + acc_consumer_state, + r_gate_bf: cute.Tensor, + r_up_bf: cute.Tensor, + work_tile_info, + warp_idx: int, + tidx, + norm_const, + rmem_sf: cute.Tensor, + rmem_sf_recompute: cute.Tensor, + real_fc1_recompute: cute.Tensor, + rmem_sf_col_output: cute.Tensor, + real_fc1_col_output: cute.Tensor, + beta: cutlass.Float32, + prob: cutlass.Float32, + epilog_sync, + d_pipeline, + d_num_stage, + ): + iket.range_push("mxfp8_dfc2_epilogue_subtile") + EN = EpilogueTileN + + r_layout = cute.make_layout((((EN,), 1),), stride=(((1,), 0),)) + atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), self.acc_dtype, + ) + r_acc = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + cute.copy(atom_t2r, t_subtile, r_acc) + + thread_in_warp = tidx % WarpThreadCount + token_row_in_cta = cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + subtile_dprob = cutlass.Float32(0.0) + + # load c from shared memory to registers + c_consumer_state = self._dfc2_load_c( + c_pipeline, + c_consumer_state, + tiled_copy_s2r, + tRS_sPre, + r_gate_bf, + r_up_bf, + ) + + # c_gate / c_up declared outside the validity guard: stmatrix (tiled_copy_r2s) + # is warp-cooperative and all threads must call it regardless of token validity. + c_shape = cute.make_layout(((1, EN,), 1, 1), stride=((0, 1,), 0, 0)).shape + c_gate = cute.make_rmem_tensor(c_shape, self.fc1_output_dtype) + c_up = cute.make_rmem_tensor(c_shape, self.fc1_output_dtype) + # c_recompute: flat MXFP8 row for token-major output staging + c_recompute = cute.make_rmem_tensor(cute.make_layout(EN).shape, self.fc1_output_dtype) + + is_valid_row = token_row_in_cta < valid_tokens + + r_gate = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + r_up = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + for j in cutlass.range_constexpr(EN): + r_gate[j] = r_gate_bf[j].to(self.acc_dtype) + r_up[j] = r_up_bf[j].to(self.acc_dtype) + # Zero invalid rows' inputs (their r_acc/r_gate/r_up are padding garbage) so the + # warp-wide column amax is not polluted and no NaN propagates into the reduction. + if token_row_in_cta >= valid_tokens: + for j in cutlass.range_constexpr(EN): + r_acc[j] = self.acc_dtype(0.0) + r_gate[j] = self.acc_dtype(0.0) + r_up[j] = self.acc_dtype(0.0) + + # dswiglu backward: acc(grad_h) x (gate, up) -> (d_gate, d_up) ---- + d_gate = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + d_up = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + if cutlass.const_expr(self._act_func == "swiglu"): + subtile_dprob = dswiglu_act( + d_gate, d_up, r_acc, r_gate, r_up, beta, prob, self._gate_up_clamp + ) + + if cutlass.const_expr(self._dfc2_col_output): + # Snapshot d_gate / d_up BEFORE quant_sfd_row mutates them in place; the col + # path col-quants these copies (quant_sfd_col mutates its input). + d_gate_col = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + d_up_col = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + for j in cutlass.range_constexpr(EN): + d_gate_col[j] = d_gate[j] + d_up_col[j] = d_up[j] + c_gate_col = cute.make_rmem_tensor(cute.make_layout(EN).shape, self.fc1_output_dtype) + c_up_col = cute.make_rmem_tensor(cute.make_layout(EN).shape, self.fc1_output_dtype) + qg_col = quant_sfd_col( + d_gate_col, c_gate_col, norm_const, + self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype, + ) + qu_col = quant_sfd_col( + d_up_col, c_up_col, norm_const, + self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype, + ) + for _k in cutlass.range_constexpr(self._cta_tile_n // EN): + if subtile_idx == cutlass.Int32(_k): + rmem_sf_col_output[2 * _k] = qg_col + rmem_sf_col_output[2 * _k + 1] = qu_col + + # quantize each half to MXFP8 + E8M0 row SF (per-thread, no warp reduction) ---- + qg = quant_sfd_row( + d_gate, c_gate, norm_const, self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype, + ) + qu = quant_sfd_row( + d_up, c_up, norm_const, self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype, + ) + + # accumulate the 2 E8M0 row SFs into rmem + for _k in cutlass.range_constexpr(self._cta_tile_n // EN): + if subtile_idx == cutlass.Int32(_k): + rmem_sf[2 * _k] = qg + rmem_sf[2 * _k + 1] = qu + + # dfc2_recompute: forward swiglu + column quantization + if cutlass.const_expr(self._dfc2_recompute): + c_recompute_f32 = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + if cutlass.const_expr(self._act_func == "swiglu"): + swiglu_act( + c_recompute_f32, + r_up, + r_gate, + prob, + self._gate_up_clamp, + ) + qc = quant_sfd_col( + c_recompute_f32, c_recompute, norm_const, + self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype, + ) + for _k in cutlass.range_constexpr(self._cta_tile_n // EN): + if subtile_idx == cutlass.Int32(_k): + rmem_sf_recompute[_k] = qc + + # BARRIER: drain PREVIOUS subtile's TMA BEFORE R2S. + if warp_idx == self._epilogue_warp_ids[0]: + d_pipeline.producer_acquire() + epilog_sync.arrive_and_wait() + + # Write d to smem. + d_outputs_per_stage = cutlass.const_expr(self._d_output_slots) + d_n_stages = cutlass.const_expr(d_num_stage // d_outputs_per_stage) + d_slot = cutlass.Int32(d_outputs_per_stage) * ( + subtile_i % cutlass.Int32(d_n_stages) + ) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + sd = thr_copy_r2s.partition_D(smem_d) + cute.copy(tiled_copy_r2s, c_gate, sd[(None, None, None, d_slot)]) + cute.copy(tiled_copy_r2s, c_up, sd[(None, None, None, d_slot + cutlass.Int32(1))]) + + # Auxiliary data planes use the public token-major ABI. + next_slot = d_slot + cutlass.Int32(2) + if cutlass.const_expr(self._dfc2_col_output): + s_col_gate = cute.slice_(smem_d, (token_row_in_cta, None, next_slot)) + s_col_up = cute.slice_( + smem_d, (token_row_in_cta, None, next_slot + cutlass.Int32(1)) + ) + self._store_aux_row_smem(c_gate_col, s_col_gate) + self._store_aux_row_smem(c_up_col, s_col_up) + next_slot = next_slot + cutlass.Int32(2) + if cutlass.const_expr(self._dfc2_recompute): + s_recompute = cute.slice_(smem_d, (token_row_in_cta, None, next_slot)) + self._store_aux_row_smem(c_recompute, s_recompute) + + iket.range_pop() + return c_consumer_state, subtile_dprob + + @cute.jit + def _stg_sf_dfc2( + self, + rmem_sf_f32: cute.Tensor, + real_fc1_output_sf: cute.Tensor, + work_tile_info, + tidx, + valid_inter, + ) -> None: + """Store the dfc2 grad_y1 E8M0 row SFs, 4 blocks per 128-col region.""" + if tidx < work_tile_info.valid_tokens_in_cta_tile: + token_idx = ( + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + tidx + ) + n_regions = (self._cta_tile_n * 2) // Fc1EpilogueOutputTileN + region_col = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n * 2) + ) + for r in cutlass.range_constexpr(n_regions): + sf_base = cute.local_tile( + real_fc1_output_sf, (1, 1, 1), + (token_idx, region_col, cutlass.Int32(0)), + ) + r_sf4_f32 = cute.make_rmem_tensor(cute.make_layout(4).shape, self.acc_dtype) + for idx in cutlass.range_constexpr(4): + r_sf4_f32[idx] = rmem_sf_f32[r * 4 + idx] + if region_col < valid_inter: + sf_ptr = cute.make_ptr( + self.sf_dtype, + sf_base.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=4, + ) + gmem_sf4 = cute.make_tensor(sf_ptr, cute.make_layout(4)) + r_sf4 = cute.make_rmem_tensor(cute.make_layout(4).shape, self.sf_dtype) + r_sf4.store(r_sf4_f32.load().to(self.sf_dtype)) + cute.autovec_copy(r_sf4, gmem_sf4) + region_col += cutlass.Int32(Fc1EpilogueOutputTileN) + + @cute.jit + def _stg_col_sf_atom_value( + self, + real_sf: cute.Tensor, + row_block, + feature, + _feature_atoms, + sf_value, + ) -> None: + """Store one SF in a 128-feature × 4-token-block atom.""" + token_atom = row_block // cutlass.Int32(4) + token_bank = row_block % cutlass.Int32(4) + feature_atom = feature // cutlass.Int32(128) + feature_bank = (feature // cutlass.Int32(32)) % cutlass.Int32(4) + feature_lane = feature % cutlass.Int32(32) + atom_byte = ( + feature_lane * cutlass.Int32(16) + + feature_bank * cutlass.Int32(4) + + token_bank + ) + if feature_atom < real_sf.shape[0] and token_atom < real_sf.shape[1]: + real_sf[feature_atom, token_atom, atom_byte] = sf_value.to(self.sf_dtype) + + @cute.jit + def _stg_sf_recompute( + self, + rmem_sf_f32: cute.Tensor, + real_fc1_recompute_sf: cute.Tensor, + work_tile_info, + tidx, + valid_inter, + valid_tokens, + ) -> None: + """Store fc1_recompute SFs in MN-major 128×4 atoms.""" + EN = EpilogueTileN # 32 + sf_vec_size = self._sf_vec_size + warp_lane_idx = tidx % cutlass.Int32(32) + warp_idx_local = tidx // cutlass.Int32(32) + hidden_atoms = (valid_inter + cutlass.Int32(127)) // cutlass.Int32(128) + + # Row-block within the M-tile: 4 warps × 1 row-block each (128 / 32). + row_blocks_per_m_tile = self._cta_tile_m // sf_vec_size + row_block = ( + work_tile_info.tile_m_idx * cutlass.Int32(row_blocks_per_m_tile) + + warp_idx_local + ) + col_base = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + ) + for s in cutlass.range_constexpr(self._cta_tile_n // EN): + col = col_base + cutlass.Int32(s * EN) + warp_lane_idx + if col < valid_inter: + self._stg_col_sf_atom_value( + real_fc1_recompute_sf, + row_block, + col, + hidden_atoms, + rmem_sf_f32[s], + ) + + @cute.jit + def _stg_sf_col_output( + self, + rmem_sf_f32: cute.Tensor, + real_fc1_col_output_sf: cute.Tensor, + work_tile_info, + tidx, + valid_inter, + valid_tokens, + ) -> None: + """Store fc1_col_output SFs in MN-major 128×4 atoms.""" + EN = EpilogueTileN # 32 + sf_vec_size = self._sf_vec_size + warp_lane_idx = tidx % cutlass.Int32(32) + warp_idx_local = tidx // cutlass.Int32(32) + hidden_atoms = (valid_inter + cutlass.Int32(127)) // cutlass.Int32(128) + + row_blocks_per_m_tile = self._cta_tile_m // sf_vec_size + row_block = ( + work_tile_info.tile_m_idx * cutlass.Int32(row_blocks_per_m_tile) + + warp_idx_local + ) + # Doubled N: cta_tile_n * 2 cols per fc1 N-tile. + col_base = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n * 2) + ) + for s in cutlass.range_constexpr(self._cta_tile_n // EN): + for gu in cutlass.range_constexpr(2): + col = col_base + cutlass.Int32((2 * s + gu) * EN) + warp_lane_idx + if col < valid_inter: + self._stg_col_sf_atom_value( + real_fc1_col_output_sf, + row_block, + col, + hidden_atoms, + rmem_sf_f32[2 * s + gu], + ) + + @cute.jit + def _tma_store_tile( + self, + smem_tile: cute.Tensor, + tma_atom: cute.CopyAtom, + gmem_tile: cute.Tensor, + ) -> None: + tma_smem_src, tma_gmem_dst = cpasync.tma_partition( + tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(smem_tile, 0, 2), + cute.group_modes(gmem_tile, 0, 2), + ) + cute.copy(tma_atom, tma_smem_src, tma_gmem_dst) + + @cute.jit + def tma_store_dfc2_outputs( + self, + smem_d_buffer: cute.Tensor, + tma_atom_grad_y1: cute.CopyAtom, + g_gate_2d: cute.Tensor, + g_up_2d: cute.Tensor, + tma_atom_col_output: cute.CopyAtom, + g_col_gate_2d, + g_col_up_2d, + tma_atom_recompute: cute.CopyAtom, + g_recompute_2d, + valid_tokens, + d_pipeline, + d_slot, + ) -> None: + """Issue one TMA store group for every dFC2 data plane.""" + tile_is_valid = valid_tokens > cutlass.Int32(0) + if tile_is_valid: + self._tma_store_tile( + cute.slice_(smem_d_buffer, (None, None, d_slot)), + tma_atom_grad_y1, + g_gate_2d, + ) + self._tma_store_tile( + cute.slice_( + smem_d_buffer, (None, None, d_slot + cutlass.Int32(1)) + ), + tma_atom_grad_y1, + g_up_2d, + ) + next_slot = d_slot + cutlass.Int32(2) + if cutlass.const_expr(self._dfc2_col_output): + self._tma_store_tile( + cute.slice_(smem_d_buffer, (None, None, next_slot)), + tma_atom_col_output, + g_col_gate_2d, + ) + self._tma_store_tile( + cute.slice_( + smem_d_buffer, + (None, None, next_slot + cutlass.Int32(1)), + ), + tma_atom_col_output, + g_col_up_2d, + ) + next_slot = next_slot + cutlass.Int32(2) + if cutlass.const_expr(self._dfc2_recompute): + self._tma_store_tile( + cute.slice_(smem_d_buffer, (None, None, next_slot)), + tma_atom_recompute, + g_recompute_2d, + ) + d_pipeline.producer_commit() + + + @cute.jit + def _subtile_dfc12_tmem_tensor( + self, + tmem_acc_tensor: cute.Tensor, + subtile_idx, + warp_idx, + ) -> cute.Tensor: + """ + Per-warp TMEM view for one fc2 subtile (EpilogueTileN=32 cols). + """ + base = tmem_acc_tensor.iterator + warp_lane_off = warp_idx * WarpThreadCount + subtile_col_off = subtile_idx * EpilogueTileN + total = (warp_lane_off << 16) + subtile_col_off + subtile_ptr = base + cute.assume(total, divby=16) + return cute.make_tensor( + subtile_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ) + + @cute.jit + def _advance_fc2_tmem_tensor( + self, + tmem_tensor: cute.Tensor, + col_offset: int, + ) -> cute.Tensor: + new_ptr = tmem_tensor.iterator + cute.assume(col_offset, divby=16) + return cute.make_tensor( + new_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ) + + @cute.jit + def _acc_pipeline_consumer_release( + self, + acc_pipeline, + acc_consumer_state, + is_release: bool, + ) -> None: + """Release the acc pipeline consumer.""" + if is_release: + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + @cute.jit + def _run_dfc1_subtile( + self, + subtile_idx, + subtile_i, + t_subtile: cute.Tensor, + real_fc2_output: cute.Tensor, + work_tile_info, + valid_hidden, + warp_idx: int, + tidx, + acc_pipeline, + acc_consumer_state, + token_comm_args=None, + rmem_sf_dfc1=None, + *, + preload_acc=None, + ) -> None: + """fc2 subtile: LDTM + fp32->bf16 + STG.""" + iket.range_push("mxfp8_fc2_epilogue_subtile") + dfc1_subtile_cnt = self._cta_tile_n // EpilogueTileN # = 8 + r_acc_layout = cute.make_layout((((EpilogueTileN,), 1),), stride=(((1,), 0),)) + atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), self.acc_dtype, + ) + r_acc = cute.make_rmem_tensor(r_acc_layout.shape, self.acc_dtype) + cute.copy(atom_t2r, t_subtile, r_acc) + + hidden_group = ( + work_tile_info.tile_n_idx * cutlass.Int32(dfc1_subtile_cnt) + subtile_idx + ) + hidden_col_start = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + + subtile_idx * cutlass.Int32(EpilogueTileN) + ) + r_bf16 = cute.make_rmem_tensor(r_acc_layout.shape, cutlass.BFloat16) + r_bf16.store(r_acc.load().to(cutlass.BFloat16)) + thread_in_warp = tidx % WarpThreadCount + token_row_in_cta = cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + is_valid = token_row_in_cta < valid_tokens and hidden_col_start < valid_hidden + + if cutlass.const_expr( + token_comm_args is not None + and not self._token_back_by_dispatch + and self._combine_mxfp8 + ): + fp8_dtype = self._combine_format.act_dtype + r_fp8 = cute.make_rmem_tensor(r_acc_layout.shape, fp8_dtype) + qpvscale = quant_sfd_row( + r_acc, r_fp8, 1.0, EpilogueTileN, + cutlass.Float8E8M0FNU, fp8_dtype, + ) + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + metadata_u32 = cute.recast_tensor( + token_comm_args.token_src_metadata, cutlass.Uint32, + ) + fc2_output_dest = Fc2OutputDest( + tensor=token_comm_args.combine_output, + metadata=metadata_u32, + peer_rank_ptr_mapper=token_comm_args.peer_rank_ptr_mapper, + ) + dest_row = fc2_output_dest.resolve_token_row(pool_token_global) + r_fp8_flat = cute.make_tensor(r_fp8.iterator, cute.make_layout(32)) + stg_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=256, + ) + dest_fp8_ptr = cute.make_ptr( + fp8_dtype, + dest_row.iterator.toint() + Int64(hidden_col_start), + cute.AddressSpace.gmem, + assumed_align=32, + ) + if is_valid: + cute.copy( + stg_fp8_atom, r_fp8_flat, + cute.make_tensor(dest_fp8_ptr, cute.make_layout(32)), + ) + self._write_sf_dfc1_buffer(rmem_sf_dfc1, subtile_idx, qpvscale) + elif cutlass.const_expr( + self._token_back_by_dispatch and self._combine_mxfp8 + ): + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + fp8_dtype = self._combine_format.act_dtype + r_fp8 = cute.make_rmem_tensor(r_acc_layout.shape, fp8_dtype) + qpvscale = quant_sfd_row( + r_acc, r_fp8, 1.0, EpilogueTileN, + cutlass.Float8E8M0FNU, fp8_dtype, + ) + fp8_byte_addr = ( + token_comm_args.fc2_output_workspace.iterator.toint() + + Int64(pool_token_global) * Int64(self._hidden_dfc1) + + Int64(hidden_col_start) + ) + stg_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=256, + ) + aligned_fp8_iter = cute.make_ptr( + fp8_dtype, + fp8_byte_addr, + cute.AddressSpace.gmem, + assumed_align=32, + ) + r_fp8_flat = cute.make_tensor(r_fp8.iterator, cute.make_layout(EpilogueTileN)) + if is_valid: + cute.copy( + stg_fp8_atom, r_fp8_flat, + cute.make_tensor(aligned_fp8_iter, cute.make_layout(EpilogueTileN)), + ) + self._write_sf_dfc1_buffer(rmem_sf_dfc1, subtile_idx, qpvscale) + else: + # BF16 path (default): fp32->bf16, two 256-bit STGs. + stg_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=256, + ) + if cutlass.const_expr( + token_comm_args is not None + and not self._token_back_by_dispatch + ): + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + metadata_u32 = cute.recast_tensor( + token_comm_args.token_src_metadata, cutlass.Uint32, + ) + fc2_output_dest = Fc2OutputDest( + tensor=token_comm_args.combine_output, + metadata=metadata_u32, + peer_rank_ptr_mapper=token_comm_args.peer_rank_ptr_mapper, + # Collapse topk -> src_topk=0 so every contribution of a + # source token resolves to the SAME combine row (red-added + # below). No-op when reduce is off. + reduce_topk_in_kernel=self._reduce_topk_in_epilogue, + ) + dest_row = fc2_output_dest.resolve_token_row(pool_token_global) + for stg_half in cutlass.range(EpilogueTileN // 16, unroll_full=True): + reg_view = cute.make_tensor( + r_bf16.iterator + stg_half * 16, + cute.make_layout(16), + ) + if cutlass.const_expr( + token_comm_args is not None + and not self._token_back_by_dispatch + ): + # epi_warps: peer-write grad_x directly to combine_output + hidden_off = hidden_col_start + cutlass.Int32(stg_half * 16) + if cutlass.const_expr(self._reduce_topk_in_epilogue): + if is_valid: + reg_u32 = cute.recast_tensor(reg_view, cutlass.Uint32) + for redg_i in cutlass.range_constexpr(16 // 4): + chunk_ptr = cute.make_ptr( + cutlass.BFloat16, + dest_row.iterator.toint() + + (hidden_off + cutlass.Int32(redg_i * 4)) + * cutlass.Int64(2), + cute.AddressSpace.gmem, + assumed_align=8, + ) + _red_add_relaxed_sys_v2_bf16x2( + chunk_ptr, + cutlass.Uint32(reg_u32[2 * redg_i]), + cutlass.Uint32(reg_u32[2 * redg_i + 1]), + ) + else: + dest_ptr = cute.make_ptr( + cutlass.BFloat16, + dest_row.iterator.toint() + hidden_off * cutlass.Int64(2), + cute.AddressSpace.gmem, + assumed_align=32, + ) + if is_valid: + cute.copy( + stg_atom, reg_view, + cute.make_tensor(dest_ptr, cute.make_layout(16)), + ) + else: + # Lean path (token_comm_args is None) OR dispatch-push + # (token_back_by_dispatch) + g_fc2_output_tile = cute.local_tile( + real_fc2_output, + (self._cta_tile_m, EpilogueTileN, 1), + (work_tile_info.tile_m_idx, hidden_group, 0), + ) + g_fc2_slice = cute.slice_(g_fc2_output_tile, (None, None, 0)) + g_thread_row = cute.local_tile( + g_fc2_slice, (1, 16), (token_row_in_cta, stg_half), + ) + g_flat = cute.coalesce(g_thread_row) + aligned_iter = cute.make_ptr( + cutlass.BFloat16, + g_flat.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=32, + ) + if is_valid: + cute.copy(stg_atom, reg_view, cute.make_tensor(aligned_iter, g_flat.layout)) + + iket.range_pop() + + @cute.jit + def _write_sf_dfc1_buffer(self, rmem_sf_dfc1, subtile_idx, qpvscale) -> None: + """Scatter one subtile's E8M0 scale into the per-tile SF buffer.""" + for j in cutlass.range_constexpr(self._cta_tile_n // EpilogueTileN): + if subtile_idx == cutlass.Int32(j): + rmem_sf_dfc1[j] = qpvscale + + @cute.jit + def _stg_sf_dfc1( + self, + rmem_sf_dfc1: cute.Tensor, + token_comm_args, + work_tile_info, + valid_hidden, + warp_idx: int, + tidx, + ) -> None: + """Flush a task tile's dfc1 E8M0 scales to local fc2_output_sf.""" + dfc1_subtile_cnt = self._cta_tile_n // EpilogueTileN + thread_in_warp = tidx % WarpThreadCount + token_row_in_cta = cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + if token_row_in_cta < work_tile_info.valid_tokens_in_cta_tile: + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + hidden_group_base = ( + work_tile_info.tile_n_idx * cutlass.Int32(dfc1_subtile_cnt) + ) + sf_byte_addr = ( + token_comm_args.fc2_output_sf.iterator.toint() + + Int64(pool_token_global) * Int64(self._dfc1_sf_block_pad) + + Int64(hidden_group_base) + ) + if cutlass.const_expr(self._dfc1_sf_batch8): + stg_e8m0x8_from_f32( + sf_byte_addr, + rmem_sf_dfc1[0], rmem_sf_dfc1[1], rmem_sf_dfc1[2], rmem_sf_dfc1[3], + rmem_sf_dfc1[4], rmem_sf_dfc1[5], rmem_sf_dfc1[6], rmem_sf_dfc1[7], + ) + else: + for j in cutlass.range_constexpr(dfc1_subtile_cnt): + block_hidden_start = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + + cutlass.Int32(j * EpilogueTileN) + ) + if block_hidden_start < valid_hidden: + stg_e8m0_from_f32(sf_byte_addr + Int64(j), rmem_sf_dfc1[j]) + + + @cute.jit + def _run_dfc1_task_tile( + self, + work_tile_info, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + acc_consumer_state, + sched_ext, + gmem_fc2_output: cute.Tensor, + valid_hidden, + warp_idx: int, + tidx, + token_comm_args=None, + ) -> None: + """fc2 (Linear2) task-tile body following fc1 pattern exactly.""" + real_fc2_output, _ = sched_ext.get_gmem_tensor( + "d", gmem_fc2_output, work_tile_info, + ) + acc_pipeline.consumer_wait(acc_consumer_state) + iket.range_push("mxfp8_dfc1_epi_tile") + + dfc1_subtile_cnt = self._cta_tile_n // EpilogueTileN # = 8 + + # Start subtile mirrors fc1: last for odd turn, first for even. + start_subtile = 0 + tmem_t = self._subtile_dfc12_tmem_tensor( + tmem_acc_tensor, cutlass.Int32(start_subtile), warp_idx, + ) + tmem_forward_cols = EpilogueTileN + + # Quantized combine: buffer the per-subtile E8M0 scales and flush them + # in one stg.64 after the loop (see _stg_sf_dfc1). + if cutlass.const_expr(self._combine_mxfp8 and token_comm_args is not None): + layout_sf_dfc1 = cute.make_layout(dfc1_subtile_cnt) + rmem_sf_dfc1 = cute.make_rmem_tensor(layout_sf_dfc1.shape, self.acc_dtype) + else: + rmem_sf_dfc1 = None + + for i in cutlass.range(0, dfc1_subtile_cnt, 1, unroll=1): + self._run_dfc1_subtile( + subtile_idx=cutlass.Int32(i), + subtile_i=i, + t_subtile=tmem_t, + real_fc2_output=real_fc2_output, + work_tile_info=work_tile_info, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + token_comm_args=token_comm_args, + rmem_sf_dfc1=rmem_sf_dfc1, + ) + + tmem_t = self._advance_fc2_tmem_tensor(tmem_t, tmem_forward_cols) + + # Release AFTER all subtile reads (never early-release for FC2). + self._acc_pipeline_consumer_release(acc_pipeline, acc_consumer_state, True) + + # Flush the buffered E8M0 scales (one stg.64 per thread when aligned). + if cutlass.const_expr(self._combine_mxfp8 and token_comm_args is not None): + self._stg_sf_dfc1( + rmem_sf_dfc1=rmem_sf_dfc1, + token_comm_args=token_comm_args, + work_tile_info=work_tile_info, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + ) + + iket.range_pop() + + + @cute.jit + def run( + self, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + sched_consumer, + sched_ext, + gmem_fc1_output: cute.Tensor, + gmem_fc1_output_sf: cute.Tensor, + tma_atom_fc1_recompute: cute.CopyAtom, + gmem_fc1_recompute: Optional[cute.Tensor], + gmem_fc1_recompute_sf: Optional[cute.Tensor], + tma_atom_fc1_col_output: cute.CopyAtom, + gmem_fc1_col_output: Optional[cute.Tensor], + gmem_fc1_col_output_sf: Optional[cute.Tensor], + smem_preact_buffer: cute.Tensor, + c_pipeline, + c_num_stage, + smem_d_buffer: cute.Tensor, + d_pipeline, + d_num_stage, + tma_atom_grad_y1: cute.CopyAtom, + gmem_topk_scores: cute.Tensor, + gmem_fc2_output: cute.Tensor, + gmem_fc1_done_counter: cute.Tensor, + warp_idx: int, + tidx, + alpha, + norm_const, + gmem_beta: cute.Tensor, + gmem_dprob: cute.Tensor, + token_comm_args=None, + ) -> None: + """ + Run the full MXFP8 dfc2+dfc1-fused (backward) epilogue task-tile loop. + """ + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._num_acc_pipeline_stages + ) + task_tile_boundary_bar = pipeline.NamedBarrier( + barrier_id=self._epilog_sync_bar_id, + num_threads=32 * len(self._epilogue_warp_ids), + ) + + valid_hidden = cutlass.Int32(gmem_fc2_output.shape[1]) + + c_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, c_num_stage + ) + + bidx, bidy, bidz = cute.arch.block_idx() + work_tile_info = sched_consumer.consume_work() + + flag_tracker = GpuReleaseFlagBatchTracker( + flag_address=Int64(0), + accumulated_flags=cutlass.Int32(0), + phase=cutlass.Int32(work_tile_info.phase), + thread_idx=tidx % (len(self._epilogue_warp_ids) * WarpThreadCount), + ) + + while work_tile_info.is_valid_tile: + acc_stage_index = acc_consumer_state.index + tmem_acc_stage_tesnor = tmem_acc_tensor[(None, None, None, acc_stage_index)] + + if work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1): + c_consumer_state = self._run_dfc2_task_tile( + work_tile_info=work_tile_info, + tmem_acc_tensor=tmem_acc_stage_tesnor, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + sched_ext=sched_ext, + gmem_fc1_output=gmem_fc1_output, + gmem_fc1_output_sf=gmem_fc1_output_sf, + tma_atom_fc1_recompute=tma_atom_fc1_recompute, + gmem_fc1_recompute=gmem_fc1_recompute, + gmem_fc1_recompute_sf=gmem_fc1_recompute_sf, + tma_atom_fc1_col_output=tma_atom_fc1_col_output, + gmem_fc1_col_output=gmem_fc1_col_output, + gmem_fc1_col_output_sf=gmem_fc1_col_output_sf, + c_pipeline=c_pipeline, + smem_preact_buffer=smem_preact_buffer, + c_consumer_state=c_consumer_state, + smem_d_buffer=smem_d_buffer, + tma_atom_grad_y1=tma_atom_grad_y1, + warp_idx=warp_idx, + tidx=tidx, + norm_const=norm_const, + gmem_topk_scores=gmem_topk_scores, + gmem_beta=gmem_beta, + gmem_dprob=gmem_dprob, + d_pipeline=d_pipeline, + d_num_stage=d_num_stage, + token_comm_args=token_comm_args, + ) + else: + self._run_dfc1_task_tile( + work_tile_info=work_tile_info, + tmem_acc_tensor=tmem_acc_stage_tesnor, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + sched_ext=sched_ext, + gmem_fc2_output=gmem_fc2_output, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + token_comm_args=token_comm_args, + ) + + acc_consumer_state.advance() + + cur_was_linear1 = work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + cur_fc1_counter_slot = ( + work_tile_info.cumulative_token_block_count + + work_tile_info.tile_m_idx // cutlass.Int32(self._atom_thr_size) + ) + cur_fc2_expert_idx = work_tile_info.expert_idx + + work_tile_info = sched_consumer.consume_work() + + if cur_was_linear1: + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.fence_proxy("async") + cute.arch.fence_acq_rel_gpu() + + task_tile_boundary_bar.arrive_and_wait() + + if cur_was_linear1: + flag_tracker = flag_tracker.accumulate( + work_tile_info.phase, + self._epi_fc1_batch, + (gmem_fc1_done_counter.iterator + cur_fc1_counter_slot).toint(), + ) + else: + if cutlass.const_expr( + self._token_back_by_dispatch or self._combine_mxfp8 + ): + # Fence before (deferred) counter release: make the fc2 + # pool-output STG writes device-visible. + cute.arch.fence_acq_rel_gpu() + fc2_flag_addr = ( + token_comm_args.fc2_done_counter.iterator + cur_fc2_expert_idx + ).toint() + else: + fc2_flag_addr = Int64(0) + no_fire: cutlass.Constexpr = not ( + self._token_back_by_dispatch or self._combine_mxfp8 + ) + flag_tracker = flag_tracker.accumulate( + work_tile_info.phase, + self._epi_fc2_batch, + fc2_flag_addr, + no_fire, + ) + + flag_tracker.fire() + + d_pipeline.producer_tail() diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py new file mode 100644 index 000000000..c6746baef --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py @@ -0,0 +1,139 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Sched extension for the fused fc1+fc2 dGLU-backward MXFP8 kernel.""" + +import dataclasses +from typing import Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import Pointer +from cutlass.cutlass_dsl import Int32, Int64, extract_mlir_values, new_from_mlir_values + +from ..fwd_glu.glu_mxfp8_fc12_extension import GluMxFp8Fc12SchedExtension +from .....schedulers.fc12_mapping import NonSwapAbFc12WorkTileInfo + + +@dataclasses.dataclass(frozen=True) +class DgluMxFp8Fc12SchedExtension(GluMxFp8Fc12SchedExtension): + """dGLU adapter for token-major auxiliary data and per-expert blocked SF.""" + + expert_token_sizes: Optional[cute.Tensor] = None + token_padding_block: int = 128 + sf_padding_block: int = 128 + + def __post_init__(self) -> None: + super().__post_init__() + if self.expert_token_sizes is None: + raise ValueError("dGLU auxiliaries require expert_token_sizes.") + if self.token_padding_block != self.sf_padding_block or self.token_padding_block % 128 != 0: + raise ValueError("dGLU auxiliaries require equal token/SF padding divisible by 128.") + + def __extract_mlir_values__(self) -> list: + values = super().__extract_mlir_values__() + values.extend(extract_mlir_values(self.expert_token_sizes)) + return values + + def __new_from_mlir_values__(self, values: list) -> "DgluMxFp8Fc12SchedExtension": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + fc1_done_counter_pointer = rebuild(self.fc1_done_counter_pointer) + fc2_spin_threshold = rebuild(self.fc2_spin_threshold) + fc1_ready_counter_pointer = ( + rebuild(self.fc1_ready_counter_pointer) if self.fc1_ready_counter_pointer is not None else None + ) + expert_token_sizes = rebuild(self.expert_token_sizes) + if value_index != len(values): + raise ValueError( + f"DgluMxFp8Fc12SchedExtension MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return type(self)( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter_pointer, + fc2_spin_threshold=fc2_spin_threshold, + fc1_ready_counter_pointer=fc1_ready_counter_pointer, + cluster_m=self.cluster_m, + expert_token_sizes=expert_token_sizes, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + ) + + @cute.jit + def _physical_token_count(self, work_tile_info: NonSwapAbFc12WorkTileInfo, padding_block: int): + expert_idx = work_tile_info.expert_idx + valid_tokens = Int32(self.expert_token_sizes[expert_idx]) + return ((valid_tokens + Int32(padding_block - 1)) // Int32(padding_block)) * Int32(padding_block) + + @cute.jit + def _aux_data_tensor( + self, + tensor: cute.Tensor, + work_tile_info: NonSwapAbFc12WorkTileInfo, + feature_extent, + ) -> cute.Tensor: + physical_tokens = self._physical_token_count(work_tile_info, self.token_padding_block) + real = cute.domain_offset( + (work_tile_info.cumulative_data_physical_row, 0, 0), + tensor, + ) + return cute.make_tensor( + real.iterator, + cute.make_layout( + (physical_tokens, feature_extent, Int32(1)), + stride=real.stride, + ), + ) + + @cute.jit + def _aux_sf_tensor( + self, + tensor: cute.Tensor, + work_tile_info: NonSwapAbFc12WorkTileInfo, + feature_padded, + ) -> cute.Tensor: + physical_tokens = self._physical_token_count(work_tile_info, self.sf_padding_block) + feature_atoms = Int32(feature_padded) // Int32(128) + token_atoms = physical_tokens // Int32(128) + element_offset = Int64(feature_padded) * ( + Int64(work_tile_info.cumulative_sf_physical_row) // Int64(self.sf_vec_size) + ) + return cute.make_tensor( + tensor.iterator + element_offset, + cute.make_layout( + (feature_atoms, token_atoms, Int32(512)), + stride=(token_atoms * Int32(512), Int32(512), Int32(1)), + ), + ) + + @cute.jit + def get_gmem_tensor( + self, + tensor_name: str, + gmem_tensor_in_moe_view: cute.Tensor, + work_tile_info: NonSwapAbFc12WorkTileInfo, + ) -> Tuple[cute.Tensor, Optional[Pointer]]: + """dGLU-backward operand views; every other name delegates to the base.""" + shape = gmem_tensor_in_moe_view.shape + + if cutlass.const_expr(tensor_name == "recompute"): + return (self._aux_data_tensor(gmem_tensor_in_moe_view, work_tile_info, shape[1]), None) + + elif cutlass.const_expr(tensor_name == "sfrecompute"): + return (self._aux_sf_tensor(gmem_tensor_in_moe_view, work_tile_info, shape[0]), None) + + elif cutlass.const_expr(tensor_name == "col_output"): + return (self._aux_data_tensor(gmem_tensor_in_moe_view, work_tile_info, shape[1]), None) + + elif cutlass.const_expr(tensor_name == "sfcol_output"): + return (self._aux_sf_tensor(gmem_tensor_in_moe_view, work_tile_info, shape[0]), None) + + return GluMxFp8Fc12SchedExtension.get_gmem_tensor( + self, tensor_name, gmem_tensor_in_moe_view, work_tile_info + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py new file mode 100644 index 000000000..142639b71 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py @@ -0,0 +1,2377 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +""" +Fused fc1+fc2 GLU MXFP8 MegaMoE kernel for SM100. +""" + +import dataclasses +from typing import Any, Literal, Optional, Tuple, Type, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute + +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +import cutlass.utils.rubin_helpers as sm107_utils +from cutlass.cute.nvgpu.tcgen05 import CollectorOp + +from ..tmem_transpose import _TmemTranspose16x32Core +from .dglu_mxfp8_fc12_epilogue import DgluMxfp8Epilogue +from .....schedulers import BlockPhase +from .....schedulers.base import WorkIdAcquisitionMode +from .....schedulers.fc12_scheduler import BlackwellFusedFc12Scheduler +from .dglu_mxfp8_fc12_extension import DgluMxFp8Fc12SchedExtension +from ......api import ImplDesc, KernelClass, ProblemDesc, StaticOrRuntimeIntegerType +from ..helpers.constants import ( + SupportedMmaTileM, + SupportedMmaTileN, +) +from ......helpers.iket_compat import iket +from ......helpers.device_workspace import DeviceWorkspace +from ......helpers.smem_workspace import SmemWorkspace +from ......helpers.dsl_helpers import spin_wait +from ......quant_def import CombineFormat, QuantKind +from ......communication.nvlink_domain.token_comm import TokenCommArgs + + +@dataclasses.dataclass(frozen=True) +class _EpilogueCommView: + """Fields the dGLU epilogue reads for cross-rank dfc1/dprob routing.""" + token_src_metadata: Any + combine_output: Any + dprob_output: Any + peer_rank_ptr_mapper: Any + fc2_output_sf: Any = None + fc2_done_counter: Any = None + fc2_output_workspace: Any = None + + +# ============================================================================= +# Sm107Mxfp8DgluDfc21Kernel +# ============================================================================= + +class Sm107Mxfp8DgluDfc21Kernel: + + # SMEM budget for buffers like mbarriers, sched, work-tile buffer, TMEM allocator state + _SmemMiscBudget = 1024 + + # Supported (ab_dtype, sf_vec_size) pairings. + # MXFP8 → Float8E4M3FN / Float8E5M2 + sf_vec_size=32 (FP8-E8M0 scales, MmaMXF8Op) + VALID_AB_DTYPE_SF_SIZE: dict = { + 32: (cutlass.Float8E4M3FN, cutlass.Float8E5M2,), + } + + # Interleave granularity for gate and up in SwiGLU / GeGlu + GateUpInterleave: int = 32 + + def __init__( + self, + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mnk: Tuple[int, int, int], + use_2cta_instrs: bool, + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + load_balance_mode: Literal["static", "atomic_counter"] = "static", + static_expert_shape: Optional[Tuple[int, int, int]] = None, + force_static_sched: bool = True, + clc_bundle_size: Optional[int] = None, + num_sched_stages: Optional[int] = None, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + sf_vec_size: int = 32, + ab_dtype: Type[cutlass.Numeric] = cutlass.Float4E2M1FN, + epi_flag_batch: Optional[Tuple[int, int]] = (1, 1), + dfc2_recompute: bool = False, + dfc2_col_output: bool = False, + fc2_in_kernel_topk_reduce: bool = False, + act_func: str = "swiglu", + gate_up_clamp: Optional[float] = None, + ) -> None: + if not force_static_sched: + raise NotImplementedError( + "v1 only implements force_static_sched=True (lean 7-warp). " + "Dynamic CLC (force_static_sched=False) is future work." + ) + + # Validate (ab_dtype, sf_vec_size) pairing. + if sf_vec_size in self.VALID_AB_DTYPE_SF_SIZE: + valid_ab = self.VALID_AB_DTYPE_SF_SIZE[sf_vec_size] + if ab_dtype not in valid_ab: + raise ValueError( + f"ab_dtype={ab_dtype.__name__} is not valid for " + f"sf_vec_size={sf_vec_size}. " + f"Expected one of: {[t.__name__ for t in valid_ab]}." + ) + else: + valid_sf_vec_sizes = tuple(self.VALID_AB_DTYPE_SF_SIZE) + raise NotImplementedError( + f"sf_vec_size must be one of {valid_sf_vec_sizes} (MXFP8); got {sf_vec_size}." + ) + + + if load_balance_mode not in ("static", "atomic_counter"): + raise ValueError( + f"load_balance_mode must be 'static' or 'atomic_counter'; " + f"got {load_balance_mode!r}." + ) + if act_func not in ("swiglu", "geglu"): + raise ValueError( + f"act_func must be 'swiglu' or 'geglu'; got {act_func!r}." + ) + if act_func != "swiglu": + raise NotImplementedError( + f"act_func={act_func!r} is not yet implemented; only " + "'swiglu' is currently supported (geglu support is planned)." + ) + + # Store ab_dtype so workspace-size helpers can use it without tensors. + self.ab_dtype = ab_dtype + self.act_func = act_func + + self.acc_dtype = acc_dtype + self.mma_tiler_mnk = mma_tiler_mnk + self.cluster_shape_mn = (cluster_shape_mnk[0], cluster_shape_mnk[1]) + self.use_2cta_instrs = use_2cta_instrs + self.force_static_sched = force_static_sched + # static_expert_shape / clc_bundle_size / num_sched_stages + self.static_expert_shape = static_expert_shape + self.clc_bundle_size = clc_bundle_size + self.num_sched_stages = num_sched_stages + + # Fused fc12 sched-side knobs + self.group_hint = group_hint + self.token_padding_block = token_padding_block + self.sf_padding_block = sf_padding_block + self.load_balance_mode = load_balance_mode + + self.sf_vec_size = sf_vec_size + self.arch = "sm_107" + self.epi_flag_batch = epi_flag_batch + self.dfc2_recompute = dfc2_recompute + self.dfc2_col_output = dfc2_col_output + self.fc2_in_kernel_topk_reduce = fc2_in_kernel_topk_reduce + self.gate_up_clamp = abs(gate_up_clamp) if gate_up_clamp is not None else None + + self._validate_mma_tiler_and_cluster_shape() + self.mma_tiler = mma_tiler_mnk + + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + # Warp specialization (9-warp / 288 thread: + dedicated preact-C load warp) + self.occupancy = 1 + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_a_warp_id = 5 + self.tma_b_warp_id = 6 + self.sched_warp_id = 7 + # Dedicated TMA-load warp for the forward pre-activation (dswiglu C), + self.c_load_warp_id = 8 + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_a_warp_id, + self.tma_b_warp_id, + self.sched_warp_id, + self.c_load_warp_id, + *self.epilogue_warp_id, + ) + ) + + # NamedBarriers. + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + self.epi_subtile_bar_ids = (4, 5, 6, 7) + + self.smem_capacity = utils.get_smem_capacity_in_bytes() + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols( + self.arch + ) + + # Warp-specialized register split. + self.epi_reg_cnt = 256 + self.task_reg_cnt = 72 + + # Token-comm (MegaMoE) + self.enable_token_comm: bool = False + self.dispatch_warp_id: Optional[Tuple[int, int, int, int]] = None + self.token_back_by_dispatch: bool = False + self.token_back_standalone: bool = False + self.token_back_warp_id: Optional[Tuple[int, int, int, int]] = None + + def _validate_mma_tiler_and_cluster_shape(self) -> None: + """Validate user-provided geometry against v1 fused-fc12 constraints.""" + m, n, k = self.mma_tiler_mnk + cm, cn = self.cluster_shape_mn + + if m not in SupportedMmaTileM: + raise ValueError( + f"mma_tiler M ({m}) must be one of {SupportedMmaTileM}" + ) + + per_cta_m = m // (2 if self.use_2cta_instrs else 1) + if per_cta_m != 128: + raise ValueError( + f"per-CTA mma_tiler M must be 128, got {per_cta_m} " + f"(mma_tiler_m={m}, use_2cta_instrs={self.use_2cta_instrs})" + ) + + for _name, _blk in ( + ("token_padding_block", self.token_padding_block), + ("sf_padding_block", self.sf_padding_block), + ): + if _blk <= 0 or _blk % self.sf_vec_size != 0: + raise ValueError( + f"{_name} ({_blk}) must be a positive multiple of " + f"sf_vec_size ({self.sf_vec_size}); the col-quant epilogue " + f"turns a per-expert row offset into a col-SF row-block " + f"index by an exact '// sf_vec_size' division." + ) + + if n not in SupportedMmaTileN: + raise ValueError( + f"mma_tiler N ({n}) must be one of {SupportedMmaTileN} in fused fc12 " + f"(N=64 SFB hack is dropped; swap-AB sched handles short-N " + f"via subtile early-exit)." + ) + + sf_k_granularity = self.sf_vec_size * 4 + if k % sf_k_granularity != 0: + raise ValueError( + f"mma_tiler K ({k}) must be a multiple of " + f"sf_vec_size * 4 = {sf_k_granularity}" + ) + + if cm % (2 if self.use_2cta_instrs else 1) != 0: + raise ValueError( + f"cluster_shape M ({cm}) must be even when use_2cta_instrs=True" + ) + + is_pow2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if cm * cn > 16 or not is_pow2(cm) or not is_pow2(cn) or cm > 4 or cn > 4: + raise ValueError( + f"Invalid cluster_shape ({cm}, {cn}): each dim must be " + f"a power of 2 and <= 4, product must be <= 16" + ) + + # v1 swap-AB requires cluster_n == 1. + if cn != 1: + raise NotImplementedError( + f"v1 fused fc12 requires cluster_n == 1 (got {cn}). " + f"cluster_n > 1 needs sentinel-style acc/ab pipeline release." + ) + + def _create_tiled_mmas(self) -> Tuple[cute.TiledMma, cute.TiledMma]: + """Return (tiled_mma, tiled_mma_sfb).""" + common = ( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + ) + # Rubin: the SM107 blockscaled FP8 MMA op hard-requires instruction + tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma( + *common, self.cta_group, + (*self.mma_inst_shape_mn, 64), + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma( + *common, tcgen05.CtaGroup.ONE, + (*self.mma_inst_shape_mn_sfb, 64), + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + return tiled_mma, tiled_mma_sfb + + def _build_scheduler( + self, *, expert_cnt, intermediate_gateup, hidden_dim, launch_cluster_count + ) -> None: + """Construct FC12 scheduler and its SMEM/device workspaces.""" + work_id_mode = "grid_stride" if self.load_balance_mode == "static" else "atomic_counter" + num_scheduler_consumer_threads = 32 * (len(self.epilogue_warp_id) + 4) + if self.static_expert_shape is not None: + expert_cnt, intermediate_gateup, hidden_dim = self.static_expert_shape + problem_desc = ProblemDesc( + { + "expert_count": expert_cnt, + "intermediate_gateup_size": intermediate_gateup, + "hidden_size": hidden_dim, + } + ) + impl_desc = ImplDesc( + { + "num_scheduler_consumer_threads": num_scheduler_consumer_threads, + "mma_tiler_mnk": self.mma_tiler, + "cluster_shape_mn": self.cluster_shape_mn, + "use_2cta_instrs": self.use_2cta_instrs, + "hint": self.group_hint, + "token_padding_block": self.token_padding_block, + "sf_padding_block": self.sf_padding_block, + "work_id_mode": work_id_mode, + "is_swap_ab": False, + "launch_cluster_count": launch_cluster_count, + } + ) + self.scheduler = BlackwellFusedFc12Scheduler(problem_desc, impl_desc) + + sched_smem_ws = SmemWorkspace() + self.scheduler.register_smem_regions(sched_smem_ws) + sched_smem_ws.finalize(max_bytes=self.smem_capacity) + self.sched_smem_ws = sched_smem_ws + + sched_device_ws = DeviceWorkspace() + self.scheduler.register_device_workspace(sched_device_ws) + sched_device_ws.finalize() + self.sched_device_ws = sched_device_ws + + def _setup_attributes(self) -> None: + """Set up MMA / cluster / tile shapes, SMEM layouts, stage counts. + + The fc12 path shares ``mma_tiler_mnk`` and SMEM layouts across phases. + """ + self.mma_inst_shape_mn = (self.mma_tiler[0], self.mma_tiler[1]) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + tiled_mma, tiled_mma_sfb = self._create_tiled_mmas() + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + assert self.mma_tiler[2] % mma_inst_shape_k == 0, ( + f"mma_tiler K ({self.mma_tiler[2]}) must be a multiple of " + f"MMA instruction K ({mma_inst_shape_k})" + ) + + # SFB-specific tiler: rounded-up MN; same K as main tiler. + self.mma_tiler_sfb = ( + self.mma_inst_shape_mn_sfb[0], + self.mma_inst_shape_mn_sfb[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + # Multicast CTA counts + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + _epi_common = dict( + mma_tiler_mnk=self.mma_tiler, + cluster_shape_mn=self.cluster_shape_mn, + use_2cta_instrs=self.use_2cta_instrs, + sf_vec_size=self.sf_vec_size, + fc1_output_dtype=self.fc1_output_dtype, + fc1_output_layout=self.fc1_output_layout, + acc_dtype=self.acc_dtype, + epilog_sync_bar_id=self.epilog_sync_bar_id, + epilogue_warp_ids=self.epilogue_warp_id, + static_expert_shape=self.static_expert_shape, + epi_flag_batch=self.epi_flag_batch, + token_back_by_dispatch=self.token_back_by_dispatch, + dfc2_recompute=self.dfc2_recompute, + dfc2_col_output=self.dfc2_col_output, + fc2_in_kernel_topk_reduce=self.fc2_in_kernel_topk_reduce, + combine_format=getattr(self, "combine_format", None), + combine_hidden=getattr(self, "hidden", None), + act_func=self.act_func, + gate_up_clamp=self.gate_up_clamp, + ) + self.epilogue = DgluMxfp8Epilogue(**_epi_common) + + if self.num_sched_stages is None: + self.num_sched_stages = 2 + + # Reserve SMEM for the preact (dswiglu C) pipeline staging buffer + self.num_c_stage = 2 + assert self.num_c_stage % 2 == 0, f"num_c_stage must be even, got {self.num_c_stage}" + self.num_c_pipe_stage = self.num_c_stage // 2 + # One PipelineTmaStore stage contains every dFC2 data output tile. + self.num_d_stage = self.epilogue.d_output_slots + c_bytes_total = self.num_c_stage * self.epilogue.preact_bytes_per_stage + d_bytes_total = self.num_d_stage * self.epilogue.d_bytes_per_stage + self.c_bytes_total = c_bytes_total + self.d_bytes_total = d_bytes_total + + ( + self.num_acc_stage, + self.num_ab_stage, + self.num_sched_stages, + ) = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.sf_dtype, + self.sf_vec_size, + self.c_bytes_total, + self.d_bytes_total, + self.smem_capacity, + self.occupancy, + self.num_sched_stages, + self._smem_misc_budget_bytes() - self._SmemMiscBudget, + ) + + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + + # Read epilogue's accumulator and scale-factor sizing decisions. + self.num_acc_pipeline_stages = self.epilogue.num_acc_pipeline_stages + self.num_acc_stage = self.epilogue.num_acc_stage + self.num_sfa_tmem_cols = self.epilogue.num_sfa_tmem_cols + self.num_sfb_tmem_cols = self.epilogue.num_sfb_tmem_cols + self.num_accumulator_tmem_cols = self.epilogue.num_accumulator_tmem_cols + + # TMA load bytes per stage (A + B + SFA + SFB). + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + self.atom_thr_size = atom_thr_size # store as Python int for use in @cute.kernel + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = ( + a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + ) * atom_thr_size + + # SMEM usage report (all sizes are per-CTA) + _ab_per_stage = a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + _misc_total = self._smem_misc_budget_bytes() + _fixed = _misc_total + self.c_bytes_total + self.d_bytes_total + _total_used = _fixed + self.num_ab_stage * _ab_per_stage + _per_cta_budget = self.smem_capacity // self.occupancy + _free = _per_cta_budget - _total_used + _extra_misc = _misc_total - self._SmemMiscBudget + print( + f"[smem] capacity={self.smem_capacity}B ({self.smem_capacity//1024}KB)" + f" occupancy={self.occupancy}" + f" per-CTA budget={_per_cta_budget}B ({_per_cta_budget//1024}KB)\n" + f" AB stages: {self.num_ab_stage} × {_ab_per_stage}B ({_ab_per_stage/1024:.1f}KB)" + f" = {self.num_ab_stage * _ab_per_stage}B" + f" [A={a_copy_size}B B={b_copy_size}B" + f" SFA={sfa_copy_size}B SFB={sfb_copy_size}B]\n" + f" fixed: misc={_misc_total}B (base={self._SmemMiscBudget}B" + f" + subclass_extra={_extra_misc}B)" + f" preact(C)={self.num_c_stage}×{self.epilogue.preact_bytes_per_stage}B" + f" sD(D)={self.num_d_stage}×{self.epilogue.d_bytes_per_stage}B" + f" used={_total_used}B ({_total_used/1024:.1f}KB)" + f" free={_free}B ({_free/1024:.1f}KB)\n" + ) + + @staticmethod + def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_bytes_total: int, + d_bytes_total: int, + smem_capacity: int, + occupancy: int, + num_sched_stages: int, + extra_misc_bytes: int = 0, + ) -> Tuple[int, int, int]: + """Compute stage counts for ACC, AB+SF, and scheduler. + """ + num_acc_stage = 2 + + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1, + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 1, + ) + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, mma_tiler_mnk, sf_vec_size, 1, + ) + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, mma_tiler_mnk, sf_vec_size, 1, + ) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + + fixed_overhead = ( + Sm107Mxfp8DgluDfc21Kernel._SmemMiscBudget + extra_misc_bytes + c_bytes_total + d_bytes_total + ) + + num_ab_stage = ( + smem_capacity // occupancy - fixed_overhead + ) // ab_bytes_per_stage + return num_acc_stage, num_ab_stage, num_sched_stages + + def get_workspace_size_in_bytes( + self, + fc1_activation_tensor, + fc1_weight_tensor, + ) -> int: + """Compute opaque workspace size for one fused dfc2+dfc1 launch.""" + sf_padding_block = self.sf_padding_block + sf_vec_size = self.sf_vec_size + + mma_tiler_n = self.mma_tiler_mnk[1] + + data_total_rows, _hidden = fc1_activation_tensor.shape + experts, _hidden_w, dfc2_weight_n = fc1_weight_tensor.shape + # grad_y1 (doubled dswiglu output) width = intermediate = 2 * inter_half. + intermediate_out = dfc2_weight_n * 2 + + # Conservative upper bound for sf_total_rows. + sf_total_rows_upper = data_total_rows + experts * sf_padding_block + + # grad_y1 byte size (MXFP8, 8-bit: 1 element per byte). + fc1_output_bytes = ( + data_total_rows * intermediate_out * self.ab_dtype.width // 8 + ) + + # grad_y1 SF sf_vec_size matches the kernel's sf_vec_size. + fc1_out_sf_vec_size = self.sf_vec_size + sf_block_cols = ( + (intermediate_out // fc1_out_sf_vec_size) + 3 + ) // 4 * 4 + fc1_output_sf_bytes = sf_total_rows_upper * sf_block_cols + + # fc1_recompute (forward-swiglu recompute): N = inter_half = intermediate_out // 2. + fc1_recompute_bytes = ( + data_total_rows * dfc2_weight_n * self.ab_dtype.width // 8 + ) + fc1_recompute_row_blocks_upper = sf_total_rows_upper // fc1_out_sf_vec_size + fc1_recompute_sf_bytes = fc1_recompute_row_blocks_upper * dfc2_weight_n + + # fc1_col_output (col-quant grad_y1): N = intermediate_out (same as + # grad_y1's row-quant fc1_output). Col-SF: row_blocks × intermediate. + fc1_col_output_bytes = fc1_output_bytes + fc1_col_output_sf_bytes = fc1_recompute_row_blocks_upper * intermediate_out + + # fc1_done_counter: one Int32 per CTA-level token block (each cluster block + # has atom_thr_size CTAs, each with its own per-CTA counter slot). + counter_slots_upper = ( + (data_total_rows + mma_tiler_n - 1) // mma_tiler_n + + experts + ) + fc1_done_counter_bytes = counter_slots_upper * 4 + + # load_balance_counter: Int32 scalar. + if self.load_balance_mode == "atomic_counter": + load_balance_counter_bytes = 4 + else: + load_balance_counter_bytes = 0 + + total = ( + fc1_output_bytes + + fc1_output_sf_bytes + + fc1_recompute_bytes + + fc1_recompute_sf_bytes + + fc1_col_output_bytes + + fc1_col_output_sf_bytes + + fc1_done_counter_bytes + + load_balance_counter_bytes + ) + + # 128B align (TMA tensor base address alignment requirement). + alignment = 128 + total = ((total + alignment - 1) // alignment) * alignment + return total + + def mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """SMEM → TMEM tiled copy + partition for SFA / SFB.""" + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t, tCsSF_compact_s2t_ + ) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + # ========================================================================= + # Token-comm hook surface (MegaMoE-only; lean base = no-op stubs) + # + # Mirrors the hook interface in ``moe_mxfp8_glu.kernel_mxfp8_glu_fc12`` + # so that ``Sm107MegaMoEMxfp8DgluKernel`` can override exactly the same + # methods. The mega wrapper realigns dispatch onto warps 8-11 (128-aligned + # for next token_comm) and relocates ``c_load_warp_id`` above the transfer + # block (warp 12 or 16); the lean base keeps c_load at warp 8. + # ========================================================================= + + def _smem_misc_budget_bytes(self) -> int: + """SMEM reserved for non-problem-tensor buffers (mbarriers, sched, TMEM state). + + MegaMoE subclass adds dispatch-warp SMEM on top via:: + + return super()._smem_misc_budget_bytes() + self._dispatch_smem_bytes() + """ + return self._SmemMiscBudget + + def token_comm_extra_smem_storage_class(self) -> type: + """Return a ``@cute.struct`` for dispatch-warp SMEM, or None.""" + return None + + def token_comm_hook_fc1_ready_counter_ptr(self, token_comm_args): + """Return dispatch->fc1 release counter pointer, or None (lean: disabled).""" + return None + + def sched_ext_fc1_peek_threshold(self) -> int: + """Return the fc1 ready-counter peek threshold for DgluMxFp8Fc12SchedExtension.""" + return 0 + + def sched_ext_fc1_counter_cumul_scale(self) -> int: + """Return the scale factor for the fc1 ready-counter slot formula.""" + return 1 + + @cute.jit + def token_comm_hook_sched_warp_pre_init_wait(self, token_comm_args): + """Sched warp: wait for dispatch barrier before reading sizes. No-op base.""" + pass + + @cute.jit + def token_comm_hook_fc1_tma_b_predispatch_spin(self, token_comm_args, work_tile_info): + """TMA-A warp: spin until dispatch-pulled tokens are resident. No-op base.""" + pass + + @cute.jit + def token_comm_hook_dispatch_warp_body( + self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx, + ): + """Body for dispatch warps 8-11 (MegaMoE-only). No-op base.""" + pass + + @cute.jit + def token_comm_hook_token_back_warp_body( + self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx, + ): + """Body for standalone token-back warps 12-15 (MegaMoE-only). No-op base.""" + pass + + @cute.jit + def token_comm_hook_kernel_tail(self, token_comm_args, *, warp_idx, lane_idx, tidx): + """All-warp kernel tail (NVLink release, etc.). No-op base.""" + pass + + @cute.jit + def __call__( + self, + activation: cute.Tensor, # (token_sum_padded, hidden) grad_out + fc1_weight: cute.Tensor, # (experts, hidden, intermediate_downproj) + activation_sf: cute.Tensor, # row-SF for activation + fc1_weight_sf: cute.Tensor, # dfc2-weight SF + fc1_output: cute.Tensor, # (token_sum_padded, intermediate_gateup) + fc1_output_sf: cute.Tensor, # (token_sum_padded_sf, intermediate_gateup_sf) + fc1_recompute: Optional[cute.Tensor], # (token_sum_padded, intermediate_downproj) + fc1_recompute_sf: Optional[cute.Tensor], # (intermediate_downproj_padded, col_sf_rows) + fc1_col_output: Optional[cute.Tensor], # (token_sum_padded, intermediate_gateup) + fc1_col_output_sf: Optional[cute.Tensor], # (intermediate_gateup_padded, col_sf_rows) + fc2_weight: cute.Tensor, # (experts, intermediate_gateup, hidden) + fc2_weight_sf: cute.Tensor, # (experts, hidden_padded * intermediate_gateup_sf_padded) + fc2_output: cute.Tensor, # (token_sum_padded, hidden) + fc1_preact: cute.Tensor, # (token_sum_padded, intermediate_gateup) BFloat16 + topk_scores: cute.Tensor, # (token_sum_padded,) Float32 + beta: cute.Tensor, # (experts,) Float32 + dprob: cute.Tensor, # (token_sum_padded,) Float32 + fc1_done_counter: cute.Tensor, # (fc1_ready_slot_count,) Int32 + offs: Optional[cute.Tensor] = None, # unsupported for dGLU; use expert_token_sizes + max_active_clusters: cutlass.Constexpr = None, + stream: cuda.CUstream = None, + norm_const_tensor: Optional[cute.Tensor] = None, + global_activation_sf: Optional[cute.Tensor] = None, + global_fc1_weight_sf: Optional[cute.Tensor] = None, + load_balance_counter: Optional[cute.Tensor] = None, + expert_token_sizes: Optional[cute.Tensor] = None, # (experts,) valid token counts + token_comm_args=None, + overflow_flag: cute.Tensor = None, + mega_peer_rank_ptr_mapper=None, + mega_local_rank: Optional[cutlass.Int32] = None, + mega_local_workspace: Optional[cute.Pointer] = None, + mega_shared_workspace: Optional[cute.Pointer] = None, + mega_activation: Optional[cute.Tensor] = None, + mega_activation_sf: Optional[cute.Tensor] = None, + mega_pre_reduced_activation: Optional[cute.Tensor] = None, + mega_pre_reduced_activation_sf: Optional[cute.Tensor] = None, + ) -> None: + """Launch the fused dfc2+dfc1 dGLU MXFP8 (backward) kernel.""" + if cutlass.const_expr(self.static_expert_shape is not None): + ( + experts_static, + intermediate_gateup_static, # inter_half = dfc2 weight N + hidden_static, + ) = self.static_expert_shape + intermediate_out_static = intermediate_gateup_static * 2 # grad_y1 / dfc1-K + + fc1_weight = cute.make_tensor( + fc1_weight.iterator, + cute.make_layout( + (experts_static, hidden_static, intermediate_gateup_static), + stride=fc1_weight.stride, + ), + ) + fc2_weight = cute.make_tensor( + fc2_weight.iterator, + cute.make_layout( + (experts_static, intermediate_out_static, hidden_static), + stride=fc2_weight.stride, + ), + ) + activation = cute.make_tensor( + activation.iterator, + cute.make_layout( + (activation.shape[0], hidden_static), + stride=activation.stride, + ), + ) + fc1_output = cute.make_tensor( + fc1_output.iterator, + cute.make_layout( + (fc1_output.shape[0], intermediate_out_static), + stride=fc1_output.stride, + ), + ) + fc1_recompute = cute.make_tensor( + fc1_recompute.iterator, + cute.make_layout( + (fc1_recompute.shape[0], intermediate_gateup_static), + stride=fc1_recompute.stride, + ), + ) + fc1_col_output = cute.make_tensor( + fc1_col_output.iterator, + cute.make_layout( + (fc1_col_output.shape[0], intermediate_out_static), + stride=fc1_col_output.stride, + ), + ) + fc1_preact = cute.make_tensor( + fc1_preact.iterator, + cute.make_layout( + (fc1_preact.shape[0], intermediate_out_static), + stride=fc1_preact.stride, + ), + ) + if cutlass.const_expr(len(fc2_output.shape) == 3): + fc2_output = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (fc2_output.shape[0], fc2_output.shape[1], hidden_static), + stride=fc2_output.stride, + ), + ) + else: + fc2_output = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (fc2_output.shape[0], hidden_static), + stride=fc2_output.stride, + ), + ) + + # ── GEMM-domain transform for fc1 phase ── + c1 = cutlass.Int32(1) + c0 = cutlass.Int32(0) + + # A_gemm (fc1 activations): (tokens_sum, hidden) -> (M=tokens, K=hidden, L=1). + tokens_sum, hidden = activation.shape + activation_gemm = cute.make_tensor( + activation.iterator, + cute.make_layout( + (tokens_sum, hidden, 1), + stride=(activation.stride[0], activation.stride[1], 0), + ), + ) + + # B_gemm (fc1 weights): (experts, hidden, intermediate_gateup) with hidden stride-1 (K-major) + # -> (N=intermediate_gateup, K=hidden, L=experts). + experts, hidden_b, intermediate_gateup = fc1_weight.shape + fc1_weight_gemm = cute.make_tensor( + fc1_weight.iterator, + cute.make_layout( + (intermediate_gateup, hidden_b, experts), + stride=(fc1_weight.stride[2], fc1_weight.stride[1], fc1_weight.stride[0]), + ), + ) + + intermediate_downproj = fc1_output.shape[1] + fc1_output_gemm = cute.make_tensor( + fc1_output.iterator, + cute.make_layout( + (tokens_sum, intermediate_downproj, 1), + stride=(fc1_output.stride[0], fc1_output.stride[1], 0), + ), + ) + + # dFC2 auxiliary data planes use the public token-major ABI. + fc1_recompute_gemm = cute.make_tensor( + fc1_recompute.iterator, + cute.make_layout( + (fc1_recompute.shape[0], fc1_recompute.shape[1], 1), + stride=(fc1_recompute.stride[0], fc1_recompute.stride[1], 0), + ), + ) + + fc1_col_output_gemm = cute.make_tensor( + fc1_col_output.iterator, + cute.make_layout( + (fc1_col_output.shape[0], fc1_col_output.shape[1], 1), + stride=(fc1_col_output.stride[0], fc1_col_output.stride[1], 0), + ), + ) + + # Forward pre-activation (gate||up) + fc1_preact_gemm = cute.make_tensor( + fc1_preact.iterator, + cute.make_layout( + (tokens_sum, intermediate_downproj, 1), + stride=(fc1_preact.stride[0], fc1_preact.stride[1], 0), + ), + ) + + # SFA / SFB scale tensors (atom-tiled) + tokens_sum_padded = activation_sf.shape[0] + hidden_padded = activation_sf.shape[1] * self.sf_vec_size + activation_sf_gemm = cute.make_tensor( + activation_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (tokens_sum_padded, hidden_padded, 1), self.sf_vec_size + ), + ) + intermediate_gateup_padded_mul_hidden_padded = fc1_weight_sf.shape[1] + intermediate_gateup_padded = ( + intermediate_gateup_padded_mul_hidden_padded * self.sf_vec_size + ) // hidden_padded + fc1_weight_sf_gemm = cute.make_tensor( + fc1_weight_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (intermediate_gateup_padded, hidden_padded, experts), + self.sf_vec_size, + ), + ) + + # GEMM-domain transform for fc2 phase ── + experts2, intermediate_downproj_b2, hidden_b2 = fc2_weight.shape + fc2_weight_gemm = cute.make_tensor( + fc2_weight.iterator, + cute.make_layout( + (hidden_b2, intermediate_downproj_b2, experts2), + stride=(fc2_weight.stride[2], fc2_weight.stride[1], fc2_weight.stride[0]), + ), + ) + + if cutlass.const_expr(len(fc2_output.shape) == 3): + fc2_hidden_out = fc2_output.shape[2] + fc2_output_gemm = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (fc2_output.shape[0], fc2_hidden_out, c1), + stride=(fc2_output.stride[0], fc2_output.stride[2], c0), + ), + ) + else: + fc2_hidden_out = fc2_output.shape[1] + fc2_output_gemm = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (tokens_sum, fc2_hidden_out, c1), + stride=(fc2_output.stride[0], fc2_output.stride[1], c0), + ), + ) + + fc1_out_sf_vec_size = self.sf_vec_size + tokens_sum_padded_sf = fc1_output_sf.shape[0] + intermediate_downproj_padded = fc1_output_sf.shape[1] * fc1_out_sf_vec_size + fc1_output_sf_gemm_for_fc2_load = cute.make_tensor( + fc1_output_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (tokens_sum_padded_sf, intermediate_downproj_padded, 1), + fc1_out_sf_vec_size, + ), + ) + + hidden_padded_fc2_mul_intermediate_downproj_padded = fc2_weight_sf.shape[1] + hidden_padded_fc2 = ( + hidden_padded_fc2_mul_intermediate_downproj_padded * self.sf_vec_size + ) // intermediate_downproj_padded + fc2_weight_sf_gemm = cute.make_tensor( + fc2_weight_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (hidden_padded_fc2, intermediate_downproj_padded, experts2), + self.sf_vec_size, + ), + ) + + expert_cnt = experts + hidden_dim = hidden + + # Infer dtypes and major modes + self.a_dtype: Type[cutlass.Numeric] = activation_gemm.element_type + self.b_dtype: Type[cutlass.Numeric] = fc1_weight_gemm.element_type + self.fc1_output_dtype: Type[cutlass.Numeric] = fc1_output_gemm.element_type + self.sf_dtype: Type[cutlass.Numeric] = activation_sf_gemm.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(activation_gemm).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(fc1_weight_gemm).mma_major_mode() + self.fc1_output_layout = utils.LayoutEnum.from_tensor(fc1_output_gemm) + + self._setup_attributes() + tiled_mma, tiled_mma_sfb = self._create_tiled_mmas() + + # fc1 TMA atoms load A1 + a_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_fc1_activation, tma_tensor_fc1_activation = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + activation_gemm, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load B1 + b_op = sm100_utils.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_fc1_weight, tma_tensor_fc1_weight = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + fc1_weight_gemm, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load SFA1 + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_fc1_activation_sf, tma_tensor_fc1_activation_sf = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + activation_sf_gemm, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64, + ) + + # TMA load SFB1 + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_fc1_weight_sf, tma_tensor_fc1_weight_sf = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + fc1_weight_sf_gemm, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Uint64, + ) + + # Coalesced TMA G2S load of the forward preact (dswiglu C input). + preact_tma_op = cpasync.CopyBulkTensorTileG2SOp() + tma_atom_fc1_preact, tma_tensor_fc1_preact = cpasync.make_tiled_tma_atom( + preact_tma_op, + fc1_preact_gemm, + self.epilogue.preact_smem_layout_one_stage, + self.epilogue.preact_epi_tile, + ) + + # Coalesced TMA S2G store of grad_y1 (dfc2 fp8 output). + grad_y1_tma_op = cpasync.CopyBulkTensorTileS2GOp() + tma_atom_grad_y1, tma_tensor_grad_y1 = cpasync.make_tiled_tma_atom( + grad_y1_tma_op, + fc1_output_gemm, + self.epilogue.d_smem_layout_one_stage, + self.epilogue.d_epi_tile, + ) + tma_atom_fc1_recompute, tma_tensor_fc1_recompute = cpasync.make_tiled_tma_atom( + grad_y1_tma_op, + fc1_recompute_gemm, + self.epilogue.d_smem_layout_one_stage, + self.epilogue.d_epi_tile, + ) + tma_atom_fc1_col_output, tma_tensor_fc1_col_output = cpasync.make_tiled_tma_atom( + grad_y1_tma_op, + fc1_col_output_gemm, + self.epilogue.d_smem_layout_one_stage, + self.epilogue.d_epi_tile, + ) + + # fc1 SFC GMEM tensor (= fc1_output_sf user view). No TMA atom; it is + # per-thread STG. + fc1_output_sf_gemm = cute.make_tensor( + fc1_output_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (tokens_sum_padded, intermediate_downproj, 1), + self.sf_vec_size, + ), + ) + + # Token-major blocked SF carriers. + fc1_recompute_sf_gemm = cute.make_tensor( + fc1_recompute_sf.iterator, + cute.make_layout( + (fc1_recompute_sf.shape[0], fc1_recompute_sf.shape[1], 1), + stride=(fc1_recompute_sf.stride[0], fc1_recompute_sf.stride[1], 0), + ), + ) + + fc1_col_output_sf_gemm = cute.make_tensor( + fc1_col_output_sf.iterator, + cute.make_layout( + (fc1_col_output_sf.shape[0], fc1_col_output_sf.shape[1], 1), + stride=(fc1_col_output_sf.stride[0], fc1_col_output_sf.stride[1], 0), + ), + ) + + # ── fc2 TMA atoms: fc1_output → A-side (M=tokens), fc2_weight → B-side (N=hidden) ── + tma_atom_fc2_activation, tma_tensor_fc2_activation = ( + cute.nvgpu.make_tiled_tma_atom_A( + a_op, + fc1_output_gemm, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + ) + tma_atom_fc2_weight, tma_tensor_fc2_weight = ( + cute.nvgpu.make_tiled_tma_atom_B( + b_op, + fc2_weight_gemm, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + ) + tma_atom_fc2_activation_sf, tma_tensor_fc2_activation_sf = ( + cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + fc1_output_sf_gemm_for_fc2_load, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64, + ) + ) + tma_atom_fc2_weight_sf, tma_tensor_fc2_weight_sf = ( + cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + fc2_weight_sf_gemm, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Uint64, + ) + ) + + # ── Scheduler params + grid + launch ── + if cutlass.const_expr(self.load_balance_mode == "atomic_counter"): + if cutlass.const_expr(load_balance_counter is None): + raise ValueError( + "load_balance_counter must be provided when " + "load_balance_mode == 'atomic_counter'" + ) + load_balance_counter_ptr = load_balance_counter.iterator + else: + load_balance_counter_ptr = None + + # dGLU auxiliary layouts require per-expert token counts; cumulative offsets + # alone are insufficient. + if cutlass.const_expr(not self.enable_token_comm): + if cutlass.const_expr(offs is not None): + raise ValueError( + "`offs` is not supported by dGLU; provide `expert_token_sizes`." + ) + if cutlass.const_expr(expert_token_sizes is None): + raise ValueError( + "`expert_token_sizes` must be provided for dGLU auxiliary layouts." + ) + + self._build_scheduler( + expert_cnt=expert_cnt, + intermediate_gateup=intermediate_gateup, + hidden_dim=hidden_dim, + launch_cluster_count=max_active_clusters, + ) + grid = self.scheduler.get_grid_shape(max_active_clusters=max_active_clusters) + + self.kernel( + tiled_mma, + tiled_mma_sfb, + # fc1 TMA atoms / tensors (A=activations, B=weights) + tma_atom_fc1_activation, + tma_tensor_fc1_activation, + tma_atom_fc1_weight, + tma_tensor_fc1_weight, + tma_atom_fc1_activation_sf, + tma_tensor_fc1_activation_sf, + tma_atom_fc1_weight_sf, + tma_tensor_fc1_weight_sf, + # fc2 TMA atoms / tensors (fc1_output→A, fc2_weight→B) + tma_atom_fc2_activation, + tma_tensor_fc2_activation, + tma_atom_fc2_weight, + tma_tensor_fc2_weight, + tma_atom_fc2_activation_sf, + tma_tensor_fc2_activation_sf, + tma_atom_fc2_weight_sf, + tma_tensor_fc2_weight_sf, + # GEMM-domain tensors (fc1) + activation_gemm, + fc1_weight_gemm, + fc1_output_gemm, + activation_sf_gemm, + fc1_weight_sf_gemm, + fc1_output_sf_gemm, + # GEMM-domain tensors (fc2) + fc2_weight_gemm, + fc2_output_gemm, + fc2_weight_sf_gemm, + fc1_output_sf_gemm_for_fc2_load, + # forward pre-activation (dswiglu input) — TMA G2S into SMEM + tma_atom_fc1_preact, + tma_tensor_fc1_preact, + tma_atom_grad_y1, + tma_tensor_grad_y1, + # token-major auxiliary data — TMA S2G stores + tma_atom_fc1_recompute, + tma_tensor_fc1_recompute, + fc1_recompute_sf_gemm, + tma_atom_fc1_col_output, + tma_tensor_fc1_col_output, + fc1_col_output_sf_gemm, + # topk / beta / dprob + cross-phase sync workspace + topk_scores, + beta, + dprob, + overflow_flag, + fc1_done_counter, + # Scheduling + offs, + expert_token_sizes, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + # SMEM layouts + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + token_comm_args, + # MegaMoE push-model token-comm inputs (None on the lean path) + mega_peer_rank_ptr_mapper, + mega_local_rank, + mega_local_workspace, + mega_shared_workspace, + mega_activation, + mega_activation_sf, + mega_pre_reduced_activation, + mega_pre_reduced_activation_sf, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=self.occupancy, + ) + + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + # fc1 TMA atoms / tensors + tma_atom_fc1_activation_1: cute.CopyAtom, + tma_tensor_fc1_activation_1: cute.Tensor, + tma_atom_weight: cute.CopyAtom, + tma_tensor_weight: cute.Tensor, + tma_atom_fc1_activation_1_sf: cute.CopyAtom, + tma_tensor_fc1_activation_1_sf: cute.Tensor, + tma_atom_fc1_weight_sf: cute.CopyAtom, + tma_tensor_fc1_weight_sf: cute.Tensor, + # fc2 TMA atoms / tensors (fc1_output→A, fc2_weight→B) + tma_atom_fc2_activation: cute.CopyAtom, + tma_tensor_fc2_activation: cute.Tensor, + tma_atom_fc2_weight: cute.CopyAtom, + tma_tensor_fc2_weight: cute.Tensor, + tma_atom_fc2_activation_sf: cute.CopyAtom, + tma_tensor_fc2_activation_sf: cute.Tensor, + tma_atom_fc2_weight_sf: cute.CopyAtom, + tma_tensor_fc2_weight_sf: cute.Tensor, + # GEMM-domain tensors (fc1) + activation_gemm: cute.Tensor, + fc1_weight_gemm: cute.Tensor, + fc1_output_gemm: cute.Tensor, + activation_sf_gemm: cute.Tensor, + fc1_weight_sf_gemm: cute.Tensor, + fc1_output_sf_gemm: cute.Tensor, + # GEMM-domain tensors (fc2) + fc2_weight_gemm: cute.Tensor, + fc2_output_gemm: cute.Tensor, + fc2_weight_sf_gemm: cute.Tensor, + fc1_output_sf_gemm_for_fc2_load: cute.Tensor, + # forward pre-activation (dswiglu input) — TMA G2S into SMEM + tma_atom_fc1_preact: cute.CopyAtom, + tma_tensor_fc1_preact: cute.Tensor, + # grad_y1 (dfc2 output) — TMA S2G store + tma_atom_grad_y1: cute.CopyAtom, + tma_tensor_grad_y1: cute.Tensor, + # token-major auxiliary data — TMA S2G stores + tma_atom_fc1_recompute: cute.CopyAtom, + tma_tensor_fc1_recompute: cute.Tensor, + fc1_recompute_sf_gemm: cute.Tensor, + tma_atom_fc1_col_output: cute.CopyAtom, + tma_tensor_fc1_col_output: cute.Tensor, + fc1_col_output_sf_gemm: cute.Tensor, + # topk / beta / dprob + cross-phase sync workspace + topk_scores: cute.Tensor, + beta: cute.Tensor, + dprob: cute.Tensor, + overflow_flag: cute.Tensor, + fc1_done_counter: cute.Tensor, + # Scheduling + offs: Optional[cute.Tensor], + expert_token_sizes: Optional[cute.Tensor], + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + # SMEM layouts + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + token_comm_args=None, + # MegaMoE push-model token-comm inputs + mega_peer_rank_ptr_mapper=None, + mega_local_rank: Optional[cutlass.Int32] = None, + mega_local_workspace: Optional[cute.Pointer] = None, + mega_shared_workspace: Optional[cute.Pointer] = None, + mega_activation: Optional[cute.Tensor] = None, + mega_activation_sf: Optional[cute.Tensor] = None, + mega_pre_reduced_activation: Optional[cute.Tensor] = None, + mega_pre_reduced_activation_sf: Optional[cute.Tensor] = None, + ): + """Device kernel for fused fc1+fc2 swap-AB GLU MXFP8 grouped GEMM.""" + a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, None, 0)) + b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, None, 0)) + sfa_smem_layout = cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)) + sfb_smem_layout = cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)) + + # MegaMoE (push model): bind the device workspace here so the fc1_ready counter + # pointer that the scheduler extension spins on (built just below) resolves. + if cutlass.const_expr(self.enable_token_comm): + self._mega_device_workspace.assign_device_members( + mega_local_workspace, mega_shared_workspace + ) + + # fc2 waits for all fc1 intermediate N-tiles in the same token block. + ext_fc2_spin_threshold = ( + fc1_weight_gemm.shape[0] + self.cta_tile_shape_mnk[1] - 1 + ) // self.cta_tile_shape_mnk[1] * self.epilogue._atom_thr_size + + if cutlass.const_expr(self.enable_token_comm): + _aux_expert_sizes = self.token_comm.local_expert_sizes( + self._mega_device_workspace, mega_local_rank + ) + else: + _aux_expert_sizes = expert_token_sizes + + ext = DgluMxFp8Fc12SchedExtension( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter.iterator, + fc2_spin_threshold=ext_fc2_spin_threshold, + # MegaMoE: peek the dispatch->fc1 ready counter (None on the lean + # path). Parity with the forward kernel's SchedExtension wiring. + fc1_ready_counter_pointer=self.token_comm_hook_fc1_ready_counter_ptr( + token_comm_args + ), + # Fold the 2 CTAs of a cluster onto one fc1_ready slot + cluster_m=self.epilogue._atom_thr_size, + expert_token_sizes=_aux_expert_sizes, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + ) + + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + tidx, _, _ = cute.arch.thread_idx() + + # MegaMoE (push model): bind token-comm device members (transfer-warp state + the + # NVLink barrier's peer mapper) before any token_in / token_back / size-wait runs. + if cutlass.const_expr(self.enable_token_comm): + _mega_token_comm_args = TokenCommArgs( + mega_activation, + mega_activation_sf, + mega_pre_reduced_activation, + mega_pre_reduced_activation_sf, + mega_peer_rank_ptr_mapper, + ) + _mega_cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] + _, _, _mega_cluster_idx = cute.arch.block_idx() + _mega_linear_cta_idx = cta_rank_in_cluster + _mega_cluster_idx * _mega_cluster_size + self.token_comm.assign_device_members( + device_workspace=self._mega_device_workspace, + token_comm_args=_mega_token_comm_args, + local_rank=mega_local_rank, + linear_cta_idx=_mega_linear_cta_idx, + ) + + # preact (dswiglu C) pipeline + num_c_stage = self.num_c_stage + num_c_pipe_stage = self.num_c_pipe_stage + num_d_stage = self.num_d_stage + + # SharedStorage (mainloop + epilogue SMEM). next's scheduler owns its own + # SMEM workspace, allocated separately below. + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_pipeline_stages * 2 + ] + c_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_c_pipe_stage * 2] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + sPre: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(self.epilogue.preact_staged_smem_layout(num_c_stage).outer), + ], + 1024, + ] + # Unified dFC2 data-output staging; slot count is compile-time gated. + sD: cute.struct.Align[ + cute.struct.MemRange[ + self.fc1_output_dtype, + cute.cosize(self.epilogue.d_staged_smem_layout(num_d_stage).outer), + ], + 1024, + ] + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # next scheduler SMEM: a self-contained workspace carved from the same + # allocator; its transport regions resolve against ``sched_smem_base``. + sched_storage = smem.allocate(self.sched_smem_ws.storage_class()) + sched_smem_base = sched_storage.buffer.data_ptr() + + # MegaMoE-only dispatch-warp SMEM (pull_buffer, mbarriers, etc.). + # Kept out of ``SharedStorage`` so the lean path never allocates it. + TokenCommStorageCls = self.token_comm_extra_smem_storage_class() + if cutlass.const_expr(TokenCommStorageCls is not None): + token_comm_storage = smem.allocate(TokenCommStorageCls) + else: + token_comm_storage = None + + # ── Pipelines: two TMA producer warps share the AB pipeline. ── + + ab_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, 2 + ) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes // 2, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = ( + len(self.epilogue_warp_id) * 32 * (2 if use_2cta_instrs else 1) + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_pipeline_stages, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # preact (dswiglu C) pipeline + c_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + c_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len(self.epilogue_warp_id) + ) + c_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.c_full_mbar_ptr.data_ptr(), + num_stages=num_c_pipe_stage, + producer_group=c_pipeline_producer_group, + consumer_group=c_pipeline_consumer_group, + tx_count=2 * self.epilogue.preact_bytes_per_stage, + defer_sync=True, + ) + # d pipeline + d_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + d_pipeline = pipeline.PipelineTmaStore.create( + num_stages=num_d_stage // self.epilogue.d_output_slots, + producer_group=d_producer_group, + ) + + + # TMEM allocator + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, + arch=self.arch, + ) + + # Sched + scheduler = self.scheduler + if cutlass.const_expr(self.enable_token_comm): + _sched_expert_sizes = self.token_comm.local_expert_sizes( + self._mega_device_workspace, mega_local_rank + ) + _sched_prefix_sum = None + # Bind the scheduler's own device workspace + self.sched_device_ws.assign_device_members( + cute.make_ptr( + cutlass.Uint8, + mega_local_workspace.toint() + + self._mega_device_workspace.offset(self.sched_work_id_region), + cute.AddressSpace.gmem, + assumed_align=16, + ), + mega_shared_workspace, + ) + else: + _sched_expert_sizes = expert_token_sizes + _sched_prefix_sum = offs + scheduler.assign_device_members( + expert_token_sizes=_sched_expert_sizes, + expert_token_prefix_sum=_sched_prefix_sum, + actual_expert_shape=None, + block_idx=cute.arch.block_idx(), + smem_workspace=self.sched_smem_ws, + smem_base=sched_smem_base, + device_workspace=self.sched_device_ws, + ) + sched_consumer = scheduler.make_consumer() + + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # SMEM tensors A / B / SFA / SFB (shared by fc1 / fc2) + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + sSFA = smem.allocate_tensor( + element_type=self.sf_dtype, + layout=sfa_smem_layout_staged, + byte_alignment=128, + ) + sSFB = smem.allocate_tensor( + element_type=self.sf_dtype, + layout=sfb_smem_layout_staged, + byte_alignment=128, + ) + + # preact (dswiglu C) staging tensor + preact_smem_layout_staged = self.epilogue.preact_staged_smem_layout( + num_c_stage + ) + sPre = storage.sPre.get_tensor( + preact_smem_layout_staged.outer, + swizzle=preact_smem_layout_staged.inner, + ) + + # Unified dFC2 data-output store staging tensor. + d_smem_layout_staged = self.epilogue.d_staged_smem_layout(num_d_stage) + sD = storage.sD.get_tensor( + d_smem_layout_staged.outer, + swizzle=d_smem_layout_staged.inner, + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + + # acc_fake layout: (MMA, MMA_M, MMA_N, STAGE) + acc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # Cluster wait before TMEM alloc. + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + mma_tiler_k = self.mma_tiler[2] + k_tile_cnt_fc1 = (fc1_weight_gemm.shape[1] + mma_tiler_k - 1) // mma_tiler_k + k_tile_cnt_fc2 = (fc2_weight_gemm.shape[1] + mma_tiler_k - 1) // mma_tiler_k + # fc2 spin threshold: number of N-tiles per CTA (per-CTA counter now). + fc2_spin_threshold = ( + (fc1_weight_gemm.shape[0] + self.cta_tile_shape_mnk[1] - 1) + // self.cta_tile_shape_mnk[1] + ) * self.epilogue._atom_thr_size + + # ════════════════════════════════════════════════════════════════════ + # Scheduler warp (warp 7) — lean path + # ════════════════════════════════════════════════════════════════════ + if warp_idx == self.sched_warp_id: + cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) + # MegaMoE: block until the Router has published this rank's per-expert sizes + self.token_comm_hook_sched_warp_pre_init_wait(token_comm_args) + work_tile = scheduler.gen_next_work() + while work_tile.is_valid_tile: + scheduler.publish_work(ext.prepare_work_tile(work_tile)) + work_tile = scheduler.gen_next_work() + # Sentinel publish (the tile is already invalid here). + scheduler.publish_work(work_tile) + scheduler.produce_tail() + + # ════════════════════════════════════════════════════════════════════ + # TMA load warps (warps 5 / 6) + # ════════════════════════════════════════════════════════════════════ + # + # TMA-A loads weights/SFA; TMA-B loads activations/SFB and waits for + # fc1 workspace readiness in fc2 phase. Both feed the same AB pipeline. + + # ── TMA-A warp (warp 5) ───────────────────────────────────────────── + if warp_idx == self.tma_a_warp_id: + cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) + a_full_mcast_mask = None + sfa_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + + b_full_mcast_mask = None + if cutlass.const_expr(self.is_b_mcast or use_2cta_instrs): + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + sfa_cta_layout = a_cta_layout + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + + work_tile_info = sched_consumer.consume_work() + + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + if is_phase_linear1: + iket.range_push("tma_weight_fc1") + # MegaMoE: spin until the dispatch (token_in) warps have pulled + ext.wait_for_input(work_tile_info) + self.token_comm_hook_fc1_tma_b_predispatch_spin( + token_comm_args, work_tile_info, + ) + + k_tile_cnt = k_tile_cnt_fc1 + real_a, desc_ptr_a = ext.get_gmem_tensor( + "fc1_activation", tma_tensor_fc1_activation_1, work_tile_info, + ) + real_sfa, desc_ptr_sfa = ext.get_gmem_tensor( + "fc1_activation_sf", tma_tensor_fc1_activation_1_sf, work_tile_info, + ) + + gA_mkl = cute.local_tile( + real_a, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + gSFA_mkl = cute.local_tile( + real_sfa, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + tCgA = thr_mma.partition_A(gA_mkl) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + + tAsA, tAgA = cpasync.tma_partition( + tma_atom_fc1_activation_1, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_fc1_activation_1_sf, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + mma_tile_m = work_tile_info.tile_m_idx // cute.size( + tiled_mma.thr_id.shape + ) + tAgA_slice = tAgA[(None, mma_tile_m, None, 0)] + tAgSFA_slice = tAgSFA[(None, mma_tile_m, None, 0)] + + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance( + peek_ab_empty_status + ) + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + cute.copy( + tma_atom_fc1_activation_1, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_a, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_fc1_activation_1_sf, + tAgSFA_slice[(None, handle.count)], + tAsSFA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfa, + mcast_mask=sfa_full_mcast_mask, + ) + iket.range_pop() + + else: + # fc2 phase A-side: load fc1_output (M=tokens) + wait for fc1 done + iket.range_push("tma_token_fc2") + counter_slot = ( + work_tile_info.cumulative_token_block_count + + work_tile_info.tile_m_idx // cutlass.Int32(self.epilogue._atom_thr_size) + ) + counter_ptr = fc1_done_counter.iterator + counter_slot + iket.range_push("tma_token_fc2_a_wait") + spin_wait( + counter_ptr, + lambda v: v >= fc2_spin_threshold, + sleep_cycles=20, + ) + iket.range_pop() + k_tile_cnt = k_tile_cnt_fc2 + real_a, desc_ptr_a = ext.get_gmem_tensor( + "fc2_activation", tma_tensor_fc2_activation, work_tile_info, + ) + real_sfa, desc_ptr_sfa = ext.get_gmem_tensor( + "fc2_activation_sf", tma_tensor_fc2_activation_sf, work_tile_info, + ) + + gA_mkl = cute.local_tile( + real_a, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + gSFA_mkl = cute.local_tile( + real_sfa, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + tCgA = thr_mma.partition_A(gA_mkl) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + + tAsA, tAgA = cpasync.tma_partition( + tma_atom_fc2_activation, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_fc2_activation_sf, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + mma_tile_m = work_tile_info.tile_m_idx // cute.size( + tiled_mma.thr_id.shape + ) + tAgA_slice = tAgA[(None, mma_tile_m, None, 0)] + tAgSFA_slice = tAgSFA[(None, mma_tile_m, None, 0)] + + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance( + peek_ab_empty_status + ) + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + cute.copy( + tma_atom_fc2_activation, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_a, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_fc2_activation_sf, + tAgSFA_slice[(None, handle.count)], + tAsSFA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfa, + mcast_mask=sfa_full_mcast_mask, + ) + + iket.range_pop() + work_tile_info = sched_consumer.consume_work() + + ab_producer.tail() + + # ── TMA-B warp (warp 6) ───────────────────────────────────────────── + if warp_idx == self.tma_b_warp_id: + cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) + b_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr(self.is_b_mcast or use_2cta_instrs): + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, + block_in_cluster_coord_sfb_vmnk, + mcast_mode=1, + ) + + # FC1: weight (B) is multicast (like original A) + a_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + sfb_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) + + work_tile_info = sched_consumer.consume_work() + + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + + if is_phase_linear1: + iket.range_push("tma_token_fc1") + + k_tile_cnt = k_tile_cnt_fc1 + real_b, desc_ptr_b = ext.get_gmem_tensor( + "fc1_weight", tma_tensor_weight, work_tile_info, + ) + real_sfb, desc_ptr_sfb = ext.get_gmem_tensor( + "fc1_weight_sf", tma_tensor_fc1_weight_sf, work_tile_info, + ) + + # N-K tiling for N-side weight (N=intermediate, K=hidden). + gB_nkl = cute.local_tile( + real_b, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + gSFB_nkl = cute.local_tile( + real_sfb, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + tCgB = thr_mma.partition_B(gB_nkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + + tBsB, tBgB = cpasync.tma_partition( + tma_atom_weight, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_fc1_weight_sf, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + tBgB_slice = tBgB[(None, work_tile_info.tile_n_idx, None, 0)] + tBgSFB_slice = tBgSFB[(None, work_tile_info.tile_n_idx, None, 0)] + + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance( + peek_ab_empty_status + ) + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + cute.copy( + tma_atom_weight, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_b, + mcast_mask=a_full_mcast_mask, # same as A-loading for weights + ) + cute.copy( + tma_atom_fc1_weight_sf, + tBgSFB_slice[(None, handle.count)], + tBsSFB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfb, + mcast_mask=sfb_full_mcast_mask, + ) + iket.range_pop() + + else: + # fc2 phase B-side: load fc2_weight (N=hidden), no counter wait + iket.range_push("tma_weight_fc2") + k_tile_cnt = k_tile_cnt_fc2 + real_b, desc_ptr_b = ext.get_gmem_tensor( + "fc2_weight", tma_tensor_fc2_weight, work_tile_info, + ) + real_sfb, desc_ptr_sfb = ext.get_gmem_tensor( + "fc2_weight_sf", tma_tensor_fc2_weight_sf, work_tile_info, + ) + + gB_nkl = cute.local_tile( + real_b, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + gSFB_nkl = cute.local_tile( + real_sfb, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + tCgB = thr_mma.partition_B(gB_nkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + + tBsB, tBgB = cpasync.tma_partition( + tma_atom_fc2_weight, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_fc2_weight_sf, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + fc2_b_hidden_tile = work_tile_info.tile_n_idx + tBgB_slice = tBgB[(None, fc2_b_hidden_tile, None, 0)] + tBgSFB_slice = tBgSFB[(None, fc2_b_hidden_tile, None, 0)] + + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance( + peek_ab_empty_status + ) + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + cute.copy( + tma_atom_fc2_weight, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_b, + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_fc2_weight_sf, + tBgSFB_slice[(None, handle.count)], + tBsSFB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfb, + mcast_mask=sfb_full_mcast_mask, + ) + iket.range_pop() + work_tile_info = sched_consumer.consume_work() + + ab_producer.tail() + + # ════════════════════════════════════════════════════════════════════ + # MMA warp (warp 4) + # ════════════════════════════════════════════════════════════════════ + # + # Both phases share tiled_mma and TMEM; only K-tile count differs. + if warp_idx == self.mma_warp_id: + cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + acc_base = cute.make_tensor(acc_tmem_ptr, acc_fake.layout) + + # SFA TMEM tensor (placed after the acc cols). + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + # SFB TMEM tensor (after acc + SFA cols). + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + + ( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t, + tCtSFA_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + ( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t, + tCtSFB_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_pipeline_stages + ) + + work_tile_info = sched_consumer.consume_work() + + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + # Prebind k_tile_cnt due to DSL AST. + k_tile_cnt = cutlass.Int32(0) + if is_phase_linear1: + k_tile_cnt = k_tile_cnt_fc1 + iket.range_push("mma_dfc2") + else: + k_tile_cnt = k_tile_cnt_fc2 + iket.range_push("mma_dfc1") + + acc_stage_index = acc_producer_state.index + + if is_leader_cta: + tCtAcc = acc_base[(None, None, None, acc_stage_index)] + + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if k_tile_cnt > 0: + peek_ab_full_status = ab_consumer.try_wait() + acc_pipeline.producer_acquire(acc_producer_state) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + iket.range_push("mma_ab_wait") + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + iket.range_pop() + + s2t_stage_coord = (None, None, None, None, handle.index) + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t[s2t_stage_coord], + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t[s2t_stage_coord], + tCtSFB_compact_s2t, + ) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + tile_crd = (None, None, None, handle.index) + cute.gemm( + tiled_mma, + tCtAcc, + [tCrA[tile_crd], tCtSFA], + [tCrB[tile_crd], tCtSFB], + tCtAcc, + ) + handle.release() + + if k_tile_cnt > 0: + acc_pipeline.producer_commit(acc_producer_state) + if k_tile_cnt > 0: + acc_producer_state.advance() + + iket.range_pop() + + work_tile_info = sched_consumer.consume_work() + + acc_pipeline.producer_tail(acc_producer_state) + + # ════════════════════════════════════════════════════════════════════ + # Dedicated preact-C TMA-load warp (c_load_warp_id) — c_pipeline PRODUCER + # ════════════════════════════════════════════════════════════════════ + # + # Mirrors the reference's epilog_load_tma warp: consume the same work + # tiles in lockstep, and for each Linear1 (dfc2) tile TMA-load gate + # (epi-tile 2*s) then up (2*s+1) into successive c_pipeline stages. + if warp_idx == self.c_load_warp_id: + cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) + + c_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, num_c_pipe_stage + ) + thr_mma_c = tiled_mma.get_slice(mma_tile_coord_v) + preact_epi_tile = self.epilogue.preact_epi_tile + c_subtile_cnt = self.cta_tile_shape_mnk[1] // 32 # 8 + + work_tile_info = sched_consumer.consume_work() + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + if is_phase_linear1: + real_preact, _ = ext.get_gmem_tensor( + "c", tma_tensor_fc1_preact, work_tile_info + ) + gC_mnl = cute.local_tile( + real_preact, cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + tCgC = thr_mma_c.partition_C(gC_mnl) + gC_epi = cute.flat_divide( + tCgC[((None, None), 0, 0, None, None, None)], preact_epi_tile + ) + bGS_sPre, bGS_gC = cpasync.tma_partition( + tma_atom_fc1_preact, 0, cute.make_layout(1), + cute.group_modes(sPre, 0, 2), + cute.group_modes(gC_epi, 0, 2), + ) + mma_m_coord = work_tile_info.tile_m_idx // cutlass.Int32(self.atom_thr_size) + mma_n_coord = work_tile_info.tile_n_idx * cutlass.Int32(2) + bGS_gC = bGS_gC[(None, None, None, mma_m_coord, mma_n_coord, 0)] + bGS_gC = cute.group_modes(bGS_gC, 1, cute.rank(bGS_gC)) + + for i in cutlass.range(0, c_subtile_cnt, 1, unroll=1): + subtile_idx = cutlass.Int32(i) + # gate (2*subtile) then up (2*subtile+1) + c_pipeline.producer_acquire(c_producer_state) + c_bar = c_pipeline.producer_get_barrier(c_producer_state) + c_slot = 2 * c_producer_state.index + cute.copy( + tma_atom_fc1_preact, + bGS_gC[(None, subtile_idx * cutlass.Int32(2) + cutlass.Int32(0))], + bGS_sPre[(None, c_slot)], + tma_bar_ptr=c_bar + ) + cute.copy( + tma_atom_fc1_preact, + bGS_gC[(None, subtile_idx * cutlass.Int32(2) + cutlass.Int32(1))], + bGS_sPre[(None, c_slot + 1)], + tma_bar_ptr=c_bar, + ) + c_producer_state.advance() + + work_tile_info = sched_consumer.consume_work() + + c_pipeline.producer_tail(c_producer_state) + + # ════════════════════════════════════════════════════════════════════ + # Epilogue warps (warps 0-3) + # ════════════════════════════════════════════════════════════════════ + # + # Fully delegated to ``self.epilogue.run(...)`` -- the epilogue owns + # the entire 2-phase task-tile loop. + if warp_idx < self.mma_warp_id: + cute.arch.warpgroup_reg_alloc(self.epi_reg_cnt) + epi_warp_idx = warp_idx + + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + acc_tensor = cute.make_tensor(acc_tmem_ptr, acc_fake.layout) + + # The epilogue is the preact c_pipeline CONSUMER (the dedicated + # c_load warp is the producer); it reads gate/up from sPre stages. + _run_kwargs = dict( + tmem_acc_tensor=acc_tensor, + acc_pipeline=acc_pipeline, + sched_consumer=sched_consumer, + sched_ext=ext, + gmem_fc1_output=tma_tensor_grad_y1, + gmem_fc1_output_sf=fc1_output_sf_gemm, + tma_atom_fc1_recompute=tma_atom_fc1_recompute, + gmem_fc1_recompute=tma_tensor_fc1_recompute, + gmem_fc1_recompute_sf=fc1_recompute_sf_gemm, + tma_atom_fc1_col_output=tma_atom_fc1_col_output, + gmem_fc1_col_output=tma_tensor_fc1_col_output, + gmem_fc1_col_output_sf=fc1_col_output_sf_gemm, + smem_preact_buffer=sPre, + c_pipeline=c_pipeline, + c_num_stage=num_c_pipe_stage, + smem_d_buffer=sD, + d_pipeline=d_pipeline, + d_num_stage=num_d_stage, + tma_atom_grad_y1=tma_atom_grad_y1, + gmem_topk_scores=topk_scores, + gmem_fc2_output=fc2_output_gemm, + gmem_fc1_done_counter=fc1_done_counter, + warp_idx=epi_warp_idx, + tidx=tidx, + alpha=cutlass.Float32(1.0), + norm_const=cutlass.Float32(1.0), + gmem_beta=beta, + gmem_dprob=dprob, + ) + + # MegaMoE: pass token_comm_args only when it is a real bundle (not + # None). Passing Python None explicitly to @cute.jit methods + # triggers a CuteDSL codegen issue; const_expr dispatch avoids any + # None-as-JIT-argument path. + if cutlass.const_expr(self.enable_token_comm): + # MegaMoE (push model): bridge next's TokenComm accessors + peer mapper into + # the dGLU epilogue's Fc2OutputDest peer-store expectations for grad_x combine. + _epi_comm = _EpilogueCommView( + token_src_metadata=self.token_comm.token_src_metadata_tensor( + self._mega_device_workspace + ), + combine_output=mega_pre_reduced_activation, + dprob_output=dprob, + peer_rank_ptr_mapper=mega_peer_rank_ptr_mapper, + fc2_output_sf=self.token_comm.fc2_activation_sf_tensor(self._mega_device_workspace), + fc2_done_counter=self.token_comm.fc2_done_counter_tensor(self._mega_device_workspace), + fc2_output_workspace=self.token_comm.fc2_activation_tensor(self._mega_device_workspace), + ) + self.epilogue.run(**_run_kwargs, token_comm_args=_epi_comm) + elif cutlass.const_expr(token_comm_args is not None): + self.epilogue.run(**_run_kwargs, token_comm_args=token_comm_args) + else: + self.epilogue.run(**_run_kwargs) + + tmem.relinquish_alloc_permit() + tmem.free(acc_tmem_ptr) + if cutlass.const_expr(self.enable_token_comm): + cute.arch.fence_acq_rel_sys() + + # ════════════════════════════════════════════════════════════════════ + # Dispatch / token_back warps hook (warps 8-11 [+ 12-15]; MegaMoE-only) + # ════════════════════════════════════════════════════════════════════ + # + # ``enable_token_comm=False`` → these warps don't exist (lean base has 9 + # warps), so the guard is const_expr-eliminated in the lean path. + # NOTE: c_load now lives ABOVE the transfer block (warp 12 or 16), so the + # gate must be UPPER-bounded at the last transfer warp — otherwise the + # c_load warp (already run above) would re-enter the dispatch body. + if cutlass.const_expr(self.enable_token_comm): + _last_transfer_warp = ( + self.token_back_warp_id[-1] + if self.token_back_standalone + else self.dispatch_warp_id[-1] + ) + if (warp_idx >= self.dispatch_warp_id[0]) & (warp_idx <= _last_transfer_warp): + cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) + lane_idx_for_dispatch = cute.arch.lane_idx() + if cutlass.const_expr(self.token_back_standalone): + if warp_idx < self.token_back_warp_id[0]: + self.token_comm_hook_dispatch_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + else: + self.token_comm_hook_token_back_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + else: + self.token_comm_hook_dispatch_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + + # ════════════════════════════════════════════════════════════════════ + # Kernel tail hook (MegaMoE-only; lean base = no-op) + # ════════════════════════════════════════════════════════════════════ + lane_idx = cute.arch.lane_idx() + self.token_comm_hook_kernel_tail( + token_comm_args, + warp_idx=warp_idx, + lane_idx=lane_idx, + tidx=tidx, + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py new file mode 100644 index 000000000..07f0bbfa5 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py @@ -0,0 +1,955 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Full MegaMoE (multi-rank) mxfp8 dGLU training-backward kernel.""" + +from types import SimpleNamespace +from typing import Any, Literal, Optional, Tuple, Type + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 + +from ......api import ImplDesc, KernelClass, ProblemDesc, StaticOrRuntimeIntegerType +from ......helpers.device_workspace import DeviceWorkspace +from ......helpers.smem_workspace import SmemWorkspace +from ......helpers.utils import ceil_div, round_up +from ......quant_def import CombineFormat, QuantKind +from ......communication.nvlink_domain.token_comm_deterministic import TokenCommDeterministic +from ..topk_reduce import TopkReduce +from ..fwd_glu.glu_mxfp8_col_requant import Mxfp8ColRequant +from .dglu_mxfp8_fc12_kernel import Sm107Mxfp8DgluDfc21Kernel + + +_AB_DTYPE_TO_QUANT_KIND = {cutlass.Float8E4M3FN: QuantKind.mxfp8_e4m3, cutlass.Float8E5M2: QuantKind.mxfp8_e5m2} +_QUANT_KIND_TO_AB_DTYPE = {str(k): d for d, k in _AB_DTYPE_TO_QUANT_KIND.items()} + +# TVM-FFI export symbol for the AOT-compiled callable. +_aot_symbol_prefix = "rubin_mega_moe_dglu_mxfp8_aot" + + +class Sm107MegaMoEMxfp8DgluKernel(Sm107Mxfp8DgluDfc21Kernel, KernelClass): + """Multi-rank MegaMoE wrapper around the lean mxfp8 dGLU kernel.""" + + # grad_y1 (dfc2 output), its SF, cross-phase counter, and internal dGLU pools. + fc1_output_region = "rubin.dglu_mxfp8.mega.fc1_output" + fc1_output_sf_region = "rubin.dglu_mxfp8.mega.fc1_output_sf" + fc1_done_counter_region = "rubin.dglu_mxfp8.mega.fc1_done_counter" + load_balance_counter_region = "rubin.dglu_mxfp8.mega.load_balance_counter" + sched_work_id_region = "rubin.dglu_mxfp8.mega.sched_work_id" + # Host-side local mirror of the token_comm shared token_src_metadata, so the + # legacy dfc2_recompute / dfc2_col_output validation (which reads it from the + # LOCAL workspace) works with the next's shared-heap metadata layout. + token_src_metadata_local_region = "rubin.dglu_mxfp8.mega.token_src_metadata_local" + fc1_preact_region = "rubin.dglu_mxfp8.mega.fc1_preact" + grad_y2_sizes_region = "rubin.dglu_mxfp8.mega.grad_y2_expert_token_sizes" + + # Reserved on top of the exact token_comm/sched SMEM to cover smem.allocate + # inter-allocation alignment padding that _compute_stages does not model. + _SMEM_ALLOC_MARGIN = 2048 + + @classmethod + def problem_desc_require(cls): + return { + "expert_count": StaticOrRuntimeIntegerType, + "intermediate_gateup_size": StaticOrRuntimeIntegerType, + "hidden_size": StaticOrRuntimeIntegerType, + "quant_kind": str, + "combine_format": CombineFormat, + "world_size": int, + "local_rank": int, + "topk": int, + "max_tokens_per_rank": int, + "max_recv_size_per_rank": int, + "gate_up_clamp": Optional[float], + } + + @classmethod + def impl_desc_require(cls): + return { + "mma_tiler_mnk": tuple, + "cluster_shape_mnk": tuple, + "use_2cta_instrs": bool, + "group_hint": int, + "token_padding_block": int, + "sf_padding_block": int, + "load_balance_mode": str, + "force_static_sched": bool, + "clc_bundle_size": Optional[int], + "num_sched_stages": Optional[int], + "acc_dtype": type, + "sf_vec_size": int, + "launch_cluster_count": int, + "drop_on_overflow": bool, + "fc2_in_kernel_topk_reduce": bool, + "token_back_mode": str, + "epi_flag_batch": tuple, + "flag_batch": int, + "act_func": str, + "dfc2_recompute": bool, + "dfc2_col_output": bool, + "enable_grad_y2_col_quant": bool, + "num_ctas_grad_y2_col_quant": int, + } + + def name(self) -> str: + return ( + f"sm107_megamoe_dglu_{self.quant_kind}_m{self.mma_tiler_mnk[0]}n{self.mma_tiler_mnk[1]}" + f"k{self.mma_tiler_mnk[2]}_e{self.expert_count}_ep{self.world_size}_topk{self.topk}_" + f"h{self.hidden_size}_i{self.intermediate_gateup_size}_combine{self.combine_format}_" + f"clamp{self.gate_up_clamp}_" + f"tokenback{self.token_back_mode}_hint{self.group_hint}_" + f"epi{self.epi_flag_batch[0]}x{self.epi_flag_batch[1]}_tif{self.flag_batch}_" + f"deterministic_mtpr{self.max_tokens_per_rank}_mrpr{self.max_recv_size_per_rank}_" + f"drop{int(self.drop_on_overflow)}_lc{self.launch_cluster_count}_" + f"recompute{int(self.dfc2_recompute)}x{int(self.dfc2_col_output)}_" + f"redtopk{int(self.reduce_topk_in_kernel)}_preactarg1" + ) + + def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): + """Compile against fake (metadata-only) inputs; ``out_path=None`` returns the in-memory callable.""" + import math + + from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_ptr + from cutlass.cute.typing import AddressSpace, sym_int64 + from cutlass.cutlass_dsl import Int32, Int64 + + from ......communication.nvlink_domain.symmetric_buffer import SymmetricBufferHost + + def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): + extents = tuple( + sym_int64(divisibility=math.gcd(int(extent), 128)) if axis in dynamic_axes else int(extent) + for axis, extent in enumerate(shape) + ) + return make_fake_compact_tensor(dtype, extents, stride_order=stride_order, assumed_align=alignment) + + tokens = self.max_tokens_per_rank + hidden = self.hidden + inter_half = self.intermediate_downproj # dfc2 weight N + gate_up = self.intermediate_gateup # grad_y1 width / dfc1 K = 2 * inter_half + experts = self.num_experts_per_rank + vec = self.sf_vec_size + activation_dtype = self.token_comm.activation_dtype # grad_out fp8 + sf_dtype = self.token_comm.activation_sf_dtype # E8M0 + # Atom-swizzled weight SF extents (to_blocked pads rows->128, cols->4). + fc1_weight_sf_columns = round_up(inter_half, 128) * round_up(hidden // vec, 4) + fc2_weight_sf_columns = round_up(hidden, 128) * round_up(gate_up // vec, 4) + aux_shapes = self.get_aux_output_shapes() + + fake_arguments = dict( + grad_out=fake_tensor(activation_dtype, (tokens, hidden), (1, 0), {0}, 16), + grad_out_sf=fake_tensor(sf_dtype, (tokens, self.token_comm.activation_sf_hidden_padded), (1, 0), {0}, 16), + topk_idx=fake_tensor(cutlass.Int32, (tokens, self.num_topk), (1, 0), {0}, 16), + topk_weights=fake_tensor(cutlass.Float32, (tokens, self.num_topk), (1, 0), {0}, 4), + fc1_weight=fake_tensor(self.ab_dtype, (experts, hidden, inter_half), (2, 0, 1), {0, 2}, 16), + fc1_weight_sf=fake_tensor(sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16), + fc2_weight=fake_tensor(self.ab_dtype, (experts, gate_up, hidden), (2, 0, 1), {0, 2}, 16), + fc2_weight_sf=fake_tensor(sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16), + beta=fake_tensor(cutlass.Float32, (experts,), (0,), {0}, 4), + fc1_preact=fake_tensor(cutlass.BFloat16, self.get_fc1_preact_shape(), (1, 0), set(), 128), + output_activation=fake_tensor(cutlass.BFloat16, (tokens, hidden), (1, 0), {0}, 16), + overflow_flag=fake_tensor(cutlass.Int32, (1,), (0,), set(), 4), + dprob=fake_tensor(cutlass.Float32, aux_shapes["dprob"], (1, 0), {0}, 16), + fc1_recompute=fake_tensor(self.ab_dtype, aux_shapes["fc1_recompute"], (1, 0), set(), 128), + fc1_recompute_sf=fake_tensor(sf_dtype, aux_shapes["fc1_recompute_sf"], (1, 0), set(), 128), + fc1_col_output=fake_tensor(self.ab_dtype, aux_shapes["fc1_col_output"], (1, 0), set(), 128), + fc1_col_output_sf=fake_tensor(sf_dtype, aux_shapes["fc1_col_output_sf"], (1, 0), set(), 128), + local_workspace=make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), + shared_workspace=make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), + peer_rank_ptr_mapper_host=SymmetricBufferHost( + base_address=Int64(0), + offsets=tuple(Int64(0) for _ in range(self.world_size)), + rank=Int32(0), + max_ranks=self.world_size, + ), + stream=make_fake_stream(), + ) + fake_arguments["grad_y2"] = fake_tensor( + self.ab_dtype, aux_shapes["grad_y2"], (0, 1), set(), 128 + ) + fake_arguments["grad_y2_sf"] = fake_tensor( + cutlass.Uint8, aux_shapes["grad_y2_sf"], (0,), set(), 16 + ) + + compiled = cute.compile[cute.EnableTVMFFI(True)](self, **fake_arguments) + if out_path is None: + return compiled + compiled.export_to_c(out_path, function_name=_aot_symbol_prefix, export_only_tvm_ffi_symbols=True) + return out_path + + @staticmethod + def load_compiled(path: str): + from cutlass.cute.runtime import load_module + + return load_module(path, enable_tvm_ffi=True)[_aot_symbol_prefix] + + @classmethod + def from_kwargs( + cls, + # Base-class (lean dfc2+dfc1) kwargs. + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mnk: Tuple[int, int, int], + use_2cta_instrs: bool, + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + load_balance_mode: str = "static", + static_expert_shape: Optional[Tuple[int, int, int]] = None, + force_static_sched: bool = True, + clc_bundle_size: Optional[int] = None, + num_sched_stages: Optional[int] = None, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + ab_dtype: Type[cutlass.Numeric] = cutlass.Float8E4M3FN, + sf_vec_size: int = 32, + *, + world_size: int, + local_rank: int, + num_topk: int, + max_tokens_per_rank: int, + max_recv_size_per_rank: int, + hidden: int, + launch_cluster_count: int, + drop_on_overflow: bool, + fc2_in_kernel_topk_reduce: bool = False, + token_back_mode: Literal["epi_warps", "standalone_warps", "reuse_dispatch_warps"] = "epi_warps", + epi_flag_batch: Optional[Tuple[int, int]] = (1, 1), + flag_batch: int = 1, + combine_format: Optional[CombineFormat] = None, + act_func: str = "swiglu", + gate_up_clamp: Optional[float] = None, + dfc2_recompute: bool = False, + dfc2_col_output: bool = False, + enable_grad_y2_col_quant: bool = False, + num_ctas_grad_y2_col_quant: int = 2368, + ) -> "Sm107MegaMoEMxfp8DgluKernel": + """Build the ``(ProblemDesc, ImplDesc)`` pair from the legacy flat signature.""" + if static_expert_shape is None: + raise NotImplementedError("Sm107MegaMoEMxfp8DgluKernel requires a static_expert_shape.") + if hidden != static_expert_shape[2]: + raise ValueError(f"hidden ({hidden}) must equal static_expert_shape[2] ({static_expert_shape[2]}).") + if ab_dtype not in _AB_DTYPE_TO_QUANT_KIND: + raise ValueError(f"ab_dtype {ab_dtype} has no mxfp8 QuantKind.") + num_experts_per_rank, intermediate_gateup, _hidden = static_expert_shape + combine_format = CombineFormat.parse("bf16" if combine_format is None else str(combine_format)) + problem_desc = ProblemDesc( + { + "expert_count": world_size * num_experts_per_rank, + "intermediate_gateup_size": intermediate_gateup, + "hidden_size": hidden, + "quant_kind": str(_AB_DTYPE_TO_QUANT_KIND[ab_dtype]), + "combine_format": combine_format, + "world_size": world_size, + "local_rank": local_rank, + "topk": num_topk, + "max_tokens_per_rank": max_tokens_per_rank, + "max_recv_size_per_rank": max_recv_size_per_rank, + "gate_up_clamp": gate_up_clamp, + } + ) + impl_desc = ImplDesc( + { + "mma_tiler_mnk": tuple(mma_tiler_mnk), + "cluster_shape_mnk": tuple(cluster_shape_mnk), + "use_2cta_instrs": use_2cta_instrs, + "group_hint": group_hint, + "token_padding_block": token_padding_block, + "sf_padding_block": sf_padding_block, + "load_balance_mode": load_balance_mode, + "force_static_sched": force_static_sched, + "clc_bundle_size": clc_bundle_size, + "num_sched_stages": num_sched_stages, + "acc_dtype": acc_dtype, + "sf_vec_size": sf_vec_size, + "launch_cluster_count": launch_cluster_count, + "drop_on_overflow": drop_on_overflow, + "fc2_in_kernel_topk_reduce": fc2_in_kernel_topk_reduce, + "token_back_mode": token_back_mode, + "epi_flag_batch": tuple(epi_flag_batch) if epi_flag_batch is not None else (1, 1), + "flag_batch": flag_batch, + "act_func": act_func, + "dfc2_recompute": dfc2_recompute, + "dfc2_col_output": dfc2_col_output, + "enable_grad_y2_col_quant": enable_grad_y2_col_quant, + "num_ctas_grad_y2_col_quant": num_ctas_grad_y2_col_quant, + } + ) + return cls(problem_desc, impl_desc) + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + # -- Extract descriptors into locals matching the legacy param names. -- + world_size = problem_desc["world_size"] + local_rank = problem_desc["local_rank"] + num_topk = problem_desc["topk"] + max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + max_recv_size_per_rank = problem_desc["max_recv_size_per_rank"] + hidden = problem_desc["hidden_size"] + combine_format = problem_desc["combine_format"] + gate_up_clamp = problem_desc["gate_up_clamp"] + _quant_kind = problem_desc["quant_kind"] + ab_dtype = _QUANT_KIND_TO_AB_DTYPE[_quant_kind] + static_expert_shape = ( + problem_desc["expert_count"] // world_size, + problem_desc["intermediate_gateup_size"], + hidden, + ) + + mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + cluster_shape_mnk = impl_desc["cluster_shape_mnk"] + use_2cta_instrs = impl_desc["use_2cta_instrs"] + group_hint = impl_desc["group_hint"] + token_padding_block = impl_desc["token_padding_block"] + sf_padding_block = impl_desc["sf_padding_block"] + load_balance_mode = impl_desc["load_balance_mode"] + force_static_sched = impl_desc["force_static_sched"] + clc_bundle_size = impl_desc["clc_bundle_size"] + num_sched_stages = impl_desc["num_sched_stages"] + acc_dtype = impl_desc["acc_dtype"] + sf_vec_size = impl_desc["sf_vec_size"] + launch_cluster_count = impl_desc["launch_cluster_count"] + drop_on_overflow = impl_desc["drop_on_overflow"] + fc2_in_kernel_topk_reduce = impl_desc["fc2_in_kernel_topk_reduce"] + token_back_mode = impl_desc["token_back_mode"] + epi_flag_batch = impl_desc["epi_flag_batch"] + flag_batch = impl_desc["flag_batch"] + act_func = impl_desc["act_func"] + dfc2_recompute = impl_desc["dfc2_recompute"] + dfc2_col_output = impl_desc["dfc2_col_output"] + self.enable_grad_y2_col_quant = impl_desc["enable_grad_y2_col_quant"] + self.num_ctas_grad_y2_col_quant = impl_desc["num_ctas_grad_y2_col_quant"] + + if hidden != static_expert_shape[2]: + raise ValueError(f"hidden ({hidden}) must equal static_expert_shape[2] ({static_expert_shape[2]}).") + token_back_by_dispatch = token_back_mode != "epi_warps" + combine_format = CombineFormat.parse("bf16" if combine_format is None else str(combine_format)) + # in-kernel topk reduce only conflicts with a QUANTIZED combine (no per-topk + # reduced-plane accumulation for quantized). It DOES work with the + # standalone/reuse_dispatch token-back modes (the token_comm token_back path + # honours token_back_reduce_topk), so those are allowed (mirrors the legacy). + if fc2_in_kernel_topk_reduce and combine_format.is_quantized: + raise ValueError("fc2_in_kernel_topk_reduce requires a non-quantized (bf16) combine.") + if token_back_mode not in ("epi_warps", "standalone_warps", "reuse_dispatch_warps"): + raise ValueError(f"unsupported token_back_mode={token_back_mode!r}.") + if ab_dtype not in _AB_DTYPE_TO_QUANT_KIND: + raise ValueError(f"ab_dtype {ab_dtype} has no mxfp8 QuantKind.") + + super().__init__( + mma_tiler_mnk=mma_tiler_mnk, + cluster_shape_mnk=cluster_shape_mnk, + use_2cta_instrs=use_2cta_instrs, + group_hint=group_hint, + token_padding_block=token_padding_block, + sf_padding_block=sf_padding_block, + load_balance_mode=load_balance_mode, + static_expert_shape=static_expert_shape, + force_static_sched=force_static_sched, + clc_bundle_size=clc_bundle_size, + num_sched_stages=num_sched_stages, + acc_dtype=acc_dtype, + ab_dtype=ab_dtype, + sf_vec_size=sf_vec_size, + epi_flag_batch=epi_flag_batch, + dfc2_recompute=dfc2_recompute, + dfc2_col_output=dfc2_col_output, + fc2_in_kernel_topk_reduce=fc2_in_kernel_topk_reduce, + act_func=act_func, + gate_up_clamp=gate_up_clamp, + ) + + # --- Warp topology (realigned for next's TokenCommDeterministic). --- + # next derives each transfer warp's transfer index as ``thread_idx % 128`` + # (token_comm.py:1421-1424 / 1911), which HARD-REQUIRES the dispatch warps -- + # and, standalone, the token_back warps -- to begin on a 4-warp / 128-thread + # boundary. So dispatch sits at warps 8-11 (thread 256 -> 256%128=0) and + # standalone token_back at 12-15 (thread 384 -> 384%128=0), mirroring the + # forward GLU. The dGLU c_load warp does NOT use the transfer index, so it + # moves ABOVE the transfer block (warp 16 iff standalone else 12), overriding + # the base kernel's default (warp 8, now occupied by dispatch). + self.enable_token_comm = True + self.dispatch_warp_id = (8, 9, 10, 11) + self.token_back_mode = token_back_mode + # Thread token_back_by_dispatch to the FC12 base (it hardcodes False) so the + # epilogue, built later in _setup_attributes(), fires the fc2_done counter for + # the standalone / reuse_dispatch token-back warps. Without this the dedicated + # token-back warps spin forever on fc2_done < target and the block-wide + # sync_threads() in kernel_tail deadlocks (M09/M10/M14/M15 hang). epi_warps is + # unaffected (it peer-writes grad_x directly and never reads fc2_done). + self.token_back_by_dispatch = token_back_by_dispatch + self.token_back_standalone = token_back_by_dispatch and token_back_mode == "standalone_warps" + self.token_back_warp_id = (12, 13, 14, 15) if self.token_back_standalone else None + num_token_back_warps = len(self.token_back_warp_id) if self.token_back_standalone else 0 + self.c_load_warp_id = 16 if self.token_back_standalone else 12 + + # Register re-balance for the mega warp layout. The base kernel sizes + # ``epi_reg_cnt`` (256) for the lean 9-warp dGLU; mega adds the 4 dispatch + # warps (+4 token-back if standalone) and the dedicated c_load warp, so the + # per-CTA register file can no longer grant 256 regs to all 4 epilogue warps + # -- the epilogue warpgroup then stalls forever inside + # ``warpgroup_reg_alloc`` and the mma/tmem barrier deadlocks. Mirror the + # legacy mega dGLU (megamoe_kernel_mxfp8_dglu.py:181-184). + self.epi_reg_cnt = 168 if self.token_back_standalone else 200 + self.threads_per_cta = 32 * ( + len(self.epilogue_warp_id) # 4 (warps 0-3) + + 1 # mma (warp 4) + + 1 # tma_a (warp 5) + + 1 # tma_b (warp 6) + + 1 # sched (warp 7) + + len(self.dispatch_warp_id) # 4 (warps 8-11) + + num_token_back_warps # 4 iff standalone_warps (warps 12-15) + + 1 # c_load (warp 12 or 16, dGLU-specific) + ) + + # --- MegaMoE constants. --- + self.world_size = world_size + self.local_rank = local_rank + self.num_topk = num_topk + self.max_tokens_per_rank = max_tokens_per_rank + self.max_recv_size_per_rank = max_recv_size_per_rank + self.hidden = hidden + self.launch_cluster_count = launch_cluster_count + self.drop_on_overflow = drop_on_overflow + self.combine_format = combine_format + self.num_experts_per_rank = static_expert_shape[0] + self.intermediate_downproj = static_expert_shape[1] + self.intermediate_gateup = self.intermediate_downproj * 2 + self.num_total_experts = world_size * self.num_experts_per_rank + self.reduce_topk_in_kernel = fc2_in_kernel_topk_reduce + self.token_back_schedule_mode = load_balance_mode if load_balance_mode == "atomic_counter" else "static" + + # --- next Router-push token communication component. --- + # dGLU dispatches raw grad_out tokens + the per-token routing prob (topk score) + # into the pool; the dfc2 epilogue folds the prob into d_gate/d_up. So the router + # ALWAYS pushes scores into the pool -> apply_topk_at_fc1=True. + mma_cta_count = 2 if use_2cta_instrs else 1 + cta_tile_m = mma_tiler_mnk[0] // mma_cta_count + cluster_m, cluster_n = self.cluster_shape_mn + tokens_per_fc1_ready_slot = cta_tile_m * cluster_m + hidden_per_fc2_cluster_tile = cta_tile_m * cluster_m + fc2_done_signals_per_token_tile = ceil_div(hidden, hidden_per_fc2_cluster_tile) * cluster_m * cluster_n + promised_launchable_sm_count = launch_cluster_count * cluster_m * cluster_n + quant_kind = _AB_DTYPE_TO_QUANT_KIND[ab_dtype] + tc_problem_desc = ProblemDesc( + { + "world_size": world_size, + "expert_count": self.num_total_experts, + "topk": num_topk, + "max_tokens_per_rank": max_tokens_per_rank, + "max_recv_size_per_rank": max_recv_size_per_rank, + "hidden_size": hidden, + "quant_kind": str(quant_kind), + "combine_format": combine_format, + "apply_topk_at_fc1": True, + } + ) + tc_impl_desc = ImplDesc( + { + "token_padding_block": token_padding_block, + "sf_padding_block": sf_padding_block, + "tokens_per_fc1_ready_slot": tokens_per_fc1_ready_slot, + "fc2_done_signals_per_token_tile": fc2_done_signals_per_token_tile, + "promised_launchable_sm_count": promised_launchable_sm_count, + "drop_on_overflow": drop_on_overflow, + "token_in_flag_batch": flag_batch, + "token_back_mode": token_back_mode, + "token_back_schedule_mode": self.token_back_schedule_mode, + "reduce_topk_in_kernel": fc2_in_kernel_topk_reduce, + } + ) + self.token_comm = TokenCommDeterministic(tc_problem_desc, tc_impl_desc) + self.pool_token_capacity = self.token_comm.worst_case_token_count + + # --- SMEM sub-buffer for the token_comm transport. --- + tc_smem_ws = SmemWorkspace() + self.token_comm.register_smem_regions(tc_smem_ws) + tc_smem_ws.finalize(max_bytes=self.smem_capacity) + self.tc_smem_ws = tc_smem_ws + self._token_comm_smem_bytes = tc_smem_ws.total_bytes + + # Build the scheduler NOW (launch_cluster_count known at construction) so its + # separately-allocated SMEM is reservable by ``_smem_misc_budget_bytes``. + _ec, _ig, _hd = static_expert_shape + self._build_scheduler( + expert_cnt=_ec, intermediate_gateup=_ig, hidden_dim=_hd, launch_cluster_count=launch_cluster_count + ) + self._sched_smem_bytes = self.sched_smem_ws.total_bytes + + # --- Post-kernel top-k reduction (skipped under in-kernel reduce). --- + self._topk_reduce = None if fc2_in_kernel_topk_reduce else TopkReduce(hidden, num_topk, combine_format) + + # --- Device workspace (next model): dGLU pools + token_comm regions. --- + self._mega_device_workspace = self._build_megamoe_device_workspace() + + # --- Bind every KernelClass schema field under its schema name. --- + self.expert_count = self.num_total_experts + self.intermediate_gateup_size = self.intermediate_downproj + self.hidden_size = hidden + self.quant_kind = _quant_kind + self.topk = num_topk + self.cluster_shape_mnk = tuple(cluster_shape_mnk) + self.mma_tiler_mnk = tuple(mma_tiler_mnk) + self.group_hint = group_hint + self.token_padding_block = token_padding_block + self.sf_padding_block = sf_padding_block + self.load_balance_mode = load_balance_mode + self.force_static_sched = force_static_sched + self.clc_bundle_size = clc_bundle_size + self.num_sched_stages = num_sched_stages + self.acc_dtype = acc_dtype + self.sf_vec_size = sf_vec_size + self.fc2_in_kernel_topk_reduce = fc2_in_kernel_topk_reduce + self.epi_flag_batch = tuple(epi_flag_batch) + self.flag_batch = flag_batch + self.act_func = act_func + self.dfc2_recompute = dfc2_recompute + self.dfc2_col_output = dfc2_col_output + self.use_2cta_instrs = use_2cta_instrs + + # Optional post-kernel token-axis MXFP8 requantization of the routed + # grad_out pool consumed as the dfc2 input. + if self.enable_grad_y2_col_quant: + col_quant_type = "mxfp8_e4m3" if ab_dtype is cutlass.Float8E4M3FN else "mxfp8_e5m2" + self.grad_y2_col_quant = Mxfp8ColRequant( + hidden=self.hidden, + num_experts=self.num_experts_per_rank, + max_total_tokens=( + self.world_size + * self.max_tokens_per_rank + * min(self.num_topk, self.num_experts_per_rank) + ), + quant_type=col_quant_type, + num_persistent_ctas=self.num_ctas_grad_y2_col_quant, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + dst_k_major=True, + ) + + def _smem_misc_budget_bytes(self) -> int: + """Reserve the token_comm transport + scheduler SMEM on top of the base misc budget.""" + _sched = getattr(self, "_sched_smem_bytes", 0) + return super()._smem_misc_budget_bytes() + self._token_comm_smem_bytes + _sched + self._SMEM_ALLOC_MARGIN + + def get_aux_output_shapes(self) -> dict: + """Shapes of the fixed-ABI dFC2 auxiliary outputs. + + Data planes are compact token-major matrices. Column-quantized scale + planes retain the WGrad 128x4 atom layout. + """ + data_token_capacity = self.token_comm.worst_case_token_count + sf_token_capacity = self.token_comm.worst_case_sf_token_count + column_sf_row_count = sf_token_capacity // self.sf_vec_size + return { + "dprob": (self.max_tokens_per_rank, self.num_topk), + "fc1_recompute": (data_token_capacity, self.intermediate_downproj), + "fc1_recompute_sf": (round_up(self.intermediate_downproj, 128), column_sf_row_count), + "fc1_col_output": (data_token_capacity, self.intermediate_gateup), + "fc1_col_output_sf": (round_up(self.intermediate_gateup, 128), column_sf_row_count), + "grad_y2": (data_token_capacity, self.hidden), + "grad_y2_sf": (sf_token_capacity * (self.hidden // self.sf_vec_size),), + } + + def get_fc1_preact_shape(self) -> Tuple[int, int]: + """Shape of the externally supplied, pool-indexed gate||up pre-activations.""" + return (self.token_comm.worst_case_token_count, self.intermediate_gateup) + + @cute.jit + def _validate_fixed_pool_tensor(self, tensor: cute.Tensor, dtype, expected_shape, expected_stride=None) -> None: + if cutlass.const_expr(tensor.element_type is not dtype): + raise TypeError("pool-domain tensor has an unexpected element type.") + if cutlass.const_expr(cute.rank(tensor.layout) != 2): + raise ValueError("pool-domain tensor must be rank 2.") + if cutlass.const_expr( + not isinstance(tensor.shape[0], int) + or not isinstance(tensor.shape[1], int) + or tensor.shape[0] != expected_shape[0] + or tensor.shape[1] != expected_shape[1] + ): + raise ValueError(f"pool-domain tensor must have static shape {expected_shape}.") + stride = (expected_shape[1], 1) if expected_stride is None else expected_stride + if cutlass.const_expr(tensor.stride[0] != stride[0] or tensor.stride[1] != stride[1]): + raise ValueError("pool-domain tensor has an unexpected stride.") + + def _build_megamoe_device_workspace(self) -> DeviceWorkspace: + """Register internal dGLU pools, counters, and token-comm regions.""" + sf_dtype = cutlass.Float8E8M0FNU + sf_vec_size = self.sf_vec_size + data_token_capacity = self.token_comm.worst_case_token_count + sf_token_capacity = self.token_comm.worst_case_sf_token_count + inter_gateup = self.intermediate_gateup # grad_y1 width (DOUBLED) + + # grad_y1 SF: row-quant, DOUBLED N columns. + sf_block_cols_back = round_up(ceil_div(inter_gateup, sf_vec_size), 4) + counter_slot_count = self.token_comm.max_fc1_ready_slot_count + + dw = DeviceWorkspace() + # grad_y1 (dfc2 output), consumed as the dfc1 (fc2) GEMM-B. + dw.register( + self.fc1_output_region, + self.ab_dtype, + (data_token_capacity, inter_gateup), + buffer_space="local", + mem_order=(1, 0), + byte_alignment=128, + ) + dw.register( + self.fc1_output_sf_region, + sf_dtype, + (sf_token_capacity, sf_block_cols_back), + buffer_space="local", + mem_order=(1, 0), + byte_alignment=128, + ) + # cross-phase fc1->fc2 done counter. + dw.register( + self.fc1_done_counter_region, + cutlass.Int32, + (counter_slot_count,), + buffer_space="local", + byte_alignment=16, + reset="tail_reset", + ) + # Dynamic load-balance atomic counter (scheduler claims work by atomic-inc). + # Always registered (cheap 1-int); the base kernel only reads it in + # atomic_counter mode, but it must be zeroed between back-to-back launches. + dw.register( + self.load_balance_counter_region, + cutlass.Int32, + (1,), + buffer_space="local", + byte_alignment=16, + reset="tail_reset", + ) + # Reserve a slot for the scheduler's atomic work-id counter (persistent-grid + # dynamic work distribution in atomic_counter mode). + dw.register( + self.sched_work_id_region, cutlass.Int32, (4,), buffer_space="local", byte_alignment=16, reset="tail_reset" + ) + # Local mirror of the shared token_src_metadata (Int64 per pool slot), filled + # host-side after the launch so the recompute/col-output validation can read it. + dw.register( + self.token_src_metadata_local_region, + cutlass.Int64, + (data_token_capacity,), + buffer_space="local", + byte_alignment=16, + ) + if self.enable_grad_y2_col_quant: + dw.register( + self.grad_y2_sizes_region, + cutlass.Int32, + (self.num_experts_per_rank,), + buffer_space="local", + byte_alignment=16, + ) + self.token_comm.register_device_workspace(dw) + dw.finalize() + return dw + + @property + def _local_offsets(self) -> dict: + """Legacy-name -> byte-offset map for inherited tester pool reads.""" + dw = self._mega_device_workspace + return { + "fc1_output": dw.offset(self.fc1_output_region), + "fc1_output_sf": dw.offset(self.fc1_output_sf_region), + "fc1_done_counter": dw.offset(self.fc1_done_counter_region), + # For the legacy dfc2_recompute / dfc2_col_output validation: + "l1_token_buffer": dw.offset(self.token_comm.fc1_activation_region), + "token_src_metadata": dw.offset(self.token_src_metadata_local_region), + } + + @property + def _shared_metadata_offset(self) -> int: + """Byte offset of the token_comm shared token_src_metadata (for the host mirror copy).""" + return self._mega_device_workspace.offset(self.token_comm._router.token_src_metadata_region) + + @property + def _local_region_by_name(self) -> dict: + """Legacy-name -> object exposing ``.nbytes`` for the inherited tester's pool reads.""" + dw = self._mega_device_workspace + name_to_region = { + "fc1_output": self.fc1_output_region, + "fc1_output_sf": self.fc1_output_sf_region, + "fc1_done_counter": self.fc1_done_counter_region, + "l1_token_buffer": self.token_comm.fc1_activation_region, + "token_src_metadata": self.token_src_metadata_local_region, + } + return {name: SimpleNamespace(nbytes=dw.nbytes(region)) for name, region in name_to_region.items()} + + def get_workspace_sizes(self) -> Tuple[int, int]: + """Return required (local, shared/symmetric) workspace bytes.""" + return self._mega_device_workspace.local_and_shared_bytes + + @property + def require_zero_workspace_leading_bytes(self) -> Tuple[int, int]: + return self._mega_device_workspace.require_zero_workspace_leading_bytes + + # ========================================================================= + # token_comm_hook_* -- filled with next's Router-push TokenCommDeterministic calls. + # ========================================================================= + + def token_comm_extra_smem_storage_class(self) -> type: + return self.tc_smem_ws.storage_class() + + def token_comm_hook_fc1_ready_counter_ptr(self, token_comm_args): + return self.token_comm.fc1_ready_counter_pointer(self._mega_device_workspace) + + def sched_ext_fc1_peek_threshold(self) -> int: # noqa: D401 + return super().sched_ext_fc1_peek_threshold() + + @cute.jit + def token_comm_hook_sched_warp_pre_init_wait(self, token_comm_args): + """The scheduler warp must wait for the Router to publish per-expert sizes.""" + self.token_comm.wait_for_sizes_ready(self._mega_device_workspace) + + @cute.jit + def token_comm_hook_fc1_tma_b_predispatch_spin(self, token_comm_args, work_tile_info): + """No-op: FC1 input readiness is enforced by the scheduler extension's fc1_ready spin.""" + pass + + @cute.jit + def token_comm_hook_dispatch_warp_body(self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx): + """Transfer warps (8-11): pull grad_out from peers into the local FC1 pool.""" + self.token_comm.token_in(self.tc_smem_ws, token_comm_storage.buffer.data_ptr()) + if cutlass.const_expr(self.token_comm.token_back_enabled and not self.token_back_standalone): + self.token_comm.token_back(self.tc_smem_ws, token_comm_storage.buffer.data_ptr()) + + @cute.jit + def token_comm_hook_token_back_warp_body(self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx): + """Standalone token-back warps (12-15): push grad_x back to source ranks.""" + self.token_comm.token_back(self.tc_smem_ws, token_comm_storage.buffer.data_ptr()) + + @cute.jit + def token_comm_hook_kernel_tail(self, token_comm_args, *, warp_idx, lane_idx, tidx): + """Cross-rank drain + workspace tail reset, performed by the transfer warps.""" + if cutlass.const_expr(self.enable_grad_y2_col_quant): + self._snapshot_grad_y2_expert_sizes(tidx) + cute.arch.sync_threads() + if (warp_idx >= self.dispatch_warp_id[0]) & (warp_idx <= self.dispatch_warp_id[-1]): + self.token_comm.reset_tail() + self.token_comm.remove_device_members() + + @cute.jit + def _snapshot_grad_y2_expert_sizes(self, tidx) -> None: + """Preserve local expert counts before token_comm tail reset.""" + dw = self._mega_device_workspace + if self.token_comm._linear_cta_idx == Int32(0): + sizes = self.token_comm.local_expert_sizes(dw, self.token_comm._local_rank) + snapshot = dw.tensor(self.grad_y2_sizes_region) + block_dim_x, _, _ = cute.arch.block_dim() + expert_idx = tidx + while expert_idx < Int32(self.num_experts_per_rank): + snapshot[expert_idx] = Int32(sizes[expert_idx]) + expert_idx = expert_idx + block_dim_x + + # ========================================================================= + # Host launch: Router kernel -> fused MegaMoE backward kernel -> top-k reduction. + # ========================================================================= + + @cute.jit + def __call__( + self, + grad_out: cute.Tensor, # (max_tokens_per_rank, hidden) fp8 + grad_out_sf: cute.Tensor, # (max_tokens_per_rank, activation_sf_hidden_padded) E8M0 + topk_idx: cute.Tensor, # (max_tokens_per_rank, num_topk) + topk_weights: cute.Tensor, # (max_tokens_per_rank, num_topk) Float32 (prob) + fc1_weight: cute.Tensor, # W2^T: (experts_per_rank, hidden, inter_downproj) + fc1_weight_sf: cute.Tensor, + fc2_weight: cute.Tensor, # W1^T: (experts_per_rank, intermediate, hidden) + fc2_weight_sf: cute.Tensor, + beta: cute.Tensor, # (experts_per_rank,) Float32 + fc1_preact: cute.Tensor, # (pool_token_capacity, intermediate_gateup) BFloat16 + output_activation: cute.Tensor, # (max_tokens_per_rank, hidden) BF16 + overflow_flag: cute.Tensor, # (1,) Int32, per-rank FC12 overflow output + dprob: cute.Tensor, # (max_tokens_per_rank, topk) Float32; symmetric, pre-zeroed + fc1_recompute: cute.Tensor, # (pool_token_capacity, inter_downproj) + fc1_recompute_sf: cute.Tensor, # (inter_padded, col_sf_rows) + fc1_col_output: cute.Tensor, # (pool_token_capacity, gateup) + fc1_col_output_sf: cute.Tensor, # (gateup_padded, col_sf_rows) + grad_y2: cute.Tensor, # (pool_token_capacity, hidden) token-axis MXFP8 + grad_y2_sf: cute.Tensor, # flat MN-major E8M0 bytes + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, # symmetric (NVLink) heap base + peer_rank_ptr_mapper_host, + stream: cuda.CUstream, + ) -> None: + """Launch the Router, then the fused backward main kernel, then (optionally) the top-k reduce.""" + + dw = self._mega_device_workspace + local_rank = peer_rank_ptr_mapper_host.rank + aux_shapes = self.get_aux_output_shapes() + self._validate_fixed_pool_tensor( + fc1_preact, cutlass.BFloat16, self.get_fc1_preact_shape() + ) + self._validate_fixed_pool_tensor( + fc1_recompute, self.ab_dtype, aux_shapes["fc1_recompute"] + ) + self._validate_fixed_pool_tensor( + fc1_recompute_sf, + self.token_comm.activation_sf_dtype, + aux_shapes["fc1_recompute_sf"], + ) + self._validate_fixed_pool_tensor( + fc1_col_output, + self.ab_dtype, + aux_shapes["fc1_col_output"], + ) + self._validate_fixed_pool_tensor( + fc1_col_output_sf, + self.token_comm.activation_sf_dtype, + aux_shapes["fc1_col_output_sf"], + ) + self._validate_fixed_pool_tensor( + grad_y2, + self.ab_dtype, + aux_shapes["grad_y2"], + (1, aux_shapes["grad_y2"][0]), + ) + if cutlass.const_expr( + grad_y2_sf.element_type is not cutlass.Uint8 + or cute.rank(grad_y2_sf.layout) != 1 + or grad_y2_sf.shape[0] != aux_shapes["grad_y2_sf"][0] + ): + raise ValueError("grad_y2_sf must be the fixed-size flat Uint8 carrier.") + self.token_comm.launch_router( + topk_indices=topk_idx, + topk_scores=topk_weights, + local_rank=local_rank, + local_workspace=local_workspace, + shared_workspace=shared_workspace, + peer_rank_ptr_mapper_host=peer_rank_ptr_mapper_host, + device_workspace=dw, + overflow_flag=overflow_flag, + stream=stream, + ) + peer_mapper = peer_rank_ptr_mapper_host.make_device_object() + dw.assign_device_members(local_workspace, shared_workspace) + + activation_pool = self.token_comm.fc1_activation_tensor(dw) + _sf_pool_atom = self.token_comm.fc1_activation_sf_tensor(dw) + activation_sf_pool = cute.make_tensor( + _sf_pool_atom.iterator, + cute.make_layout( + (self.token_comm.worst_case_sf_token_count, self.hidden // self.sf_vec_size), + stride=(self.token_comm.activation_sf_hidden_padded, 1), + ), + ) + fc1_output = dw.tensor(self.fc1_output_region) + fc1_output_sf = dw.tensor(self.fc1_output_sf_region) + fc1_done_counter = dw.tensor(self.fc1_done_counter_region) + load_balance_counter = ( + dw.tensor(self.load_balance_counter_region) if self.token_back_schedule_mode == "atomic_counter" else None + ) + pool_topk_scores = self.token_comm.fc1_topk_scores_tensor(dw) + + if cutlass.const_expr(self.reduce_topk_in_kernel): + # In-kernel top-k reduce (epi_warps + bf16 combine): the epilogue red-adds each + # topk grad_x contribution straight into the (pre-zeroed) 2D output. + pre_reduced = cute.make_tensor( + output_activation.iterator, + cute.make_layout( + (output_activation.shape[0], 1, output_activation.shape[1]), + stride=(output_activation.stride[0], output_activation.stride[0], output_activation.stride[1]), + ), + ) + pre_reduced_sf = None + else: + pre_reduced = self.token_comm.pre_reduced_activation_tensor(dw) + pre_reduced_sf = self.token_comm.pre_reduced_activation_sf_tensor(dw) + + if cutlass.const_expr(self.token_comm.token_back_push_data): + # token_back-by-dispatch: the epilogue writes grad_x to the LOCAL fc2_activation + # pool (same pool token_back reads), in its native (tokens, 1, hidden) shape. + fc2_output = self.token_comm.fc2_activation_tensor(dw) + else: + # epi_warps: the epilogue peer-writes grad_x directly (combine_output = pre_reduced), + _combine_hidden = pre_reduced.shape[2] + fc2_output = cute.make_tensor( + pre_reduced.iterator, + cute.make_layout( + (pre_reduced.shape[0] * pre_reduced.shape[1], _combine_hidden), stride=(_combine_hidden, 1) + ), + ) + + # dprob is a source-domain combine plane. Add a singleton value mode so + # the epilogue can reuse Fc2OutputDest's (token, topk, value) resolver. + dprob_combine = cute.make_tensor( + dprob.iterator, + cute.make_layout( + (dprob.shape[0], dprob.shape[1], 1), + stride=(dprob.stride[0], dprob.stride[1], 0), + ), + ) + + super().__call__( + activation_pool, + fc1_weight, + activation_sf_pool, + fc1_weight_sf, + fc1_output, + fc1_output_sf, + fc1_recompute, + fc1_recompute_sf, + fc1_col_output, + fc1_col_output_sf, + fc2_weight, + fc2_weight_sf, + fc2_output, + fc1_preact, + pool_topk_scores, + beta, + dprob_combine, + fc1_done_counter, + offs=None, + load_balance_counter=load_balance_counter, + max_active_clusters=self.launch_cluster_count, + stream=stream, + overflow_flag=overflow_flag, + mega_peer_rank_ptr_mapper=peer_mapper, + mega_local_rank=local_rank, + mega_local_workspace=local_workspace, + mega_shared_workspace=shared_workspace, + mega_activation=grad_out, + mega_activation_sf=grad_out_sf, + mega_pre_reduced_activation=pre_reduced, + mega_pre_reduced_activation_sf=pre_reduced_sf, + ) + + # Post-kernel top-k reduction: dequant + K-sum into the final output. + if cutlass.const_expr(not self.reduce_topk_in_kernel): + self._topk_reduce(pre_reduced, pre_reduced_sf, output_activation, None, stream) + + # Export the routed dfc2 input in token-axis MXFP8 form. The source + # grad_out pool and its row-wise SF remain resident after reset_tail. + if cutlass.const_expr(self.enable_grad_y2_col_quant): + lw = local_workspace + data_offset = dw.offset(self.token_comm.fc1_activation_region) + sf_offset = dw.offset(self.token_comm.fc1_activation_sf_region) + sizes_offset = dw.offset(self.grad_y2_sizes_region) + sf_pool_bytes = self.token_comm.worst_case_sf_token_count * (self.hidden // self.sf_vec_size) + src_data = cute.make_tensor( + cute.make_ptr( + self.ab_dtype, lw.toint() + Int64(data_offset), AddressSpace.gmem, assumed_align=128 + ), + cute.make_layout( + (self.token_comm.worst_case_token_count, self.hidden), stride=(self.hidden, 1) + ), + ) + src_sf_u8 = cute.make_tensor( + cute.make_ptr(cutlass.Uint8, lw.toint() + Int64(sf_offset), AddressSpace.gmem, assumed_align=16), + cute.make_layout((sf_pool_bytes,)), + ) + expert_sizes = cute.make_tensor( + cute.make_ptr(cutlass.Int32, lw.toint() + Int64(sizes_offset), AddressSpace.gmem, assumed_align=16), + cute.make_layout((self.num_experts_per_rank,)), + ) + self.grad_y2_col_quant( + src_data, + src_sf_u8, + expert_sizes, + grad_y2, + grad_y2_sf, + stream, + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/__init__.py new file mode 100644 index 000000000..4e9e25577 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Rubin training MegaMoE (mxfp8 GLU) kernel components.""" + +from .glu_mxfp8_fc12_epilogue import Fc2OutputDest, GluMxfp8Epilogue +from .glu_mxfp8_fc12_extension import GluMxFp8Fc12SchedExtension, TensorRole +from .glu_mxfp8_fc12_kernel import Sm107Mxfp8GluFc12Kernel +from .glu_mxfp8_mega_moe_kernel import Sm107MegaMoEMxfp8GluKernel + + +__all__ = [ + "Fc2OutputDest", + "GluMxFp8Fc12SchedExtension", + "GluMxfp8Epilogue", + "Sm107MegaMoEMxfp8GluKernel", + "Sm107Mxfp8GluFc12Kernel", + "TensorRole", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py new file mode 100644 index 000000000..cf745ee35 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py @@ -0,0 +1,1562 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Warp-specialized, persistent, TMA MXFP8 row-scale -> column-scale requant. + +Data uses the padded row-major dispatch pool. Destination SF is concatenated +per expert in ``[hidden_atom][token_atom]`` order. +""" + +from typing import Literal + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.cpasync as cpasync +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cutlass_dsl import ( + Float32, + Int32, + Int64, + T, + Uint8, + dsl_user_op, +) +from cutlass._mlir.dialects import arith, llvm +from cutlass.cute.typing import AddressSpace + +from ......helpers.constants import ( + Fp8E4M3FNMax, + Fp8E5M2Max, +) + +# The next/ helpers do not export these two block-size constants; they are fixed +# by the MXFP8 spec / the dispatch pool's 32x4x4 SF atom layout, so define them +# locally. +Mxfp8BlockSize = 32 +SfPaddingBlock = 128 + + +def _lcm(a: int, b: int) -> int: + from math import gcd + + return a * b // gcd(a, b) + + +def _smem_capacity() -> int: + """Max dynamic SMEM per CTA, from CUTLASS's per-arch table.""" + try: + from cutlass.utils import get_smem_capacity_in_bytes + + return int(get_smem_capacity_in_bytes()) + except Exception: + return 227 * 1024 + + +# ptxas 13.2 accepts ``scaled::n1`` only on sm_107a -- not sm_107, sm_107f, nor +# the later sm_110a / sm_120a -- so this is an exact set, not a floor. +_SCALED_CVT_ARCHS = frozenset({(10, 7)}) + + +def _target_arch_tuple() -> "tuple[int, int, str]": + """``(major, minor, suffix)`` of the active cuTeDSL compilation target.""" + from cutlass.cutlass_dsl import CuTeDSL + + arch = CuTeDSL._get_dsl().get_arch_enum() + return int(arch.major), int(arch.minor), (getattr(arch, "suffix", "") or "") + + +def _scaled_cvt_available() -> bool: + """Can this target assemble ``cvt...scaled::n1::ue8m0.e4m3x2.bf16x2``?""" + major, minor, suffix = _target_arch_tuple() + if (major, minor) not in _SCALED_CVT_ARCHS: + return False + if suffix != "a": + raise ValueError( + f"MXFP8 column requant targets sm_{major}{minor}{suffix}, but its " + f"block-scaled requant instruction " + f"'cvt.rn.satfinite.scaled::n1::ue8m0.e4m3x2.bf16x2' is accepted by " + f"ptxas only for the 'a' architecture variant; sm_{major}{minor} and " + f"sm_{major}{minor}f both fail with \"Arguments mismatch for " + f"instruction 'cvt'\". Compile for sm_{major}{minor}a, or pass " + f"scaled_cvt=False to select the portable requant path." + ) + return True + + +_SM_COUNT_CACHE: "list[int | None]" = [None] + + +def _resolve_sm_count(default: int) -> int: + """SM count of the current device.""" + if _SM_COUNT_CACHE[0] is None: + n = 0 + try: + import ctypes + + lib = ctypes.CDLL("libcuda.so.1") + lib.cuInit(0) + dev = ctypes.c_int() + if lib.cuDeviceGet(ctypes.byref(dev), 0) == 0: + val = ctypes.c_int() + # CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT + if lib.cuDeviceGetAttribute(ctypes.byref(val), 16, dev) == 0: + n = int(val.value) + except Exception: + n = 0 + _SM_COUNT_CACHE[0] = n + return _SM_COUNT_CACHE[0] or default + + +def _address_value(pointer_or_address, *, loc=None, ip=None): + if isinstance(pointer_or_address, Int64): + return pointer_or_address.ir_value() + return pointer_or_address.toint(loc=loc, ip=ip).ir_value() + + +@dsl_user_op +def tma_load_1d( + destination_smem, source_gmem, mbarrier_smem, num_bytes, *, loc=None, ip=None, +) -> None: + """Issue a 1D GMEM-to-SMEM bulk copy.""" + llvm.inline_asm( + None, + [ + destination_smem.toint(loc=loc, ip=ip).ir_value(), + _address_value(source_gmem, loc=loc, ip=ip), + num_bytes.ir_value(), + mbarrier_smem.toint(loc=loc, ip=ip).ir_value(), + ], + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1], $2, [$3];", + "r,l,r,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cp_async_bulk_s2g(destination_gmem, source_smem, num_bytes, *, loc=None, ip=None) -> None: + """Issue a 1D SMEM-to-GMEM bulk copy; the caller commits the group.""" + llvm.inline_asm( + None, + [ + _address_value(destination_gmem, loc=loc, ip=ip), + source_smem.toint(loc=loc, ip=ip).ir_value(), + num_bytes.ir_value(), + ], + "cp.async.bulk.global.shared::cta.bulk_group [$0], [$1], $2;", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _fp8x2_mnemonic(fp8_type) -> str: + if fp8_type is cutlass.Float8E4M3FN: + return "e4m3x2" + if fp8_type is cutlass.Float8E5M2: + return "e5m2x2" + raise TypeError(f"unsupported FP8 type {fp8_type}") + + +@dsl_user_op +def cvt_scaled_up_bf16x2(pair_b32, scale_b32, half: int, fp8_type, *, loc=None, ip=None) -> Int32: + """Two FP8 values + their E8M0 scale -> BF16x2, in one SASS instruction.""" + mn = _fp8x2_mnemonic(fp8_type) + asm = ( + "{\n" + " .reg .b16 a0,a1,s0,s1;\n" + " mov.b32 {a0,a1}, $1;\n" + " mov.b32 {s0,s1}, $2;\n" + f" cvt.rn.scaled::n2::ue8m0.bf16x2.{mn} $0, a{half}, s0;\n" + "}" + ) + return Int32( + llvm.inline_asm( + T.i32(), + [Int32(pair_b32).ir_value(loc=loc, ip=ip), Int32(scale_b32).ir_value(loc=loc, ip=ip)], + asm, + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def cvt_scaled_dn_fp8x2(v_bf16x2, raw_b32, fp8_type, *, loc=None, ip=None) -> Int32: + """BF16x2 + one E8M0 scale -> two FP8 bytes, in one SASS instruction.""" + mn = _fp8x2_mnemonic(fp8_type) + asm = ( + "{\n" + " .reg .b16 q, s16, junk;\n" + " .reg .b8 sb;\n" + " mov.b32 {s16, junk}, $2;\n" + " cvt.u8.u16 sb, s16;\n" + f" cvt.rn.satfinite.scaled::n1::ue8m0.{mn}.bf16x2 q, $1, sb;\n" + " cvt.u32.u16 $0, q;\n" + "}" + ) + return Int32( + llvm.inline_asm( + T.i32(), + [Int32(v_bf16x2).ir_value(loc=loc, ip=ip), Int32(raw_b32).ir_value(loc=loc, ip=ip)], + asm, + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def cvt_dn_fp8x2_portable(v_bf16x2, inv_lo, inv_hi, fp8_type, *, loc=None, ip=None) -> Int32: + """Two hidden columns of one token, as BF16x2, plus those two columns' + exact FP32 reciprocal scales -> two FP8 bytes in bits [15:0] (byte 0 from + the low BF16 half, byte 1 from the high half). + + Same rounding as ``cvt_scaled_dn_fp8x2``, not an approximation of it: BF16 + -> FP32 is an exact left shift, the reciprocal is an exact power of two, and + ``cvt.rn.satfinite`` is the RNE-and-saturate the hardware instruction + applies. Taking a scale per half is what lets the caller skip the transpose + that the one-scale-per-pair hardware instruction forces. + """ + mn = _fp8x2_mnemonic(fp8_type) + asm = ( + "{\n" + " .reg .b32 a, b;\n" + " .reg .b16 q;\n" + " shl.b32 a, $1, 16;\n" + " and.b32 b, $1, 0xffff0000;\n" + " mul.f32 a, a, $2;\n" + " mul.f32 b, b, $3;\n" + # ``cvt d, a, b`` yields d[15:8] = cvt(a) and d[7:0] = cvt(b), so the + # HIGH column has to be the first source for byte 0 to be the low one. + f" cvt.rn.satfinite.{mn}.f32 q, b, a;\n" + " cvt.u32.u16 $0, q;\n" + "}" + ) + return Int32( + llvm.inline_asm( + T.i32(), + [ + Int32(v_bf16x2).ir_value(loc=loc, ip=ip), + Float32(inv_lo).ir_value(loc=loc, ip=ip), + Float32(inv_hi).ir_value(loc=loc, ip=ip), + ], + asm, + "=r,r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def max_xorsign_abs_bf16x2(a, b, *, loc=None, ip=None) -> Int32: + """Packed magnitude max of two BF16x2.""" + return Int32( + llvm.inline_asm( + T.i32(), + [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)], + "max.xorsign.abs.bf16x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def e8m0_raw_from_bf16(bf16_bits, limit_exponent: int, *, loc=None, ip=None) -> Int32: + """E8M0 raw byte for a non-negative BF16 magnitude.""" + bits = Int32(bf16_bits) << Int32(16) + biased = (bits + Int32(0x1FFFFF - (limit_exponent << 23))) >> Int32(23) + return Int32( + arith.select( + (bits >= Int32(0x7F800000)).ir_value(loc=loc, ip=ip), + Int32(254).ir_value(loc=loc, ip=ip), + cutlass.max(Int32(0), biased).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def bits_f32(x: Int32, *, loc=None, ip=None) -> Float32: + return Float32(llvm.bitcast(T.f32(), Int32(x).ir_value(loc=loc, ip=ip), loc=loc, ip=ip)) + + +def fp8_limit_exponent(limit: float) -> int: + """Exponent ``k`` of an FP8 limit written as ``1.75 * 2**k``.""" + scaled = limit / 1.75 + exponent = int(scaled).bit_length() - 1 + if float(1 << exponent) != scaled: + raise ValueError( + f"FP8 limit {limit} is not of the form 1.75 * 2**k, which the integer E8M0 encoding assumes." + ) + return exponent + + +class Mxfp8ColRequant: + """Warp-specialized persistent TMA launcher for token-axis MXFP8 requant.""" + + # --- fixed MXFP8 / SF-atom geometry (do not retune) --------------------- + TokensPerBlock: int = Mxfp8BlockSize # 32: one E8M0 scale covers 32 values + SfAtomBytes: int = 512 + SfAtomNonK: int = 128 + SfAtomKBanks: int = 4 + + TokenPaddingBlocks: tuple = (128, 256) + + # --- this kernel's tile ------------------------------------------------ + TILE_TOK: int = 128 # == SfAtomNonK, so a tile is one whole SF atom row + NSTAGE: int = 2 # 2 is what leaves shared memory for 2 CTAs per SM + KMajorStages: int = 1 # Short-lived K-major CTAs do not amortize stage two. + + # Tile width and per-lane access width, per requant path. Both are tuned: + # the two paths have different shared-memory budgets and different + # instruction mixes, and the pairs below are the measured optima. + TileHidScaled: int = 512 + ColsPerLaneScaled: int = 8 + TileHidPortable: int = 256 + ColsPerLanePortable: int = 4 + + # Sub-tile width of the consumer, in hidden columns. See _sp_tile_body. + # Must divide ColsPerLane and divide 32. + SP_LW: int = 4 + + # Declared upper bound on the block size. ptxas budgets + # 512//ceil(T/128) registers from it, so 576 yields 96 registers per + # thread, which is what two resident CTAs need. + MaxNTidTarget: int = 576 + + # Padding between consecutive SF atoms in shared memory, in bytes; it is + # what keeps the consumer's strided scale accesses off one bank. + SfInPad: int = 16 + SfOutPad: int = 16 + + # Grid depth, in resident waves. The curve is broad and flat above this. + GridWaves: int = 24 + KMajorGridWaves: int = 13 + + ProducerWarps: int = 1 + # ConsumerWarps is derived: + # (TILE_TOK / TokensPerBlock) * (TILE_HID / (32 * ColsPerLane)) + + SmCount: int = 148 # only a fallback; _resolve_sm_count queries the driver + # Upper bound on the grid: below this many tiles per CTA the + # O(num_experts) prologue dominates. Stated per 256 experts. + MinTilesPerCta: int = 4 + MinTilesRefExperts: int = 256 + + ConsumerBarrierId: int = 1 + + @classmethod + def _require_token_padding_block(cls, block) -> int: + """Refuse any token padding block this kernel is not built for.""" + if isinstance(block, bool) or not isinstance(block, int): + raise ValueError( + f"token_padding_block must be an int, got {block!r} " + f"({type(block).__name__})." + ) + if block not in cls.TokenPaddingBlocks: + raise ValueError( + f"token_padding_block must be one of " + f"{tuple(cls.TokenPaddingBlocks)}, got {block}. The work tile is " + f"TILE_TOK = SfAtomNonK = {cls.SfAtomNonK} tokens, i.e. one " + f"{cls.SfAtomBytes} B SF atom of {cls.SfAtomNonK}x{cls.SfAtomNonK} " + f"(token x hidden), so an expert's padded extent has to be a whole " + f"number of {cls.SfAtomNonK}-token tiles for a tile to belong to " + f"exactly one expert; {cls.TokenPaddingBlocks} are the only blocks " + f"the dispatch pool emits and the only ones validated." + ) + return int(block) + + def __init__( + self, + hidden: int, + num_experts: int, + max_total_tokens: int, + quant_type: Literal["mxfp8_e4m3", "mxfp8_e5m2"], + num_persistent_ctas: int = -1, + token_padding_block: int = SfPaddingBlock, + sf_padding_block: int = SfPaddingBlock, + *, + scaled_cvt: "bool | None" = None, + dst_k_major: bool = False, + ) -> None: + """``scaled_cvt`` selects the requant path: ``None`` asks the + compilation target, ``True`` and ``False`` force the block-scaled and + the portable path so that either can be exercised on one machine. + Everything else is derived from the problem.""" + self.hidden = int(hidden) + self.num_experts = int(num_experts) + self.max_total_tokens = int(max_total_tokens) + self.quant_type = quant_type + self.token_padding_block = int(token_padding_block) + self.sf_padding_block = int(sf_padding_block) + self.dst_k_major = bool(dst_k_major) + + self._require_token_padding_block(self.token_padding_block) + if self.sf_padding_block != self.SfAtomNonK: + raise ValueError( + f"sf_padding_block must be {self.SfAtomNonK} for the 32x4x4 atom layout, " + f"got {self.sf_padding_block}." + ) + + if scaled_cvt is None: + self.scaled_cvt = _scaled_cvt_available() + elif scaled_cvt: + if not _scaled_cvt_available(): + _major, _minor, _suffix = _target_arch_tuple() + raise ValueError( + f"scaled_cvt=True forces the block-scaled cvt consumer, but " + f"the target is sm_{_major}{_minor}{_suffix} and " + f"'cvt.rn.satfinite.scaled::n1::ue8m0.e4m3x2.bf16x2' is not " + f"available there. Pass scaled_cvt=None to let the target " + f"choose." + ) + self.scaled_cvt = True + else: + self.scaled_cvt = False + + if self.hidden <= 0: + raise ValueError(f"hidden must be positive, got {self.hidden}.") + + # The tile has to be a multiple of lcm(SF atom, 32*C) and divide hidden, + # so a shape the tuned pair cannot tile falls back to a narrower access. + if self.scaled_cvt: + _cols_choices = (self.ColsPerLaneScaled, self.ColsPerLanePortable) + _preferred_tile_hid = self.TileHidScaled + # The extra K-major staging tile makes 256 columns faster on Rubin: + # it preserves two resident CTAs and wins despite twice as many + # hidden groups (84 us versus 101 us for the DS3 production case). + if self.dst_k_major: + _preferred_tile_hid = min( + _preferred_tile_hid, self.TileHidPortable + ) + else: + _cols_choices = (self.ColsPerLanePortable,) + _preferred_tile_hid = self.TileHidPortable + _picked = None + for _c in _cols_choices: + _grain = _lcm(self.SfAtomNonK, 32 * _c) + _choices = [ + t + for t in range(_grain, _preferred_tile_hid + 1, _grain) + if self.hidden % t == 0 + ] + if _choices: + _picked = (_c, _choices) + break + if _picked is None: + raise ValueError( + f"hidden={self.hidden} is not supported: it must be a multiple " + f"of {_lcm(self.SfAtomNonK, 32 * _cols_choices[-1])}." + ) + self.ColsPerLane, self._tile_hid_choices = _picked + self.TILE_HID = self._tile_hid_choices[-1] + + # Narrow shapes can end up with fewer columns per lane than the sub-tile + # width, so clamp instead of refusing; the split is then trivial. + self.sp_lw = min(self.SP_LW, self.ColsPerLane) + if self.num_experts <= 0: + raise ValueError(f"num_experts must be positive, got {self.num_experts}.") + if self.max_total_tokens <= 0: + raise ValueError(f"max_total_tokens must be positive, got {self.max_total_tokens}.") + + if quant_type == "mxfp8_e4m3": + self.quant_dtype = cutlass.Float8E4M3FN + self._data_limit_exponent = fp8_limit_exponent(float(Fp8E4M3FNMax)) + elif quant_type == "mxfp8_e5m2": + self.quant_dtype = cutlass.Float8E5M2 + self._data_limit_exponent = fp8_limit_exponent(float(Fp8E5M2Max)) + else: + raise ValueError(f"Unsupported quant_type: {quant_type!r}") + self.sf_dtype = cutlass.Float8E8M0FNU + + self.sf_in_pad = self.SfInPad + self.sf_out_pad = self.SfOutPad + self.smem_capacity = _smem_capacity() + self.NumStages = self.KMajorStages if self.dst_k_major else self.NSTAGE + while ( + self._smem_bytes_for(self.TILE_HID, self.NumStages) > self.smem_capacity + and len(self._tile_hid_choices) > 1 + ): + self._tile_hid_choices.pop() + self.TILE_HID = self._tile_hid_choices[-1] + + self.TmaBoxHidU32 = self.TILE_HID // 4 + + self._hidden_atoms = self.hidden // self.SfAtomNonK + self.HidAtomsPerTile = self.TILE_HID // self.SfAtomNonK + self.HidSegs = self.TILE_HID // (32 * self.ColsPerLane) + self.TokBlocks = self.TILE_TOK // self.TokensPerBlock # 4 + self.ConsumerWarps = self.HidSegs * self.TokBlocks + self.SfInStride = self.SfAtomBytes + self.sf_in_pad + self.SfTileBytes = self.HidAtomsPerTile * self.SfInStride + self.SfTileXferBytes = self.HidAtomsPerTile * self.SfAtomBytes + + if (self.ProducerWarps + self.ConsumerWarps) * 32 > 1024: + raise ValueError( + f"hidden={self.hidden} needs " + f"{(self.ProducerWarps + self.ConsumerWarps) * 32} threads per " + f"CTA, over the 1024-thread hardware limit." + ) + self.WarpsPerCta = self.ProducerWarps + self.ConsumerWarps + self.ThreadsPerCta = self.WarpsPerCta * 32 + # ``.maxntid`` is an upper bound, so it can never be below the launch. + self.MaxNTid = max(self.MaxNTidTarget, self.ThreadsPerCta) + + # --- SMEM ------------------------------------------------------------ + self.smem_data_bytes = self.NumStages * self.TILE_TOK * self.TILE_HID + # K-major output staging is separate from the input pipeline. + self.smem_data_out_bytes = ( + self.TILE_TOK * self.TILE_HID if self.dst_k_major else 0 + ) + self.smem_sf_in_bytes = self.NumStages * self.SfTileBytes + self.SfOutStride = self.SfAtomBytes + self.sf_out_pad + self.smem_sf_out_bytes = self.HidAtomsPerTile * self.SfOutStride + self.smem_table_bytes = 3 * (self.num_experts + 1) * 4 + self.smem_bytes = ( + self.smem_data_bytes + + self.smem_data_out_bytes + + self.smem_sf_in_bytes + + self.smem_sf_out_bytes + + self.smem_table_bytes + + 2 * self.NumStages * 8 + + (1024 if self.dst_k_major else 256) + ) + if self.smem_bytes > self.smem_capacity: + raise ValueError( + f"hidden={self.hidden} needs {self.smem_bytes} B of shared " + f"memory per CTA, over this target's " + f"{self.smem_capacity} B limit." + ) + + self.hidden_groups = self.hidden // self.TILE_HID + + # --- grid ------------------------------------------------------------ + # The grid quantum is RES = (resident CTAs) = CtasPerSm * SM_count: a + # grid that is not a multiple of RES leaves a fractional resident wave, + # which is a large loss. Across multiples of RES the curve is broad and + # flat, so one tuned wave count serves every production size. + self.SmCount = _resolve_sm_count(type(self).SmCount) + # Both gates: SMEM, and the 8-warps-per-scheduler cap (warps go to the + # 4 schedulers round robin, so one CTA occupies ceil(warps/4) slots). + _warp_gate = 8 // -(-self.WarpsPerCta // 4) + self.CtasPerSm = max(1, min(2, self.smem_capacity // self.smem_bytes, _warp_gate)) + self.ResidentCtas = self.CtasPerSm * self.SmCount + # Every CTA pays an O(num_experts) prefix-table prologue, so the grid has + # an upper bound, and that bound rises in proportion to the expert count: + # with too few tiles per CTA the prologue dominates the tile work. + _min_tiles = self.MinTilesPerCta * max( + 1, -(-self.num_experts // self.MinTilesRefExperts) + ) + # The wave count is tuned for 2 resident CTAs. A shape that gets only + # one keeps a single-wave grid rather than extrapolating that tuning + # point outside the regime it was taken in. + if self.CtasPerSm >= 2: + _want = ( + self.KMajorGridWaves if self.dst_k_major else self.GridWaves + ) + else: + _want = 1 + _max_tiles = -(-self.max_total_tokens // self.TILE_TOK) * self.hidden_groups + _waves = max(1, min(_want, _max_tiles // (_min_tiles * self.ResidentCtas))) + if num_persistent_ctas > 0: + self.num_persistent_ctas = int(num_persistent_ctas) + else: + self.num_persistent_ctas = _waves * self.ResidentCtas + self.grid = self.num_persistent_ctas + + # Reported by the runner's PASS line; this kernel has a fixed split. + self.HiddenPerCta = self.TILE_HID + self.hidden_tiles_per_work = 1 + + # Binary-search ladder over the valid-token prefix table: the powers of + # two below num_experts, largest first. + steps = [] + span = 1 + while span < self.num_experts: + span <<= 1 + span >>= 1 + while span >= 1: + steps.append(span) + span >>= 1 + self._search_steps = tuple(steps) + self._search_needs_guard = (self.num_experts & (self.num_experts - 1)) != 0 + self._experts_per_lane = (self.num_experts + 31) // 32 + + # ------------------------------------------------------------------ host + def _k_major_tma_smem_layout(self): + """Canonical swizzled FP8 hidden-by-token TMA staging tile.""" + staged = sm100_utils.make_smem_layout_epi( + self.quant_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self.TILE_HID, self.TILE_TOK), + 1, + ) + return cute.select(staged, mode=[0, 1]) + + @cute.jit + def __call__( + self, + src_data: cute.Tensor, + src_sf_u8: cute.Tensor, + expert_token_sizes: cute.Tensor, + dst_data: cute.Tensor, + dst_sf_u8: cute.Tensor, + cuda_stream: cuda.CUstream, + token_padding_block: cutlass.Constexpr = None, + ) -> None: + TOKPAD = cutlass.const_expr( + self.token_padding_block if token_padding_block is None else token_padding_block + ) + self._require_token_padding_block(TOKPAD) + if cutlass.const_expr(TOKPAD != self.token_padding_block): + raise ValueError( + f"token_padding_block passed to __call__ ({TOKPAD}) disagrees with " + f"the one the Mxfp8ColRequant was constructed with " + f"({self.token_padding_block}); the launch geometry is derived from " + f"the constructor value, so construct a new instance instead." + ) + if cutlass.const_expr(src_data.element_type is not self.quant_dtype): + raise TypeError(f"src_data must use {self.quant_dtype}, got {src_data.element_type}.") + if cutlass.const_expr(dst_data.element_type is not self.quant_dtype): + raise TypeError(f"dst_data must use {self.quant_dtype}, got {dst_data.element_type}.") + + HID_U32 = cutlass.const_expr(self.hidden // 4) + BOX_H = cutlass.const_expr(self.TmaBoxHidU32) + BOX_T = cutlass.const_expr(self.TILE_TOK) + src_u32 = cute.make_tensor( + cute.recast_ptr(src_data.iterator, dtype=cutlass.Uint32), + cute.make_layout((src_data.shape[0], HID_U32), stride=(HID_U32, 1)), + ) + tma_atom, tma_tensor = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + src_u32, + cute.make_layout((BOX_T, BOX_H), stride=(BOX_H, 1)), + (BOX_T, BOX_H), + ) + + if cutlass.const_expr(self.dst_k_major): + k_major_smem_layout = self._k_major_tma_smem_layout() + dst_u8 = cute.make_tensor( + cute.recast_ptr(dst_data.iterator, dtype=cutlass.Uint8), + cute.make_layout( + (dst_data.shape[1], dst_data.shape[0]), + stride=(dst_data.shape[0], 1), + ), + ) + tma_atom_st, tma_tensor_st = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + dst_u8, + k_major_smem_layout, + (self.TILE_HID, BOX_T), + ) + else: + dst_u32 = cute.make_tensor( + cute.recast_ptr(dst_data.iterator, dtype=cutlass.Uint32), + cute.make_layout((dst_data.shape[0], HID_U32), stride=(HID_U32, 1)), + ) + tma_atom_st, tma_tensor_st = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + dst_u32, + cute.make_layout((BOX_T, BOX_H), stride=(BOX_H, 1)), + (BOX_T, BOX_H), + ) + k = self.ws_kernel( + src_data, src_sf_u8, expert_token_sizes, dst_data, dst_sf_u8, + tma_atom, tma_tensor, tma_atom_st, tma_tensor_st, TOKPAD, + ) + # Validated in __init__, not here: __call__ is DSL-preprocessed and a + # plain if/raise in a traced body is rejected at trace time. + _ntid = cutlass.const_expr(self.MaxNTid) + k.launch( + grid=[self.grid, 1, 1], + block=[self.ThreadsPerCta, 1, 1], + max_number_threads=[_ntid, 1, 1], + stream=cuda_stream, + ) + + # ------------------------------------------------------- prefix tables + @cute.jit + def _pad_up(self, count, block: int): + if cutlass.const_expr(block > 0 and (block & (block - 1)) == 0): + return (count + Int32(block - 1)) & Int32(-block) + return ((count + Int32(block - 1)) // Int32(block)) * Int32(block) + + @cute.jit + def warp_prefix_sum(self, value, lane_idx): + acc = Int32(value) + for shift in cutlass.range_constexpr(0, 5, 1): + step = 1 << shift + other = Int32(cute.arch.shuffle_sync_up(acc, Int32(step), mask_and_clamp=0)) + if lane_idx >= Int32(step): + acc = acc + other + return acc + + @cute.jit + def build_prefix_tables( + self, expert_token_sizes, tbl_vend, tbl_data, tbl_sf, lane_idx, + token_padding_block: cutlass.Constexpr = None, + ): + """Exclusive prefixes of the padded-data / padded-SF row counts.""" + E = cutlass.const_expr(self.num_experts) + EPL = cutlass.const_expr(self._experts_per_lane) + TOKPAD = cutlass.const_expr( + self.token_padding_block if token_padding_block is None else token_padding_block + ) + self._require_token_padding_block(TOKPAD) + + counts = cute.make_rmem_tensor((EPL,), cutlass.Int32) + local_data = Int32(0) + local_sf = Int32(0) + for slot in cutlass.range_constexpr(0, EPL, 1): + expert = lane_idx * Int32(EPL) + Int32(slot) + count = Int32(0) + if expert < Int32(E): + count = Int32(expert_token_sizes[expert]) + counts[slot] = count + local_data = local_data + self._pad_up(count, TOKPAD) + local_sf = local_sf + self._pad_up(count, self.sf_padding_block) + + base_data = self.warp_prefix_sum(local_data, lane_idx) - local_data + base_sf = self.warp_prefix_sum(local_sf, lane_idx) - local_sf + + for slot in cutlass.range_constexpr(0, EPL, 1): + expert = lane_idx * Int32(EPL) + Int32(slot) + count = Int32(counts[slot]) + if expert <= Int32(E): + tbl_vend[expert] = base_data + count + tbl_data[expert] = base_data + tbl_sf[expert] = base_sf + base_data = base_data + self._pad_up(count, TOKPAD) + base_sf = base_sf + self._pad_up(count, self.sf_padding_block) + + # When 32 * EPL == E no lane owns index E, so it is written separately. + if (lane_idx + Int32(1)) * Int32(EPL) == Int32(E): + tbl_vend[E] = base_data + tbl_data[E] = base_data + tbl_sf[E] = base_sf + + @cute.jit + def find_expert(self, tbl, key): + E = cutlass.const_expr(self.num_experts) + lo = Int32(0) + for step in self._search_steps: + probe = lo + Int32(step) + if cutlass.const_expr(self._search_needs_guard): + if probe < Int32(E) and Int32(tbl[probe]) <= key: + lo = probe + else: + if Int32(tbl[probe]) <= key: + lo = probe + return lo + + # ---------------------------------------------------------------- kernel + @cute.kernel + def ws_kernel( + self, + src_data: cute.Tensor, + src_sf_u8: cute.Tensor, + expert_token_sizes: cute.Tensor, + dst_data: cute.Tensor, + dst_sf_u8: cute.Tensor, + tma_atom=None, + tma_tensor=None, + tma_atom_st=None, + tma_tensor_st=None, + token_padding_block: cutlass.Constexpr = None, + ) -> None: + # Compile-time constant, folded before a single instruction is emitted. + TOKPAD = cutlass.const_expr( + self.token_padding_block if token_padding_block is None else token_padding_block + ) + self._require_token_padding_block(TOKPAD) + TOK = cutlass.const_expr(self.TILE_TOK) + W = cutlass.const_expr(self.TILE_HID) + S = cutlass.const_expr(self.NumStages) + SFB = cutlass.const_expr(self.SfTileBytes) + NCONS = cutlass.const_expr(self.ConsumerWarps) + table_len = cutlass.const_expr(self.num_experts + 1) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + grid_dim_x, _, _ = cute.arch.grid_dim() + warp_idx = tidx // Int32(32) + lane_idx = tidx % Int32(32) + + smem = cutlass.utils.SmemAllocator() + mbar_full = smem.allocate_array(cutlass.Int64, S) + mbar_empty = smem.allocate_array(cutlass.Int64, S) + tbl_vend = smem.allocate_tensor(cutlass.Int32, cute.make_layout((table_len,)), 16) + tbl_data = smem.allocate_tensor(cutlass.Int32, cute.make_layout((table_len,)), 16) + tbl_sf = smem.allocate_tensor(cutlass.Int32, cute.make_layout((table_len,)), 16) + smem_sf_in = smem.allocate_array(self.sf_dtype, S * SFB, byte_alignment=128) + smem_sf_out = smem.allocate_array( + self.sf_dtype, cutlass.const_expr(self.smem_sf_out_bytes), byte_alignment=128 + ) + smem_data = smem.allocate_array( + self.quant_dtype, self.NumStages * TOK * W, byte_alignment=128 + ) + smem_data_out = None + if cutlass.const_expr(self.dst_k_major): + smem_data_out = smem.allocate_array( + self.quant_dtype, TOK * W, byte_alignment=1024 + ) + + if tidx == Int32(0): + for s in cutlass.range_constexpr(0, S, 1): + cute.arch.mbarrier_init(mbar_full + s, 1) + cute.arch.mbarrier_init(mbar_empty + s, NCONS) + cute.arch.mbarrier_init_fence() + + if warp_idx == Int32(0): + self.build_prefix_tables( + expert_token_sizes, tbl_vend, tbl_data, tbl_sf, lane_idx, TOKPAD + ) + cute.arch.sync_threads() + + total_tiles = Int32(tbl_data[self.num_experts]) // Int32(TOK) + + smem_data_base = smem_data.toint() + smem_data_out_base = None + if cutlass.const_expr(self.dst_k_major): + smem_data_out_base = smem_data_out.toint() + smem_sf_in_base = smem_sf_in.toint() + smem_sf_out_base = smem_sf_out.toint() + src_sf_base = src_sf_u8.iterator.toint() + dst_sf_base = dst_sf_u8.iterator.toint() + + if warp_idx < Int32(self.ProducerWarps): + self.produce( + smem_data_base, smem_sf_in_base, mbar_full, mbar_empty, + tbl_data, tbl_sf, src_sf_base, + bidx, grid_dim_x, total_tiles, lane_idx, + tma_atom, tma_tensor, TOKPAD, + ) + else: + self.consume_scaled( + smem_data_base, smem_data_out_base, smem_sf_in_base, + smem_sf_out_base, mbar_full, mbar_empty, + tbl_vend, tbl_data, tbl_sf, dst_sf_base, + bidx, grid_dim_x, total_tiles, + warp_idx - Int32(self.ProducerWarps), lane_idx, + tma_atom_st, tma_tensor_st, TOKPAD, + ) + + # -------------------------------------------------------------- producer + @cute.jit + def produce( + self, smem_data_base, smem_sf_in_base, mbar_full, mbar_empty, + tbl_data, tbl_sf, src_sf_base, + bidx, grid_dim_x, total_tiles, lane_idx, + tma_atom=None, tma_tensor=None, + token_padding_block: cutlass.Constexpr = None, + ): + TOKPAD = cutlass.const_expr( + self.token_padding_block if token_padding_block is None else token_padding_block + ) + self._require_token_padding_block(TOKPAD) + TOK = cutlass.const_expr(self.TILE_TOK) + W = cutlass.const_expr(self.TILE_HID) + S = cutlass.const_expr(self.NumStages) + SFB = cutlass.const_expr(self.SfTileBytes) + # data bytes + SF bytes, both fenced by bar_full[stage] + EXPECT = cutlass.const_expr(TOK * W + self.SfTileXferBytes) + SF_PREFIX_DIFFERS = cutlass.const_expr(TOKPAD != self.sf_padding_block) + + BOX_H = cutlass.const_expr(self.TmaBoxHidU32) + sD = cute.make_tensor( + cute.make_ptr( + cutlass.Uint32, smem_data_base, AddressSpace.smem, assumed_align=128, + ), + cute.make_layout((TOK, BOX_H, S), stride=(BOX_H, 1, TOK * BOX_H)), + ) + gD = cute.group_modes( + cute.local_tile(tma_tensor, (TOK, BOX_H), (None, None)), 0, 2 + ) + tDsD, tDgD = cpasync.tma_partition( + tma_atom, 0, cute.make_layout(1), cute.group_modes(sD, 0, 2), gD, + ) + cpasync.prefetch_descriptor(tma_atom) + + t = Int32(0) + work_idx = Int32(bidx) + total_work = total_tiles * Int32(self.hidden_groups) + while work_idx < total_work: + stage = t % Int32(S) + token_tile = work_idx // Int32(self.hidden_groups) + hid_begin = (work_idx % Int32(self.hidden_groups)) * Int32(W) + + data_row0 = token_tile * Int32(TOK) + sf_row0 = data_row0 + if cutlass.const_expr(SF_PREFIX_DIFFERS): + owner = self.find_expert(tbl_data, data_row0) + sf_row0 = Int32(tbl_sf[owner]) + (data_row0 - Int32(tbl_data[owner])) + sf_row0 = cutlass.min( + sf_row0, Int32(tbl_sf[owner + Int32(1)]) - Int32(self.SfAtomNonK) + ) + + if t >= Int32(S): + cute.arch.mbarrier_wait(mbar_empty + stage, ((t // Int32(S)) - Int32(1)) % Int32(2)) + + if lane_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(mbar_full + stage, Int32(EXPECT)) + cute.arch.sync_warp() + + if lane_idx == Int32(0): + sf_src = ( + Int64(src_sf_base) + + Int64(sf_row0 // Int32(self.SfAtomNonK)) * Int64(self._hidden_atoms * self.SfAtomBytes) + + Int64(hid_begin // Int32(self.SfAtomNonK)) * Int64(self.SfAtomBytes) + ) + # One copy per SF atom: the atoms are padded apart in shared + # memory, so they are not one contiguous run. + for a in cutlass.range_constexpr(0, cutlass.const_expr(self.HidAtomsPerTile), 1): + tma_load_1d( + cute.make_ptr( + self.sf_dtype, + smem_sf_in_base + stage * Int32(SFB) + + Int32(a * self.SfInStride), + AddressSpace.smem, assumed_align=16, + ), + sf_src + Int64(a * self.SfAtomBytes), + mbar_full + stage, + Int32(self.SfAtomBytes), + ) + + cute.copy( + tma_atom, + tDgD[(None, token_tile, hid_begin // Int32(W))], + tDsD[(None, stage)], + tma_bar_ptr=mbar_full + stage, + ) + + t = t + Int32(1) + work_idx = work_idx + grid_dim_x + + # ------------------------------------------------- consumer (scaled cvt) + @cute.jit + def consume_scaled( + self, smem_data_base, smem_data_out_base, smem_sf_in_base, + smem_sf_out_base, mbar_full, mbar_empty, + tbl_vend, tbl_data, tbl_sf, dst_sf_base, + bidx, grid_dim_x, total_tiles, cw, lane_idx, + tma_atom_st=None, tma_tensor_st=None, + token_padding_block: cutlass.Constexpr = None, + ): + """The single-pass all-BF16 consumer, shared by every target.""" + TOKPAD = cutlass.const_expr( + self.token_padding_block if token_padding_block is None else token_padding_block + ) + self._require_token_padding_block(TOKPAD) + TOK = cutlass.const_expr(self.TILE_TOK) + W = cutlass.const_expr(self.TILE_HID) + S = cutlass.const_expr(self.NumStages) + SFB = cutlass.const_expr(self.SfTileBytes) + NB = cutlass.const_expr(self.TokensPerBlock) + CONS_THREADS = cutlass.const_expr(self.ConsumerWarps * 32) + HATOMS = cutlass.const_expr(self.HidAtomsPerTile) + SF_PREFIX_DIFFERS = cutlass.const_expr(TOKPAD != self.sf_padding_block) + C = cutlass.const_expr(self.ColsPerLane) # hidden columns per lane + SEGW = cutlass.const_expr(32 * C) # columns per consumer segment + + tb = cw // Int32(self.HidSegs) + seg = cw % Int32(self.HidSegs) + + LW = cutlass.const_expr(self.sp_lw) + ldsw = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=LW * 8 + ) + data_lane_off = tb * Int32(NB * W) + seg * Int32(SEGW) + lane_idx * Int32(LW) + + hb0 = seg * Int32(C) + (lane_idx * Int32(LW)) // Int32(32) + sf_lane_off = tb * Int32(4) + + if cutlass.const_expr(self.dst_k_major): + k_major_smem_layout = self._k_major_tma_smem_layout() + sDoT = cute.make_tensor( + cute.make_ptr( + cutlass.Uint8, + smem_data_out_base, + AddressSpace.smem, + assumed_align=128, + ), + k_major_smem_layout, + ) + gDo = cute.group_modes( + cute.local_tile(tma_tensor_st, (W, TOK), (None, None)), 0, 2 + ) + tDsDo, tDgDo = cpasync.tma_partition( + tma_atom_st, + 0, + cute.make_layout(1), + cute.group_modes(sDoT, 0, 2), + gDo, + ) + cpasync.prefetch_descriptor(tma_atom_st) + else: + BOX_HS = cutlass.const_expr(self.TmaBoxHidU32) + sDo = cute.make_tensor( + cute.make_ptr( + cutlass.Uint32, smem_data_base, AddressSpace.smem, assumed_align=128, + ), + cute.make_layout((TOK, BOX_HS, S), stride=(BOX_HS, 1, TOK * BOX_HS)), + ) + gDo = cute.group_modes( + cute.local_tile(tma_tensor_st, (TOK, BOX_HS), (None, None)), 0, 2 + ) + tDsDo, tDgDo = cpasync.tma_partition( + tma_atom_st, 0, cute.make_layout(1), cute.group_modes(sDo, 0, 2), gDo, + ) + cpasync.prefetch_descriptor(tma_atom_st) + + t = Int32(0) + work_idx = Int32(bidx) + total_work = total_tiles * Int32(self.hidden_groups) + while work_idx < total_work: + stage = t % Int32(S) + token_tile = work_idx // Int32(self.hidden_groups) + hid_begin = (work_idx % Int32(self.hidden_groups)) * Int32(W) + + data_row0 = token_tile * Int32(TOK) + owner = self.find_expert(tbl_data, data_row0) + valid_rows = cutlass.min(Int32(TOK), Int32(tbl_vend[owner]) - data_row0) + + # Destination SF is concatenated by expert, with token atoms + # contiguous inside each hidden atom. + sf_expert_token_atom = Int32(tbl_sf[owner]) // Int32(TOK) + sf_token_atom = ( + data_row0 - Int32(tbl_data[owner]) + ) // Int32(TOK) + sf_token_atoms = ( + Int32(tbl_sf[owner + Int32(1)]) - Int32(tbl_sf[owner]) + ) // Int32(TOK) + sf_live = Int32(1) + if cutlass.const_expr(SF_PREFIX_DIFFERS): + sf_live = cutlass.min(Int32(1), cutlass.max(Int32(0), valid_rows)) + + cute.arch.mbarrier_wait( + mbar_full + stage, (t // Int32(S)) % Int32(2) + ) + stage_data = ( + smem_data_base + stage * Int32(TOK * W) + data_lane_off + ) + stage_sf = ( + smem_sf_in_base + stage * Int32(SFB) + sf_lane_off + ) + sfout = smem_sf_out_base + + self._sp_tile_body( + stage_data, + stage_sf, + sfout, + smem_data_out_base, + ldsw, + hb0, + seg, + tb, + lane_idx, + valid_rows, + ) + + # Cross-proxy ordering, and it is NOT optional. The consumer warps + # wrote this tile's SMEM through the generic proxy; the store below + # reads it through the async proxy. Without this fence the store may + # observe stale bytes, and the kernel becomes nondeterministic: the + # same input yields different outputs from run to run. + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.barrier(barrier_id=self.ConsumerBarrierId, number_of_threads=CONS_THREADS) + + # The whole tile goes out, padding rows included; those were + # neutralised to zero in shared memory, which is what the pool + # expects to find there. + if cutlass.const_expr(self.dst_k_major): + if cw == Int32(0): + cute.copy( + tma_atom_st, + tDsDo, + tDgDo[ + ( + None, + hid_begin // Int32(W), + token_tile, + ) + ], + ) + else: + if cw == Int32(0): + cute.copy( + tma_atom_st, + tDsDo[(None, stage)], + tDgDo[(None, token_tile, hid_begin // Int32(W))], + ) + if cw == Int32(0) and lane_idx >= Int32(16) and lane_idx < Int32(16) + Int32(HATOMS) * sf_live: + atom = lane_idx - Int32(16) + cp_async_bulk_s2g( + Int64(dst_sf_base) + + ( + Int64(sf_expert_token_atom) * Int64(self._hidden_atoms) + + ( + Int64(hid_begin // Int32(self.SfAtomNonK)) + + Int64(atom) + ) + * Int64(sf_token_atoms) + + Int64(sf_token_atom) + ) + * Int64(self.SfAtomBytes), + cute.make_ptr( + self.sf_dtype, + sfout + atom * Int32(self.SfOutStride), + AddressSpace.smem, + assumed_align=16, + ), + Int32(self.SfAtomBytes), + ) + # Only consumer warp 0 issues async stores. Its wait followed by + # the CTA barrier makes completion visible to every consumer before + # either the output buffer or the producer stage is reused. + if cw == Int32(0): + cute.arch.cp_async_bulk_commit_group() + if cutlass.const_expr(self.dst_k_major): + # K-major has a separate output tile: release the input stage + # while its async store drains, overlapping the next TMA load. + if lane_idx == Int32(0): + cute.arch.mbarrier_arrive(mbar_empty + stage) + if cw == Int32(0): + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.barrier( + barrier_id=self.ConsumerBarrierId, + number_of_threads=CONS_THREADS, + ) + else: + # Row-major stores directly from the input stage, so it cannot + # be released until the async store has finished reading it. + if cw == Int32(0): + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.barrier( + barrier_id=self.ConsumerBarrierId, + number_of_threads=CONS_THREADS, + ) + if lane_idx == Int32(0): + cute.arch.mbarrier_arrive(mbar_empty + stage) + + t = t + Int32(1) + work_idx = work_idx + grid_dim_x + + @cute.jit + def _sp_tile_body( + self, + stage_data, + stage_sf_base, + sfout, + smem_data_out_base, + ldsw, + hb0, + seg, + tb, + lane_idx, + valid_rows, + ): + """The single-pass arithmetic for one lane's share of one tile.""" + TOK = cutlass.const_expr(self.TILE_TOK) + W = cutlass.const_expr(self.TILE_HID) + NB = cutlass.const_expr(self.TokensPerBlock) + C = cutlass.const_expr(self.ColsPerLane) + LW = cutlass.const_expr(self.sp_lw) + NWc = cutlass.const_expr(LW // 4) # 4-byte words per sub-tile row + NPc = cutlass.const_expr(LW // 2) # BF16x2 registers per sub-tile row + NCH = cutlass.const_expr(C // LW) # sub-tiles per lane per tile + SEGW = cutlass.const_expr(32 * C) # columns per consumer segment + QT = self.quant_dtype + + # Dead rows ARE reachable: token_padding_block constrains an expert's + # PADDED extent, not its valid count, so counts like 127,127,127 leave a + # padded tail in every expert's last tile. Neutralising them here needs + # no barrier: a consumer thread only reads the bytes it just wrote. + zeros = cute.make_rmem_tensor((NWc,), cutlass.Int32) + for w in cutlass.range_constexpr(0, NWc, 1): + zeros[w] = Int32(0) + + for ch in cutlass.range_constexpr(0, NCH, 1): + base = stage_data + Int32(ch * 32 * LW) + hbc = hb0 + Int32(ch * LW) + sfb = ( + stage_sf_base + + (hbc // Int32(4)) * Int32(self.SfInStride) + + hbc % Int32(4) + ) + + # The dead row's SCALE has to be neutralised as well as its data. + # ``cvt.rn.scaled::n2::ue8m0.bf16x2`` turns a NaN scale (raw 0xFF) + # into a NaN BF16 even from a zero payload. The amax survives -- + # ``max.xorsign.abs`` returns the non-NaN operand -- so the payload + # stays right and only the PADDING row is written back as 0x7F + # instead of 0x00. The fc1 pool really does leave 0xFF in padding + # scale bytes. Writing 127 is the same value the masked arm + # substitutes, paid once per tile instead of once per read. + dead_tt = cutlass.min( + Int32(NB), + cutlass.max(Int32(0), valid_rows - tb * Int32(NB)), + ) + while dead_tt < Int32(NB): + self._store_words( + base + dead_tt * Int32(W), ldsw, zeros, NWc, LW + ) + sf_t = cute.make_tensor( + cute.make_ptr( + cutlass.Uint8, sfb + dead_tt * Int32(16), + AddressSpace.smem, assumed_align=1, + ), + cute.make_layout((1,)), + ) + sf_t[0] = Uint8(127) + dead_tt = dead_tt + Int32(1) + + # ---- the one scan: unpack-with-scale, keep BF16, accumulate amax -- + d = [[None] * NPc for _ in range(NB)] + acc = [Int32(0)] * NPc + for tt in cutlass.range_constexpr(0, NB, 1): + # Both loads are issued before either is consumed, so the scale + # load overlaps the data load. + words = self._load_words( + base + Int32(tt * W), ldsw, NWc, LW + ) + raw_sf = self._src_scale_raw(sfb + Int32(tt * 16)) + s16 = raw_sf | (raw_sf << Int32(8)) + for w in cutlass.range_constexpr(0, NWc, 1): + qw = Int32(words[w]) + lo = cvt_scaled_up_bf16x2(qw, s16, 0, QT) + hi = cvt_scaled_up_bf16x2(qw, s16, 1, QT) + # Kept LIVE across the amax -- this is the single pass. + d[tt][2 * w] = lo + d[tt][2 * w + 1] = hi + acc[2 * w] = max_xorsign_abs_bf16x2(acc[2 * w], lo) + acc[2 * w + 1] = max_xorsign_abs_bf16x2(acc[2 * w + 1], hi) + + # max.xorsign.abs leaves junk in every sign bit. + raws = [None] * LW + for k in cutlass.range_constexpr(0, NPc, 1): + a = acc[k] & Int32(0x7FFF7FFF) + raws[2 * k] = e8m0_raw_from_bf16(a & Int32(0xFFFF), self._data_limit_exponent) + raws[2 * k + 1] = e8m0_raw_from_bf16( + (a >> Int32(16)) & Int32(0xFFFF), self._data_limit_exponent + ) + scs = raws + if cutlass.const_expr(self.scaled_cvt): + invs = None + else: + # The portable down-convert scales by an exact FP32 reciprocal + # instead of handing an E8M0 byte to the hardware. + invs = [ + bits_f32( + cutlass.max( + (Int32(254) - raws[j]) << Int32(23), Int32(0x400000) + ) + ) + for j in range(LW) + ] + + col0 = seg * Int32(SEGW) + Int32(ch * 32 * LW) + lane_idx * Int32(LW) + for j in cutlass.range_constexpr(0, LW, 1): + col = col0 + Int32(j) + off = ( + sfout + + (col // Int32(128)) * Int32(self.SfOutStride) + + (col % Int32(32)) * Int32(16) + + ((col % Int32(128)) // Int32(32)) * Int32(4) + + tb + ) + out_t = cute.make_tensor( + cute.make_ptr(cutlass.Uint8, off, AddressSpace.smem, assumed_align=1), + cute.make_layout((1,)), + ) + out_t[0] = Uint8(raws[j]) + + if cutlass.const_expr(self.dst_k_major): + for tt in cutlass.range_constexpr(0, NB, 16): + # Keep the two-token result column-major and stage it + # directly. The previous implementation first stored + # out0/out1 back to row-major SMEM and then reread every + # byte in a separate transpose pass. + rot = (lane_idx >> Int32(1)) & Int32(3) + swap_adjacent = (rot & Int32(1)) != Int32(0) + swap_halves = (rot & Int32(2)) != Int32(0) + adjacent_rotated = [ + cute.make_rmem_tensor((4,), cutlass.Int32) + for _ in range(LW) + ] + + # Produce only two token pairs at a time and immediately + # combine them into one four-token word. This avoids + # keeping all eight pair tensors live until a later pass. + for q in cutlass.range_constexpr(0, 4, 1): + pair0 = self._requant_token_pair_kmajor( + d[tt + 4 * q], + d[tt + 4 * q + 1], + scs, + invs, + NPc, + NWc, + LW, + ) + pair1 = self._requant_token_pair_kmajor( + d[tt + 4 * q + 2], + d[tt + 4 * q + 3], + scs, + invs, + NPc, + NWc, + LW, + ) + packed = [None] * LW + for j in cutlass.range_constexpr(0, LW, 1): + packed[j] = Int32( + cute.arch.prmt( + pair0[j], + pair1[j], + Int32(0x5410), + ) + ) + for j in cutlass.range_constexpr(0, LW, 1): + adjacent_rotated[j][q] = Int32( + arith.select( + swap_adjacent.ir_value(), + packed[j ^ 1].ir_value(), + packed[j].ir_value(), + ) + ) + + token = tb * Int32(NB) + Int32(tt) + # The adjacent-column exchange above applies rot bit 0. + # This final exchange applies rot bit 1, preserving the + # same XOR-rotated store order with one select per word. + for phase in cutlass.range_constexpr(0, LW, 1): + j_rot = Int32(phase) ^ rot + token_words = cute.make_rmem_tensor( + (4,), cutlass.Int32 + ) + for q in cutlass.range_constexpr(0, 4, 1): + token_words[q] = Int32( + arith.select( + swap_halves.ir_value(), + Int32( + adjacent_rotated[phase ^ 2][q] + ).ir_value(), + Int32( + adjacent_rotated[phase][q] + ).ir_value(), + ) + ) + self._store_kmajor_token_16( + smem_data_out_base, + col0 + j_rot, + token, + token_words, + ) + else: + for tt in cutlass.range_constexpr(0, NB, 2): + out0, out1 = self._requant_token_pair( + d[tt], d[tt + 1], scs, invs, NPc, NWc + ) + self._store_words(base + Int32(tt * W), ldsw, out0, NWc, LW) + self._store_words( + base + Int32((tt + 1) * W), ldsw, out1, NWc, LW + ) + + def _requant_token_pair_kmajor(self, d0, d1, scs, invs, NPc, NWc, LW): + """Return one packed token pair per hidden column for K-major staging.""" + QT = self.quant_dtype + pairs = cute.make_rmem_tensor((LW,), cutlass.Int32) + if cutlass.const_expr(self.scaled_cvt): + for k in range(0, NPc, 1): + lo = Int32(cute.arch.prmt(d0[k], d1[k], Int32(0x5410))) + hi = Int32(cute.arch.prmt(d0[k], d1[k], Int32(0x7632))) + pairs[2 * k] = cvt_scaled_dn_fp8x2(lo, scs[2 * k], QT) + pairs[2 * k + 1] = cvt_scaled_dn_fp8x2( + hi, scs[2 * k + 1], QT + ) + else: + for w in range(0, NWc, 1): + k0 = 2 * w + k1 = 2 * w + 1 + p00 = cvt_dn_fp8x2_portable( + d0[k0], invs[2 * k0], invs[2 * k0 + 1], QT + ) + p01 = cvt_dn_fp8x2_portable( + d0[k1], invs[2 * k1], invs[2 * k1 + 1], QT + ) + p10 = cvt_dn_fp8x2_portable( + d1[k0], invs[2 * k0], invs[2 * k0 + 1], QT + ) + p11 = cvt_dn_fp8x2_portable( + d1[k1], invs[2 * k1], invs[2 * k1 + 1], QT + ) + pair01 = Int32(cute.arch.prmt(p00, p10, Int32(0x5140))) + pair23 = Int32(cute.arch.prmt(p01, p11, Int32(0x5140))) + pairs[4 * w] = pair01 & Int32(0xFFFF) + pairs[4 * w + 1] = (pair01 >> Int32(16)) & Int32(0xFFFF) + pairs[4 * w + 2] = pair23 & Int32(0xFFFF) + pairs[4 * w + 3] = (pair23 >> Int32(16)) & Int32(0xFFFF) + return pairs + + @cute.jit + def _store_kmajor_token_16( + self, smem_data_out_base, col, token, token_words + ): + """Store 16 adjacent tokens with one 128-bit R2S copy.""" + linear = col * Int32(self.TILE_TOK) + token + offset = linear ^ ((linear & Int32(0x380)) >> Int32(3)) + st128 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Int32, + num_bits_per_copy=128, + ) + dst = cute.make_tensor( + cute.make_ptr( + cutlass.Int32, + smem_data_out_base + offset, + AddressSpace.smem, + assumed_align=16, + ), + cute.make_layout((4,)), + ) + cute.copy(st128, token_words, dst) + + def _requant_token_pair(self, d0, d1, scs, invs, NPc, NWc): + """Requantise one lane's two consecutive tokens for row-major output. + + ``d0``/``d1`` are lists of NPc Int32, each a token-major BF16x2: register + ``k`` is hidden columns ``(2k, 2k+1)`` of one token. ``out0``/``out1`` + come back as NWc packed FP8 words per token, in the byte order + ``_load_words`` read, so ``_store_words`` can put them straight back. + ``scs[j]`` is column ``j``'s output E8M0 raw byte and ``invs[j]`` its + exact FP32 reciprocal; each arm reads only the one it needs. + + Deliberately NOT ``@cute.jit``: it has to inline into the caller's trace. + A jitted callee would fail to marshal the Python lists and would change + the emitted IR. + + The transpose belongs inside this function: it is not a layout + preference but a consequence of ``scaled::n1`` taking a single scale for + both halves, which forces the two elements to share a hidden column + while SMEM is token-major. + """ + QT = self.quant_dtype + if cutlass.const_expr(self.scaled_cvt): + # Allocate the rmem tensors AFTER the cvt chain: the allocation + # order reaches the IR, so moving them changes the emitted code. + o = [None] * (2 * NPc) + # token-major -> column-major: o[j] = (col j of t0, t1) + for k in range(0, NPc, 1): + lo = Int32(cute.arch.prmt(d0[k], d1[k], Int32(0x5410))) + hi = Int32(cute.arch.prmt(d0[k], d1[k], Int32(0x7632))) + o[2 * k] = cvt_scaled_dn_fp8x2(lo, scs[2 * k], QT) + o[2 * k + 1] = cvt_scaled_dn_fp8x2(hi, scs[2 * k + 1], QT) + + # column-major -> token-major, one 4-byte word per token/w + out0 = cute.make_rmem_tensor((NWc,), cutlass.Int32) + out1 = cute.make_rmem_tensor((NWc,), cutlass.Int32) + for w in range(0, NWc, 1): + b = 4 * w + m01 = (o[b] & Int32(0xFFFF)) | (o[b + 1] << Int32(16)) + m23 = (o[b + 2] & Int32(0xFFFF)) | (o[b + 3] << Int32(16)) + out0[w] = Int32(cute.arch.prmt(m01, m23, Int32(0x6420))) + out1[w] = Int32(cute.arch.prmt(m01, m23, Int32(0x7531))) + else: + # Word w is columns 4w..4w+3; d[2w] is (4w, 4w+1) and d[2w+1] is + # (4w+2, 4w+3), and each helper returns byte 0 = its low column, so + # the OR below reproduces the byte order _load_words saw. + out0 = cute.make_rmem_tensor((NWc,), cutlass.Int32) + out1 = cute.make_rmem_tensor((NWc,), cutlass.Int32) + for w in range(0, NWc, 1): + k0 = 2 * w + k1 = 2 * w + 1 + p00 = cvt_dn_fp8x2_portable(d0[k0], invs[2 * k0], invs[2 * k0 + 1], QT) + p01 = cvt_dn_fp8x2_portable(d0[k1], invs[2 * k1], invs[2 * k1 + 1], QT) + p10 = cvt_dn_fp8x2_portable(d1[k0], invs[2 * k0], invs[2 * k0 + 1], QT) + p11 = cvt_dn_fp8x2_portable(d1[k1], invs[2 * k1], invs[2 * k1 + 1], QT) + out0[w] = (p00 & Int32(0xFFFF)) | (p01 << Int32(16)) + out1[w] = (p10 & Int32(0xFFFF)) | (p11 << Int32(16)) + return out0, out1 + + @cute.jit + def _load_words(self, byte_addr, atom, NW, C): + """One LDS of C raw bytes (no FP8 decode -- the scaled cvt does that).""" + regs = cute.make_rmem_tensor((NW,), cutlass.Int32) + cute.copy( + atom, + cute.make_tensor( + cute.make_ptr(cutlass.Int32, byte_addr, AddressSpace.smem, assumed_align=C), + cute.make_layout((NW,)), + ), + regs, + ) + return regs + + @cute.jit + def _store_words(self, byte_addr, atom, regs, NW, C): + cute.copy( + atom, + regs, + cute.make_tensor( + cute.make_ptr(cutlass.Int32, byte_addr, AddressSpace.smem, assumed_align=C), + cute.make_layout((NW,)), + ), + ) + + def _smem_bytes_for(self, tile_hid: int, stages: int) -> int: + """SMEM a (tile_hid, stages) pair would need, in bytes.""" + hid_atoms = tile_hid // self.SfAtomNonK + return ( + stages * self.TILE_TOK * tile_hid + + (self.TILE_TOK * tile_hid if self.dst_k_major else 0) + + stages * hid_atoms * (self.SfAtomBytes + self.sf_in_pad) + + hid_atoms * (self.SfAtomBytes + self.sf_out_pad) + + 3 * (self.num_experts + 1) * 4 + + 2 * stages * 8 + + (1024 if self.dst_k_major else 256) + ) + + @cute.jit + def _src_scale_raw(self, byte_addr): + """Raw source E8M0 byte, forced to 0..255. + + The mask is load-bearing: the pointer type says unsigned, but the + emitted load sign-extends, and the caller packs this into both halves + of a word with `raw | (raw << 8)`. Without the mask a scale byte >= + 0x80 poisons the E8M0 pair. Removing it took the mega suite from 28/28 + to 0/28. + """ + return ( + Int32( + cute.make_tensor( + cute.make_ptr(cutlass.Uint8, byte_addr, AddressSpace.smem, assumed_align=1), + cute.make_layout((1,)), + )[0] + ) + & Int32(0xFF) + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py new file mode 100644 index 000000000..bca96e800 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py @@ -0,0 +1,1648 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Autonomous epilogue for the fused fc1+fc2 swap-AB MegaMoE kernel. + +Component boundaries use ``TensorWithContract`` to keep per-thread RMEM layout +semantics explicit. See ``megamoe_design.md`` for the epilogue dataflow. +""" + +from typing import Optional, Tuple, Type, Union, Any +import dataclasses + +import cutlass +import cutlass.cute as cute + +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import AddressSpace +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils + +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith as _arith +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import dsl_user_op, Int32 as _epi_Int32, Int64 + +from cutlass.cute.typing import Float32 + +from ......helpers.iket_compat import iket +from ......helpers.flag_batch import GpuReleaseFlagBatchTracker +from ......helpers.ptx_helpers import ( + cp_async_bulk_s2g, + red_add_relaxed_sys_v2_bf16x2 as _red_add_relaxed_sys_v2_bf16x2, + stg_e8m0_from_f32, + stg_e8m0x8_from_f32, +) +from ..helpers.utils import quant_sfd_row, swiglu_act +from ......quant_def import CombineFormat +from ......communication.token_protocol import TokenSrcMetadata +from .....schedulers import BlockPhase + +# The 16x32 TMEM transpose core is architecture-neutral; reuse it via the local +# source-copy shim rather than re-porting the transpose math or reaching into the +# inference deliverable's directory. +from ..tmem_transpose import _TmemTranspose16x32Core + +Fc1GateUpInterleave = 32 +EpilogueTileN = 32 +Fc1EpilogueOutputTileM = 256 +Fc1EpilogueOutputTileN = 128 +WarpThreadCount = 32 +EpiWarpCount = 4 +Fc1CTMAStages = 1 + + +# ============================================================================= +# Fc2OutputDest (MegaMoE fc2 STG destination resolver, non-swap MXFP8 path) +# ============================================================================= + + +@dataclasses.dataclass(frozen=True) +class Fc2OutputDest: + """fc2 output destination in MoE-domain ``(token_max, topk, hidden)`` layout.""" + + tensor: cute.Tensor + metadata: Optional[cute.Tensor] = None + peer_rank_ptr_mapper: Any = None + reduce_topk_in_kernel: bool = False + + def __post_init__(self) -> None: + if (self.metadata is None) != (self.peer_rank_ptr_mapper is None): + raise ValueError( + "Fc2OutputDest: ``metadata`` and ``peer_rank_ptr_mapper`` must be " + "both None (direct mode) or both non-None (MegaMoE / indirect " + "mode). Got metadata=" + f"{'set' if self.metadata is not None else 'None'}, " + f"peer_rank_ptr_mapper=" + f"{'set' if self.peer_rank_ptr_mapper is not None else 'None'}." + ) + + @cute.jit + def resolve_token_row(self, pool_token_global) -> cute.Tensor: + """Return the ``(hidden,)`` BF16 GMEM row this pool token's STG lands on.""" + if cutlass.const_expr(self.metadata is None): + # Int64 token coord: the (token, topk, hidden) row offset is token*(topk*hidden), + # which overflows int32 once max_tokens*topk*hidden > 2^31 (mirrors inference). + return cute.slice_(self.tensor, (Int64(pool_token_global), 0, None)) + + md = TokenSrcMetadata.load( + self.metadata.iterator.toint() + + Int64(pool_token_global) * Int64(TokenSrcMetadata.nbytes) + ) + src_rank = md.src_rank + src_token = md.src_token + if cutlass.const_expr(self.reduce_topk_in_kernel): + src_topk = cutlass.Int32(0) + else: + src_topk = md.src_topk + # Int64 token coord: src_token*(topk*hidden) overflows int32 once + local_row = cute.slice_(self.tensor, (Int64(src_token), src_topk, None)) + # next's SymmetricBufferDevice exposes ``map_pointer`` + peer_iter = self.peer_rank_ptr_mapper.map_pointer( + local_row.iterator, src_rank, + ) + return cute.make_tensor(peer_iter, local_row.layout) + + +# ============================================================================= +# GluMxfp8Epilogue +# ============================================================================= + +class GluMxfp8Epilogue: + + _SubtileBarIdBase = 4 + # Named barrier for cross-warp sync during raw-C TMA stores. + _CStoreBarId = 10 + + def __init__( + self, + *, + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + use_2cta_instrs: bool, + sf_vec_size: int, + fc1_output_dtype: Type[cutlass.Numeric], + fc1_output_layout: utils.LayoutEnum, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + sf_dtype: Type[cutlass.Numeric] = cutlass.Float8E8M0FNU, + c_dtype: Type[cutlass.Numeric] = cutlass.BFloat16, + glu_clamp: Optional[float] = None, + epilog_sync_bar_id: int = 1, + epilogue_warp_ids: Tuple[int, ...] = (0, 1, 2, 3), + static_expert_shape: Optional[Tuple[int, int, int]] = None, + fc2_in_kernel_topk_reduce: bool = False, + token_back_by_dispatch: bool = False, + epi_flag_batch: Optional[Tuple[int, int]] = (1, 1), + apply_topk_in_fc1: bool = False, + generate_c: bool = False, + use_stg_fc1: bool = False, + combine_format: Optional[Any] = None, + act_func: str = "swiglu", + fc2_use_bulk: bool = False, + fc2_tma_stages: Optional[int] = None, + ) -> None: + self._act_func = act_func + self.fc1_output_dtype = fc1_output_dtype + self.fc1_output_layout = fc1_output_layout + self.acc_dtype = acc_dtype + self.sf_dtype = sf_dtype + self._sf_vec_size = sf_vec_size + self._c_dtype = c_dtype + self._epilog_sync_bar_id = epilog_sync_bar_id + self._epilogue_warp_ids = epilogue_warp_ids + self._use_2cta_instrs = use_2cta_instrs + + self._atom_thr_size = 2 if use_2cta_instrs else 1 + self._cta_tile_m = mma_tiler_mnk[0] // self._atom_thr_size + self._cta_tile_n = mma_tiler_mnk[1] + self._mma_tiler_k = mma_tiler_mnk[2] + self._cta_tile_n_sfb = ((mma_tiler_mnk[1] + 127) // 128) * 128 + self._static_expert_shape = static_expert_shape + if ( + static_expert_shape is not None + and static_expert_shape[2] % (self._cta_tile_m * cluster_shape_mn[0]) == 0 + ): + self._fc2_stg_needs_predicate: bool = False + else: + self._fc2_stg_needs_predicate: bool = True + + # TMA tile is (EpilogueTileN tokens, Fc1EpilogueOutputTileN intermediates) + self._epi_tile = (EpilogueTileN, Fc1EpilogueOutputTileN) + self._subtile_cnt = self._cta_tile_n // 2 // EpilogueTileN + + self._num_acc_stage = 2 + self._num_acc_pipeline_stages = self._num_acc_stage + + k = self._mma_tiler_k + self._num_sfa_tmem_cols = self._cta_tile_m * k // sf_vec_size * 4 // 4 // 128 + self._num_sfb_tmem_cols = ( + self._cta_tile_n_sfb * k // sf_vec_size * 4 // 4 // 128 + ) + + self._num_accumulator_tmem_cols = self._cta_tile_n * self._num_acc_stage + + self._fc2_in_kernel_topk_reduce = fc2_in_kernel_topk_reduce + self._token_back_by_dispatch = token_back_by_dispatch + self._apply_topk_in_fc1 = apply_topk_in_fc1 + self._generate_c = generate_c + self._use_stg_fc1 = use_stg_fc1 + # combine_format determines fc2 output encoding: bf16 (default) or quantized. + if combine_format is None: + combine_format = CombineFormat.parse("bf16") + self._combine_format = combine_format + self._combine_mxfp8 = combine_format.is_quantized + # sf_block_pad for fc2 MXFP8 combine + if self._combine_mxfp8 and static_expert_shape is not None: + _hidden_fc2 = static_expert_shape[2] + _sf_blocks_fc2 = _hidden_fc2 // EpilogueTileN + self._fc2_sf_block_pad = ((_sf_blocks_fc2 + 15) // 16) * 16 + self._hidden_fc2 = _hidden_fc2 + else: + self._fc2_sf_block_pad = 0 + self._hidden_fc2 = 0 + # batching stg.64 SF + self._fc2_sf_batch8 = ( + self._combine_mxfp8 + and self._hidden_fc2 > 0 + and (self._hidden_fc2 % self._cta_tile_n == 0) + and (self._cta_tile_n // EpilogueTileN == 8) + ) + self._epi_tile_c = (self._cta_tile_m, 2 * Fc1GateUpInterleave) + self._epi_fc1_batch = max(1, epi_flag_batch[0]) + self._epi_fc2_batch = max(1, epi_flag_batch[1]) + + self.glu_clamp = ( + cutlass.Float32(glu_clamp) if glu_clamp is not None else None + ) + + self._fc2_subtile_cnt = self._cta_tile_n // EpilogueTileN # = 8 + self._fc2_use_tma = ( + bool(fc2_use_bulk) and token_back_by_dispatch and self._combine_mxfp8 + ) + _ublk_hidden_ok = ( + self._combine_mxfp8 + and static_expert_shape is not None + and self._hidden_fc2 % self._cta_tile_n == 0 + ) + self._fc2_use_ublk = ( + bool(fc2_use_bulk) and (not token_back_by_dispatch) and _ublk_hidden_ok + ) + self._fc2_needs_staging = self._fc2_use_tma or self._fc2_use_ublk + if fc2_tma_stages is not None and not 1 <= fc2_tma_stages <= self._fc2_subtile_cnt: + raise ValueError( + f"fc2_tma_stages must be in [1, {self._fc2_subtile_cnt}], got {fc2_tma_stages}." + ) + if self._fc2_needs_staging: + self._fc2_tma_stages = ( + fc2_tma_stages if fc2_tma_stages is not None else min(2, self._fc2_subtile_cnt) + ) + else: + self._fc2_tma_stages = 0 + self._fc2_wire_dtype = ( + self._combine_format.act_dtype if self._combine_mxfp8 else cutlass.BFloat16 + ) + self._fc2_tma_stage_bytes = ( + self._cta_tile_m * EpilogueTileN * self._fc2_wire_dtype.width // 8 + ) + self._fc2_tma_staging_bytes = self._fc2_tma_stages * self._fc2_tma_stage_bytes + + self._fc2_reduce_coalesce = bool(self._fc2_in_kernel_topk_reduce) + # bf16 staging tile: (cta_tile_m tokens x EpilogueTileN hidden), 2 B/elem. + self._fc2_reduce_staging_bytes = ( + self._cta_tile_m * EpilogueTileN * cutlass.BFloat16.width // 8 + if self._fc2_reduce_coalesce + else 0 + ) + + # -- Codegen-time queries -- + + @property + def epi_tile(self) -> Tuple[int, int]: + return self._epi_tile + + @property + def num_acc_pipeline_stages(self) -> int: + return self._num_acc_pipeline_stages + + @property + def num_acc_stage(self) -> int: + return self._num_acc_stage + + @property + def subtile_cnt(self) -> int: + return self._subtile_cnt + + @property + def cta_tile_n(self) -> int: + return self._cta_tile_n + + @property + def num_sfa_tmem_cols(self) -> int: + return self._num_sfa_tmem_cols + + @property + def num_sfb_tmem_cols(self) -> int: + return self._num_sfb_tmem_cols + + @property + def num_accumulator_tmem_cols(self) -> int: + return self._num_accumulator_tmem_cols + + def staged_smem_layout( + self, + n_stages: int, + ) -> Union[cute.Layout, cute.ComposedLayout]: + return sm100_utils.make_smem_layout_epi( + self.fc1_output_dtype, + self.fc1_output_layout, + self._epi_tile, + n_stages, + ) + + @property + def smem_layout_one_stage(self) -> Union[cute.Layout, cute.ComposedLayout]: + staged = self.staged_smem_layout(1) + return cute.select(staged, mode=[0, 1]) + + @property + def bytes_per_stage(self) -> int: + return cute.size_in_bytes(self.fc1_output_dtype, self.smem_layout_one_stage) + + @property + def epi_tile_c(self) -> Tuple[int, int]: + """TMA tile for raw gate+up output: (cta_tile_m=128, 2*Fc1GateUpInterleave=64) fp32.""" + return self._epi_tile_c + + def staged_c_smem_layout(self, n_stages: int): + """SMEM layout for n_stages of raw gate+up (Float32, row-major, epi_tile_c).""" + return sm100_utils.make_smem_layout_epi( + self._c_dtype, + self.fc1_output_layout, # same N-major direction as fc1 output + self._epi_tile_c, + n_stages, + ) + + @property + def c_smem_layout_one_stage(self): + return cute.select(self.staged_c_smem_layout(1), mode=[0, 1]) + + @property + def c_bytes_per_stage(self) -> int: + return cute.size_in_bytes(self._c_dtype, self.c_smem_layout_one_stage) + + # ── FC2 TMASTG staging ─────────────────────────────────────────────── + @property + def fc2_use_tma(self) -> bool: + return self._fc2_use_tma + + @property + def fc2_use_ublk(self) -> bool: + return self._fc2_use_ublk + + @property + def fc2_needs_staging(self) -> bool: + """True when the FC2 store stages into the SMEM tile: either the dispatch + TMASTG path (``fc2_use_tma``) or the epi_warps UBLK path (``fc2_use_ublk``).""" + return self._fc2_needs_staging + + @property + def fc2_tma_stages(self) -> int: + return self._fc2_tma_stages + + @property + def fc2_tma_staging_bytes(self) -> int: + return self._fc2_tma_staging_bytes + + @property + def fc2_tma_tile(self) -> Tuple[int, int]: + """TMA store tile: (cta_tile_m=128 tokens, EpilogueTileN=32 hidden).""" + return (self._cta_tile_m, EpilogueTileN) + + def fc2_tma_staged_smem_layout(self, n_stages: int): + """Row-major (128 tokens, 32 hidden, n_stages) staging tile for the FC2 + bulk store. No swizzle: each thread owns one token row and writes its 32 + contiguous wire-dtype elements as a single 256-bit STS, and a 32-byte + innermost box is a valid (unswizzled) TMA tile. ``select(mode=[0,1])`` + yields the single-stage tile used to build the TMA atom.""" + cta_tile_m = self._cta_tile_m + stage_stride = cta_tile_m * EpilogueTileN + layout = cute.make_layout( + (cta_tile_m, EpilogueTileN, n_stages), + stride=(EpilogueTileN, 1, stage_stride if n_stages > 1 else 0), + ) + return layout + + @property + def fc2_tma_smem_layout_one_stage(self): + return cute.select(self.fc2_tma_staged_smem_layout(1), mode=[0, 1]) + + # ── FC2 in-kernel reduce coalescing ────────────────────────────────── + @property + def fc2_reduce_coalesce(self) -> bool: + return self._fc2_reduce_coalesce + + @property + def fc2_reduce_staging_bytes(self) -> int: + return self._fc2_reduce_staging_bytes + + def fc2_reduce_smem_layout(self): + """Row-major (cta_tile_m tokens, EpilogueTileN hidden) bf16 transpose tile. + Each epi thread writes its token's contiguous hidden row (64 B STS); the + coalesced-issue re-partition reads 4-elem chunks hidden-major.""" + return cute.make_layout( + (self._cta_tile_m, EpilogueTileN), stride=(EpilogueTileN, 1) + ) + + @staticmethod + @cute.jit + def tma_store_fc1_output( + warp_idx, + sC, + store_idx, + tma_atom_fc1_output: cute.CopyAtom, + g_fc1_output_subtile_view: cute.Tensor, + valid_tokens, + ) -> None: + """Per-warp TMA store for FC1 output.""" + cute.arch.fence_proxy("async.shared", space="cta") + sC_stage = cute.slice_(sC, (None, None, store_idx)) + g_fc1_output_2d = cute.slice_(g_fc1_output_subtile_view, (None, None, 0)) + bSG_sC, bSG_g = cpasync.tma_partition( + tma_atom_fc1_output, + 0, + cute.make_layout(1), + cute.group_modes(sC_stage, 0, 2), + cute.group_modes(g_fc1_output_2d, 0, 2), + ) + + leader_warp = store_idx + tile_has_valid = ( + store_idx * cutlass.Int32(EpilogueTileN) < valid_tokens + ) + + bar_id = store_idx + cutlass.Int32(GluMxfp8Epilogue._SubtileBarIdBase) + bar = pipeline.NamedBarrier( + barrier_id=bar_id, + num_threads=EpiWarpCount * WarpThreadCount, + ) + if warp_idx == leader_warp: + bar.arrive_and_wait() + # TMA bulk-tensor stores are all-or-nothing per issue (no per-element + # `pred` mask, no scalar predicate arg), so guard the whole copy with + # a runtime `if`. Skips fully-padding token-tiles that would alias + # the next expert's region. + if tile_has_valid: + cute.copy(tma_atom_fc1_output, bSG_sC, bSG_g) + else: + bar.arrive() + + @cute.jit + def _store_fc1_c_subtile( + self, + r_gate: cute.Tensor, + r_up: cute.Tensor, + smem_c_buffer: cute.Tensor, + tma_atom_c: cute.CopyAtom, + gmem_c_subtile_view: cute.Tensor, + c_buffer_idx, + work_tile_info, + warp_idx: int, + tidx, + c_pipeline, + ) -> None: + """Store pre-SwiGLU gate/up accumulators to the global C tensor via SMEM staging.""" + r_layout = cute.make_layout((((Fc1GateUpInterleave,), 1),), stride=(((1,), 0),)) + + # Cast acc_dtype (Float32) → c_dtype (e.g. BFloat16) before R2S. + r_gate_c = cute.make_rmem_tensor(r_layout.shape, self._c_dtype) + r_up_c = cute.make_rmem_tensor(r_layout.shape, self._c_dtype) + r_gate_c.store(r_gate.load().to(self._c_dtype)) + r_up_c.store(r_up.load().to(self._c_dtype)) + + r2s_c_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self._c_dtype, num_bits_per_copy=128, + ) + thread_in_warp_c = tidx % cutlass.Int32(WarpThreadCount) + c_row = cutlass.Int32(warp_idx * EpilogueTileN) + thread_in_warp_c + sC_raw_stage = cute.slice_(smem_c_buffer, (None, None, c_buffer_idx)) + c_gate_smem = cute.local_tile( + sC_raw_stage, (1, Fc1GateUpInterleave), (c_row, cutlass.Int32(0)), + ) + cute.copy(r2s_c_atom, cute.coalesce(r_gate_c), cute.coalesce(c_gate_smem)) + c_up_smem = cute.local_tile( + sC_raw_stage, (1, Fc1GateUpInterleave), (c_row, cutlass.Int32(1)), + ) + cute.copy(r2s_c_atom, cute.coalesce(r_up_c), cute.coalesce(c_up_smem)) + + # Fence + barrier: ensure all warps have written before TMA issue. + cute.arch.fence_proxy("async.shared", space="cta") + c_store_bar = pipeline.NamedBarrier( + barrier_id=self._CStoreBarId, + num_threads=EpiWarpCount * WarpThreadCount, + ) + c_store_bar.arrive_and_wait() + + # Warp 0 issues TMA S2G, commits, then pre-acquires the next stage so + # smem is guaranteed free before the next call's R2S writes. + if warp_idx == 0: + if work_tile_info.valid_tokens_in_cta_tile > cutlass.Int32(0): + g_c_2d = cute.slice_(gmem_c_subtile_view, (None, None, 0)) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, 0, cute.make_layout(1), + cute.group_modes(sC_raw_stage, 0, 2), + cute.group_modes(g_c_2d, 0, 2), + ) + cute.copy(tma_atom_c, bSG_sC, bSG_gC) + c_pipeline.producer_commit() + + def _subtile_local_tmem_tensor_pair( + self, + tmem_acc_tensor: cute.Tensor, + subtile_idx, + warp_idx, + ) -> cute.Tensor: + """ + Build a (gate, up) pair of TMEM tensor views for the MXFP8 fc1 epilogue. + """ + base = tmem_acc_tensor.iterator + warp_lane_off = warp_idx * WarpThreadCount + subtile_col_off = subtile_idx * EpilogueTileN * 2 + total = (warp_lane_off << 16) + subtile_col_off + subtile_gate_ptr = base + cute.assume(total, divby=16) + subtile_up_ptr = base + cute.assume(total + Fc1GateUpInterleave, divby=16) + return ( + cute.make_tensor( + subtile_gate_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ), + cute.make_tensor( + subtile_up_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ) + ) + + def _subtile_forward_tmem_tensor( + self, + tmem_gate_tensor: cute.Tensor, + tmem_up_tensor: cute.Tensor, + col_offset: int, + ) -> (cute.Tensor, cute.Tensor): + """Move the tmem tensor to the correct position.""" + tmem_gate_ptr = tmem_gate_tensor.iterator + cute.assume(col_offset, divby=16) + tmem_up_ptr = tmem_up_tensor.iterator + cute.assume(col_offset, divby=16) + return ( + cute.make_tensor( + tmem_gate_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ), + cute.make_tensor( + tmem_up_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ) + ) + + # -- fc1 subtile: SM100 path -- + @cute.jit + def _run_fc1_task_tile( + self, + work_tile_info, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + acc_consumer_state, + smem_fc1_output_buffer: cute.Tensor, + tma_atom_fc1_output: cute.CopyAtom, + sched_ext, + gmem_fc1_output: cute.Tensor, + gmem_fc1_output_sf: cute.Tensor, + gmem_topk_scores: cute.Tensor, + warp_idx: int, + tidx, + alpha, + norm_const, + smem_c_buffer: cute.Tensor, + tma_atom_c: cute.CopyAtom, + gmem_c: cute.Tensor, + c_pipeline, + ) -> None: + """MXFP8 fc1 task-tile: TmemTranspose16x32 TMEM loading + cross-warp E8M0 exchange.""" + real_fc1_output, _ = sched_ext.get_gmem_tensor("d", gmem_fc1_output, work_tile_info) + if cutlass.const_expr(self.fc1_output_dtype.width == 8): + real_fc1_output_sf, _ = sched_ext.get_gmem_tensor("sfd", gmem_fc1_output_sf, work_tile_info) + else: + real_fc1_output_sf = None + + real_topk_scores = gmem_topk_scores + if cutlass.const_expr(self._apply_topk_in_fc1): + real_topk_scores, _ = sched_ext.get_gmem_tensor( + "topk", gmem_topk_scores, work_tile_info + ) + + if cutlass.const_expr(self._generate_c): + real_c, _ = sched_ext.get_gmem_tensor("c", gmem_c, work_tile_info) + c_n_base = work_tile_info.tile_n_idx * cutlass.Int32(self._subtile_cnt) + + acc_pipeline.consumer_wait(acc_consumer_state) + if warp_idx == 0: + iket.range_push("fc1_epi_tile") + + subtile_cnt = self._subtile_cnt + tmem_gate, tmem_up = self._subtile_local_tmem_tensor_pair( + tmem_acc_tensor, 0, warp_idx, + ) + tmem_forward_cols = Fc1GateUpInterleave * 2 + + if cutlass.const_expr(self.fc1_output_dtype.width == 8): + layout_sf = cute.make_layout(4) + rmem_sf = cute.make_rmem_tensor(layout_sf.shape, self.acc_dtype) + else: + rmem_sf = None + + # Set up RMEM→SMEM copy atom (direct CopyUniversalOp) + r2s_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.fc1_output_dtype, + num_bits_per_copy=128, + ) + tRS_sC = None + + for i in cutlass.range(0, subtile_cnt, 1, unroll=1): + subtile_idx = cutlass.Int32(i) + + if cutlass.const_expr(self._generate_c): + c_buffer_idx = cutlass.Int32(i % Fc1CTMAStages) + g_c_subtile = cute.local_tile( + real_c, + (self._cta_tile_m, 2 * Fc1GateUpInterleave, 1), + (work_tile_info.tile_m_idx, c_n_base + subtile_idx, cutlass.Int32(0)), + ) + smem_c_buf_arg = smem_c_buffer + tma_atom_c_arg = tma_atom_c + gmem_c_subtile_arg = g_c_subtile + else: + c_buffer_idx = cutlass.Int32(0) + smem_c_buf_arg = smem_fc1_output_buffer + tma_atom_c_arg = tma_atom_fc1_output + gmem_c_subtile_arg = real_fc1_output + + self._run_fc1_subtile( + subtile_idx=subtile_idx, + tmem_gate_tensor=tmem_gate, + tmem_up_tensor=tmem_up, + real_fc1_output=real_fc1_output, + real_fc1_output_sf=real_fc1_output_sf, + real_topk_scores=real_topk_scores, + work_tile_info=work_tile_info, + smem_fc1_output_buffer=smem_fc1_output_buffer, + tma_atom_fc1_output=tma_atom_fc1_output, + r2s_copy_atom=r2s_copy_atom, + warp_idx=warp_idx, + tidx=tidx, + alpha=alpha, + norm_const=norm_const, + rmem_sf=rmem_sf, + smem_c_buffer=smem_c_buf_arg, + tma_atom_c=tma_atom_c_arg, + gmem_c_subtile_view=gmem_c_subtile_arg, + c_buffer_idx=c_buffer_idx, + c_pipeline=c_pipeline, + ) + + tmem_gate, tmem_up = self._subtile_forward_tmem_tensor(tmem_gate, tmem_up, tmem_forward_cols) + + self._acc_pipeline_consumer_release(acc_pipeline, acc_consumer_state, True) + + if cutlass.const_expr(self.fc1_output_dtype.width == 8): + self._stg_sf_fc1(rmem_sf, real_fc1_output_sf, work_tile_info, tidx) + + # TMA store: 4 stores (one per warp group) after all subtiles + if cutlass.const_expr(not self._use_stg_fc1): + base_token_tile = ( + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m // EpilogueTileN) + ) + for idx in cutlass.range_constexpr(EpiWarpCount): + g_fc1_output_warp_view = cute.local_tile( + real_fc1_output, + (EpilogueTileN, Fc1EpilogueOutputTileN, 1), + (base_token_tile + idx, work_tile_info.tile_n_idx, 0), + ) + GluMxfp8Epilogue.tma_store_fc1_output( + warp_idx, smem_fc1_output_buffer, idx, + tma_atom_fc1_output, g_fc1_output_warp_view, + work_tile_info.valid_tokens_in_cta_tile, + ) + + if warp_idx == 0: + iket.range_pop() + + @cute.jit + def _swiglu_act( + self, + t_swiglu: cute.Tensor, + t_up: cute.Tensor, + t_gate: cute.Tensor, + prob: Optional[Float32] = None, + ) -> None: + """SwiGLU hook consumed by ``_run_fc1_subtile`` """ + swiglu_act(t_swiglu, t_up, t_gate, prob) + + @cute.jit + def _run_fc1_subtile( + self, + subtile_idx, + tmem_gate_tensor: cute.Tensor, + tmem_up_tensor: cute.Tensor, + real_fc1_output: cute.Tensor, + real_fc1_output_sf: cute.Tensor, + real_topk_scores: cute.Tensor, + work_tile_info, + smem_fc1_output_buffer: cute.Tensor, + tma_atom_fc1_output: cute.CopyAtom, + r2s_copy_atom: cute.CopyAtom, + warp_idx: int, + tidx, + alpha, + norm_const, + rmem_sf: cute.Tensor, + smem_c_buffer: cute.Tensor, + tma_atom_c: cute.CopyAtom, + gmem_c_subtile_view: cute.Tensor, + c_buffer_idx, + c_pipeline, + ) -> None: + """MXFP8 fc1 subtile: GLU + E8M0 SF + fp8 R2S.""" + if warp_idx == 0: + iket.range_push("fc1_epilogue_subtile") + + r_layout = cute.make_layout((((Fc1GateUpInterleave,), 1),), stride=(((1,), 0),)) + r_gate = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + r_up = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + + atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), self.acc_dtype, + ) + cute.copy(atom_t2r, tmem_gate_tensor, r_gate) + cute.copy(atom_t2r, tmem_up_tensor, r_up) + + # ── generate_c: store raw gate+up to GMEM C tensor via SMEM staging ── + if cutlass.const_expr(self._generate_c): + self._store_fc1_c_subtile( + r_gate=r_gate, + r_up=r_up, + smem_c_buffer=smem_c_buffer, + tma_atom_c=tma_atom_c, + gmem_c_subtile_view=gmem_c_subtile_view, + c_buffer_idx=c_buffer_idx, + work_tile_info=work_tile_info, + warp_idx=warp_idx, + tidx=tidx, + c_pipeline=c_pipeline, + ) + + if cutlass.const_expr(self.glu_clamp is not None): + for i in cutlass.range_constexpr(cute.size(r_up)): + r_gate[i] = cute.arch.fmin(r_gate[i], self.glu_clamp) + r_up[i] = cute.arch.fmin(r_up[i], self.glu_clamp) + r_up[i] = cute.arch.fmax(r_up[i], -self.glu_clamp) + + topk = None + if cutlass.const_expr(self._apply_topk_in_fc1): + thread_in_warp = tidx % cutlass.Int32(WarpThreadCount) + token_in_tile = ( + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + ) + topk = Float32(real_topk_scores[token_in_tile]) + + swiglu = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + if cutlass.const_expr(self._act_func == "swiglu"): + self._swiglu_act(swiglu, r_up, r_gate, topk) + + c = cute.make_rmem_tensor(r_layout.shape, self.fc1_output_dtype) + if cutlass.const_expr(self.fc1_output_dtype.width == 8): + # Quantized hand-off: fp8 data + E8M0 block scale. + qpvscale = quant_sfd_row(swiglu, c, norm_const, self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype) + if subtile_idx == 0: + rmem_sf[0] = qpvscale + elif subtile_idx == 1: + rmem_sf[1] = qpvscale + elif subtile_idx == 2: + rmem_sf[2] = qpvscale + elif subtile_idx == 3: + rmem_sf[3] = qpvscale + else: + # Plain-data hand-off: direct cast to the fc1 output dtype (the + # fc1_output workspace is reloaded as fc2's A operand). + c.store(swiglu.load().to(self.fc1_output_dtype)) + + thread_in_warp = tidx % WarpThreadCount + if cutlass.const_expr(self._use_stg_fc1): + # Direct STG.256 to GMEM — no SMEM staging or TMA store needed. + token_in_tile = cutlass.Int32(warp_idx * EpilogueTileN) + thread_in_warp + if token_in_tile < work_tile_info.valid_tokens_in_cta_tile: + abs_token = ( + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + cutlass.Int32(warp_idx * EpilogueTileN) + thread_in_warp + ) + # absolute column start (element index in the intermediate axis) + col_elem = ( + (work_tile_info.tile_n_idx * cutlass.Int32(self._subtile_cnt) + subtile_idx) + * cutlass.Int32(Fc1GateUpInterleave) + ) + g_base = cute.local_tile( + real_fc1_output, + (1, 1, 1), + (abs_token, col_elem, cutlass.Int32(0)), + ) + stg_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.fc1_output_dtype, num_bits_per_copy=256, + ) + # col_elem is always a multiple of Fc1GateUpInterleave=32 (FP8 elements) + aligned_iter = cute.make_ptr( + self.fc1_output_dtype, + g_base.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=32, + ) + g_vec = cute.make_tensor(aligned_iter, cute.make_layout(Fc1GateUpInterleave)) + cute.copy(stg_atom, cute.coalesce(c), g_vec) + else: + sC_stage = cute.slice_(smem_fc1_output_buffer, (None, None, warp_idx)) + sC_thread_row = cute.local_tile( + sC_stage, (1, Fc1GateUpInterleave), (thread_in_warp, subtile_idx) + ) + cute.copy(r2s_copy_atom, cute.coalesce(c), cute.coalesce(sC_thread_row)) + + + if cutlass.const_expr(self._generate_c): + if warp_idx == 0: + c_pipeline.producer_acquire() + c_store_bar = pipeline.NamedBarrier( + barrier_id=self._CStoreBarId, + num_threads=EpiWarpCount * WarpThreadCount, + ) + c_store_bar.arrive_and_wait() + + if warp_idx == 0: + iket.range_pop() + + @cute.jit + def _subtile_fc2_tmem_tensor( + self, + tmem_acc_tensor: cute.Tensor, + subtile_idx, + warp_idx, + ) -> cute.Tensor: + """ + Per-warp TMEM view for one fc2 subtile (EpilogueTileN=32 cols). + """ + base = tmem_acc_tensor.iterator + warp_lane_off = warp_idx * WarpThreadCount + subtile_col_off = subtile_idx * EpilogueTileN + total = (warp_lane_off << 16) + subtile_col_off + subtile_ptr = base + cute.assume(total, divby=16) + return cute.make_tensor( + subtile_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ) + + @cute.jit + def _advance_fc2_tmem_tensor( + self, + tmem_tensor: cute.Tensor, + col_offset: int, + ) -> cute.Tensor: + """Advance the fc2 TMEM tensor by col_offset cols (mirrors _subtile_forward_tmem_tensor).""" + new_ptr = tmem_tensor.iterator + cute.assume(col_offset, divby=16) + return cute.make_tensor( + new_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ) + + @cute.jit + def _acc_pipeline_consumer_release( + self, + acc_pipeline, + acc_consumer_state, + is_release: bool, + ) -> None: + """Release the acc pipeline consumer.""" + if is_release: + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + @cute.jit + def _run_fc2_subtile( + self, + subtile_idx, + tmem_subtile_tensor: cute.Tensor, + real_fc2_output: cute.Tensor, + work_tile_info, + valid_hidden, + warp_idx: int, + tidx, + token_comm_args=None, + rmem_sf_fc2=None, + smem_fc2_tma_buffer=None, + tma_atom_fc2_output=None, + gmem_fc2_tma_output=None, + smem_fc2_reduce_buffer=None, + ) -> None: + """fc2 subtile: LDTM + encode + STG.""" + if warp_idx == 0: + iket.range_push("fc2_epi_subtile") + + fc2_subtile_cnt = self._cta_tile_n // EpilogueTileN # = 8 + hidden_group = ( + work_tile_info.tile_n_idx * cutlass.Int32(fc2_subtile_cnt) + subtile_idx + ) + hidden_col_start = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + + subtile_idx * cutlass.Int32(EpilogueTileN) + ) + r_acc_layout = cute.make_layout((((EpilogueTileN,), 1),), stride=(((1,), 0),)) + atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), self.acc_dtype, + ) + r_acc = cute.make_rmem_tensor(r_acc_layout.shape, self.acc_dtype) + cute.copy(atom_t2r, tmem_subtile_tensor, r_acc) + + thread_in_warp = tidx % WarpThreadCount + token_row_in_cta = cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + if token_row_in_cta < valid_tokens and hidden_col_start < valid_hidden: + if cutlass.const_expr( + token_comm_args is not None + and not self._token_back_by_dispatch + and self._combine_mxfp8 + ): + # MegaMoE Form A, quantized combine: + # 1. Quantize fp32 → fp8 + compute E8M0 block scale. + # 2. STG fp8 data to peer's combine_output. + # 3. Write E8M0 scale to local fc2_output_sf for token-back push. + fp8_dtype = self._combine_format.act_dtype + r_fp8 = cute.make_rmem_tensor(r_acc_layout.shape, fp8_dtype) + qpvscale = quant_sfd_row( + r_acc, r_fp8, 1.0, EpilogueTileN, + cutlass.Float8E8M0FNU, fp8_dtype, + ) + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + metadata_u32 = cute.recast_tensor( + token_comm_args.token_src_metadata, cutlass.Uint32, + ) + fc2_output_dest = Fc2OutputDest( + tensor=token_comm_args.combine_output, + metadata=metadata_u32, + peer_rank_ptr_mapper=token_comm_args.peer_rank_ptr_mapper, + ) + dest_row = fc2_output_dest.resolve_token_row(pool_token_global) + r_fp8_flat = cute.make_tensor( + r_fp8.iterator, cute.make_layout(EpilogueTileN) + ) + dest_fp8_ptr = cute.make_ptr( + fp8_dtype, + dest_row.iterator.toint() + Int64(hidden_col_start), + cute.AddressSpace.gmem, + assumed_align=32, + ) + if cutlass.const_expr(self._fc2_use_ublk): + # UBLK: stage this token's 32 fp8 elems into SMEM + stage_idx = subtile_idx % cutlass.Int32(self._fc2_tma_stages) + smem_row = cute.slice_( + smem_fc2_tma_buffer, (token_row_in_cta, None, stage_idx) + ) + sts_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=256, + ) + cute.copy(sts_fp8_atom, r_fp8_flat, smem_row) + cute.arch.fence_proxy("async.shared", space="cta") + cp_async_bulk_s2g( + dest_fp8_ptr, + smem_row.iterator, + cutlass.Int32(EpilogueTileN * fp8_dtype.width // 8), + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group( + self._fc2_tma_stages - 1, read=True + ) + else: + # STG 32 fp8 elements = 256 bits in one shot. + stg_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, + num_bits_per_copy=256, + ) + cute.copy( + stg_fp8_atom, r_fp8_flat, + cute.make_tensor(dest_fp8_ptr, cute.make_layout(EpilogueTileN)), + ) + # Buffer the E8M0 scale; the whole task tile's 8 scales are + # flushed together by _stg_sf_fc2 (single stg.64 when aligned). + self._write_sf_fc2_buffer(rmem_sf_fc2, subtile_idx, qpvscale) + elif cutlass.const_expr( + self._token_back_by_dispatch and self._combine_mxfp8 + ): + # MegaMoE token-back-by-dispatch + quantized combine: + # Epilogue writes fp8 data to local pool; dispatch warps push + # both data (fc2_output_workspace) and SF (fc2_output_sf) to peers. + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + fp8_dtype = self._combine_format.act_dtype + r_fp8 = cute.make_rmem_tensor(r_acc_layout.shape, fp8_dtype) + qpvscale = quant_sfd_row( + r_acc, r_fp8, 1.0, EpilogueTileN, + cutlass.Float8E8M0FNU, fp8_dtype, + ) + r_fp8_flat = cute.make_tensor(r_fp8.iterator, cute.make_layout(EpilogueTileN)) + if cutlass.const_expr(self._fc2_use_tma): + stage_idx = subtile_idx % cutlass.Int32(self._fc2_tma_stages) + smem_row = cute.slice_( + smem_fc2_tma_buffer, (token_row_in_cta, None, stage_idx) + ) + sts_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=256, + ) + cute.copy(sts_fp8_atom, r_fp8_flat, smem_row) + else: + # Write 32 fp8 elements to local fc2_output_workspace pool. + fp8_byte_addr = ( + token_comm_args.fc2_output_workspace.iterator.toint() + + Int64(pool_token_global) * Int64(self._hidden_fc2) + + Int64(hidden_col_start) + ) + stg_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, + num_bits_per_copy=256, + ) + aligned_fp8_iter = cute.make_ptr( + fp8_dtype, + fp8_byte_addr, + cute.AddressSpace.gmem, + assumed_align=32, + ) + cute.copy( + stg_fp8_atom, r_fp8_flat, + cute.make_tensor(aligned_fp8_iter, cute.make_layout(EpilogueTileN)), + ) + # Buffer the E8M0 scale; flushed together by _stg_sf_fc2 after + # the subtile loop (single stg.64 when hidden-aligned). + self._write_sf_fc2_buffer(rmem_sf_fc2, subtile_idx, qpvscale) + elif cutlass.const_expr( + token_comm_args is None and self._combine_mxfp8 + ): + # Lean path (no token-comm), quantized fc2 output + fp8_dtype = self._combine_format.act_dtype + r_fp8 = cute.make_rmem_tensor(r_acc_layout.shape, fp8_dtype) + qpvscale = quant_sfd_row( + r_acc, r_fp8, 1.0, EpilogueTileN, + cutlass.Float8E8M0FNU, fp8_dtype, + ) + g_fc2_output_tile = cute.local_tile( + real_fc2_output, + (self._cta_tile_m, EpilogueTileN, 1), + (work_tile_info.tile_m_idx, hidden_group, 0), + ) + g_fc2_slice = cute.slice_(g_fc2_output_tile, (None, None, 0)) + g_thread_row = cute.local_tile( + g_fc2_slice, (1, EpilogueTileN), (token_row_in_cta, 0), + ) + g_flat = cute.coalesce(g_thread_row) + aligned_iter = cute.make_ptr( + fp8_dtype, + g_flat.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=32, + ) + stg_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, + num_bits_per_copy=256, + ) + r_fp8_flat = cute.make_tensor( + r_fp8.iterator, cute.make_layout(EpilogueTileN) + ) + cute.copy( + stg_fp8_atom, r_fp8_flat, + cute.make_tensor(aligned_iter, cute.make_layout(EpilogueTileN)), + ) + # Buffer the E8M0 scale; flushed together by _stg_sf_fc2 after + # the subtile loop (single stg.64 when hidden-aligned). + self._write_sf_fc2_buffer(rmem_sf_fc2, subtile_idx, qpvscale) + else: + # BF16 path (default): fp32->bf16, two 256-bit STGs. + r_bf16 = cute.make_rmem_tensor(r_acc_layout.shape, cutlass.BFloat16) + r_bf16.store(r_acc.load().to(cutlass.BFloat16)) + stg_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=256, + ) + if cutlass.const_expr(self._fc2_reduce_coalesce): + # In-kernel reduce: stage the full bf16 token row token-major into + # SMEM (row start is 64 B aligned). + cute.autovec_copy( + cute.make_tensor(r_bf16.iterator, cute.make_layout(EpilogueTileN)), + cute.slice_(smem_fc2_reduce_buffer, (token_row_in_cta, None)), + ) + if cutlass.const_expr( + token_comm_args is not None and not self._token_back_by_dispatch + ): + metadata_u32 = cute.recast_tensor( + token_comm_args.token_src_metadata, cutlass.Uint32, + ) + fc2_output_dest = Fc2OutputDest( + tensor=token_comm_args.combine_output, + metadata=metadata_u32, + peer_rank_ptr_mapper=token_comm_args.peer_rank_ptr_mapper, + reduce_topk_in_kernel=self._fc2_in_kernel_topk_reduce, + ) + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + dest_row = fc2_output_dest.resolve_token_row(pool_token_global) + for stg_half in cutlass.range_constexpr(EpilogueTileN // 16): + reg_view = cute.make_tensor( + r_bf16.iterator + stg_half * 16, + cute.make_layout(16), + ) + if cutlass.const_expr( + token_comm_args is not None and not self._token_back_by_dispatch + ): + hidden_off = hidden_col_start + cutlass.Int32(stg_half * 16) + dest_ptr = cute.make_ptr( + cutlass.BFloat16, + dest_row.iterator.toint() + hidden_off * cutlass.Int64(2), + cute.AddressSpace.gmem, + assumed_align=32, + ) + # Reduce path staged to SMEM above; only the non-reduce Form-A + # combine does the direct peer STG here. + if cutlass.const_expr(not self._fc2_reduce_coalesce): + cute.copy( + stg_atom, reg_view, + cute.make_tensor(dest_ptr, cute.make_layout(16)), + ) + else: + g_fc2_output_tile = cute.local_tile( + real_fc2_output, + (self._cta_tile_m, EpilogueTileN, 1), + (work_tile_info.tile_m_idx, hidden_group, 0), + ) + g_fc2_slice = cute.slice_(g_fc2_output_tile, (None, None, 0)) + g_thread_row = cute.local_tile( + g_fc2_slice, (1, 16), (token_row_in_cta, stg_half), + ) + g_flat = cute.coalesce(g_thread_row) + aligned_iter = cute.make_ptr( + cutlass.BFloat16, + g_flat.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=32, + ) + cute.copy(stg_atom, reg_view, cute.make_tensor(aligned_iter, g_flat.layout)) + + if cutlass.const_expr(self._fc2_use_tma): + self._issue_fc2_tma_store( + subtile_idx=subtile_idx, + smem_fc2_tma_buffer=smem_fc2_tma_buffer, + tma_atom_fc2_output=tma_atom_fc2_output, + gmem_fc2_tma_output=gmem_fc2_tma_output, + work_tile_info=work_tile_info, + warp_idx=warp_idx, + ) + + # In-kernel reduce: the token-major SMEM tile is now populated; all 128 epi + # threads re-partition hidden-major and issue coalesced peer red.add. + if cutlass.const_expr(self._fc2_reduce_coalesce): + self._issue_fc2_reduce_coalesced( + subtile_idx=subtile_idx, + smem_fc2_reduce_buffer=smem_fc2_reduce_buffer, + token_comm_args=token_comm_args, + work_tile_info=work_tile_info, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + ) + + if warp_idx == 0: + iket.range_pop() + + + @cute.jit + def _issue_fc2_tma_store( + self, + subtile_idx, + smem_fc2_tma_buffer: cute.Tensor, + tma_atom_fc2_output: cute.CopyAtom, + gmem_fc2_tma_output: cute.Tensor, + work_tile_info, + warp_idx: int, + ) -> None: + """Issue the FC2 bulk-tensor store for one subtile's staged SMEM tile.""" + stage_idx = subtile_idx % cutlass.Int32(self._fc2_tma_stages) + smem_stage = cute.slice_(smem_fc2_tma_buffer, (None, None, stage_idx)) + token_base = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + ) + hidden_col = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + + subtile_idx * cutlass.Int32(EpilogueTileN) + ) + token_tile_idx = token_base // cutlass.Int32(self._cta_tile_m) + hidden_tile_idx = hidden_col // cutlass.Int32(EpilogueTileN) + tiled_output = cute.flat_divide( + gmem_fc2_tma_output, (self._cta_tile_m, EpilogueTileN) + ) + gmem_subtile = tiled_output[None, None, token_tile_idx, hidden_tile_idx] + bSG_s, bSG_g = cpasync.tma_partition( + tma_atom_fc2_output, + 0, + cute.make_layout(1), + cute.group_modes(smem_stage, 0, 2), + cute.group_modes(gmem_subtile, 0, 2), + ) + cute.arch.fence_proxy("async.shared", space="cta") + bar = pipeline.NamedBarrier( + barrier_id=self._epilog_sync_bar_id, + num_threads=EpiWarpCount * WarpThreadCount, + ) + bar.arrive_and_wait() + if warp_idx == 0: + cute.copy(tma_atom_fc2_output, bSG_s, bSG_g) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(self._fc2_tma_stages - 1, read=True) + bar.arrive_and_wait() + + @cute.jit + def _issue_fc2_reduce_coalesced( + self, + subtile_idx, + smem_fc2_reduce_buffer: cute.Tensor, + token_comm_args, + work_tile_info, + valid_hidden, + warp_idx: int, + tidx, + ) -> None: + """Coalesced peer red.add of one subtile's staged reduce tile.""" + redg_width: cutlass.Constexpr[int] = 4 # bf16 per red.v2.bf16x2 + lanes_per_token: cutlass.Constexpr[int] = EpilogueTileN // redg_width # 8 + tokens_per_warp: cutlass.Constexpr[int] = WarpThreadCount // lanes_per_token # 4 + tokens_per_pass: cutlass.Constexpr[int] = tokens_per_warp * EpiWarpCount # 16 + passes: cutlass.Constexpr[int] = self._cta_tile_m // tokens_per_pass # 8 + + # Rendezvous: make every thread's staged SMEM row visible before the reads. + bar = pipeline.NamedBarrier( + barrier_id=self._epilog_sync_bar_id, + num_threads=EpiWarpCount * WarpThreadCount, + ) + bar.arrive_and_wait() + + thread_in_warp = tidx % WarpThreadCount + token_in_warp = thread_in_warp // cutlass.Int32(lanes_per_token) # 0..tokens_per_warp-1 + chunk = thread_in_warp % cutlass.Int32(lanes_per_token) # 0..lanes_per_token-1 + hidden_off = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + + subtile_idx * cutlass.Int32(EpilogueTileN) + + chunk * cutlass.Int32(redg_width) + ) + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + metadata_u32 = cute.recast_tensor(token_comm_args.token_src_metadata, cutlass.Uint32) + fc2_output_dest = Fc2OutputDest( + tensor=token_comm_args.combine_output, + metadata=metadata_u32, + peer_rank_ptr_mapper=token_comm_args.peer_rank_ptr_mapper, + reduce_topk_in_kernel=True, + ) + for p in cutlass.range_constexpr(passes): + token_in_cta = ( + cutlass.Int32(p * tokens_per_pass) + + cutlass.Int32(warp_idx) * cutlass.Int32(tokens_per_warp) + + token_in_warp + ) + if token_in_cta < valid_tokens and hidden_off < valid_hidden: + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_in_cta + ) + dest_row = fc2_output_dest.resolve_token_row(pool_token_global) + # LDS this lane's 4 staged bf16 (contiguous hidden chunk). + smem_chunk_ptr = cute.slice_( + smem_fc2_reduce_buffer, (token_in_cta, None) + ).iterator + chunk * cutlass.Int32(redg_width) + r4 = cute.make_rmem_tensor((redg_width,), cutlass.BFloat16) + cute.autovec_copy( + cute.make_tensor(smem_chunk_ptr, cute.make_layout(redg_width)), r4 + ) + r4_u32 = cute.recast_tensor(r4, cutlass.Uint32) # 2 u32 = 4 bf16 + dest_ptr = cute.make_ptr( + cutlass.BFloat16, + dest_row.iterator.toint() + Int64(hidden_off) * Int64(2), + cute.AddressSpace.gmem, + assumed_align=8, + ) + _red_add_relaxed_sys_v2_bf16x2( + dest_ptr, + cutlass.Uint32(r4_u32[0]), + cutlass.Uint32(r4_u32[1]), + ) + bar.arrive_and_wait() + + @cute.jit + def _write_sf_fc2_buffer(self, rmem_sf_fc2, subtile_idx, qpvscale) -> None: + """Scatter one subtile's E8M0 scale into the per-tile SF buffer.""" + for j in cutlass.range_constexpr(self._cta_tile_n // EpilogueTileN): + if subtile_idx == cutlass.Int32(j): + rmem_sf_fc2[j] = qpvscale + + @cute.jit + def _stg_sf_fc2( + self, + rmem_sf_fc2: cute.Tensor, + sf_base_addr, + sf_row_stride, + work_tile_info, + valid_hidden, + warp_idx: int, + tidx, + ) -> None: + """Flush a task tile's fc2 E8M0 scales to the local ``fc2_output_sf``.""" + fc2_subtile_cnt = self._cta_tile_n // EpilogueTileN + thread_in_warp = tidx % WarpThreadCount + token_row_in_cta = cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + if token_row_in_cta < work_tile_info.valid_tokens_in_cta_tile: + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + hidden_group_base = ( + work_tile_info.tile_n_idx * cutlass.Int32(fc2_subtile_cnt) + ) + sf_byte_addr = ( + sf_base_addr + + Int64(pool_token_global) * Int64(sf_row_stride) + + Int64(hidden_group_base) + ) + if cutlass.const_expr(self._fc2_sf_batch8): + stg_e8m0x8_from_f32( + sf_byte_addr, + rmem_sf_fc2[0], rmem_sf_fc2[1], rmem_sf_fc2[2], rmem_sf_fc2[3], + rmem_sf_fc2[4], rmem_sf_fc2[5], rmem_sf_fc2[6], rmem_sf_fc2[7], + ) + else: + for j in cutlass.range_constexpr(fc2_subtile_cnt): + block_hidden_start = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + + cutlass.Int32(j * EpilogueTileN) + ) + if block_hidden_start < valid_hidden: + stg_e8m0_from_f32(sf_byte_addr + Int64(j), rmem_sf_fc2[j]) + + @cute.jit + def _run_fc2_task_tile( + self, + work_tile_info, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + acc_consumer_state, + sched_ext, + gmem_fc2_output: cute.Tensor, + valid_hidden, + warp_idx: int, + tidx, + token_comm_args=None, + gmem_fc2_output_sf=None, + smem_fc2_tma_buffer=None, + tma_atom_fc2_output=None, + gmem_fc2_tma_output=None, + smem_fc2_reduce_buffer=None, + ) -> None: + """FC2 (Linear2) task-tile body using two-stage TMEM accumulation.""" + real_fc2_output, _ = sched_ext.get_gmem_tensor( + "d", gmem_fc2_output, work_tile_info, + ) + acc_pipeline.consumer_wait(acc_consumer_state) + if warp_idx == 0: + iket.range_push("fc2_epi_tile") + + fc2_subtile_cnt = self._cta_tile_n // EpilogueTileN # = 8 + + tmem_t = self._subtile_fc2_tmem_tensor( + tmem_acc_tensor, cutlass.Int32(0), warp_idx, + ) + + tmem_forward_cols = EpilogueTileN + + # Quantized combine: buffer the per-subtile E8M0 scales and flush them in + # one stg.64 after the loop (see _stg_sf_fc2). Indexed by subtile_idx, so + # the reversed odd-turn walk fills the same slots. Both the token-comm + # combine plane and the lean stand-alone fc2_output_sf plane need it. + if cutlass.const_expr(self._combine_mxfp8): + layout_sf_fc2 = cute.make_layout(fc2_subtile_cnt) + rmem_sf_fc2 = cute.make_rmem_tensor(layout_sf_fc2.shape, self.acc_dtype) + else: + rmem_sf_fc2 = None + + for i in cutlass.range(0, fc2_subtile_cnt, 1, unroll=1): + subtile_idx = cutlass.Int32(i) + + self._run_fc2_subtile( + subtile_idx=subtile_idx, + tmem_subtile_tensor=tmem_t, + real_fc2_output=real_fc2_output, + work_tile_info=work_tile_info, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + token_comm_args=token_comm_args, + rmem_sf_fc2=rmem_sf_fc2, + smem_fc2_tma_buffer=smem_fc2_tma_buffer, + tma_atom_fc2_output=tma_atom_fc2_output, + gmem_fc2_tma_output=gmem_fc2_tma_output, + smem_fc2_reduce_buffer=smem_fc2_reduce_buffer, + ) + + tmem_t = self._advance_fc2_tmem_tensor(tmem_t, tmem_forward_cols) + + self._acc_pipeline_consumer_release(acc_pipeline, acc_consumer_state, True) + + # Flush the buffered E8M0 scales + if cutlass.const_expr(self._combine_mxfp8 and token_comm_args is not None): + self._stg_sf_fc2( + rmem_sf_fc2=rmem_sf_fc2, + sf_base_addr=token_comm_args.fc2_output_sf.iterator.toint(), + sf_row_stride=Int64(self._fc2_sf_block_pad), + work_tile_info=work_tile_info, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + ) + elif cutlass.const_expr( + self._combine_mxfp8 and gmem_fc2_output_sf is not None + ): + self._stg_sf_fc2( + rmem_sf_fc2=rmem_sf_fc2, + sf_base_addr=gmem_fc2_output_sf.iterator.toint(), + sf_row_stride=Int64(gmem_fc2_output_sf.stride[0]), + work_tile_info=work_tile_info, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + ) + + if warp_idx == 0: + iket.range_pop() + + + @cute.jit + def _stg_sf_fc1( + self, + rmem_sf_f32: cute.Tensor, + real_fc1_output_sf: cute.Tensor, + work_tile_info, + tidx, + ) -> None: + """Compute gmem SF tile coords and store fc1 scale factors to gmem.""" + bx, _, _ = cute.arch.block_idx() + sf_idx = work_tile_info.tile_n_idx + token_idx = ( + work_tile_info.tile_m_idx * self._cta_tile_m + + tidx + ) + if tidx < work_tile_info.valid_tokens_in_cta_tile: + sf_base = cute.local_tile( + real_fc1_output_sf, + (1, 1, 1), + (token_idx, sf_idx * cutlass.Int32(Fc1EpilogueOutputTileN), cutlass.Int32(0)), + ) + sf_ptr = cute.make_ptr( + self.sf_dtype, + sf_base.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=4, + ) + gmem_sf_f8 = cute.make_tensor(sf_ptr, cute.make_layout(4)) + sf_layout = cute.make_layout(4) + r_sf_f8 = cute.make_rmem_tensor(sf_layout.shape, self.sf_dtype) + r_sf_f8.store(rmem_sf_f32.load().to(self.sf_dtype)) + cute.autovec_copy(r_sf_f8, gmem_sf_f8) + + + @cute.jit + def run( + self, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + sched_consumer, + sched_ext, + smem_fc1_output_buffer: Optional[cute.Tensor], + tma_atom_fc1_output: cute.CopyAtom, + gmem_fc1_output: cute.Tensor, + gmem_topk_scores: cute.Tensor, + gmem_fc2_output: cute.Tensor, + gmem_fc1_done_counter: cute.Tensor, + warp_idx: int, + tidx, + gmem_fc1_output_sf: Optional[cute.Tensor] = None, + alpha=None, + norm_const=None, + token_comm_args=None, + gmem_fc2_output_sf: Optional[cute.Tensor] = None, + smem_c_buffer: cute.Tensor = None, + tma_atom_c: cute.CopyAtom = None, + gmem_c: cute.Tensor = None, + smem_fc2_tma_buffer: Optional[cute.Tensor] = None, + tma_atom_fc2_output: cute.CopyAtom = None, + gmem_fc2_tma_output: cute.Tensor = None, + smem_fc2_reduce_buffer: Optional[cute.Tensor] = None, + ) -> None: + """ + Run the full fc1+fc2-fused epilogue task-tile loop. + """ + assert self._use_stg_fc1 or smem_fc1_output_buffer is not None, ( + "smem_fc1_output_buffer=None requires use_stg_fc1=True (the TMA " + "fc1-output store path consumes the sD staging buffer)" + ) + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._num_acc_pipeline_stages + ) + + if cutlass.const_expr(self._generate_c): + _c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=Fc1CTMAStages, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + EpiWarpCount * WarpThreadCount, + ), + ) + else: + _c_pipeline = None + + task_tile_boundary_bar = pipeline.NamedBarrier( + barrier_id=self._epilog_sync_bar_id, + num_threads=32 * len(self._epilogue_warp_ids), + ) + + valid_hidden = cutlass.Int32(gmem_fc2_output.shape[1]) + + bidx, bidy, bidz = cute.arch.block_idx() + work_tile_info = sched_consumer.consume_work() + + flag_tracker = GpuReleaseFlagBatchTracker( + flag_address=Int64(0), + accumulated_flags=cutlass.Int32(0), + phase=cutlass.Int32(work_tile_info.phase), + thread_idx=tidx % (len(self._epilogue_warp_ids) * WarpThreadCount), + ) + + while work_tile_info.is_valid_tile: + acc_stage_index = acc_consumer_state.index + tmem_acc_stage_tesnor = tmem_acc_tensor[(None, None, None, acc_stage_index)] + + if work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1): + if cutlass.const_expr(self._generate_c): + _smem_c_buf = smem_c_buffer + _tma_atom_c = tma_atom_c + _gmem_c = gmem_c + else: + _smem_c_buf = smem_fc1_output_buffer + _tma_atom_c = tma_atom_fc1_output + _gmem_c = gmem_fc1_output + self._run_fc1_task_tile( + work_tile_info=work_tile_info, + tmem_acc_tensor=tmem_acc_stage_tesnor, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + smem_fc1_output_buffer=smem_fc1_output_buffer, + tma_atom_fc1_output=tma_atom_fc1_output, + sched_ext=sched_ext, + gmem_fc1_output=gmem_fc1_output, + gmem_fc1_output_sf=gmem_fc1_output_sf, + gmem_topk_scores=gmem_topk_scores, + warp_idx=warp_idx, + tidx=tidx, + alpha=alpha, + norm_const=norm_const, + smem_c_buffer=_smem_c_buf, + tma_atom_c=_tma_atom_c, + gmem_c=_gmem_c, + c_pipeline=_c_pipeline, + ) + else: + self._run_fc2_task_tile( + work_tile_info=work_tile_info, + tmem_acc_tensor=tmem_acc_stage_tesnor, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + sched_ext=sched_ext, + gmem_fc2_output=gmem_fc2_output, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + token_comm_args=token_comm_args, + gmem_fc2_output_sf=gmem_fc2_output_sf, + smem_fc2_tma_buffer=smem_fc2_tma_buffer, + tma_atom_fc2_output=tma_atom_fc2_output, + gmem_fc2_tma_output=gmem_fc2_tma_output, + smem_fc2_reduce_buffer=smem_fc2_reduce_buffer, + ) + + acc_consumer_state.advance() + + cur_was_linear1 = work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + cur_fc1_counter_slot = ( + work_tile_info.cumulative_token_block_count + + work_tile_info.tile_m_idx // cutlass.Int32(self._atom_thr_size) + ) + cur_fc2_expert_idx = work_tile_info.expert_idx + + work_tile_info = sched_consumer.consume_work() + + # Drain in-flight bulk stores before publishing the done counter. + if cur_was_linear1 or cutlass.const_expr( + self._fc2_use_tma or self._fc2_use_ublk + ): + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.fence_proxy("async") + cute.arch.fence_acq_rel_gpu() + + task_tile_boundary_bar.arrive_and_wait() + + if cur_was_linear1: + flag_tracker = flag_tracker.accumulate( + work_tile_info.phase, + self._epi_fc1_batch, + (gmem_fc1_done_counter.iterator + cur_fc1_counter_slot).toint(), + ) + else: + _fire_fc2_counter: cutlass.Constexpr = ( + (self._token_back_by_dispatch or self._combine_mxfp8) + and token_comm_args is not None + ) + if cutlass.const_expr(_fire_fc2_counter): + # Fence before (deferred) counter release: make the fc2 + # pool-output STG writes device-visible. + cute.arch.fence_acq_rel_gpu() + fc2_flag_addr = ( + token_comm_args.fc2_done_counter.iterator + cur_fc2_expert_idx + ).toint() + else: + fc2_flag_addr = Int64(0) + no_fire: cutlass.Constexpr = not _fire_fc2_counter + flag_tracker = flag_tracker.accumulate( + work_tile_info.phase, + self._epi_fc2_batch, + fc2_flag_addr, + no_fire, + ) + + flag_tracker.fire() + diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_extension.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_extension.py new file mode 100644 index 000000000..4dfd4cd29 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_extension.py @@ -0,0 +1,232 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Scheduling adapter for the MXFP8 GLU FC12 kernel.""" + +import dataclasses +from typing import ClassVar, List, Literal, Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass.cute.typing import Pointer +from cutlass.cutlass_dsl import Int32, extract_mlir_values, new_from_mlir_values +from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF + +from ......helpers.dsl_helpers import spin_peek, spin_wait +from .....schedulers.fc12_mapping import BlockPhase, NonSwapAbFc12WorkTileInfo, peek_ready_bit + + +# Forward GLU tensor roles: FC1 activation is token-indexed (M), FC1 weight is +# expert-indexed; "d"/"sfd" are the FC1 fp8 output + E8M0 plane, "c" the raw gate/up. +TensorRole = Literal[ + "fc1_activation", + "fc1_weight", + "fc1_activation_sf", + "fc1_weight_sf", + "c", + "d", + "sfd", + "topk", + "fc2_activation", + "fc2_activation_sf", + "fc2_weight", + "fc2_weight_sf", +] + + +@cute.jit +def _rewrite_tensor_shape(tensor: cute.Tensor, new_shape: Tuple) -> cute.Tensor: + return cute.make_tensor(tensor.iterator, cute.make_layout(new_shape, stride=tensor.stride)) + + +@dataclasses.dataclass(frozen=True) +class GluMxFp8Fc12SchedExtension: + """Kernel-owned work-tile preparation and GMEM view adapter (non-swap MXFP8).""" + + work_tile_type: ClassVar[type] = NonSwapAbFc12WorkTileInfo + + sf_vec_size: int + fc1_done_counter_pointer: Pointer + fc2_spin_threshold: Int32 + fc1_ready_counter_pointer: Optional[Pointer] = None + cluster_m: int = 1 + + def __post_init__(self) -> None: + if self.sf_vec_size <= 0: + raise ValueError(f"sf_vec_size must be positive, got {self.sf_vec_size}.") + if self.cluster_m <= 0: + raise ValueError(f"cluster_m must be positive, got {self.cluster_m}.") + object.__setattr__(self, "fc2_spin_threshold", Int32(self.fc2_spin_threshold)) + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + values.extend(extract_mlir_values(self.fc1_done_counter_pointer)) + values.extend(extract_mlir_values(self.fc2_spin_threshold)) + if self.fc1_ready_counter_pointer is not None: + values.extend(extract_mlir_values(self.fc1_ready_counter_pointer)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GluMxFp8Fc12SchedExtension": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + fc1_done_counter_pointer = rebuild(self.fc1_done_counter_pointer) + fc2_spin_threshold = rebuild(self.fc2_spin_threshold) + fc1_ready_counter_pointer = ( + rebuild(self.fc1_ready_counter_pointer) if self.fc1_ready_counter_pointer is not None else None + ) + if value_index != len(values): + raise ValueError( + f"GluMxFp8Fc12SchedExtension MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return type(self)( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter_pointer, + fc2_spin_threshold=fc2_spin_threshold, + fc1_ready_counter_pointer=fc1_ready_counter_pointer, + cluster_m=self.cluster_m, + ) + + @cute.jit + def _counter_slot(self, work_tile: NonSwapAbFc12WorkTileInfo) -> Int32: + # Cluster-granular token-block slot: dispatch_pull increments one counter per + # cluster-level token block, so M-direction tiles fold by cluster_m. + return work_tile.cumulative_token_block_count + work_tile.tile_m_idx // Int32(self.cluster_m) + + @cute.jit + def prepare_work_tile(self, work_tile: NonSwapAbFc12WorkTileInfo) -> NonSwapAbFc12WorkTileInfo: + """Pack kernel readiness observations into the published tile flags.""" + phase_and_flags = work_tile.phase_and_flags + if work_tile.is_valid_tile: + counter_slot = self._counter_slot(work_tile) + is_fc1 = work_tile.phase == Int32(BlockPhase.Linear1) + is_fc2 = work_tile.phase == Int32(BlockPhase.Linear2) + + if cutlass.const_expr(self.fc1_ready_counter_pointer is not None): + if is_fc1: + counter_pointer = self.fc1_ready_counter_pointer + counter_slot + peek_flag = Int32(0) + if spin_peek(counter_pointer, lambda value: value >= work_tile.valid_tokens_in_cluster_tile): + peek_flag = Int32(peek_ready_bit) + phase_and_flags = work_tile.phase_and_flags | peek_flag + + if is_fc2: + counter_pointer = self.fc1_done_counter_pointer + counter_slot + peek_flag = Int32(0) + if spin_peek(counter_pointer, lambda value: value >= self.fc2_spin_threshold): + peek_flag = Int32(peek_ready_bit) + phase_and_flags = work_tile.phase_and_flags | peek_flag + + return NonSwapAbFc12WorkTileInfo( + expert_idx=work_tile.expert_idx, + tile_m_idx=work_tile.tile_m_idx, + tile_n_idx=work_tile.tile_n_idx, + cumulative_data_physical_row=work_tile.cumulative_data_physical_row, + cumulative_sf_physical_row=work_tile.cumulative_sf_physical_row, + cumulative_token_block_count=work_tile.cumulative_token_block_count, + valid_tokens_in_cta_cluster_tile=work_tile.valid_tokens_in_cta_cluster_tile, + phase_and_flags=phase_and_flags, + ) + + @cute.jit + def wait_for_input(self, work_tile: NonSwapAbFc12WorkTileInfo) -> None: + """Wait until this FC1 input tile's cluster-level token count has arrived.""" + if cutlass.const_expr(self.fc1_ready_counter_pointer is not None): + counter_pointer = self.fc1_ready_counter_pointer + self._counter_slot(work_tile) + spin_wait( + counter_pointer, + lambda value: value >= work_tile.valid_tokens_in_cluster_tile, + peek_status=work_tile.peek_ready, + ) + + @cute.jit + def get_gmem_tensor( + self, + tensor_name: TensorRole, + gmem_tensor_in_moe_view: cute.Tensor, + work_tile_info: NonSwapAbFc12WorkTileInfo, + ) -> Tuple[cute.Tensor, Optional[Pointer]]: + """Phase-invariant GMEM slice for the operands.""" + expert_idx = work_tile_info.expert_idx + data_token_offset = work_tile_info.cumulative_data_physical_row + sf_token_offset = work_tile_info.cumulative_sf_physical_row + + shape = gmem_tensor_in_moe_view.shape + stride = gmem_tensor_in_moe_view.stride + c1 = cutlass.Int32(1) + sf_vec_size = self.sf_vec_size + + if cutlass.const_expr(tensor_name == "fc1_activation"): + real = cute.domain_offset((data_token_offset, 0, 0), gmem_tensor_in_moe_view) + return (_rewrite_tensor_shape(real, (shape[0], shape[1], c1)), None) + + elif cutlass.const_expr(tensor_name == "fc1_weight"): + real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view) + return (_rewrite_tensor_shape(real, (shape[0], shape[1], c1)), None) + + elif cutlass.const_expr(tensor_name == "fc1_activation_sf"): + real = cute.domain_offset((sf_token_offset, 0, 0), gmem_tensor_in_moe_view) + per_expert_shape = (shape[0], shape[1], c1) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size) + real = cute.make_tensor(real.iterator, cute.make_layout(sf_layout.shape, stride=stride)) + return (real, None) + + elif cutlass.const_expr(tensor_name == "fc1_weight_sf"): + real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view) + per_expert_shape = (shape[0], shape[1], c1) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size) + real = cute.make_tensor(real.iterator, cute.make_layout(sf_layout.shape, stride=stride)) + return (real, None) + + elif cutlass.const_expr(tensor_name == "c"): + # Raw fc1 accumulator output (gate+up FP32, pre-SwiGLU): token-indexed. + real = cute.domain_offset((data_token_offset, 0, 0), gmem_tensor_in_moe_view) + return (_rewrite_tensor_shape(real, (shape[0], shape[1], c1)), None) + + elif cutlass.const_expr(tensor_name == "d"): + real = cute.domain_offset((data_token_offset, 0, 0), gmem_tensor_in_moe_view) + return (_rewrite_tensor_shape(real, (shape[0], shape[1], c1)), None) + + elif cutlass.const_expr(tensor_name == "sfd"): + real = cute.domain_offset((sf_token_offset, 0, 0), gmem_tensor_in_moe_view) + per_expert_shape = (shape[0], shape[1], c1) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size) + real = cute.make_tensor(real.iterator, cute.make_layout(sf_layout.shape, stride=stride)) + return (real, None) + + elif cutlass.const_expr(tensor_name == "topk"): + real = cute.domain_offset((data_token_offset,), gmem_tensor_in_moe_view) + return (real, None) + + elif cutlass.const_expr(tensor_name == "fc2_activation"): + real = cute.domain_offset((data_token_offset, 0, 0), gmem_tensor_in_moe_view) + return (_rewrite_tensor_shape(real, (shape[0], shape[1], c1)), None) + + elif cutlass.const_expr(tensor_name == "fc2_activation_sf"): + real = cute.domain_offset((sf_token_offset, 0, 0), gmem_tensor_in_moe_view) + per_expert_shape = (shape[0], shape[1], c1) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size) + real = cute.make_tensor(real.iterator, cute.make_layout(sf_layout.shape, stride=stride)) + return (real, None) + + elif cutlass.const_expr(tensor_name == "fc2_weight"): + real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view) + return (_rewrite_tensor_shape(real, (shape[0], shape[1], c1)), None) + + elif cutlass.const_expr(tensor_name == "fc2_weight_sf"): + real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view) + per_expert_shape = (shape[0], shape[1], c1) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size) + real = cute.make_tensor(real.iterator, cute.make_layout(sf_layout.shape, stride=stride)) + return (real, None) + + raise ValueError(f"Unknown tensor_name: {tensor_name!r}.") + + +__all__ = ["GluMxFp8Fc12SchedExtension", "TensorRole"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_kernel.py new file mode 100644 index 000000000..73f42e30b --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_kernel.py @@ -0,0 +1,2372 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +""" +Fused fc1+fc2 GLU MXFP8 MegaMoE kernel for SM100. +""" + +import dataclasses +from typing import Any, Literal, Optional, Tuple, Type, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, tcgen05, OperandMajorMode +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +import cutlass.utils.rubin_helpers as sm107_utils +from cutlass.cute.nvgpu.tcgen05 import CollectorOp + + +from ..helpers.constants import SupportedMmaTileM, SupportedMmaTileN +from ......helpers.iket_compat import iket +from ......api import ImplDesc, KernelClass, ProblemDesc, StaticOrRuntimeIntegerType +from ......helpers.device_workspace import DeviceWorkspace +from ......helpers.smem_workspace import SmemWorkspace +from ......helpers.dsl_helpers import spin_wait +from ......quant_def import CombineFormat, QuantKind +from ......communication.nvlink_domain.token_comm import TokenCommArgs +from .....schedulers import BlockPhase +from .....schedulers.base import WorkIdAcquisitionMode +from .....schedulers.fc12_scheduler import BlackwellFusedFc12Scheduler +from .glu_mxfp8_fc12_epilogue import GluMxfp8Epilogue +from .glu_mxfp8_fc12_extension import GluMxFp8Fc12SchedExtension + + +@dataclasses.dataclass(frozen=True) +class _EpilogueCommView: + """The exact fields the GLU epilogue reads for cross-rank FC2 routing. + + Built inside the device kernel (so it carries no MLIR-marshaled scalars -- its fields + are already-traced device views) and passed to ``epilogue.run(token_comm_args=...)`` on + the MegaMoE path. Bridges next's Router-push component (``token_src_metadata`` / + ``fc2_done_counter`` come from ``TokenComm`` accessors; ``combine_output`` is the symmetric + ``pre_reduced_activation``) to the epilogue's ``Fc2OutputDest`` peer-store expectations. + """ + + token_src_metadata: Any + combine_output: Any + peer_rank_ptr_mapper: Any + fc2_output_sf: Any = None + fc2_done_counter: Any = None + fc2_output_workspace: Any = None + + +# ============================================================================= +# Sm107Mxfp8GluFc12Kernel +# ============================================================================= + +class Sm107Mxfp8GluFc12Kernel: + + # SMEM budget for all "non-problem-tensor" buffers (mbarriers, sched + # work-tile buffer, TMEM allocator state). Reserved at host side in + # ``_compute_stages``. Bump if ``SharedStorage`` over-allocates SMEM. + _SmemMiscBudget = 1024 + + # Supported (ab_dtype, sf_vec_size) pairings. + # MXFP8 → Float8E4M3FN / Float8E5M2 + sf_vec_size=32 (FP8-E8M0 scales, MmaMXF8Op) + VALID_AB_DTYPE_SF_SIZE: dict = { + 32: (cutlass.Float8E4M3FN, cutlass.Float8E5M2,), + } + + # Interleave granularity for gate and up in SwiGLU / GeGlu + GateUpInterleave: int = 32 + + def __init__( + self, + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mnk: Tuple[int, int, int], + use_2cta_instrs: bool, + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + load_balance_mode: Literal["static", "atomic_counter"] = "static", + static_expert_shape: Optional[Tuple[int, int, int]] = None, + force_static_sched: bool = True, + clc_bundle_size: Optional[int] = None, + num_sched_stages: Optional[int] = None, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + sf_vec_size: int = 32, + ab_dtype: Type[cutlass.Numeric] = cutlass.Float4E2M1FN, + fc2_in_kernel_topk_reduce: bool = False, + token_back_by_dispatch: bool = False, + epi_flag_batch: Tuple[int, int] = (1, 1), + gate_up_clamp: Optional[float] = None, + apply_topk_in_fc1: bool = False, + generate_c: bool = False, + use_stg_fc1: bool = False, + act_func: str = "swiglu", + combine_format: Optional[Any] = None, + fc2_use_bulk: bool = False, + fc2_tma_stages: Optional[int] = None, + ) -> None: + if not force_static_sched: + raise NotImplementedError( + "v1 only implements force_static_sched=True (lean 7-warp). " + "Dynamic CLC (force_static_sched=False) is future work." + ) + + # Validate (ab_dtype, sf_vec_size) pairing. + if sf_vec_size in self.VALID_AB_DTYPE_SF_SIZE: + valid_ab = self.VALID_AB_DTYPE_SF_SIZE[sf_vec_size] + if ab_dtype not in valid_ab: + raise ValueError( + f"ab_dtype={ab_dtype.__name__} is not valid for " + f"sf_vec_size={sf_vec_size}. " + f"Expected one of: {[t.__name__ for t in valid_ab]}." + ) + else: + valid_sf_vec_sizes = tuple(self.VALID_AB_DTYPE_SF_SIZE) + raise NotImplementedError( + f"sf_vec_size must be one of {valid_sf_vec_sizes} (MXFP8); got {sf_vec_size}." + ) + + + if load_balance_mode not in ("static", "atomic_counter"): + raise ValueError( + f"load_balance_mode must be 'static' or 'atomic_counter'; " + f"got {load_balance_mode!r}." + ) + if act_func not in ("swiglu", "geglu"): + raise ValueError( + f"act_func must be 'swiglu' or 'geglu'; got {act_func!r}." + ) + if act_func != "swiglu": + raise NotImplementedError( + f"act_func={act_func!r} is not yet implemented; only " + "'swiglu' is currently supported (geglu support is planned)." + ) + + # Only (M=256, N=256) with 2-CTA instructions is validated now. + m, n, _k = mma_tiler_mnk + if (m, n) != (256, 256) or not use_2cta_instrs: + raise ValueError( + "Sm107Mxfp8GluFc12Kernel only supports mma_tiler (M, N) = " + "(256, 256) with use_2cta_instrs=True; " + f"got mma_tiler_mnk={mma_tiler_mnk}, use_2cta_instrs={use_2cta_instrs}." + ) + + # Store ab_dtype so workspace-size helpers can use it without tensors. + self.ab_dtype = ab_dtype + self.c_dtype = cutlass.BFloat16 + self.act_func = act_func + self.combine_format = combine_format + + self.fc2_in_kernel_topk_reduce = fc2_in_kernel_topk_reduce + self.token_back_by_dispatch = token_back_by_dispatch + self.epi_flag_batch = epi_flag_batch + self.apply_topk_in_fc1 = apply_topk_in_fc1 + self.generate_c = generate_c + self.use_stg_fc1 = use_stg_fc1 + self.fc2_use_bulk = fc2_use_bulk + self._fc2_tma_stages_arg = fc2_tma_stages + self.fc2_tma_stages = fc2_tma_stages if fc2_tma_stages is not None else 0 + self.gate_up_clamp = ( + abs(gate_up_clamp) if gate_up_clamp is not None else None + ) + + self.acc_dtype = acc_dtype + self.mma_tiler_mnk = mma_tiler_mnk + self.cluster_shape_mn = (cluster_shape_mnk[0], cluster_shape_mnk[1]) + self.use_2cta_instrs = use_2cta_instrs + self.force_static_sched = force_static_sched + self.static_expert_shape = static_expert_shape + self.clc_bundle_size = clc_bundle_size + self.num_sched_stages = num_sched_stages + + # Fused fc12 sched-side knobs + self.group_hint = group_hint + self.token_padding_block = token_padding_block + self.sf_padding_block = sf_padding_block + self.load_balance_mode = load_balance_mode + + self.sf_vec_size = sf_vec_size + self.arch = "sm_107" + + self._validate_mma_tiler_and_cluster_shape() + self.mma_tiler = mma_tiler_mnk + + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + # Warp specialization (lean 8-warp / 256 thread) + self.occupancy = 1 + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_a_warp_id = 5 + self.tma_b_warp_id = 6 + self.sched_warp_id = 7 + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_a_warp_id, + self.tma_b_warp_id, + self.sched_warp_id, + *self.epilogue_warp_id, + ) + ) + + # NamedBarrier IDs. + # + # Per-subtile rotated-leader scheme lives inside the epilogue; this + # kernel only owns the four reserved IDs and forwards + # them via ``self.epilog_sync_bar_id`` to the epilogue ctor. + # IDs 8 and 9 are reserved for MXFP8 warp-pair absmax exchange. + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + self.epi_subtile_bar_ids = (4, 5, 6, 7) + + self.enable_token_comm: bool = False + self.dispatch_warp_id: Optional[Tuple[int, int, int, int]] = None + self.token_back_warp_id: Optional[Tuple[int, int, int, int]] = None + self.token_back_standalone: bool = False + + self.smem_capacity = utils.get_smem_capacity_in_bytes() + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols( + self.arch + ) + + def _validate_mma_tiler_and_cluster_shape(self) -> None: + """Validate user-provided geometry against v1 fused-fc12 constraints. + + ``mma_tiler_n`` is restricted to {128, 256}. Short-N is handled by + the swap-AB scheduler via subtile-level early-exit. + """ + m, n, k = self.mma_tiler_mnk + cm, cn = self.cluster_shape_mn + + if m not in SupportedMmaTileM: + raise ValueError( + f"mma_tiler M ({m}) must be one of {SupportedMmaTileM}" + ) + + per_cta_m = m // (2 if self.use_2cta_instrs else 1) + if per_cta_m != 128: + raise ValueError( + f"per-CTA mma_tiler M must be 128, got {per_cta_m} " + f"(mma_tiler_m={m}, use_2cta_instrs={self.use_2cta_instrs})" + ) + + if n not in SupportedMmaTileN: + raise ValueError( + f"mma_tiler N ({n}) must be one of {SupportedMmaTileN} in fused fc12 " + f"(N=64 SFB hack is dropped; swap-AB sched handles short-N " + f"via subtile early-exit)." + ) + + sf_k_granularity = self.sf_vec_size * 4 + if k % sf_k_granularity != 0: + raise ValueError( + f"mma_tiler K ({k}) must be a multiple of " + f"sf_vec_size * 4 = {sf_k_granularity}" + ) + + if cm % (2 if self.use_2cta_instrs else 1) != 0: + raise ValueError( + f"cluster_shape M ({cm}) must be even when use_2cta_instrs=True" + ) + + is_pow2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if cm * cn > 16 or not is_pow2(cm) or not is_pow2(cn) or cm > 4 or cn > 4: + raise ValueError( + f"Invalid cluster_shape ({cm}, {cn}): each dim must be " + f"a power of 2 and <= 4, product must be <= 16" + ) + + if cn > 2: + raise NotImplementedError( + f"cluster_n={cn} is not supported yet (deadlocks in the " + f"mainloop, likely the 4-way scale-factor multicast). " + f"cluster_n in {{1, 2}} is supported and verified." + ) + + if cm > 2: + raise NotImplementedError( + f"cluster_m={cm} is not supported yet (residual fc1->fc2 race " + f"gives ~3-4% mismatch). cluster_m in {{1, 2}} is supported " + f"and verified." + ) + + if cn > 1 and self.static_expert_shape is not None: + cta_tile_n = n # N is not split across the 2-CTA (M) pair + cluster_tile_n = cta_tile_n * cn + _experts, intermediate_gateup, hidden = self.static_expert_shape + if intermediate_gateup % cluster_tile_n != 0: + raise ValueError( + f"cluster_n={cn}: fc1 intermediate_gateup " + f"({intermediate_gateup}) must be a multiple of " + f"cta_tile_n * cluster_n (= {cta_tile_n} * {cn} = " + f"{cluster_tile_n}) to avoid ragged N-peers. Ragged-N " + f"(per-peer N-store predication) is not yet supported." + ) + if hidden % cluster_tile_n != 0: + raise ValueError( + f"cluster_n={cn}: fc2 hidden ({hidden}) must be a multiple " + f"of cta_tile_n * cluster_n (= {cta_tile_n} * {cn} = " + f"{cluster_tile_n}) to avoid ragged N-peers. Ragged-N " + f"(per-peer N-store predication) is not yet supported." + ) + + def _create_tiled_mmas(self) -> Tuple[cute.TiledMma, cute.TiledMma]: + common = ( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + ) + # Rubin: the SM107 blockscaled FP8 MMA op instruction + tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma( + *common, self.cta_group, + (*self.mma_inst_shape_mn, 64), + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma( + *common, tcgen05.CtaGroup.ONE, + (*self.mma_inst_shape_mn_sfb, 64), + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + return tiled_mma, tiled_mma_sfb + + def _build_scheduler( + self, *, expert_cnt, intermediate_gateup, hidden_dim, launch_cluster_count + ) -> None: + """Construct FC12 scheduler and its SMEM/device workspaces.""" + work_id_mode = "grid_stride" if self.load_balance_mode == "static" else "atomic_counter" + # Consumer group = every warp that calls ``consume_work`` (tma_a, tma_b, mma, epilogue). + num_scheduler_consumer_threads = 32 * (len(self.epilogue_warp_id) + 3) + if self.static_expert_shape is not None: + expert_cnt, intermediate_gateup, hidden_dim = self.static_expert_shape + problem_desc = ProblemDesc( + { + "expert_count": expert_cnt, + "intermediate_gateup_size": intermediate_gateup, + "hidden_size": hidden_dim, + } + ) + impl_desc = ImplDesc( + { + "num_scheduler_consumer_threads": num_scheduler_consumer_threads, + "mma_tiler_mnk": self.mma_tiler, + "cluster_shape_mn": self.cluster_shape_mn, + "use_2cta_instrs": self.use_2cta_instrs, + "hint": self.group_hint, + "token_padding_block": self.token_padding_block, + "sf_padding_block": self.sf_padding_block, + "work_id_mode": work_id_mode, + "is_swap_ab": False, + "launch_cluster_count": launch_cluster_count, + } + ) + self.scheduler = BlackwellFusedFc12Scheduler(problem_desc, impl_desc) + + sched_smem_ws = SmemWorkspace() + self.scheduler.register_smem_regions(sched_smem_ws) + sched_smem_ws.finalize(max_bytes=self.smem_capacity) + self.sched_smem_ws = sched_smem_ws + + sched_device_ws = DeviceWorkspace() + self.scheduler.register_device_workspace(sched_device_ws) + sched_device_ws.finalize() + self.sched_device_ws = sched_device_ws + + def _setup_attributes(self) -> None: + """Set up MMA / cluster / tile shapes, SMEM layouts, stage counts. + + The fc12 path shares ``mma_tiler_mnk`` and SMEM layouts across phases. + """ + if self.enable_token_comm: + self.dispatch_warp_id = (8, 9, 10, 11) + num_token_back_warps = ( + len(self.token_back_warp_id) if self.token_back_standalone else 0 + ) + self.threads_per_cta = 32 * ( + len(self.epilogue_warp_id) + + 1 # mma + + 1 # tma_a + + 1 # tma_b + + 1 # sched + + len(self.dispatch_warp_id) + + num_token_back_warps + ) + + self.mma_inst_shape_mn = (self.mma_tiler[0], self.mma_tiler[1]) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + tiled_mma, tiled_mma_sfb = self._create_tiled_mmas() + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + assert self.mma_tiler[2] % mma_inst_shape_k == 0, ( + f"mma_tiler K ({self.mma_tiler[2]}) must be a multiple of " + f"MMA instruction K ({mma_inst_shape_k})" + ) + + # SFB-specific tiler: rounded-up MN; same K as main tiler. + self.mma_tiler_sfb = ( + self.mma_inst_shape_mn_sfb[0], + self.mma_inst_shape_mn_sfb[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + # Multicast CTA counts + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + # Epilogue owns all epi-side decisions (acc stages, subtile dispatch, + # TMA commit/drain, and piggyback red.add). + _epi_common = dict( + mma_tiler_mnk=self.mma_tiler, + cluster_shape_mn=self.cluster_shape_mn, + use_2cta_instrs=self.use_2cta_instrs, + sf_vec_size=self.sf_vec_size, + fc1_output_dtype=self.fc1_output_dtype, + fc1_output_layout=self.fc1_output_layout, + acc_dtype=self.acc_dtype, + epilog_sync_bar_id=self.epilog_sync_bar_id, + epilogue_warp_ids=self.epilogue_warp_id, + static_expert_shape=self.static_expert_shape, + fc2_in_kernel_topk_reduce=self.fc2_in_kernel_topk_reduce, + token_back_by_dispatch=self.token_back_by_dispatch, + epi_flag_batch=self.epi_flag_batch, + glu_clamp=self.gate_up_clamp, + apply_topk_in_fc1=self.apply_topk_in_fc1, + generate_c=self.generate_c, + use_stg_fc1=self.use_stg_fc1, + combine_format=getattr(self, "combine_format", None), + act_func=self.act_func, + fc2_use_bulk=self.fc2_use_bulk, + fc2_tma_stages=self._fc2_tma_stages_arg, + ) + self.epilogue = GluMxfp8Epilogue(**_epi_common) + + if self.num_sched_stages is None: + self.num_sched_stages = 2 + + self.num_d_stage = self.epilogue.subtile_cnt + # fc1 output (fp8 quantised) SMEM — always present. + d_bytes_total = self.epilogue.bytes_per_stage * self.num_d_stage + # Raw gate+up SMEM (BF16, ping-pong) — only when generate_c=True. + c_bytes_total = 0 + if self.generate_c: + from .glu_mxfp8_fc12_epilogue import Fc1CTMAStages + self.num_c_raw_stage = Fc1CTMAStages + c_bytes_total += self.epilogue.c_bytes_per_stage * self.num_c_raw_stage + else: + self.num_c_raw_stage = 0 + + # FC2 TMASTG staging SMEM + fc2_tma_bytes_total = ( + self.epilogue.fc2_tma_staging_bytes + self.epilogue.fc2_reduce_staging_bytes + ) + + ( + self.num_acc_stage, + self.num_a_stage, + self.num_b_stage, + self.num_sched_stages, + ) = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.sf_dtype, + self.sf_vec_size, + c_bytes_total + d_bytes_total + fc2_tma_bytes_total, + self.smem_capacity, + self.occupancy, + self.num_sched_stages, + ) + + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_a_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_b_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_a_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_b_stage, + ) + self.d_smem_layout_staged = self.epilogue.staged_smem_layout( + self.num_d_stage, + ) + # Raw gate+up SMEM layout (only meaningful when generate_c=True; pass as + # dummy to kernel when False). + if self.generate_c: + self.c_smem_layout_staged = self.epilogue.staged_c_smem_layout( + self.num_c_raw_stage + ) + else: + self.c_smem_layout_staged = None + + # Read epilogue's accumulator and scale-factor sizing decisions. + self.num_acc_pipeline_stages = self.epilogue.num_acc_pipeline_stages + self.num_acc_stage = self.epilogue.num_acc_stage + self.num_sfa_tmem_cols = self.epilogue.num_sfa_tmem_cols + self.num_sfb_tmem_cols = self.epilogue.num_sfb_tmem_cols + self.num_accumulator_tmem_cols = self.epilogue.num_accumulator_tmem_cols + + # TMA load bytes per stage (A + B + SFA + SFB). + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + self.atom_thr_size = atom_thr_size # store as Python int for use in @cute.kernel + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_a_bytes = (a_copy_size + sfa_copy_size) * atom_thr_size + self.num_tma_load_b_bytes = (b_copy_size + sfb_copy_size) * atom_thr_size + + def _smem_misc_budget_bytes(self) -> int: + """Per-CTA SMEM reserved outside the ABC-stage pipeline.""" + return self._SmemMiscBudget + + def _compute_stages( + self, + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_bytes_total: int, + smem_capacity: int, + occupancy: int, + num_sched_stages: int, + ) -> Tuple[int, int, int]: + """Compute stage counts for ACC, AB+SF, and scheduler.""" + num_acc_stage = 2 + + a_smem_layout_staged_one = sm100_utils.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1, + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 1, + ) + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, mma_tiler_mnk, sf_vec_size, 1, + ) + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, mma_tiler_mnk, sf_vec_size, 1, + ) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_staged_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + b_bytes_per_stage = ( + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + + fixed_overhead = ( + self._smem_misc_budget_bytes() + c_bytes_total + ) + + num_ab_stage = ( + smem_capacity // occupancy - fixed_overhead + ) // ab_bytes_per_stage + num_a_stage = num_ab_stage + num_b_stage = num_ab_stage + + smem_per_cta = smem_capacity // occupancy + unused_smem = smem_per_cta - fixed_overhead - num_ab_stage * ab_bytes_per_stage + if unused_smem > b_bytes_per_stage: + num_b_stage = num_b_stage + 1 + unused_smem = unused_smem - b_bytes_per_stage + print( + f"[fc12 stages] num_ab_stage={num_a_stage, num_b_stage} " + f"ab_bytes_per_stage={ab_bytes_per_stage} " + f"num_acc_stage={num_acc_stage} " + f"misc_budget={self._smem_misc_budget_bytes()} " + f"c_bytes_total={c_bytes_total} " + f"smem_cap={smem_capacity} " + f"unused_smem={unused_smem}" + ) + + return num_acc_stage, num_a_stage, num_b_stage, num_sched_stages + + def get_workspace_size_in_bytes( + self, + fc1_activation_tensor, + fc1_weight_tensor, + ) -> int: + """Compute opaque workspace size for one fused fc1+fc2 launch.""" + sf_padding_block = self.sf_padding_block + sf_vec_size = self.sf_vec_size + + mma_tiler_n = self.mma_tiler_mnk[1] + + data_total_rows, _hidden = fc1_activation_tensor.shape + experts, _hidden_w, intermediate_gateup = fc1_weight_tensor.shape + intermediate_downproj = intermediate_gateup // 2 + + # Conservative upper bound for sf_total_rows. + sf_total_rows_upper = data_total_rows + experts * sf_padding_block + + fc1_output_bytes = ( + data_total_rows * intermediate_downproj * self.ab_dtype.width // 8 + ) + + # fc1_output_sf sf_vec_size matches the kernel's sf_vec_size. + fc1_out_sf_vec_size = self.sf_vec_size + sf_block_cols = ( + (intermediate_downproj // fc1_out_sf_vec_size) + 3 + ) // 4 * 4 + fc1_output_sf_bytes = sf_total_rows_upper * sf_block_cols + + # fc1_done_counter: one Int32 per global token block, plus expert slack. + counter_slots_upper = ( + (data_total_rows + mma_tiler_n - 1) // mma_tiler_n + + experts + ) + fc1_done_counter_bytes = counter_slots_upper * 4 + + # load_balance_counter: Int32 scalar. + if self.load_balance_mode == "atomic_counter": + load_balance_counter_bytes = 4 + else: + load_balance_counter_bytes = 0 + + total = ( + fc1_output_bytes + + fc1_output_sf_bytes + + fc1_done_counter_bytes + + load_balance_counter_bytes + ) + + # 128B align (TMA tensor base address alignment requirement). + alignment = 128 + total = ((total + alignment - 1) // alignment) * alignment + return total + + def mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """SMEM → TMEM tiled copy + partition for SFA / SFB.""" + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t, tCsSF_compact_s2t_ + ) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + def token_comm_extra_smem_storage_class(self) -> type: + """Return a ``@cute.struct`` for dispatch-warp SMEM, or None.""" + return None + + def token_comm_hook_fc1_ready_counter_ptr(self, token_comm_args): + """Return dispatch->fc1 release counter pointer, or None (lean: disabled).""" + return None + + def sched_ext_fc1_peek_threshold(self) -> int: + """Return the fc1 ready-counter peek threshold for GluMxFp8Fc12SchedExtension. + + Must match the spin threshold in ``token_comm_hook_fc1_tma_b_predispatch_spin`` + so that an early peek hit does not skip the spin and expose stale pool rows. + Default 0 → use ``valid_tokens_in_tile`` (no cluster, base class behaviour). + MegaMoE overrides to return ``cluster_tile_tokens`` to match the cluster spin. + """ + return 0 + + def sched_ext_fc1_counter_cumul_scale(self) -> int: + """Return the scale factor for the fc1 ready-counter slot formula. + + Slot = scale * (cumul + fc1_counter_index) + tile_m_idx % scale. + Default 1 = cluster-level granularity (slot = cumul + cluster_token_block_idx). + A MegaMoE subclass can override to cluster_m for per-CTA granularity. + """ + return 1 + + @cute.jit + def token_comm_hook_sched_warp_pre_init_wait(self, token_comm_args): + """Sched warp: wait for dispatch barrier before reading sizes. No-op base.""" + pass + + @cute.jit + def token_comm_hook_fc1_tma_b_predispatch_spin(self, token_comm_args, work_tile_info): + """TMA warp: spin until dispatch-pulled tokens are resident. No-op base.""" + pass + + @cute.jit + def token_comm_hook_dispatch_warp_body( + self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx, + ): + """Body for dispatch warps 8-11 (MegaMoE-only). No-op base.""" + pass + + @cute.jit + def token_comm_hook_token_back_warp_body( + self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx, + ): + """Body for standalone token-back warps 12-15 (MegaMoE-only). No-op base.""" + pass + + @cute.jit + def token_comm_hook_kernel_tail(self, token_comm_args, *, warp_idx, lane_idx, tidx): + """All-warp kernel tail (NVLink release, etc.). No-op base.""" + pass + + @cute.jit + def __call__( + self, + activation: cute.Tensor, # (token_sum_padded, hidden) + fc1_weight: cute.Tensor, # (experts, hidden, intermediate_gateup) + activation_sf: cute.Tensor, # (token_sum_padded_sf, hidden / sf_vec_size) + fc1_weight_sf: cute.Tensor, # (experts, intermediate_gateup_padded * hidden / sf_vec_size) + fc1_output: cute.Tensor, # (token_sum_padded, intermediate_downproj) + fc1_output_sf: cute.Tensor, # (token_sum_padded_sf, intermediate_downproj / sf_vec_size) + fc2_weight: cute.Tensor, # (experts, intermediate_downproj, hidden) + fc2_weight_sf: cute.Tensor, # (experts, hidden_padded * intermediate_downproj / sf_vec_size) + fc2_output: cute.Tensor, # (token_sum_padded, hidden) BFloat16, hidden stride-1 + topk_scores: cute.Tensor, # (token_sum_padded,) Float32 + fc1_done_counter: cute.Tensor, # (max_token_block_per_rank,) Int32 + offs: Optional[cute.Tensor] = None, # (experts,) Int32 cumulative end offsets + max_active_clusters: cutlass.Constexpr = None, + stream: cuda.CUstream = None, + norm_const_tensor: Optional[cute.Tensor] = None, + global_activation_sf: Optional[cute.Tensor] = None, + global_fc1_weight_sf: Optional[cute.Tensor] = None, + load_balance_counter: Optional[cute.Tensor] = None, + expert_token_sizes: Optional[cute.Tensor] = None, + token_comm_args=None, + fc1_c: Optional[cute.Tensor] = None, + # ── Per-rank FC12 overflow output ──────────────────────────────── + overflow_flag: cute.Tensor = None, + fc2_output_sf: Optional[cute.Tensor] = None, + mega_peer_rank_ptr_mapper=None, + mega_local_rank: Optional[cutlass.Int32] = None, + mega_local_workspace: Optional[cute.Pointer] = None, + mega_shared_workspace: Optional[cute.Pointer] = None, + mega_activation: Optional[cute.Tensor] = None, + mega_activation_sf: Optional[cute.Tensor] = None, + mega_pre_reduced_activation: Optional[cute.Tensor] = None, + mega_pre_reduced_activation_sf: Optional[cute.Tensor] = None, + ) -> None: + """Launch the fused fc1+fc2 GLU MXFP8 kernel.""" + + if cutlass.const_expr(self.static_expert_shape is not None): + ( + experts_static, + intermediate_gateup_static, + hidden_static, + ) = self.static_expert_shape + intermediate_downproj_static = intermediate_gateup_static // 2 + + fc1_weight = cute.make_tensor( + fc1_weight.iterator, + cute.make_layout( + (experts_static, hidden_static, intermediate_gateup_static), + stride=fc1_weight.stride, + ), + ) + fc2_weight = cute.make_tensor( + fc2_weight.iterator, + cute.make_layout( + (experts_static, intermediate_downproj_static, hidden_static), + stride=fc2_weight.stride, + ), + ) + activation = cute.make_tensor( + activation.iterator, + cute.make_layout( + (activation.shape[0], hidden_static), + stride=activation.stride, + ), + ) + fc1_output = cute.make_tensor( + fc1_output.iterator, + cute.make_layout( + (fc1_output.shape[0], intermediate_downproj_static), + stride=fc1_output.stride, + ), + ) + # fc2_output is 2D (tokens, hidden) on the lean path and 3D + # (max_tokens, topk, hidden) on the MegaMoE path. + if cutlass.const_expr(len(fc2_output.shape) == 3): + fc2_output = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (fc2_output.shape[0], fc2_output.shape[1], hidden_static), + stride=fc2_output.stride, + ), + ) + else: + fc2_output = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (fc2_output.shape[0], hidden_static), + stride=fc2_output.stride, + ), + ) + + # GEMM-domain transform for fc1 phase + c1 = cutlass.Int32(1) + c0 = cutlass.Int32(0) + + # A_gemm (fc1 activations): (tokens_sum, hidden) -> (M=tokens, K=hidden, L=1). + tokens_sum, hidden = activation.shape + activation_gemm = cute.make_tensor( + activation.iterator, + cute.make_layout( + (tokens_sum, hidden, 1), + stride=(activation.stride[0], activation.stride[1], 0), + ), + ) + + # B_gemm (fc1 weights): (experts, hidden, intermediate_gateup) with hidden stride-1 (K-major) + # -> (N=intermediate_gateup, K=hidden, L=experts). + experts, hidden_b, intermediate_gateup = fc1_weight.shape + fc1_weight_gemm = cute.make_tensor( + fc1_weight.iterator, + cute.make_layout( + (intermediate_gateup, hidden_b, experts), + stride=(fc1_weight.stride[2], fc1_weight.stride[1], fc1_weight.stride[0]), + ), + ) + + # D_gemm is a user-view output tensor; epilogue owns its store path. + intermediate_downproj = fc1_output.shape[1] + fc1_output_gemm = cute.make_tensor( + fc1_output.iterator, + cute.make_layout( + (tokens_sum, intermediate_downproj, 1), + stride=(fc1_output.stride[0], fc1_output.stride[1], 0), + ), + ) + + # SFA / SFB scale tensors (atom-tiled) — fc1 phase. + # SFA (mma M-side) = activation_sf (activation scales, A-side) + # SFB (mma N-side) = fc1_weight_sf (weight scales, B-side) + tokens_sum_padded = activation_sf.shape[0] + hidden_padded = activation_sf.shape[1] * self.sf_vec_size + activation_sf_gemm = cute.make_tensor( + activation_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (tokens_sum_padded, hidden_padded, 1), self.sf_vec_size + ), + ) + intermediate_gateup_padded_mul_hidden_padded = fc1_weight_sf.shape[1] + intermediate_gateup_padded = ( + intermediate_gateup_padded_mul_hidden_padded * self.sf_vec_size + ) // hidden_padded + fc1_weight_sf_gemm = cute.make_tensor( + fc1_weight_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (intermediate_gateup_padded, hidden_padded, experts), + self.sf_vec_size, + ), + ) + + # ── GEMM-domain transform for fc2 phase ── + # + # fc2 roles: M=hidden, N=tokens_sum, K=intermediate_downproj. + + # A_gemm (fc2 weights): (experts, intermediate_downproj, hidden) + # -> (M=hidden, K=intermediate_downproj, L=experts). + experts2, intermediate_downproj_b2, hidden_b2 = fc2_weight.shape + fc2_weight_gemm = cute.make_tensor( + fc2_weight.iterator, + cute.make_layout( + (hidden_b2, intermediate_downproj_b2, experts2), + stride=(fc2_weight.stride[2], fc2_weight.stride[1], fc2_weight.stride[0]), + ), + ) + + if cutlass.const_expr(len(fc2_output.shape) == 3): + fc2_hidden_out = fc2_output.shape[2] + fc2_output_gemm = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (fc2_output.shape[0], fc2_hidden_out, c1), + stride=(fc2_output.stride[0], fc2_output.stride[2], c0), + ), + ) + else: + fc2_hidden_out = fc2_output.shape[1] + fc2_output_gemm = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (tokens_sum, fc2_hidden_out, c1), + stride=(fc2_output.stride[0], fc2_output.stride[1], c0), + ), + ) + + # SFA / SFB for fc2: + # SFA (mma M-side) = fc2_weight_sf (fc2 weight scales, sf_vec_size) + # SFB (mma N-side) = fc1_output_sf (fc1 epilogue SFs, uses sf_vec_size) + # Both paths produce SFs with self.sf_vec_size: + fc1_out_sf_vec_size = self.sf_vec_size + tokens_sum_padded_sf = fc1_output_sf.shape[0] + intermediate_downproj_padded = fc1_output_sf.shape[1] * fc1_out_sf_vec_size + fc1_output_sf_gemm_for_fc2_load = cute.make_tensor( + fc1_output_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (tokens_sum_padded_sf, intermediate_downproj_padded, 1), + fc1_out_sf_vec_size, + ), + ) + + hidden_padded_fc2_mul_intermediate_downproj_padded = fc2_weight_sf.shape[1] + hidden_padded_fc2 = ( + hidden_padded_fc2_mul_intermediate_downproj_padded * self.sf_vec_size + ) // intermediate_downproj_padded + fc2_weight_sf_gemm = cute.make_tensor( + fc2_weight_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (hidden_padded_fc2, intermediate_downproj_padded, experts2), + self.sf_vec_size, + ), + ) + + expert_cnt = experts + hidden_dim = hidden + + # ── Infer dtypes and major modes ── + self.a_dtype: Type[cutlass.Numeric] = activation_gemm.element_type + self.b_dtype: Type[cutlass.Numeric] = fc1_weight_gemm.element_type + self.fc1_output_dtype: Type[cutlass.Numeric] = fc1_output_gemm.element_type + self.sf_dtype: Type[cutlass.Numeric] = activation_sf_gemm.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(activation_gemm).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(fc1_weight_gemm).mma_major_mode() + self.fc1_output_layout = utils.LayoutEnum.from_tensor(fc1_output_gemm) + + self._setup_attributes() + tiled_mma, tiled_mma_sfb = self._create_tiled_mmas() + + # ── fc1 TMA atoms ── + + # TMA load A1 + a_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_fc1_activation, tma_tensor_fc1_activation = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + activation_gemm, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load B1 + b_op = sm100_utils.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_fc1_weight, tma_tensor_fc1_weight = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + fc1_weight_gemm, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load SFA1 + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_fc1_activation_sf, tma_tensor_fc1_activation_sf = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + activation_sf_gemm, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64, + ) + + # TMA load SFB1 + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_fc1_weight_sf, tma_tensor_fc1_weight_sf = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + fc1_weight_sf_gemm, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Uint64, + ) + + # TMA store for fc1 MXFP8 output. + fc1_output_tma_op = cpasync.CopyBulkTensorTileS2GOp() + tma_atom_fc1_output, tma_tensor_fc1_output = cpasync.make_tiled_tma_atom( + fc1_output_tma_op, + fc1_output_gemm, + self.epilogue.smem_layout_one_stage, + self.epilogue.epi_tile, + ) + + # TMA store for raw fc1 accumulator + if cutlass.const_expr(self.generate_c): + c_gemm = cute.make_tensor( + fc1_c.iterator, + cute.make_layout( + (tokens_sum, intermediate_gateup, 1), + stride=(fc1_c.stride[0], fc1_c.stride[1], 0), + ), + ) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c_gemm, + self.epilogue.c_smem_layout_one_stage, + self.epilogue.epi_tile_c, + ) + else: + tma_atom_c = tma_atom_fc1_output + tma_tensor_c = tma_tensor_fc1_output + + # fc1 SFC GMEM tensor (= fc1_output_sf user view). + fc1_output_sf_gemm = cute.make_tensor( + fc1_output_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (tokens_sum_padded, intermediate_downproj, 1), + self.sf_vec_size, + ), + ) + + # ── fc2 TMA atoms: fc1_output → A-side (M=tokens), fc2_weight → B-side (N=hidden) ── + tma_atom_fc2_activation, tma_tensor_fc2_activation = ( + cute.nvgpu.make_tiled_tma_atom_A( + a_op, + fc1_output_gemm, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + ) + tma_atom_fc2_weight, tma_tensor_fc2_weight = ( + cute.nvgpu.make_tiled_tma_atom_B( + b_op, + fc2_weight_gemm, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + ) + tma_atom_fc2_activation_sf, tma_tensor_fc2_activation_sf = ( + cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + fc1_output_sf_gemm_for_fc2_load, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64, + ) + ) + tma_atom_fc2_weight_sf, tma_tensor_fc2_weight_sf = ( + cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + fc2_weight_sf_gemm, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Uint64, + ) + ) + + # ── Scheduler params + grid + launch ── + if cutlass.const_expr(self.load_balance_mode == "atomic_counter"): + if cutlass.const_expr(load_balance_counter is None): + raise ValueError( + "load_balance_counter must be provided when " + "load_balance_mode == 'atomic_counter'" + ) + load_balance_counter_ptr = load_balance_counter.iterator + else: + load_balance_counter_ptr = None + + # On the MegaMoE path the per-expert sizes come from the Router (device-side), so the + # caller supplies neither offs nor expert_token_sizes. + if cutlass.const_expr(not self.enable_token_comm): + if cutlass.const_expr((offs is None) == (expert_token_sizes is None)): + raise ValueError( + "Exactly one of `offs` / `expert_token_sizes` must be " + "provided; got " + f"offs={'set' if offs is not None else 'None'}, " + f"expert_token_sizes=" + f"{'set' if expert_token_sizes is not None else 'None'}." + ) + + self._build_scheduler( + expert_cnt=expert_cnt, + intermediate_gateup=intermediate_gateup, + hidden_dim=hidden_dim, + launch_cluster_count=max_active_clusters, + ) + grid = self.scheduler.get_grid_shape(max_active_clusters=max_active_clusters) + + # FC2 TMASTG (fc2_use_bulk) store atom (host-side) + if cutlass.const_expr(self.enable_token_comm and self.epilogue.fc2_use_tma): + self._mega_device_workspace.assign_device_members( + mega_local_workspace, mega_shared_workspace + ) + _fc2_pool3d = self.token_comm.fc2_activation_tensor(self._mega_device_workspace) + _fc2_pool2d = cute.make_tensor( + _fc2_pool3d.iterator, + cute.make_layout( + (_fc2_pool3d.shape[0], _fc2_pool3d.shape[2]), + stride=(_fc2_pool3d.stride[0], _fc2_pool3d.stride[2]), + ), + ) + tma_atom_fc2_output, fc2_tma_output = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + _fc2_pool2d, + self.epilogue.fc2_tma_smem_layout_one_stage, + self.epilogue.fc2_tma_tile, + ) + else: + tma_atom_fc2_output = None + fc2_tma_output = None + + self.kernel( + tiled_mma, + tiled_mma_sfb, + # fc1 TMA atoms / tensors (A=activations, B=weights) + tma_atom_fc1_activation, + tma_tensor_fc1_activation, + tma_atom_fc1_weight, + tma_tensor_fc1_weight, + tma_atom_fc1_activation_sf, + tma_tensor_fc1_activation_sf, + tma_atom_fc1_weight_sf, + tma_tensor_fc1_weight_sf, + tma_atom_fc1_output, + tma_tensor_fc1_output, + # fc2 TMA atoms / tensors (fc1_output→A, fc2_weight→B) + tma_atom_fc2_activation, + tma_tensor_fc2_activation, + tma_atom_fc2_weight, + tma_tensor_fc2_weight, + tma_atom_fc2_activation_sf, + tma_tensor_fc2_activation_sf, + tma_atom_fc2_weight_sf, + tma_tensor_fc2_weight_sf, + # GEMM-domain tensors (fc1) + activation_gemm, + fc1_weight_gemm, + fc1_output_gemm, + activation_sf_gemm, + fc1_weight_sf_gemm, + fc1_output_sf_gemm, + # GEMM-domain tensors (fc2) + fc2_weight_gemm, + fc2_output_gemm, + fc2_weight_sf_gemm, + fc1_output_sf_gemm_for_fc2_load, + # topk + cross-phase sync workspace + topk_scores, + fc1_done_counter, + # Scheduling + offs, + expert_token_sizes, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + # SMEM layouts + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.d_smem_layout_staged, + self.c_smem_layout_staged, + tma_atom_c, + tma_tensor_c, + overflow_flag, + token_comm_args, + fc2_output_sf, + mega_peer_rank_ptr_mapper, + mega_local_rank, + mega_local_workspace, + mega_shared_workspace, + mega_activation, + mega_activation_sf, + mega_pre_reduced_activation, + mega_pre_reduced_activation_sf, + tma_atom_fc2_output, + fc2_tma_output, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=self.occupancy, + ) + + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + # fc1 TMA atoms / tensors + tma_atom_fc1_activation_1: cute.CopyAtom, + tma_tensor_fc1_activation_1: cute.Tensor, + tma_atom_weight: cute.CopyAtom, + tma_tensor_weight: cute.Tensor, + tma_atom_fc1_activation_1_sf: cute.CopyAtom, + tma_tensor_fc1_activation_1_sf: cute.Tensor, + tma_atom_fc1_weight_sf: cute.CopyAtom, + tma_tensor_fc1_weight_sf: cute.Tensor, + tma_atom_fc1_output: cute.CopyAtom, + tma_tensor_fc1_output: cute.Tensor, + # fc2 TMA atoms / tensors (fc1_output→A, fc2_weight→B) + tma_atom_fc2_activation: cute.CopyAtom, + tma_tensor_fc2_activation: cute.Tensor, + tma_atom_fc2_weight: cute.CopyAtom, + tma_tensor_fc2_weight: cute.Tensor, + tma_atom_fc2_activation_sf: cute.CopyAtom, + tma_tensor_fc2_activation_sf: cute.Tensor, + tma_atom_fc2_weight_sf: cute.CopyAtom, + tma_tensor_fc2_weight_sf: cute.Tensor, + # GEMM-domain tensors (fc1) + activation_gemm: cute.Tensor, + fc1_weight_gemm: cute.Tensor, + fc1_output_gemm: cute.Tensor, + activation_sf_gemm: cute.Tensor, + fc1_weight_sf_gemm: cute.Tensor, + fc1_output_sf_gemm: cute.Tensor, + # GEMM-domain tensors (fc2) + fc2_weight_gemm: cute.Tensor, + fc2_output_gemm: cute.Tensor, + fc2_weight_sf_gemm: cute.Tensor, + fc1_output_sf_gemm_for_fc2_load: cute.Tensor, + # topk + cross-phase sync workspace + topk_scores: cute.Tensor, + fc1_done_counter: cute.Tensor, + # debug: (total_tokens, intermediate_half*3) fp32 for swiglu comparison + # Scheduling + offs: Optional[cute.Tensor], + expert_token_sizes: Optional[cute.Tensor], + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + # SMEM layouts + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + d_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout], + c_smem_layout_staged: Optional[Union[cute.Layout, cute.ComposedLayout]] = None, + tma_atom_c: Optional[cute.CopyAtom] = None, + tma_tensor_c: Optional[cute.Tensor] = None, + overflow_flag: cute.Tensor = None, + token_comm_args=None, + fc2_output_sf: Optional[cute.Tensor] = None, + # MegaMoE push-model token-comm inputs + mega_peer_rank_ptr_mapper=None, + mega_local_rank: Optional[cutlass.Int32] = None, + mega_local_workspace: Optional[cute.Pointer] = None, + mega_shared_workspace: Optional[cute.Pointer] = None, + mega_activation: Optional[cute.Tensor] = None, + mega_activation_sf: Optional[cute.Tensor] = None, + mega_pre_reduced_activation: Optional[cute.Tensor] = None, + mega_pre_reduced_activation_sf: Optional[cute.Tensor] = None, + # FC2 TMASTG (fc2_use_bulk): store atom + token-major pool view (None off-path) + tma_atom_fc2_output: Optional[cute.CopyAtom] = None, + fc2_tma_output: Optional[cute.Tensor] = None, + ): + """Device kernel for fused fc1+fc2 swap-AB GLU MXFP8 grouped GEMM.""" + a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, None, 0)) + b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, None, 0)) + sfa_smem_layout = cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)) + sfb_smem_layout = cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)) + + # MegaMoE (push model) + if cutlass.const_expr(self.enable_token_comm): + self._mega_device_workspace.assign_device_members( + mega_local_workspace, mega_shared_workspace + ) + + # fc2 waits for all fc1 intermediate N-tiles in the same token block. + # Each N-tile is processed by atom_thr_size CTAs (both CTA0 and CTA1 increment + # the counter), so the threshold must account for both CTAs' contributions. + ext_fc2_spin_threshold = ( + fc1_weight_gemm.shape[0] + self.cta_tile_shape_mnk[1] - 1 + ) // self.cta_tile_shape_mnk[1] * self.epilogue._atom_thr_size + + ext = GluMxFp8Fc12SchedExtension( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter.iterator, + fc2_spin_threshold=ext_fc2_spin_threshold, + fc1_ready_counter_pointer=self.token_comm_hook_fc1_ready_counter_ptr( + token_comm_args + ), + cluster_m=self.epilogue._atom_thr_size, + ) + + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + tidx, _, _ = cute.arch.thread_idx() + + # MegaMoE (push model): bind token-comm device members (transfer-warp state + the + # NVLink barrier's peer mapper) before any token_in / token_back / size-wait runs. + if cutlass.const_expr(self.enable_token_comm): + _mega_token_comm_args = TokenCommArgs( + mega_activation, + mega_activation_sf, + mega_pre_reduced_activation, + mega_pre_reduced_activation_sf, + mega_peer_rank_ptr_mapper, + ) + _mega_cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] + _, _, _mega_cluster_idx = cute.arch.block_idx() + _mega_linear_cta_idx = cta_rank_in_cluster + _mega_cluster_idx * _mega_cluster_size + self.token_comm.assign_device_members( + device_workspace=self._mega_device_workspace, + token_comm_args=_mega_token_comm_args, + local_rank=mega_local_rank, + linear_cta_idx=_mega_linear_cta_idx, + ) + + # SharedStorage (mainloop + epilogue SMEM). next's scheduler owns its own + # SMEM workspace, allocated separately below. + @cute.struct + class SharedStorage: + a_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_a_stage * 2] + b_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_b_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_pipeline_stages * 2 + ] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # next scheduler SMEM: a self-contained workspace carved from the same + # allocator; its transport regions resolve against ``sched_smem_base``. + sched_storage = smem.allocate(self.sched_smem_ws.storage_class()) + sched_smem_base = sched_storage.buffer.data_ptr() + + # MegaMoE-only dispatch-warp SMEM (pull_buffer, mbarriers, etc.). + # Kept out of ``SharedStorage`` so the lean path never allocates it. + TokenCommStorageCls = self.token_comm_extra_smem_storage_class() + if cutlass.const_expr(TokenCommStorageCls is not None): + token_comm_storage = smem.allocate(TokenCommStorageCls) + else: + token_comm_storage = None + + # ── Pipelines: separate producer/consumer groups for A and B. ── + + a_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, 1 + ) + a_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mcast_ctas_a + ) + a_producer, a_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.a_full_mbar_ptr.data_ptr(), + num_stages=self.num_a_stage, + producer_group=a_pipeline_producer_group, + consumer_group=a_pipeline_consumer_group, + tx_count=self.num_tma_load_a_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + mcast_mode_mn=(1, 0), + defer_sync=True, + ).make_participants() + b_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, 1 + ) + b_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mcast_ctas_b + ) + b_producer, b_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.b_full_mbar_ptr.data_ptr(), + num_stages=self.num_b_stage, + producer_group=b_pipeline_producer_group, + consumer_group=b_pipeline_consumer_group, + tx_count=self.num_tma_load_b_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + mcast_mode_mn=(0, 1), + defer_sync=True, + ).make_participants() + + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = ( + len(self.epilogue_warp_id) * 32 * (2 if use_2cta_instrs else 1) + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_pipeline_stages, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # TMEM allocator + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, + arch=self.arch, + ) + + scheduler = self.scheduler + if cutlass.const_expr(self.enable_token_comm): + _sched_expert_sizes = self.token_comm.local_expert_sizes( + self._mega_device_workspace, mega_local_rank + ) + _sched_prefix_sum = None + else: + _sched_expert_sizes = expert_token_sizes + _sched_prefix_sum = offs + scheduler.assign_device_members( + expert_token_sizes=_sched_expert_sizes, + expert_token_prefix_sum=_sched_prefix_sum, + actual_expert_shape=None, + block_idx=cute.arch.block_idx(), + smem_workspace=self.sched_smem_ws, + smem_base=sched_smem_base, + device_workspace=self.sched_device_ws, + ) + sched_consumer = scheduler.make_consumer() + + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # ── SMEM tensors A / B / SFA / SFB (shared by fc1 / fc2) ── + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + sSFA = smem.allocate_tensor( + element_type=self.sf_dtype, + layout=sfa_smem_layout_staged, + byte_alignment=128, + ) + sSFB = smem.allocate_tensor( + element_type=self.sf_dtype, + layout=sfb_smem_layout_staged, + byte_alignment=128, + ) + + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + + # acc_fake layout: (MMA, MMA_M, MMA_N, STAGE). + acc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # Cluster wait before TMEM alloc. + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + mma_tiler_k = self.mma_tiler[2] + k_tile_cnt_fc1 = (fc1_weight_gemm.shape[1] + mma_tiler_k - 1) // mma_tiler_k + k_tile_cnt_fc2 = (fc2_weight_gemm.shape[1] + mma_tiler_k - 1) // mma_tiler_k + + cluster_n = self.cluster_shape_mn[1] + _fc1_n_cluster_tile = self.cta_tile_shape_mnk[1] * cluster_n + fc2_spin_threshold = ( + ( + (fc1_weight_gemm.shape[0] + _fc1_n_cluster_tile - 1) + // _fc1_n_cluster_tile + ) + * cluster_n + * self.epilogue._atom_thr_size + ) + + # ════════════════════════════════════════════════════════════════════ + # Scheduler warp (warp 7) — lean path + # ════════════════════════════════════════════════════════════════════ + if warp_idx == self.sched_warp_id: + # MegaMoE: block until the Router has published this rank's per-expert sizes + # (cross-rank), so the lazy size walk in ``gen_next_work`` reads valid counts. + # No-op on the lean path. + self.token_comm_hook_sched_warp_pre_init_wait(token_comm_args) + work_tile = scheduler.gen_next_work() + while work_tile.is_valid_tile: + scheduler.publish_work(ext.prepare_work_tile(work_tile)) + work_tile = scheduler.gen_next_work() + # Sentinel publish (the tile is already invalid here). + scheduler.publish_work(work_tile) + scheduler.produce_tail() + + # ════════════════════════════════════════════════════════════════════ + # TMA load warps (warps 5 / 6) + # ════════════════════════════════════════════════════════════════════ + # + # TMA-A loads activations/SFA into the A pipeline. + # TMA-B loads weights/SFB into the B pipeline and waits for + # fc1 workspace readiness in the fc2 phase. + + # ── TMA-A warp (warp 5) ───────────────────────────────────────────── + if warp_idx == self.tma_a_warp_id: + _iket_active = (tidx == cutlass.Int32(160)) + a_full_mcast_mask = None + sfa_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + + b_full_mcast_mask = None + if cutlass.const_expr(self.is_b_mcast or use_2cta_instrs): + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + sfa_cta_layout = a_cta_layout + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + + work_tile_info = sched_consumer.consume_work() + + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + if is_phase_linear1: + # fc1 phase A-side + if _iket_active: + iket.range_push("tma_token_fc1") + ext.wait_for_input(work_tile_info) + self.token_comm_hook_fc1_tma_b_predispatch_spin( + token_comm_args, work_tile_info, + ) + + k_tile_cnt = k_tile_cnt_fc1 + real_a, desc_ptr_a = ext.get_gmem_tensor( + "fc1_activation", tma_tensor_fc1_activation_1, work_tile_info, + ) + real_sfa, desc_ptr_sfa = ext.get_gmem_tensor( + "fc1_activation_sf", tma_tensor_fc1_activation_1_sf, work_tile_info, + ) + + gA_mkl = cute.local_tile( + real_a, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + gSFA_mkl = cute.local_tile( + real_sfa, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + tCgA = thr_mma.partition_A(gA_mkl) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + + tAsA, tAgA = cpasync.tma_partition( + tma_atom_fc1_activation_1, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_fc1_activation_1_sf, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + mma_tile_m = work_tile_info.tile_m_idx // cute.size( + tiled_mma.thr_id.shape + ) + tAgA_slice = tAgA[(None, mma_tile_m, None, 0)] + tAgSFA_slice = tAgSFA[(None, mma_tile_m, None, 0)] + + a_producer.reset() + peek_a_empty_status = a_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Producer-side backpressure: TMA-A blocked waiting for the + # MMA to free an A SMEM slot. First k-tile only (later tiles + # overlap via peek-ahead). Complements mma_ab_operand_wait. + if _iket_active: + iket.range_push("tma_a_buf_acquire_wait") + handle = a_producer.acquire_and_advance( + peek_a_empty_status + ) + if _iket_active: + iket.range_pop() # tma_a_buf_acquire_wait + peek_a_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_a_empty_status = a_producer.try_acquire() + cute.copy( + tma_atom_fc1_activation_1, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_a, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_fc1_activation_1_sf, + tAgSFA_slice[(None, handle.count)], + tAsSFA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfa, + mcast_mask=sfa_full_mcast_mask, + ) + else: + # fc2 phase A-side: load fc1_output (M=tokens) + wait for fc1 done + if _iket_active: + iket.range_push("tma_token_fc2") + counter_slot = ( + work_tile_info.cumulative_token_block_count + + work_tile_info.tile_m_idx // cutlass.Int32(self.epilogue._atom_thr_size) + ) + counter_ptr = fc1_done_counter.iterator + counter_slot + if _iket_active: + iket.range_push("tma_token_fc2_a_wait") + spin_wait( + counter_ptr, + lambda v: v >= fc2_spin_threshold, + sleep_cycles=20, + ) + if _iket_active: + iket.range_pop() + cute.arch.load(counter_ptr, counter_ptr.dtype, sem="acquire", scope="gpu") + cute.arch.fence_proxy("async") + cute.arch.fence_proxy("async.global") + + if ( + tidx == cutlass.Int32(32 * self.tma_a_warp_id) + and work_tile_info.tile_n_idx == cutlass.Int32(0) + ): + counter_val_post = cute.arch.load( + counter_ptr, counter_ptr.dtype, cop="cg" + ) + fc1_byte_offset = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx + // cutlass.Int32(self.epilogue._atom_thr_size) + * cutlass.Int32(self.epilogue._cta_tile_m) + ) * fc1_output_gemm.stride[0] + fc1_probe_ptr = cute.make_ptr( + cutlass.Int32, + fc1_output_gemm.iterator.toint() + fc1_byte_offset, + cute.AddressSpace.gmem, + ) + fc1_first_i32 = cute.arch.load(fc1_probe_ptr, cutlass.Int32, cop="cg") + + k_tile_cnt = k_tile_cnt_fc2 + real_a, desc_ptr_a = ext.get_gmem_tensor( + "fc2_activation", tma_tensor_fc2_activation, work_tile_info, + ) + real_sfa, desc_ptr_sfa = ext.get_gmem_tensor( + "fc2_activation_sf", tma_tensor_fc2_activation_sf, work_tile_info, + ) + + gA_mkl = cute.local_tile( + real_a, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + gSFA_mkl = cute.local_tile( + real_sfa, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + tCgA = thr_mma.partition_A(gA_mkl) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + + tAsA, tAgA = cpasync.tma_partition( + tma_atom_fc2_activation, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_fc2_activation_sf, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + # fc2 A-side = fc1_output (M=tokens). + mma_tile_m = work_tile_info.tile_m_idx // cute.size( + tiled_mma.thr_id.shape + ) + tAgA_slice = tAgA[(None, mma_tile_m, None, 0)] + tAgSFA_slice = tAgSFA[(None, mma_tile_m, None, 0)] + + a_producer.reset() + peek_a_empty_status = a_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Producer-side backpressure: TMA-A blocked waiting for the + # MMA to free an A SMEM slot. First k-tile only (later tiles + # overlap via peek-ahead). Complements mma_ab_operand_wait. + if _iket_active: + iket.range_push("tma_a_buf_acquire_wait") + handle = a_producer.acquire_and_advance( + peek_a_empty_status + ) + if _iket_active: + iket.range_pop() # tma_a_buf_acquire_wait + peek_a_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_a_empty_status = a_producer.try_acquire() + cute.copy( + tma_atom_fc2_activation, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_a, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_fc2_activation_sf, + tAgSFA_slice[(None, handle.count)], + tAsSFA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfa, + mcast_mask=sfa_full_mcast_mask, + ) + + if _iket_active: + iket.range_pop() + work_tile_info = sched_consumer.consume_work() + + a_producer.tail() + + # ── TMA-B warp (warp 6) ───────────────────────────────────────────── + if warp_idx == self.tma_b_warp_id: + _iket_active = (tidx == cutlass.Int32(192)) + b_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr(self.is_b_mcast or use_2cta_instrs): + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, + block_in_cluster_coord_sfb_vmnk, + mcast_mode=1, + ) + + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + sfb_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) + + work_tile_info = sched_consumer.consume_work() + + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + + if is_phase_linear1: + # fc1 phase B-side (fc1_weight, N-side GEMM-B) + if _iket_active: + iket.range_push("tma_weight_fc1") + + k_tile_cnt = k_tile_cnt_fc1 + real_b, desc_ptr_b = ext.get_gmem_tensor( + "fc1_weight", tma_tensor_weight, work_tile_info, + ) + real_sfb, desc_ptr_sfb = ext.get_gmem_tensor( + "fc1_weight_sf", tma_tensor_fc1_weight_sf, work_tile_info, + ) + + # N-K tiling for N-side weight (N=intermediate, K=hidden). + gB_nkl = cute.local_tile( + real_b, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + gSFB_nkl = cute.local_tile( + real_sfb, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + + tCgB = thr_mma.partition_B(gB_nkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + + tBsB, tBgB = cpasync.tma_partition( + tma_atom_weight, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_fc1_weight_sf, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # Use tile_n_idx for N-side weight: invariant across token blocks. + tBgB_slice = tBgB[(None, work_tile_info.tile_n_idx, None, 0)] + tBgSFB_slice = tBgSFB[(None, work_tile_info.tile_n_idx, None, 0)] + + b_producer.reset() + peek_b_empty_status = b_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Producer-side backpressure: TMA-B blocked waiting for the + # MMA to free a B SMEM slot. First k-tile only (later tiles + # overlap via peek-ahead). Complements mma_ab_operand_wait. + if _iket_active: + iket.range_push("tma_b_buf_acquire_wait") + handle = b_producer.acquire_and_advance( + peek_b_empty_status + ) + if _iket_active: + iket.range_pop() # tma_b_buf_acquire_wait + peek_b_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_b_empty_status = b_producer.try_acquire() + cute.copy( + tma_atom_weight, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_b, + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_fc1_weight_sf, + tBgSFB_slice[(None, handle.count)], + tBsSFB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfb, + mcast_mask=sfb_full_mcast_mask, + ) + else: + # fc2 phase B-side: load fc2_weight (N=hidden) + if _iket_active: + iket.range_push("tma_weight_fc2") + k_tile_cnt = k_tile_cnt_fc2 + real_b, desc_ptr_b = ext.get_gmem_tensor( + "fc2_weight", tma_tensor_fc2_weight, work_tile_info, + ) + real_sfb, desc_ptr_sfb = ext.get_gmem_tensor( + "fc2_weight_sf", tma_tensor_fc2_weight_sf, work_tile_info, + ) + + gB_nkl = cute.local_tile( + real_b, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + gSFB_nkl = cute.local_tile( + real_sfb, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + tCgB = thr_mma.partition_B(gB_nkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + + tBsB, tBgB = cpasync.tma_partition( + tma_atom_fc2_weight, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_fc2_weight_sf, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # fc2 B-side = fc2_weight (N=hidden) + fc2_b_hidden_tile = work_tile_info.tile_n_idx + tBgB_slice = tBgB[(None, fc2_b_hidden_tile, None, 0)] + tBgSFB_slice = tBgSFB[(None, fc2_b_hidden_tile, None, 0)] + + b_producer.reset() + peek_b_empty_status = b_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Producer-side backpressure: TMA-B blocked waiting for the + # MMA to free a B SMEM slot. First k-tile only (later tiles + # overlap via peek-ahead). Complements mma_ab_operand_wait. + if _iket_active: + if k_tile == 0: + iket.range_push("tma_b_buf_acquire_wait") + handle = b_producer.acquire_and_advance( + peek_b_empty_status + ) + if _iket_active: + if k_tile == 0: + iket.range_pop() # tma_b_buf_acquire_wait + peek_b_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_b_empty_status = b_producer.try_acquire() + cute.copy( + tma_atom_fc2_weight, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_b, + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_fc2_weight_sf, + tBgSFB_slice[(None, handle.count)], + tBsSFB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfb, + mcast_mask=sfb_full_mcast_mask, + ) + if _iket_active: + iket.range_pop() + work_tile_info = sched_consumer.consume_work() + + b_producer.tail() + + # ════════════════════════════════════════════════════════════════════ + # MMA warp (warp 4) + # ════════════════════════════════════════════════════════════════════ + # + # Both phases share tiled_mma and TMEM; only K-tile count differs. + if warp_idx == self.mma_warp_id: + _iket_active = (tidx == cutlass.Int32(128)) + + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + acc_base = cute.make_tensor(acc_tmem_ptr, acc_fake.layout) + + # SFA TMEM tensor (placed after the acc cols). + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + # SFB TMEM tensor (after acc + SFA cols). + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + + ( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t, + tCtSFA_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + ( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t, + tCtSFB_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_pipeline_stages + ) + + work_tile_info = sched_consumer.consume_work() + + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + # Prebind k_tile_cnt due to DSL AST. + k_tile_cnt = cutlass.Int32(0) + if is_phase_linear1: + k_tile_cnt = k_tile_cnt_fc1 + if _iket_active: + iket.range_push("mma_fc1") + else: + k_tile_cnt = k_tile_cnt_fc2 + if _iket_active: + iket.range_push("mma_fc2") + + acc_stage_index = acc_producer_state.index + + if is_leader_cta: + tCtAcc = acc_base[(None, None, None, acc_stage_index)] + + if _iket_active: + iket.range_push("mma_acc_acquire") + a_consumer.reset() + b_consumer.reset() + peek_a_full_status = cutlass.Boolean(1) + peek_b_full_status = cutlass.Boolean(1) + if k_tile_cnt > 0: + peek_a_full_status = a_consumer.try_wait() + peek_b_full_status = b_consumer.try_wait() + acc_pipeline.producer_acquire(acc_producer_state) + if _iket_active: + iket.range_pop() + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Isolate the real AB-operand stall: the initial k-tile wait + # for A/B to arrive from the TMA warps. Later k-tiles overlap + # with MMA compute via the peek-ahead pipeline, so only the + # first tile's wait reflects true operand-arrival latency. + if _iket_active: + iket.range_push("mma_ab_operand_wait") + handle_a = a_consumer.wait_and_advance(peek_a_full_status) + handle_b = b_consumer.wait_and_advance(peek_b_full_status) + if _iket_active: + iket.range_pop() # mma_ab_operand_wait + peek_a_full_status = cutlass.Boolean(1) + peek_b_full_status = cutlass.Boolean(1) + if handle_a.count + 1 < k_tile_cnt: + peek_a_full_status = a_consumer.try_wait() + peek_b_full_status = b_consumer.try_wait() + + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t[(None, None, None, None, handle_a.index)], + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t[(None, None, None, None, handle_b.index)], + tCtSFB_compact_s2t, + ) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + cute.gemm( + tiled_mma, + tCtAcc, + [tCrA[(None, None, None, handle_a.index)], tCtSFA], + [tCrB[(None, None, None, handle_b.index)], tCtSFB], + tCtAcc, + ) + handle_a.release() + handle_b.release() + + if k_tile_cnt > 0: + acc_pipeline.producer_commit(acc_producer_state) + if k_tile_cnt > 0: + acc_producer_state.advance() + + if _iket_active: + iket.range_pop() + + work_tile_info = sched_consumer.consume_work() + + acc_pipeline.producer_tail(acc_producer_state) + + # ── sD SMEM (fc1 output staging; fc2 doesn't use it) ── + sD = smem.allocate_tensor( + element_type=self.fc1_output_dtype, + layout=d_smem_layout_staged.outer, + byte_alignment=128, + swizzle=d_smem_layout_staged.inner, + ) + + # ── sC SMEM (raw gate+up Float32, ping-pong; only when generate_c=True) ── + if cutlass.const_expr(self.generate_c): + sC = smem.allocate_tensor( + element_type=self.epilogue._c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + + # ── sFC2 SMEM (FC2 bulk-store staging; fc2_use_bulk only) ── + if cutlass.const_expr(self.epilogue.fc2_needs_staging): + sFC2 = smem.allocate_tensor( + element_type=self.epilogue._fc2_wire_dtype, + layout=self.epilogue.fc2_tma_staged_smem_layout( + self.epilogue.fc2_tma_stages + ), + byte_alignment=128, + ) + + # ── sRED SMEM (in-kernel reduce coalescing transpose; reduce_topk only) ── + if cutlass.const_expr(self.epilogue.fc2_reduce_coalesce): + sRED = smem.allocate_tensor( + element_type=cutlass.BFloat16, + layout=self.epilogue.fc2_reduce_smem_layout(), + byte_alignment=128, + ) + + # ════════════════════════════════════════════════════════════════════ + # Epilogue warps (warps 0-3) + # ════════════════════════════════════════════════════════════════════ + # + # Fully delegated to ``self.epilogue.run(...)`` -- the epilogue owns + # the entire 2-phase task-tile loop. + if warp_idx < self.mma_warp_id: + epi_warp_idx = warp_idx + + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + acc_tensor = cute.make_tensor(acc_tmem_ptr, acc_fake.layout) + + #acc_tensor = cute.make_tensor( + # acc_tmem_ptr, + # cute.make_layout( + # (((128, self.cta_tile_shape_mnk[1]), 1),), + # stride=(((1 << 16, 1), 0),), + # ), + #) + + # Build common kwargs shared by both epilogue flavours. + if cutlass.const_expr(self.generate_c): + _smem_c_raw_arg = sC + _tma_atom_c_arg = tma_atom_c + _gmem_c_arg = tma_tensor_c + else: + _smem_c_raw_arg = None + _tma_atom_c_arg = None + _gmem_c_arg = None + if cutlass.const_expr(self.use_stg_fc1): + _gmem_fc1_output_arg = fc1_output_gemm + else: + _gmem_fc1_output_arg = tma_tensor_fc1_output + _run_kwargs = dict( + tmem_acc_tensor=acc_tensor, + acc_pipeline=acc_pipeline, + sched_consumer=sched_consumer, + sched_ext=ext, + smem_fc1_output_buffer=sD, + tma_atom_fc1_output=tma_atom_fc1_output, + gmem_fc1_output=_gmem_fc1_output_arg, + gmem_fc1_output_sf=fc1_output_sf_gemm, + gmem_topk_scores=topk_scores, + gmem_fc2_output=fc2_output_gemm, + gmem_fc1_done_counter=fc1_done_counter, + smem_c_buffer=_smem_c_raw_arg, + tma_atom_c=_tma_atom_c_arg, + gmem_c=_gmem_c_arg, + warp_idx=epi_warp_idx, + tidx=tidx, + alpha=cutlass.Float32(1.0), + norm_const=cutlass.Float32(1.0), + ) + + _run_extra = {} + if cutlass.const_expr(fc2_output_sf is not None): + _run_extra["gmem_fc2_output_sf"] = fc2_output_sf + if cutlass.const_expr(self.epilogue.fc2_needs_staging): + # Both bulk paths stage into sFC2; only the dispatch TMASTG path + # also needs the TMA atom + local-pool tensor (UBLK peer-writes). + _run_extra["smem_fc2_tma_buffer"] = sFC2 + if cutlass.const_expr(self.epilogue.fc2_use_tma): + _run_extra["tma_atom_fc2_output"] = tma_atom_fc2_output + _run_extra["gmem_fc2_tma_output"] = fc2_tma_output + if cutlass.const_expr(self.epilogue.fc2_reduce_coalesce): + _run_extra["smem_fc2_reduce_buffer"] = sRED + if cutlass.const_expr(self.enable_token_comm): + # MegaMoE (push model): bridge next's TokenComm accessors + peer mapper into + # the epilogue's Fc2OutputDest peer-store expectations. + _epi_comm = _EpilogueCommView( + token_src_metadata=self.token_comm.token_src_metadata_tensor( + self._mega_device_workspace + ), + combine_output=mega_pre_reduced_activation, + peer_rank_ptr_mapper=mega_peer_rank_ptr_mapper, + fc2_output_sf=self.token_comm.fc2_activation_sf_tensor(self._mega_device_workspace), + fc2_done_counter=self.token_comm.fc2_done_counter_tensor(self._mega_device_workspace), + fc2_output_workspace=self.token_comm.fc2_activation_tensor(self._mega_device_workspace), + ) + self.epilogue.run(**_run_kwargs, **_run_extra, token_comm_args=_epi_comm) + elif cutlass.const_expr(token_comm_args is not None): + self.epilogue.run( + **_run_kwargs, **_run_extra, token_comm_args=token_comm_args + ) + else: + self.epilogue.run(**_run_kwargs, **_run_extra) + + tmem.relinquish_alloc_permit() + tmem.free(acc_tmem_ptr) + if cutlass.const_expr(self.enable_token_comm): + cute.arch.fence_acq_rel_sys() + + # ════════════════════════════════════════════════════════════════════ + # Dispatch warps hook (warps 8-11; MegaMoE-only) + # ════════════════════════════════════════════════════════════════════ + # + # ``enable_token_comm=False`` → warps 8-11 don't exist (threads_per_cta + # = 256), so the guard is const_expr-eliminated in the lean path. + if cutlass.const_expr(self.enable_token_comm): + if warp_idx >= self.dispatch_warp_id[0]: + lane_idx_for_dispatch = cute.arch.lane_idx() + if cutlass.const_expr(self.token_back_standalone): + if warp_idx < self.token_back_warp_id[0]: + self.token_comm_hook_dispatch_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + else: + self.token_comm_hook_token_back_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + else: + self.token_comm_hook_dispatch_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + + # ════════════════════════════════════════════════════════════════════ + # Kernel tail hook (MegaMoE-only; lean base = no-op) + # ════════════════════════════════════════════════════════════════════ + lane_idx = cute.arch.lane_idx() + self.token_comm_hook_kernel_tail( + token_comm_args, + warp_idx=warp_idx, + lane_idx=lane_idx, + tidx=tidx, + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py new file mode 100644 index 000000000..3d7493d89 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py @@ -0,0 +1,870 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Full MegaMoE (multi-rank) mxfp8 GLU training-forward kernel.""" + +from typing import Any, Literal, Optional, Tuple, Type + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace + +from ......api import ImplDesc, KernelClass, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from ......helpers.device_workspace import DeviceWorkspace +from ......helpers.smem_workspace import SmemWorkspace +from ......helpers.utils import ceil_div, round_up +from ......quant_def import CombineFormat, QuantKind +from ......communication.nvlink_domain.token_comm_deterministic import TokenCommDeterministic +from ..topk_reduce import TopkReduce +from .glu_mxfp8_col_requant import Mxfp8ColRequant +from .glu_mxfp8_fc12_kernel import Sm107Mxfp8GluFc12Kernel + + +_AB_DTYPE_TO_QUANT_KIND = {cutlass.Float8E4M3FN: QuantKind.mxfp8_e4m3, cutlass.Float8E5M2: QuantKind.mxfp8_e5m2} +_QUANT_KIND_TO_AB_DTYPE = {str(k): d for d, k in _AB_DTYPE_TO_QUANT_KIND.items()} + +# TVM-FFI export symbol for the AOT-compiled callable (consumed by ``tester.compiler``). +_aot_symbol_prefix = "rubin_mega_moe_glu_mxfp8_aot" + + +class Sm107MegaMoEMxfp8GluKernel(Sm107Mxfp8GluFc12Kernel, KernelClass): + """Multi-rank MegaMoE wrapper around the lean mxfp8 GLU FC12 kernel.""" + + fc1_output_region = "rubin.glu_mxfp8.mega.fc1_output" + fc1_output_sf_region = "rubin.glu_mxfp8.mega.fc1_output_sf" + fc1_done_counter_region = "rubin.glu_mxfp8.mega.fc1_done_counter" + col_quant_sizes_region = "rubin.glu_mxfp8.mega.col_quant_expert_token_sizes" + + # Reserved on top of the exact token_comm/sched SMEM to cover smem.allocate inter-allocation + # alignment padding that _compute_stages does not model (see _smem_misc_budget_bytes). + _SMEM_ALLOC_MARGIN = 2048 + + @classmethod + def problem_desc_require(cls): + return { + "expert_count": StaticOrRuntimeIntegerType, + "intermediate_gateup_size": StaticOrRuntimeIntegerType, + "hidden_size": StaticOrRuntimeIntegerType, + "quant_kind": str, + "combine_format": CombineFormat, + "world_size": int, + "local_rank": int, + "topk": int, + "max_tokens_per_rank": int, + "max_recv_size_per_rank": int, + "gate_up_clamp": Optional[float], + } + + @classmethod + def impl_desc_require(cls): + return { + "mma_tiler_mnk": tuple, + "cluster_shape_mnk": tuple, + "use_2cta_instrs": bool, + "group_hint": int, + "token_padding_block": int, + "sf_padding_block": int, + "load_balance_mode": str, + "force_static_sched": bool, + "clc_bundle_size": Optional[int], + "num_sched_stages": Optional[int], + "acc_dtype": type, + "sf_vec_size": int, + "launch_cluster_count": int, + "drop_on_overflow": bool, + "fc2_in_kernel_topk_reduce": bool, + "token_back_mode": str, + "epi_flag_batch": tuple, + "flag_batch": int, + "generate_c": bool, + "use_stg_fc1": bool, + "act_func": str, + "fc2_use_bulk": bool, + "fc2_tma_stages": OptionalRequirement(int), + "enable_col_quant": OptionalRequirement(bool), + "col_quant_num_ctas": OptionalRequirement(int), + } + + def name(self) -> str: + return ( + f"sm107_megamoe_glu_{self.quant_kind}_m{self.mma_tiler_mnk[0]}n{self.mma_tiler_mnk[1]}" + f"k{self.mma_tiler_mnk[2]}_e{self.expert_count}_ep{self.world_size}_topk{self.topk}_" + f"h{self.hidden_size}_i{self.intermediate_gateup_size}_combine{self.combine_format}_" + f"clamp{self.gate_up_clamp}_" + f"tokenback{self.token_back_mode}_hint{self.group_hint}_" + f"epi{self.epi_flag_batch[0]}x{self.epi_flag_batch[1]}_tif{self.flag_batch}_" + f"deterministic_mtpr{self.max_tokens_per_rank}_mrpr{self.max_recv_size_per_rank}_" + f"drop{int(self.drop_on_overflow)}_lc{self.launch_cluster_count}_" + f"genc{int(self.generate_c)}_topkfc11_" + f"fc2bulk{int(self.fc2_use_bulk)}x{self.fc2_tma_stages}_" + f"redtopk{int(self.reduce_topk_in_kernel)}" + ) + + def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): + """Compile against fake (metadata-only) inputs; ``out_path=None`` returns the in-memory callable.""" + import math + + from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_ptr + from cutlass.cute.typing import AddressSpace, sym_int64 + from cutlass.cutlass_dsl import Int32, Int64 + + from ......communication.nvlink_domain.symmetric_buffer import SymmetricBufferHost + + def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): + extents = tuple( + sym_int64(divisibility=math.gcd(int(extent), 128)) if axis in dynamic_axes else int(extent) + for axis, extent in enumerate(shape) + ) + return make_fake_compact_tensor(dtype, extents, stride_order=stride_order, assumed_align=alignment) + + tokens = self.max_tokens_per_rank + aux_shapes = self.get_aux_output_shapes() + hidden = self.hidden_size + intermediate_gateup = self.intermediate_gateup_size + intermediate_downproj = intermediate_gateup // 2 + experts = self.num_experts_per_rank + sf_vec_size = self.sf_vec_size + fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up(hidden // sf_vec_size, 4) + fc2_weight_sf_columns = round_up(hidden, 128) * round_up(intermediate_downproj // sf_vec_size, 4) + output_dtype = cutlass.BFloat16 + # Weight SF and activation SF share the E8M0 block-scale dtype for mxfp8. + weight_sf_dtype = self.token_comm.activation_sf_dtype + + fake_arguments = dict( + activation=fake_tensor(self.token_comm.activation_dtype, (tokens, hidden), (1, 0), {0}, 16), + activation_sf=fake_tensor( + self.token_comm.activation_sf_dtype, + (tokens, self.token_comm.activation_sf_hidden_padded), + (1, 0), + {0}, + 16, + ), + topk_indices=fake_tensor(cutlass.Int32, (tokens, self.topk), (1, 0), {0}, 16), + topk_scores=fake_tensor(cutlass.Float32, (tokens, self.topk), (1, 0), {0}, 4), + fc1_weight=fake_tensor(self.ab_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16), + fc1_weight_sf=fake_tensor(weight_sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16), + fc2_weight=fake_tensor(self.ab_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16), + fc2_weight_sf=fake_tensor(weight_sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16), + output_activation=fake_tensor(output_dtype, (tokens, hidden), (1, 0), {0}, 16), + overflow_flag=fake_tensor(cutlass.Int32, (1,), (0,), set(), 4), + local_workspace=make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), + shared_workspace=make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), + peer_rank_ptr_mapper_host=SymmetricBufferHost( + base_address=Int64(0), + offsets=tuple(Int64(0) for _ in range(self.world_size)), + rank=Int32(0), + max_ranks=self.world_size, + ), + stream=make_fake_stream(), + ) + if self.generate_c: + fake_arguments["fc1_c"] = fake_tensor(output_dtype, aux_shapes["fc1_c"], (1, 0), set(), 16) + else: + fake_arguments["fc1_c"] = None + + if self.enable_col_quant: + # col_quant_data is logically (token, hidden) with token unit stride; + # col_quant_sf is flat concat_e [hidden_atom][token_atom] E8M0 bytes. + fake_arguments["col_quant_data"] = fake_tensor( + self.ab_dtype, aux_shapes["col_quant_data"], (0, 1), set(), 16 + ) + fake_arguments["col_quant_sf"] = fake_tensor( + cutlass.Uint8, aux_shapes["col_quant_sf"], (0,), set(), 16 + ) + else: + fake_arguments["col_quant_data"] = None + fake_arguments["col_quant_sf"] = None + + compiled = cute.compile[cute.EnableTVMFFI(True)](self, **fake_arguments) + if out_path is None: + return compiled + compiled.export_to_c(out_path, function_name=_aot_symbol_prefix, export_only_tvm_ffi_symbols=True) + return out_path + + @staticmethod + def load_compiled(path: str): + from cutlass.cute.runtime import load_module + + return load_module(path, enable_tvm_ffi=True)[_aot_symbol_prefix] + + @classmethod + def from_kwargs( + cls, + # Base-class (lean FC12) kwargs. + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mnk: Tuple[int, int, int], + use_2cta_instrs: bool, + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + load_balance_mode: str = "static", + static_expert_shape: Optional[Tuple[int, int, int]] = None, + force_static_sched: bool = True, + clc_bundle_size: Optional[int] = None, + num_sched_stages: Optional[int] = None, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + ab_dtype: Type[cutlass.Numeric] = cutlass.Float8E4M3FN, + sf_vec_size: int = 32, + *, + world_size: int, + local_rank: int, + num_topk: int, + max_tokens_per_rank: int, + max_recv_size_per_rank: int, + hidden: int, + launch_cluster_count: int, + drop_on_overflow: bool, + fc2_in_kernel_topk_reduce: bool = False, + token_back_mode: Literal["epi_warps", "standalone_warps", "reuse_dispatch_warps"] = "epi_warps", + epi_flag_batch: Optional[Tuple[int, int]] = (4, 2), + flag_batch: int = 1, + gate_up_clamp: Optional[float] = None, + generate_c: bool = False, + use_stg_fc1: bool = False, + combine_format: Optional[CombineFormat] = None, + act_func: str = "swiglu", + fc2_use_bulk: bool = False, + fc2_tma_stages: Optional[int] = None, + enable_col_quant: bool = False, + # -1 = let Mxfp8ColRequant derive the grid from the resident-CTA quantum. + # A hardcoded CTA count is not generally a multiple of that quantum, which + # leaves a fractional resident wave. + col_quant_num_ctas: int = -1, + ) -> "Sm107MegaMoEMxfp8GluKernel": + """Build the ``(ProblemDesc, ImplDesc)`` pair from the legacy flat signature.""" + if static_expert_shape is None: + raise NotImplementedError("Sm107MegaMoEMxfp8GluKernel requires a static_expert_shape.") + if hidden != static_expert_shape[2]: + raise ValueError(f"hidden ({hidden}) must equal static_expert_shape[2] ({static_expert_shape[2]}).") + if ab_dtype not in _AB_DTYPE_TO_QUANT_KIND: + raise ValueError(f"ab_dtype {ab_dtype} has no mxfp8 QuantKind.") + num_experts_per_rank, intermediate_gateup, _hidden = static_expert_shape + combine_format = CombineFormat.parse("bf16" if combine_format is None else str(combine_format)) + problem_desc = ProblemDesc( + { + "expert_count": world_size * num_experts_per_rank, + "intermediate_gateup_size": intermediate_gateup, + "hidden_size": hidden, + "quant_kind": str(_AB_DTYPE_TO_QUANT_KIND[ab_dtype]), + "combine_format": combine_format, + "world_size": world_size, + "local_rank": local_rank, + "topk": num_topk, + "max_tokens_per_rank": max_tokens_per_rank, + "max_recv_size_per_rank": max_recv_size_per_rank, + "gate_up_clamp": gate_up_clamp, + } + ) + impl_desc = ImplDesc( + { + "mma_tiler_mnk": tuple(mma_tiler_mnk), + "cluster_shape_mnk": tuple(cluster_shape_mnk), + "use_2cta_instrs": use_2cta_instrs, + "group_hint": group_hint, + "token_padding_block": token_padding_block, + "sf_padding_block": sf_padding_block, + "load_balance_mode": load_balance_mode, + "force_static_sched": force_static_sched, + "clc_bundle_size": clc_bundle_size, + "num_sched_stages": num_sched_stages, + "acc_dtype": acc_dtype, + "sf_vec_size": sf_vec_size, + "launch_cluster_count": launch_cluster_count, + "drop_on_overflow": drop_on_overflow, + "fc2_in_kernel_topk_reduce": fc2_in_kernel_topk_reduce, + "token_back_mode": token_back_mode, + "epi_flag_batch": tuple(epi_flag_batch) if epi_flag_batch is not None else (1, 1), + "flag_batch": flag_batch, + "generate_c": generate_c, + "use_stg_fc1": use_stg_fc1, + "act_func": act_func, + "fc2_use_bulk": fc2_use_bulk, + # OptionalRequirement: present only when set (absent == None). + **({"fc2_tma_stages": fc2_tma_stages} if fc2_tma_stages is not None else {}), + # Col-quant keys present only when enabled + **( + { + "enable_col_quant": True, + "col_quant_num_ctas": col_quant_num_ctas, + } + if enable_col_quant + else {} + ), + } + ) + return cls(problem_desc, impl_desc) + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + # -- Extract descriptors into locals matching the legacy param names so the body below + # is unchanged; derive the base-class flat inputs (static_expert_shape, ab_dtype). -- + world_size = problem_desc["world_size"] + local_rank = problem_desc["local_rank"] + num_topk = problem_desc["topk"] + max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + max_recv_size_per_rank = problem_desc["max_recv_size_per_rank"] + hidden = problem_desc["hidden_size"] + gate_up_clamp = problem_desc["gate_up_clamp"] + combine_format = problem_desc["combine_format"] + _quant_kind = problem_desc["quant_kind"] + ab_dtype = _QUANT_KIND_TO_AB_DTYPE[_quant_kind] + static_expert_shape = ( + problem_desc["expert_count"] // world_size, + problem_desc["intermediate_gateup_size"], + hidden, + ) + + mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + cluster_shape_mnk = impl_desc["cluster_shape_mnk"] + use_2cta_instrs = impl_desc["use_2cta_instrs"] + group_hint = impl_desc["group_hint"] + token_padding_block = impl_desc["token_padding_block"] + sf_padding_block = impl_desc["sf_padding_block"] + load_balance_mode = impl_desc["load_balance_mode"] + force_static_sched = impl_desc["force_static_sched"] + clc_bundle_size = impl_desc["clc_bundle_size"] + num_sched_stages = impl_desc["num_sched_stages"] + acc_dtype = impl_desc["acc_dtype"] + sf_vec_size = impl_desc["sf_vec_size"] + launch_cluster_count = impl_desc["launch_cluster_count"] + drop_on_overflow = impl_desc["drop_on_overflow"] + fc2_in_kernel_topk_reduce = impl_desc["fc2_in_kernel_topk_reduce"] + token_back_mode = impl_desc["token_back_mode"] + epi_flag_batch = impl_desc["epi_flag_batch"] + flag_batch = impl_desc["flag_batch"] + generate_c = impl_desc["generate_c"] + use_stg_fc1 = impl_desc["use_stg_fc1"] + act_func = impl_desc["act_func"] + fc2_use_bulk = impl_desc["fc2_use_bulk"] + fc2_tma_stages = impl_desc.get("fc2_tma_stages") + self.enable_col_quant = bool(impl_desc.get("enable_col_quant") or False) + self._col_quant_num_ctas = int(impl_desc.get("col_quant_num_ctas") or -1) + + if static_expert_shape is None: + raise NotImplementedError("Sm107MegaMoEMxfp8GluKernel requires a static_expert_shape.") + if hidden != static_expert_shape[2]: + raise ValueError(f"hidden ({hidden}) must equal static_expert_shape[2] ({static_expert_shape[2]}).") + token_back_by_dispatch = token_back_mode != "epi_warps" + + combine_format = CombineFormat.parse("bf16" if combine_format is None else str(combine_format)) + if fc2_in_kernel_topk_reduce and (token_back_by_dispatch or combine_format.is_quantized): + raise ValueError("fc2_in_kernel_topk_reduce requires epi_warps + non-quantized (bf16) combine.") + if token_back_mode not in ("epi_warps", "standalone_warps", "reuse_dispatch_warps"): + raise ValueError(f"unsupported token_back_mode={token_back_mode!r}.") + if ab_dtype not in _AB_DTYPE_TO_QUANT_KIND: + raise ValueError(f"ab_dtype {ab_dtype} has no mxfp8 QuantKind.") + # FC2 bulk store + if fc2_use_bulk and not combine_format.is_quantized: + raise ValueError("fc2_use_bulk currently supports only a quantized (mxfp8) combine format.") + if fc2_tma_stages is not None and not fc2_use_bulk: + raise ValueError("fc2_tma_stages requires fc2_use_bulk=True.") + + super().__init__( + mma_tiler_mnk=mma_tiler_mnk, + cluster_shape_mnk=cluster_shape_mnk, + use_2cta_instrs=use_2cta_instrs, + group_hint=group_hint, + token_padding_block=token_padding_block, + sf_padding_block=sf_padding_block, + load_balance_mode=load_balance_mode, + static_expert_shape=static_expert_shape, + force_static_sched=force_static_sched, + clc_bundle_size=clc_bundle_size, + num_sched_stages=num_sched_stages, + acc_dtype=acc_dtype, + ab_dtype=ab_dtype, + sf_vec_size=sf_vec_size, + fc2_in_kernel_topk_reduce=fc2_in_kernel_topk_reduce, + token_back_by_dispatch=token_back_by_dispatch, + epi_flag_batch=epi_flag_batch, + gate_up_clamp=gate_up_clamp, + apply_topk_in_fc1=True, + generate_c=generate_c, + use_stg_fc1=use_stg_fc1, + act_func=act_func, + fc2_use_bulk=fc2_use_bulk, + fc2_tma_stages=fc2_tma_stages, + ) + + # --- Warp topology: expand to 12 warps (or 16 for standalone token-back). --- + self.enable_token_comm = True + self.dispatch_warp_id = (8, 9, 10, 11) + self.token_back_mode = token_back_mode + self.token_back_standalone = token_back_by_dispatch and token_back_mode == "standalone_warps" + self.token_back_warp_id = (12, 13, 14, 15) if self.token_back_standalone else None + num_token_back_warps = len(self.token_back_warp_id) if self.token_back_standalone else 0 + self.threads_per_cta = 32 * (len(self.epilogue_warp_id) + 4 + len(self.dispatch_warp_id) + num_token_back_warps) + + # --- MegaMoE constants. --- + self.world_size = world_size + self.local_rank = local_rank + self.num_topk = num_topk + self.max_tokens_per_rank = max_tokens_per_rank + self.max_recv_size_per_rank = max_recv_size_per_rank + self.hidden = hidden + self.launch_cluster_count = launch_cluster_count + self.drop_on_overflow = drop_on_overflow + self.combine_format = combine_format + self.num_experts_per_rank = static_expert_shape[0] + self.intermediate_gateup = static_expert_shape[1] + self.intermediate_downproj = self.intermediate_gateup // 2 + self.num_total_experts = world_size * self.num_experts_per_rank + self.reduce_topk_in_kernel = fc2_in_kernel_topk_reduce + self.token_back_schedule_mode = load_balance_mode if load_balance_mode == "atomic_counter" else "static" + + # --- next Router-push token communication component. --- + mma_cta_count = 2 if use_2cta_instrs else 1 + cta_tile_m = mma_tiler_mnk[0] // mma_cta_count + cluster_m, cluster_n = self.cluster_shape_mn + tokens_per_fc1_ready_slot = cta_tile_m * cluster_m + hidden_per_fc2_cluster_tile = cta_tile_m * cluster_m + fc2_done_signals_per_token_tile = ceil_div(hidden, hidden_per_fc2_cluster_tile) * cluster_m * cluster_n + promised_launchable_sm_count = launch_cluster_count * cluster_m * cluster_n + quant_kind = _AB_DTYPE_TO_QUANT_KIND[ab_dtype] + tc_problem_desc = ProblemDesc( + { + "world_size": world_size, + "expert_count": self.num_total_experts, + "topk": num_topk, + "max_tokens_per_rank": max_tokens_per_rank, + "max_recv_size_per_rank": max_recv_size_per_rank, + "hidden_size": hidden, + "quant_kind": str(quant_kind), + "combine_format": combine_format, + "apply_topk_at_fc1": True, + } + ) + tc_impl_desc = ImplDesc( + { + "token_padding_block": token_padding_block, + "sf_padding_block": sf_padding_block, + "tokens_per_fc1_ready_slot": tokens_per_fc1_ready_slot, + "fc2_done_signals_per_token_tile": fc2_done_signals_per_token_tile, + "promised_launchable_sm_count": promised_launchable_sm_count, + "drop_on_overflow": drop_on_overflow, + "token_in_flag_batch": flag_batch, + "token_back_mode": token_back_mode, + "token_back_schedule_mode": self.token_back_schedule_mode, + "reduce_topk_in_kernel": fc2_in_kernel_topk_reduce, + } + ) + self.token_comm = TokenCommDeterministic(tc_problem_desc, tc_impl_desc) + self.pool_token_capacity = self.token_comm.worst_case_token_count + + # --- SMEM sub-buffer for the token_comm transport (allocated in the device kernel). --- + tc_smem_ws = SmemWorkspace() + self.token_comm.register_smem_regions(tc_smem_ws) + tc_smem_ws.finalize(max_bytes=self.smem_capacity) + self.tc_smem_ws = tc_smem_ws + self._token_comm_smem_bytes = tc_smem_ws.total_bytes + + # Build the scheduler + _ec, _ig, _hd = static_expert_shape + self._build_scheduler( + expert_cnt=_ec, intermediate_gateup=_ig, hidden_dim=_hd, launch_cluster_count=launch_cluster_count + ) + self._sched_smem_bytes = self.sched_smem_ws.total_bytes + + # --- Post-kernel top-k reduction (skipped under in-kernel REDG reduce). --- + self._topk_reduce = None if fc2_in_kernel_topk_reduce else TopkReduce(hidden, num_topk, combine_format) + + # --- Device workspace (next model): fc1 pool/output + token_comm regions. --- + self._mega_device_workspace = self._build_megamoe_device_workspace() + + self.expert_count = self.num_total_experts + self.intermediate_gateup_size = self.intermediate_gateup + self.hidden_size = hidden + self.quant_kind = _quant_kind + self.topk = num_topk + self.cluster_shape_mnk = tuple(cluster_shape_mnk) + self.mma_tiler_mnk = tuple(mma_tiler_mnk) + self.group_hint = group_hint + self.token_padding_block = token_padding_block + self.sf_padding_block = sf_padding_block + self.load_balance_mode = load_balance_mode + self.force_static_sched = force_static_sched + self.clc_bundle_size = clc_bundle_size + self.num_sched_stages = num_sched_stages + self.acc_dtype = acc_dtype + self.sf_vec_size = sf_vec_size + self.fc2_in_kernel_topk_reduce = fc2_in_kernel_topk_reduce + self.epi_flag_batch = tuple(epi_flag_batch) + self.flag_batch = flag_batch + self.generate_c = generate_c + self.use_stg_fc1 = use_stg_fc1 + self.act_func = act_func + self.use_2cta_instrs = use_2cta_instrs + self.gate_up_clamp = gate_up_clamp + + # Optional standalone token-axis (column) MXFP8 requantization. + if self.enable_col_quant: + col_quant_type = "mxfp8_e4m3" if ab_dtype is cutlass.Float8E4M3FN else "mxfp8_e5m2" + self.col_quant = Mxfp8ColRequant( + hidden=self.hidden, + num_experts=self.num_experts_per_rank, + max_total_tokens=( + self.world_size + * self.max_tokens_per_rank + * min(self.num_topk, self.num_experts_per_rank) + ), + quant_type=col_quant_type, + num_persistent_ctas=self._col_quant_num_ctas, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + dst_k_major=True, + ) + + def get_aux_output_shapes(self) -> dict: + """Return receiver-domain auxiliary output shapes.""" + return { + "fc1_c": (self.pool_token_capacity, self.intermediate_gateup), + "col_quant_data": (self.pool_token_capacity, self.hidden), + "col_quant_sf": ( + self.token_comm.worst_case_sf_token_count + * (self.hidden // self.sf_vec_size), + ), + } + + @cute.jit + def _validate_fixed_matrix( + self, tensor: cute.Tensor, dtype, expected_shape, expected_stride=None + ) -> None: + if cutlass.const_expr(tensor.element_type is not dtype): + raise TypeError("pool-domain matrix has an unexpected element type.") + if cutlass.const_expr(cute.rank(tensor.layout) != 2): + raise ValueError("pool-domain matrix must be rank 2.") + if cutlass.const_expr( + not isinstance(tensor.shape[0], int) + or not isinstance(tensor.shape[1], int) + or tensor.shape[0] != expected_shape[0] + or tensor.shape[1] != expected_shape[1] + ): + raise ValueError(f"pool-domain matrix must have static shape {expected_shape}.") + stride = (expected_shape[1], 1) if expected_stride is None else expected_stride + if cutlass.const_expr(tensor.stride[0] != stride[0] or tensor.stride[1] != stride[1]): + raise ValueError("pool-domain matrix has an unexpected stride.") + + @cute.jit + def _validate_fixed_vector(self, tensor: cute.Tensor, dtype, expected_size: int) -> None: + if cutlass.const_expr(tensor.element_type is not dtype): + raise TypeError("pool-domain vector has an unexpected element type.") + if cutlass.const_expr(cute.rank(tensor.layout) != 1): + raise ValueError("pool-domain vector must be rank 1.") + if cutlass.const_expr( + not isinstance(tensor.shape[0], int) or tensor.shape[0] != expected_size + ): + raise ValueError(f"pool-domain vector must have static size {expected_size}.") + if cutlass.const_expr(tensor.stride[0] != 1): + raise ValueError("pool-domain vector must be contiguous.") + + def _smem_misc_budget_bytes(self) -> int: + """Reserve the token_comm transport SMEM on top of the base misc budget.""" + _sched = getattr(self, "_sched_smem_bytes", 0) + return super()._smem_misc_budget_bytes() + self._token_comm_smem_bytes + _sched + self._SMEM_ALLOC_MARGIN + + def _build_megamoe_device_workspace(self) -> DeviceWorkspace: + """Register the FC1 output/pool + fc1_done_counter + all token_comm regions.""" + sf_dtype = cutlass.Float8E8M0FNU + sf_column_count = round_up(ceil_div(self.intermediate_downproj, self.sf_vec_size), 4) + max_sf_rows = self.token_comm.worst_case_sf_token_count + counter_slot_count = self.token_comm.max_fc1_ready_slot_count + + device_workspace = DeviceWorkspace() + device_workspace.register( + self.fc1_output_region, + self.ab_dtype, + (self.pool_token_capacity, self.intermediate_downproj), + buffer_space="local", + mem_order=(1, 0), + byte_alignment=128, + ) + device_workspace.register( + self.fc1_output_sf_region, + sf_dtype, + (max_sf_rows, sf_column_count), + buffer_space="local", + mem_order=(1, 0), + byte_alignment=128, + ) + device_workspace.register( + self.fc1_done_counter_region, + cutlass.Int32, + (counter_slot_count,), + buffer_space="local", + byte_alignment=16, + reset="tail_reset", + ) + if self.enable_col_quant: + # Persistent per-expert token-count snapshot for the post-kernel col-quant launch. + device_workspace.register( + self.col_quant_sizes_region, + cutlass.Int32, + (self.num_experts_per_rank,), + buffer_space="local", + byte_alignment=16, + ) + self.token_comm.register_device_workspace(device_workspace) + device_workspace.finalize() + return device_workspace + + def get_workspace_sizes(self) -> Tuple[int, int]: + """Return required (local, shared/symmetric) workspace bytes.""" + return self._mega_device_workspace.local_and_shared_bytes + + @property + def require_zero_workspace_leading_bytes(self) -> Tuple[int, int]: + return self._mega_device_workspace.require_zero_workspace_leading_bytes + + # ========================================================================= + # token_comm_hook_* -- the lean device method's integration seams, here + # filled with next's Router-push TokenCommDeterministic calls. + # ========================================================================= + + def token_comm_extra_smem_storage_class(self) -> type: + """SMEM struct for the token_comm transport overlay (allocated in the device kernel).""" + return self.tc_smem_ws.storage_class() + + def token_comm_hook_fc1_ready_counter_ptr(self, token_comm_args): + """Pointer the FC1 scheduler/extension spins on; token_in increments it per ready slot.""" + return self.token_comm.fc1_ready_counter_pointer(self._mega_device_workspace) + + def sched_ext_fc1_peek_threshold(self) -> int: # noqa: D401 - lean hook override point + return super().sched_ext_fc1_peek_threshold() + + @cute.jit + def token_comm_hook_sched_warp_pre_init_wait(self, token_comm_args): + """The scheduler warp must wait for the Router to publish per-expert sizes.""" + self.token_comm.wait_for_sizes_ready(self._mega_device_workspace) + + @cute.jit + def token_comm_hook_fc1_tma_b_predispatch_spin(self, token_comm_args, work_tile_info): + """No-op: FC1 input readiness is enforced by the scheduler extension's fc1_ready spin.""" + pass + + @cute.jit + def token_comm_hook_dispatch_warp_body(self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx): + """Transfer warps (8-11): pull activation from peers into the local FC1 pool.""" + self.token_comm.token_in(self.tc_smem_ws, token_comm_storage.buffer.data_ptr()) + if cutlass.const_expr(self.token_comm.token_back_enabled and not self.token_back_standalone): + self.token_comm.token_back(self.tc_smem_ws, token_comm_storage.buffer.data_ptr()) + + @cute.jit + def token_comm_hook_token_back_warp_body(self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx): + """Standalone token-back warps (12-15): push FC2 output back to source ranks.""" + self.token_comm.token_back(self.tc_smem_ws, token_comm_storage.buffer.data_ptr()) + + @cute.jit + def token_comm_hook_tail_reset_shared_counters(self, token_comm_args, *, warp_idx, lane_idx, tidx): + """Absorbed into reset_tail (kernel_tail hook).""" + pass + + @cute.jit + def token_comm_hook_kernel_tail(self, token_comm_args, *, warp_idx, lane_idx, tidx): + """Cross-rank drain + workspace tail reset, performed by the transfer warps. + + The whole-CTA barrier is REQUIRED (mirrors inference's mainloop kernel-tail + sync_threads): it forces every compute warp (scheduler / tma / mma / epilogue) to + finish its consume loop -- including any ``ext.wait_for_input`` spin on ``fc1_ready`` + -- before the transfer warps run ``reset_tail``. ``reset_tail`` tail-resets the + (local, GPU-wide) ``fc1_ready`` / ``fc2_done`` counters; without this barrier a fast + CTA's transfer warps can zero a counter that a slower compute warp is still spinning + on, which non-deterministically deadlocks (~1/5 runs). + """ + # Preserve the per-expert dispatch token counts before reset_tail zeros the + # shared sizes_region. + if cutlass.const_expr(self.enable_col_quant): + self._snapshot_col_quant_expert_sizes(tidx) + cute.arch.sync_threads() + # reset_tail must run on EXACTLY the 4 transfer/token_in warps (8-11). For + # standalone_warps (16 warps) the token_back warps (12-15) must be EXCLUDED: their + # thread_idx aliases (thread_idx % transfer_thread_count) back onto transfer warps 0-3, + # so including them double-counts the reset_tail NVLink grid barrier -> deadlock. + if (warp_idx >= self.dispatch_warp_id[0]) & (warp_idx <= self.dispatch_warp_id[-1]): + self.token_comm.reset_tail() + self.token_comm.remove_device_members() + + @cute.jit + def _snapshot_col_quant_expert_sizes(self, tidx) -> None: + from cutlass.cutlass_dsl import Int32, Int64 + """CTA 0 copies this rank's per-expert token counts into the persistent + col-quant region before token_comm's tail_reset zeros ``sizes_region``.""" + dw = self._mega_device_workspace + if self.token_comm._linear_cta_idx == Int32(0): + sizes = self.token_comm.local_expert_sizes(dw, self.token_comm._local_rank) + snapshot = dw.tensor(self.col_quant_sizes_region) + block_dim_x, _, _ = cute.arch.block_dim() + expert_idx = tidx + while expert_idx < Int32(self.num_experts_per_rank): + snapshot[expert_idx] = Int32(sizes[expert_idx]) + expert_idx = expert_idx + block_dim_x + + # ========================================================================= + # Host launch: Router kernel -> fused MegaMoE main kernel -> top-k reduction. + # ========================================================================= + + @cute.jit + def __call__( + self, + activation: cute.Tensor, # (max_tokens_per_rank, hidden) raw per-rank, symmetric heap + activation_sf: cute.Tensor, # (max_tokens_per_rank, hidden // sf_vec_size), symmetric + topk_indices: cute.Tensor, # (max_tokens_per_rank, topk) + topk_scores: cute.Tensor, # (max_tokens_per_rank, topk) Float32 + fc1_weight: cute.Tensor, # (experts_per_rank, hidden, intermediate_gateup) + fc1_weight_sf: cute.Tensor, + fc2_weight: cute.Tensor, # (experts_per_rank, intermediate_downproj, hidden) + fc2_weight_sf: cute.Tensor, + output_activation: cute.Tensor, # (max_tokens_per_rank, topk, hidden) final combined output + fc1_c: Optional[cute.Tensor], # (pool_token_capacity, intermediate_gateup) when generate_c=True + col_quant_data: Optional[cute.Tensor], # dispatch-pool-strided fp8 segments + col_quant_sf: Optional[cute.Tensor], # flat concat_e [hidden_atom][token_atom] E8M0 bytes + overflow_flag: cute.Tensor, # (1,) Int32, per-rank router receive-overflow output + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, # symmetric (NVLink) heap base + peer_rank_ptr_mapper_host, + stream: cuda.CUstream, + ) -> None: + """Launch the Router, then the fused main kernel, then (optionally) the top-k reduce.""" + from cutlass.cutlass_dsl import Int64 + + dw = self._mega_device_workspace + local_rank = peer_rank_ptr_mapper_host.rank + aux_shapes = self.get_aux_output_shapes() + if cutlass.const_expr(self.generate_c): + if cutlass.const_expr(fc1_c is None): + raise ValueError("generate_c=True requires a receiver-domain fc1_c tensor.") + self._validate_fixed_matrix( + fc1_c, cutlass.BFloat16, aux_shapes["fc1_c"] + ) + if cutlass.const_expr(self.enable_col_quant): + if cutlass.const_expr(col_quant_data is None or col_quant_sf is None): + raise ValueError("enable_col_quant=True requires data and scale outputs.") + self._validate_fixed_matrix( + col_quant_data, + self.ab_dtype, + aux_shapes["col_quant_data"], + (1, aux_shapes["col_quant_data"][0]), + ) + self._validate_fixed_vector( + col_quant_sf, cutlass.Uint8, aux_shapes["col_quant_sf"][0] + ) + self.token_comm.launch_router( + topk_indices=topk_indices, + topk_scores=topk_scores, + local_rank=local_rank, + local_workspace=local_workspace, + shared_workspace=shared_workspace, + peer_rank_ptr_mapper_host=peer_rank_ptr_mapper_host, + device_workspace=dw, + overflow_flag=overflow_flag, + stream=stream, + ) + peer_mapper = peer_rank_ptr_mapper_host.make_device_object() + dw.assign_device_members(local_workspace, shared_workspace) + + activation_pool = self.token_comm.fc1_activation_tensor(dw) + _sf_pool_atom = self.token_comm.fc1_activation_sf_tensor(dw) + activation_sf_pool = cute.make_tensor( + _sf_pool_atom.iterator, + cute.make_layout( + (self.token_comm.worst_case_sf_token_count, self.hidden // self.sf_vec_size), + stride=(self.token_comm.activation_sf_hidden_padded, 1), + ), + ) + fc1_output = dw.tensor(self.fc1_output_region) + fc1_output_sf = dw.tensor(self.fc1_output_sf_region) + fc1_done_counter = dw.tensor(self.fc1_done_counter_region) + pool_topk_scores = self.token_comm.fc1_topk_scores_tensor(dw) + + if cutlass.const_expr(self.reduce_topk_in_kernel): + # In-kernel top-k reduce (epi_warps + bf16 combine) + pre_reduced = cute.make_tensor( + output_activation.iterator, + cute.make_layout( + (output_activation.shape[0], 1, output_activation.shape[1]), + stride=(output_activation.stride[0], output_activation.stride[0], output_activation.stride[1]), + ), + ) + pre_reduced_sf = None + else: + pre_reduced = self.token_comm.pre_reduced_activation_tensor(dw) + pre_reduced_sf = self.token_comm.pre_reduced_activation_sf_tensor(dw) + + if cutlass.const_expr(self.token_comm.token_back_push_data): + # token_back-by-dispatch (standalone_warps / reuse_dispatch_warps) + fc2_output = self.token_comm.fc2_activation_tensor(dw) + else: + # epi_warps: the epilogue peer-writes FC2 directly. + _combine_hidden = pre_reduced.shape[2] + fc2_output = cute.make_tensor( + pre_reduced.iterator, + cute.make_layout( + (pre_reduced.shape[0] * pre_reduced.shape[1], _combine_hidden), stride=(_combine_hidden, 1) + ), + ) + + super().__call__( + activation_pool, + fc1_weight, + activation_sf_pool, + fc1_weight_sf, + fc1_output, + fc1_output_sf, + fc2_weight, + fc2_weight_sf, + fc2_output, + pool_topk_scores, + fc1_done_counter, + offs=None, + max_active_clusters=self.launch_cluster_count, + stream=stream, + fc1_c=fc1_c, + overflow_flag=overflow_flag, + mega_peer_rank_ptr_mapper=peer_mapper, + mega_local_rank=local_rank, + mega_local_workspace=local_workspace, + mega_shared_workspace=shared_workspace, + mega_activation=activation, + mega_activation_sf=activation_sf, + mega_pre_reduced_activation=pre_reduced, + mega_pre_reduced_activation_sf=pre_reduced_sf, + ) + + # Top-k weights were already applied before FC1 quantization, so the + # post-kernel reduction is a plain dequantized sum. + if cutlass.const_expr(not self.reduce_topk_in_kernel): + self._topk_reduce(pre_reduced, pre_reduced_sf, output_activation, None, stream) + + # Post-kernel token-axis MXFP8 requantization of the preserved dispatch-local + # FC1 activation pool, stored by the tail snapshot of per-expert token counts. + if cutlass.const_expr(self.enable_col_quant): + lw = local_workspace + data_offset = dw.offset(self.token_comm.fc1_activation_region) + sf_offset = dw.offset(self.token_comm.fc1_activation_sf_region) + sizes_offset = dw.offset(self.col_quant_sizes_region) + sf_pool_bytes = self.token_comm.worst_case_sf_token_count * (self.hidden // self.sf_vec_size) + src_data = cute.make_tensor( + cute.make_ptr( + self.ab_dtype, lw.toint() + Int64(data_offset), AddressSpace.gmem, assumed_align=128 + ), + cute.make_layout( + (self.token_comm.worst_case_token_count, self.hidden), stride=(self.hidden, 1) + ), + ) + src_sf_u8 = cute.make_tensor( + cute.make_ptr(cutlass.Uint8, lw.toint() + Int64(sf_offset), AddressSpace.gmem, assumed_align=16), + cute.make_layout((sf_pool_bytes,)), + ) + expert_sizes = cute.make_tensor( + cute.make_ptr(cutlass.Int32, lw.toint() + Int64(sizes_offset), AddressSpace.gmem, assumed_align=16), + cute.make_layout((self.num_experts_per_rank,)), + ) + self.col_quant( + src_data, + src_sf_u8, + expert_sizes, + col_quant_data, + col_quant_sf, + stream, + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/__init__.py new file mode 100644 index 000000000..82b548c79 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Training mega helpers (SwiGLU / mxfp8 register-level quant primitives).""" + +from .constants import ( + SupportedMmaTileM, + SupportedMmaTileN, +) +from .utils import dswiglu_act, quant_sfd_col, quant_sfd_row, swiglu_act + +__all__ = [ + "SupportedMmaTileM", + "SupportedMmaTileN", + "dswiglu_act", + "quant_sfd_col", + "quant_sfd_row", + "swiglu_act", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/constants.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/constants.py new file mode 100644 index 000000000..1bd314338 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/constants.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Training-mega-specific numeric constants (MMA tiler extents). + +These are consumed only by the training GLU/dGLU FC12 kernels, so they live in the +training mega ``helpers`` package rather than the cross-inference/training +``next/sources/helpers/constants.py``. +""" + + +# MMA tiler GLU FC12 kernels accept along M and N. +SupportedMmaTileM = (128, 256) +SupportedMmaTileN = (128, 256) + + +__all__ = [ + "SupportedMmaTileM", + "SupportedMmaTileN", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/utils.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/utils.py new file mode 100644 index 000000000..a015fe2a5 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/utils.py @@ -0,0 +1,322 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +from typing import Optional + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass._mlir.dialects import math as _math +from cutlass.cutlass_dsl import Float32, T, dsl_user_op + +from ......helpers.constants import Fp32Max, Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Log2E +from ......helpers.ptx_helpers import cvt_f32_to_fp8_to_f32, cvt_f32x4_to_f8x4_pack_i32 + + +@dsl_user_op +def zero_unless_equal( + value: Float32, + raw: Float32, + clamped: Float32, + *, + loc=None, + ip=None, +) -> Float32: + """Return ``value`` if ``raw == clamped``, using branch-free PTX.""" + return Float32( + llvm.inline_asm( + T.f32(), + [ + Float32(value).ir_value(loc=loc, ip=ip), + Float32(raw).ir_value(loc=loc, ip=ip), + Float32(clamped).ir_value(loc=loc, ip=ip), + ], + "{\n" + " .reg .pred in_range;\n" + " setp.eq.f32 in_range, $2, $3;\n" + " selp.f32 $0, $1, 0f00000000, in_range;\n" + "}", + "=f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def swiglu_act( + t_swiglu: cute.Tensor, + t_up: cute.Tensor, + t_gate: cute.Tensor, + prob: Optional[Float32] = None, + gate_up_clamp: Optional[Float32] = None, +) -> None: + """SwiGLU with optional gate-upper/up-symmetric clamp.""" + for i in cutlass.range_constexpr(0, cute.size(t_swiglu), 2): + gate = (t_gate[i], t_gate[i + 1]) + up = (t_up[i], t_up[i + 1]) + if cutlass.const_expr(gate_up_clamp is not None): + gate = ( + cute.arch.fmin(gate[0], gate_up_clamp), + cute.arch.fmin(gate[1], gate_up_clamp), + ) + up = ( + cute.arch.fmax(cute.arch.fmin(up[0], gate_up_clamp), -gate_up_clamp), + cute.arch.fmax(cute.arch.fmin(up[1], gate_up_clamp), -gate_up_clamp), + ) + up_gate = cute.arch.mul_packed_f32x2( + up, + gate, + rnd="rn", + ftz=False, + ) + gate_log2e = cute.arch.mul_packed_f32x2( + gate, (-Log2E, -Log2E), rnd="rn", ftz=False, + ) + one_plus_exp = cute.arch.add_packed_f32x2( + ( + cute.math.exp2(gate_log2e[0], fastmath=True), + cute.math.exp2(gate_log2e[1], fastmath=True), + ), + (1.0, 1.0), + ) + sigmoid = (cute.arch.rcp_approx(one_plus_exp[0]), cute.arch.rcp_approx(one_plus_exp[1])) + (t_swiglu[i], t_swiglu[i + 1]) = cute.arch.mul_packed_f32x2( + up_gate, sigmoid, rnd="rn", ftz=False, + ) + if cutlass.const_expr(prob is not None): + (t_swiglu[i], t_swiglu[i + 1]) = cute.arch.mul_packed_f32x2( + (t_swiglu[i], t_swiglu[i + 1]), + (prob, prob), + rnd="rn", + ftz=False, + ) + + +@cute.jit +def quant_sfd_row( + src: cute.Tensor, + dst: cute.Tensor, + norm_const, + sf_vec_size, + sf_dtype, + d_dtype, +): + """Quantize the ``sf_vec_size`` values in ``src`` to ``d_dtype`` with one block scale.""" + rcp_limit = Fp8E4M3RcpLimit if d_dtype == cutlass.Float8E4M3FN else Fp8E5M2RcpLimit + acc_frg = src.load() + abs_acc_frg_ir = _math.absf(acc_frg.ir_value()) + abs_acc_frg = type(acc_frg)(abs_acc_frg_ir, acc_frg.shape, acc_frg.dtype) + # Fuse the two loop-invariant constants into one multiply + rcp_limit_norm = rcp_limit * norm_const + avg_fp32 = abs_acc_frg.reduce(cute.ReductionOp.MAX, Float32(0.0), 0) * rcp_limit_norm + qpvscale_up = cvt_f32_to_fp8_to_f32(avg_fp32, sf_dtype) + acc_scale = norm_const * cute.arch.rcp_approx(qpvscale_up) + acc_scale = cute.arch.fmin(acc_scale, Fp32Max, nan=True) + for ei in cutlass.range_constexpr(0, sf_vec_size, 2): + src[ei], src[ei + 1] = cute.arch.mul_packed_f32x2( + (src[ei], src[ei + 1]), (acc_scale, acc_scale), rnd="rn", ftz=False, + ) + dst_i32 = cute.recast_tensor(dst, cutlass.Int32) + for ei in cutlass.range_constexpr(0, sf_vec_size, 4): + fp32x4 = cute.make_rmem_tensor(4, Float32) + fp32x4[0] = src[ei + 0] + fp32x4[1] = src[ei + 1] + fp32x4[2] = src[ei + 2] + fp32x4[3] = src[ei + 3] + fp8x4_i32 = cvt_f32x4_to_f8x4_pack_i32(fp32x4, d_dtype) + dst_i32[ei // 4] = cutlass.Int32(fp8x4_i32) + return qpvscale_up + + +@cute.jit +def dswiglu_act( + t_dgate: cute.Tensor, + t_dup: cute.Tensor, + t_acc: cute.Tensor, + t_gate: cute.Tensor, + t_up: cute.Tensor, + beta_val: Float32, + prob: Float32, + gate_up_clamp: Optional[Float32] = None, +) -> Float32: + """SwiGLU backward with optional clamp, beta/prob scaling, and dprob. + + Given upstream gradient ``acc``, per-expert scalar ``beta_val``, per-token routing + probability ``prob``, and forward pre-activations ``gate``/``up``:: + + gate_raw = gate * beta_val + up_raw = up * beta_val + gate_b = min(gate_raw, clamp) + up_b = clamp(up_raw, -clamp, clamp) + sig = sigmoid(gate_b) + swish = gate_b * sig + + dprob += acc * up_b * swish (returned to the caller) + d_up = acc * prob * swish * I[-clamp <= up_raw <= clamp] + d_gate = acc * prob * up_b * silu'(gate_b) * I[gate_raw <= clamp] + + The clamp is skipped when ``gate_up_clamp`` is ``None``. Boundary values retain + their gradient, matching ``torch.clamp``. + """ + dprob_acc = Float32(0.0) + for i in cutlass.range_constexpr(0, cute.size(t_acc), 2): + gate_raw = cute.arch.mul_packed_f32x2( + (t_gate[i], t_gate[i + 1]), (beta_val, beta_val), rnd="rn", ftz=False, + ) + up_raw = cute.arch.mul_packed_f32x2( + (t_up[i], t_up[i + 1]), (beta_val, beta_val), rnd="rn", ftz=False, + ) + gate_b = gate_raw + up_b = up_raw + if cutlass.const_expr(gate_up_clamp is not None): + gate_b = ( + cute.arch.fmin(gate_raw[0], gate_up_clamp), + cute.arch.fmin(gate_raw[1], gate_up_clamp), + ) + up_b = ( + cute.arch.fmax(cute.arch.fmin(up_raw[0], gate_up_clamp), -gate_up_clamp), + cute.arch.fmax(cute.arch.fmin(up_raw[1], gate_up_clamp), -gate_up_clamp), + ) + + # sig = 1 / (1 + exp(-gate_b)); exp(-x) = exp2(-Log2E * x) + sig_rcp = cute.arch.mul_packed_f32x2( + gate_b, (-Log2E, -Log2E), rnd="rn", ftz=False, + ) + (sig0, sig1) = cute.arch.add_packed_f32x2( + ( + cute.math.exp2(sig_rcp[0], fastmath=True), + cute.math.exp2(sig_rcp[1], fastmath=True), + ), + (1.0, 1.0), + ) + sig0 = cute.arch.rcp_approx(sig0) + sig1 = cute.arch.rcp_approx(sig1) + + # swish = gate_b * sig + swish = cute.arch.mul_packed_f32x2(gate_b, (sig0, sig1), rnd="rn", ftz=False) + + # dprob += acc * up_b * swish (both lanes into the running scalar) + dp = cute.arch.mul_packed_f32x2( + (t_acc[i], t_acc[i + 1]), (up_b[0], up_b[1]), rnd="rn", ftz=False, + ) + dp = cute.arch.mul_packed_f32x2(dp, swish, rnd="rn", ftz=False) + dprob_acc = dprob_acc + dp[0] + dp[1] + + # acc * prob (shared factor for d_up and d_gate) + acc_prob = cute.arch.mul_packed_f32x2( + (t_acc[i], t_acc[i + 1]), (prob, prob), rnd="rn", ftz=False, + ) + + # d_up = acc * prob * swish + (t_dup[i], t_dup[i + 1]) = cute.arch.mul_packed_f32x2( + acc_prob, swish, rnd="rn", ftz=False, + ) + + # d_gate = acc * prob * up_b * sig * (1 + gate_b * (1 - sig)) + one_minus_sig = cute.arch.add_packed_f32x2( + (1.0, 1.0), (-sig0, -sig1), rnd="rn", ftz=False, + ) + dsig = cute.arch.mul_packed_f32x2(gate_b, one_minus_sig, rnd="rn", ftz=False) + term = cute.arch.add_packed_f32x2( + (dsig[0], dsig[1]), (1.0, 1.0), rnd="rn", ftz=False, + ) + dgate = cute.arch.mul_packed_f32x2( + acc_prob, (up_b[0], up_b[1]), rnd="rn", ftz=False, + ) + dgate = cute.arch.mul_packed_f32x2(dgate, (sig0, sig1), rnd="rn", ftz=False) + (t_dgate[i], t_dgate[i + 1]) = cute.arch.mul_packed_f32x2( + dgate, term, rnd="rn", ftz=False, + ) + if cutlass.const_expr(gate_up_clamp is not None): + t_dgate[i] = zero_unless_equal(t_dgate[i], gate_raw[0], gate_b[0]) + t_dgate[i + 1] = zero_unless_equal(t_dgate[i + 1], gate_raw[1], gate_b[1]) + t_dup[i] = zero_unless_equal(t_dup[i], up_raw[0], up_b[0]) + t_dup[i + 1] = zero_unless_equal(t_dup[i + 1], up_raw[1], up_b[1]) + + return dprob_acc + + +@cute.jit +def quant_sfd_col( + src: cute.Tensor, + dst: cute.Tensor, + norm_const, + sf_vec_size, + sf_dtype, + d_dtype, +): + """Column (cross-thread) block-scale quantize: the amax is a warp reduction.""" + rcp_limit = Fp8E4M3RcpLimit if d_dtype == cutlass.Float8E4M3FN else Fp8E5M2RcpLimit + acc_frg = src.load() + abs_acc_frg_ir = _math.absf(acc_frg.ir_value()) + acc_frg = type(acc_frg)(abs_acc_frg_ir, acc_frg.shape, acc_frg.dtype) + + qpvscale_up = Float32(0.0) + tidx, _, _ = cute.arch.thread_idx() + scale = rcp_limit * norm_const + + for vi in cutlass.range_constexpr(0, sf_vec_size, 4): + # Warp-wide MAX across the 32 rows for each of the 4 lanes. + max_value0 = Float32(cute.arch.warp_redux_sync(acc_frg[vi], "fmax", nan=True)) + max_value1 = Float32(cute.arch.warp_redux_sync(acc_frg[vi + 1], "fmax", nan=True)) + max_value2 = Float32(cute.arch.warp_redux_sync(acc_frg[vi + 2], "fmax", nan=True)) + max_value3 = Float32(cute.arch.warp_redux_sync(acc_frg[vi + 3], "fmax", nan=True)) + + (max_value0, max_value1) = cute.arch.mul_packed_f32x2( + (max_value0, max_value1), (scale, scale), rnd="rn", ftz=False, + ) + (max_value2, max_value3) = cute.arch.mul_packed_f32x2( + (max_value2, max_value3), (scale, scale), rnd="rn", ftz=False, + ) + + max_value0 = cvt_f32_to_fp8_to_f32(max_value0, sf_dtype) + max_value1 = cvt_f32_to_fp8_to_f32(max_value1, sf_dtype) + max_value2 = cvt_f32_to_fp8_to_f32(max_value2, sf_dtype) + max_value3 = cvt_f32_to_fp8_to_f32(max_value3, sf_dtype) + + # Each thread keeps its assigned column's pre-round-trip scale. + if tidx % 32 == vi: + qpvscale_up = max_value0 + if tidx % 32 == vi + 1: + qpvscale_up = max_value1 + if tidx % 32 == vi + 2: + qpvscale_up = max_value2 + if tidx % 32 == vi + 3: + qpvscale_up = max_value3 + + max_value_rcp0 = cute.arch.fmin(cute.arch.rcp_approx(max_value0), Fp32Max, nan=True) + max_value_rcp1 = cute.arch.fmin(cute.arch.rcp_approx(max_value1), Fp32Max, nan=True) + max_value_rcp2 = cute.arch.fmin(cute.arch.rcp_approx(max_value2), Fp32Max, nan=True) + max_value_rcp3 = cute.arch.fmin(cute.arch.rcp_approx(max_value3), Fp32Max, nan=True) + + (acc_scale_col0, acc_scale_col1) = cute.arch.mul_packed_f32x2( + (norm_const, norm_const), (max_value_rcp0, max_value_rcp1), rnd="rn", ftz=False, + ) + (acc_scale_col2, acc_scale_col3) = cute.arch.mul_packed_f32x2( + (norm_const, norm_const), (max_value_rcp2, max_value_rcp3), rnd="rn", ftz=False, + ) + + (src[vi], src[vi + 1]) = cute.arch.mul_packed_f32x2( + (src[vi], src[vi + 1]), (acc_scale_col0, acc_scale_col1), rnd="rn", ftz=False, + ) + (src[vi + 2], src[vi + 3]) = cute.arch.mul_packed_f32x2( + (src[vi + 2], src[vi + 3]), (acc_scale_col2, acc_scale_col3), rnd="rn", ftz=False, + ) + + dst_i32 = cute.recast_tensor(dst, cutlass.Int32) + for ei in cutlass.range_constexpr(0, sf_vec_size, 4): + fp32x4 = cute.make_rmem_tensor(4, Float32) + fp32x4[0] = src[ei + 0] + fp32x4[1] = src[ei + 1] + fp32x4[2] = src[ei + 2] + fp32x4[3] = src[ei + 3] + fp8x4_i32 = cvt_f32x4_to_f8x4_pack_i32(fp32x4, d_dtype) + dst_i32[ei // 4] = cutlass.Int32(fp8x4_i32) + return qpvscale_up + + +__all__ = ["dswiglu_act", "quant_sfd_col", "quant_sfd_row", "swiglu_act"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/tmem_transpose.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/tmem_transpose.py new file mode 100644 index 000000000..12767ec26 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/tmem_transpose.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Rubin-training source-copy shim for the 16x32 TMEM transpose core. + +``_TmemTranspose16x32Core`` is the register-level transpose helper shared with +the Blackwell swap-AB epilogue. It is arch-compatible (identical math), so we +re-export it through a marked import rather than re-porting the transpose, and +rather than reaching into another kernel product's directory at port time -- +the kernel_export script inlines the source here. +""" + +# <<>> +from ....blackwell.inference.mega.block_scaled_swap_ab_fc12_epilogue import ( + _TmemTranspose16x32Core, +) + + +__all__ = ["_TmemTranspose16x32Core"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/topk_reduce.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/topk_reduce.py new file mode 100644 index 000000000..5eb1d6102 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/topk_reduce.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Rubin-training source-copy shim for the compatible Blackwell TopK reduction. + +Identical implementation to the Blackwell / inference ``TopkReduce``; kept as a +marked import so ``rubin.training.mega`` stays a self-contained deliverable and +the kernel_export script can inline the source instead of pulling in another +kernel product's directory. +""" + +# <<>> +from ....blackwell.inference.mega.topk_reduce import TopkReduce + + +__all__ = ["TopkReduce"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/__init__.py new file mode 100644 index 000000000..19945fd80 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/__init__.py @@ -0,0 +1,32 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Scheduler protocols and implementations.""" + +from .base import SchedulerBase, SchedulerConsumer, SchedulerWorkTileBase, WorkIdAcquisitionMode +from .fc12_mapping import ( + BlockPhase, + Fc12WorkTileState, + NonSwapAbFc12WorkTileInfo, + SwapAbFc12WorkTileInfo, + peek_ready_bit, +) +from .fc12_scheduler import BlackwellFusedFc12Scheduler, PhaseInterleavedFc12Scheduler +from .non_clc_mixed_cga import NonClcMixedCgaConfig, NonClcMixedCgaSchedulerWorker + + +__all__ = [ + "BlackwellFusedFc12Scheduler", + "BlockPhase", + "Fc12WorkTileState", + "NonSwapAbFc12WorkTileInfo", + "NonClcMixedCgaConfig", + "NonClcMixedCgaSchedulerWorker", + "PhaseInterleavedFc12Scheduler", + "SchedulerBase", + "SchedulerConsumer", + "SchedulerWorkTileBase", + "SwapAbFc12WorkTileInfo", + "WorkIdAcquisitionMode", + "peek_ready_bit", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/base.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/base.py new file mode 100644 index 000000000..ceffeb665 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/base.py @@ -0,0 +1,209 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Scheduler façade and architecture-independent work-tile transport.""" + +from abc import ABC, abstractmethod +from typing import Any, ClassVar, Literal, Optional, Type + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass.cutlass_dsl import extract_mlir_values, new_from_mlir_values + +from ...api import ImplDesc, KernelComponent, ProblemDesc +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.smem_workspace import SmemWorkspace + + +WorkIdAcquisitionMode = Literal["grid_stride", "atomic_counter", "cluster_launch_control"] + + +class SchedulerWorkTileBase(ABC): + """Register ABI for one work tile transported through scheduler SMEM.""" + + storage_dtype: ClassVar[type] = cutlass.Int32 + storage_field_count: ClassVar[int] + + @property + @abstractmethod + def is_valid_tile(self): + """Return whether this tile names executable work.""" + ... + + @abstractmethod + def to_rmem(self) -> cute.Tensor: + """Serialize this tile into its one-dimensional register ABI.""" + ... + + @classmethod + @abstractmethod + def from_rmem(cls, registers: cute.Tensor) -> "SchedulerWorkTileBase": + """Deserialize one tile from its register ABI.""" + ... + + +class SchedulerConsumer: + """Per-consumer state for the common scheduler SMEM transport.""" + + def __init__( + self, + scheduler_pipeline: pipeline.PipelineAsync, + smem_buffer: cute.Tensor, + num_stages: int, + work_tile_type: Type[SchedulerWorkTileBase], + ) -> None: + self._pipeline = scheduler_pipeline + self._smem_buffer = smem_buffer + self._work_tile_type = work_tile_type + self._consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, num_stages) + + def __extract_mlir_values__(self) -> list: + return extract_mlir_values(self._consumer_state) + + def __new_from_mlir_values__(self, values: list) -> "SchedulerConsumer": + expected_value_count = len(extract_mlir_values(self._consumer_state)) + if len(values) != expected_value_count: + raise ValueError( + f"SchedulerConsumer MLIR value count mismatch: expected {expected_value_count}, got {len(values)}." + ) + result = type(self).__new__(type(self)) + result._pipeline = self._pipeline + result._smem_buffer = self._smem_buffer + result._work_tile_type = self._work_tile_type + result._consumer_state = new_from_mlir_values(self._consumer_state, values) + return result + + @cute.jit + def consume_work(self) -> SchedulerWorkTileBase: + """Block until the next work tile is available.""" + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self._work_tile_type.storage_dtype, num_bits_per_copy=128 + ) + self._pipeline.consumer_wait(self._consumer_state) + registers = cute.make_rmem_tensor( + (self._work_tile_type.storage_field_count,), self._work_tile_type.storage_dtype + ) + cute.copy(copy_atom, self._smem_buffer[(None, self._consumer_state.index)], registers) + work_tile = self._work_tile_type.from_rmem(registers) + cute.arch.fence_acq_rel_cta() + self._pipeline.consumer_release(self._consumer_state) + self._consumer_state.advance() + return work_tile + + +class SchedulerBase(KernelComponent): + """Common work-tile transport and façade protocol for schedulers.""" + + pipeline_mbarriers_region = "scheduler.pipeline_mbarriers" + work_tiles_region = "scheduler.work_tiles" + num_scheduler_stages = 2 + + @classmethod + def problem_desc_require(cls) -> dict: + return {} + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return {"num_scheduler_consumer_threads": int} + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + self.num_scheduler_consumer_threads = impl_desc["num_scheduler_consumer_threads"] + if self.num_scheduler_consumer_threads <= 0: + raise ValueError("num_scheduler_consumer_threads must be positive.") + + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: + """Register the common work-tile transport regions.""" + if not hasattr(self, "work_tile_type"): + raise AttributeError( + f"{type(self).__name__} must bind work_tile_type before registering scheduler SMEM regions." + ) + work_tile_type = self.work_tile_type + work_tile_field_count = work_tile_type.storage_field_count + smem_workspace.register_mbarrier(self.pipeline_mbarriers_region, self.num_scheduler_stages * 2) + smem_workspace.register_tensor( + self.work_tiles_region, + work_tile_type.storage_dtype, + (work_tile_field_count, self.num_scheduler_stages), + stride=(1, work_tile_field_count), + byte_alignment=16, + ) + + def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: + """Register scheduler-specific GMEM regions when needed.""" + pass + + @cute.jit + def create_scheduler_pipelines(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Create CTA-lifetime scheduler pipelines and transport state.""" + self._pipeline = pipeline.PipelineAsync.create( + num_stages=self.num_scheduler_stages, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_scheduler_consumer_threads), + barrier_storage=smem_workspace.ptr(self.pipeline_mbarriers_region, smem_base), + defer_sync=True, + ) + self._smem_buffer = smem_workspace.tensor(self.work_tiles_region, smem_base) + self._producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_scheduler_stages + ) + + def make_consumer(self) -> SchedulerConsumer: + """Create a consumer with an independent pipeline state.""" + return SchedulerConsumer( + scheduler_pipeline=self._pipeline, + smem_buffer=self._smem_buffer, + num_stages=self.num_scheduler_stages, + work_tile_type=self.work_tile_type, + ) + + @cute.jit + def publish_work(self, work_tile: SchedulerWorkTileBase) -> None: + """Publish one work tile through the common transport pipeline.""" + copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), work_tile.storage_dtype, num_bits_per_copy=128) + self._pipeline.producer_acquire(self._producer_state) + cute.copy(copy_atom, work_tile.to_rmem(), self._smem_buffer[(None, self._producer_state.index)]) + cute.arch.fence_proxy("async.shared", space="cta") + self._pipeline.producer_commit(self._producer_state) + self._producer_state.advance() + + @cute.jit + def produce_tail(self) -> None: + """Wait until every published work tile has been consumed.""" + self._pipeline.producer_tail(self._producer_state) + + def __extract_mlir_values__(self) -> list: + return extract_mlir_values(self._producer_state) + + def __new_from_mlir_values__(self, values: list) -> "SchedulerBase": + expected_value_count = len(extract_mlir_values(self._producer_state)) + if len(values) != expected_value_count: + raise ValueError( + f"SchedulerBase MLIR value count mismatch: expected {expected_value_count}, got {len(values)}." + ) + result = type(self).__new__(type(self)) + result.num_scheduler_consumer_threads = self.num_scheduler_consumer_threads + result.work_tile_type = self.work_tile_type + result._pipeline = self._pipeline + result._smem_buffer = self._smem_buffer + result._producer_state = new_from_mlir_values(self._producer_state, values) + return result + + @abstractmethod + def get_grid_shape(self, *, max_active_clusters: Optional[int] = None, problem_desc: Any = None): + """Return the launch grid from static scheduler policy.""" + ... + + @abstractmethod + def assign_device_members(self, *args, **kwargs) -> None: + """Initialize device members whose ownership spans one CTA lifetime.""" + ... + + @abstractmethod + def gen_next_work(self) -> SchedulerWorkTileBase: + """Claim, map, and return the next work tile.""" + ... + + +__all__ = ["SchedulerBase", "SchedulerConsumer", "SchedulerWorkTileBase", "WorkIdAcquisitionMode"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py new file mode 100644 index 000000000..ab3797072 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py @@ -0,0 +1,1173 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""FC12 work-tile ABI and grouped or phase-interleaved task mapping.""" + +import dataclasses +from enum import IntEnum +from typing import List, Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass.cutlass_dsl import Boolean, Int32, extract_mlir_values, new_from_mlir_values + +from ...helpers.iket_compat import iket +from ...helpers.utils import padded_expert_rows +from .base import SchedulerWorkTileBase + + +phase_bits = 16 +phase_mask = (1 << phase_bits) - 1 +peek_ready_bit = 1 << phase_bits + + +class Fc12WorkTileState(IntEnum): + """Sentinel values carried in the expert index field.""" + + Done = -1 + + +class BlockPhase(IntEnum): + """FC1/FC2 phase encoded in one fused work tile.""" + + None_ = 0 + Linear1 = 1 + Linear2 = 2 + + +@dataclasses.dataclass(frozen=True) +class SwapAbFc12WorkTileInfo(SchedulerWorkTileBase): + """Eight-field work tile for the swap-AB FC12 orientation.""" + + storage_field_count = 8 + + expert_idx: Int32 + tile_m_idx: Int32 + tile_n_idx: Int32 + cumulative_data_physical_row: Int32 + cumulative_sf_physical_row: Int32 + cumulative_token_block_count: Int32 + valid_tokens_in_cta_tile: Int32 + phase_and_flags: Int32 + + @property + def is_valid_tile(self): + return self.expert_idx >= Int32(0) + + @property + def phase(self) -> Int32: + return self.phase_and_flags & Int32(phase_mask) + + @property + def peek_ready(self): + return self.phase_and_flags >= Int32(peek_ready_bit) + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in ( + self.expert_idx, + self.tile_m_idx, + self.tile_n_idx, + self.cumulative_data_physical_row, + self.cumulative_sf_physical_row, + self.cumulative_token_block_count, + self.valid_tokens_in_cta_tile, + self.phase_and_flags, + ): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "SwapAbFc12WorkTileInfo": + if len(values) != self.storage_field_count: + raise ValueError( + f"SwapAbFc12WorkTileInfo expects {self.storage_field_count} MLIR values, got {len(values)}." + ) + fields = ( + self.expert_idx, + self.tile_m_idx, + self.tile_n_idx, + self.cumulative_data_physical_row, + self.cumulative_sf_physical_row, + self.cumulative_token_block_count, + self.valid_tokens_in_cta_tile, + self.phase_and_flags, + ) + rebuilt = [new_from_mlir_values(field, [value]) for field, value in zip(fields, values)] + return type(self)(*rebuilt) + + def to_rmem(self) -> cute.Tensor: + registers = cute.make_rmem_tensor((self.storage_field_count,), cutlass.Int32) + registers[0] = self.expert_idx + registers[1] = self.tile_m_idx + registers[2] = self.tile_n_idx + registers[3] = self.cumulative_data_physical_row + registers[4] = self.cumulative_sf_physical_row + registers[5] = self.cumulative_token_block_count + registers[6] = self.valid_tokens_in_cta_tile + registers[7] = self.phase_and_flags + return registers + + @classmethod + def from_rmem(cls, registers: cute.Tensor) -> "SwapAbFc12WorkTileInfo": + return cls( + expert_idx=registers[0], + tile_m_idx=registers[1], + tile_n_idx=registers[2], + cumulative_data_physical_row=registers[3], + cumulative_sf_physical_row=registers[4], + cumulative_token_block_count=registers[5], + valid_tokens_in_cta_tile=registers[6], + phase_and_flags=registers[7], + ) + + +@dataclasses.dataclass(frozen=True) +class NonSwapAbFc12WorkTileInfo(SchedulerWorkTileBase): + """Eight-field work tile for the non-swap-AB FC12 orientation.""" + + storage_field_count = 8 + + expert_idx: Int32 + tile_m_idx: Int32 + tile_n_idx: Int32 + cumulative_data_physical_row: Int32 + cumulative_sf_physical_row: Int32 + cumulative_token_block_count: Int32 + valid_tokens_in_cta_cluster_tile: Int32 + phase_and_flags: Int32 + + @property + def is_valid_tile(self): + return self.expert_idx >= Int32(0) + + @property + def phase(self) -> Int32: + return self.phase_and_flags & Int32(phase_mask) + + @property + def peek_ready(self): + return self.phase_and_flags >= Int32(peek_ready_bit) + + @property + def valid_tokens_in_cta_tile(self) -> Int32: + return self.valid_tokens_in_cta_cluster_tile >> Int32(16) + + @property + def valid_tokens_in_cluster_tile(self) -> Int32: + return self.valid_tokens_in_cta_cluster_tile & Int32(0xFFFF) + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in ( + self.expert_idx, + self.tile_m_idx, + self.tile_n_idx, + self.cumulative_data_physical_row, + self.cumulative_sf_physical_row, + self.cumulative_token_block_count, + self.valid_tokens_in_cta_cluster_tile, + self.phase_and_flags, + ): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "NonSwapAbFc12WorkTileInfo": + if len(values) != self.storage_field_count: + raise ValueError( + f"NonSwapAbFc12WorkTileInfo expects {self.storage_field_count} MLIR values, got {len(values)}." + ) + fields = ( + self.expert_idx, + self.tile_m_idx, + self.tile_n_idx, + self.cumulative_data_physical_row, + self.cumulative_sf_physical_row, + self.cumulative_token_block_count, + self.valid_tokens_in_cta_cluster_tile, + self.phase_and_flags, + ) + rebuilt = [new_from_mlir_values(field, [value]) for field, value in zip(fields, values)] + return type(self)(*rebuilt) + + def to_rmem(self) -> cute.Tensor: + registers = cute.make_rmem_tensor((self.storage_field_count,), cutlass.Int32) + registers[0] = self.expert_idx + registers[1] = self.tile_m_idx + registers[2] = self.tile_n_idx + registers[3] = self.cumulative_data_physical_row + registers[4] = self.cumulative_sf_physical_row + registers[5] = self.cumulative_token_block_count + registers[6] = self.valid_tokens_in_cta_cluster_tile + registers[7] = self.phase_and_flags + return registers + + @classmethod + def from_rmem(cls, registers: cute.Tensor) -> "NonSwapAbFc12WorkTileInfo": + return cls( + expert_idx=registers[0], + tile_m_idx=registers[1], + tile_n_idx=registers[2], + cumulative_data_physical_row=registers[3], + cumulative_sf_physical_row=registers[4], + cumulative_token_block_count=registers[5], + valid_tokens_in_cta_cluster_tile=registers[6], + phase_and_flags=registers[7], + ) + + +class _Fc12TaskCursorState: + """Register-resident cursor for the FC12 group/phase/expert state machine.""" + + def __init__( + self, + current_group_first_expert: Int32, + current_group_last_expert_exclusive: Int32, + current_phase: Int32, + current_expert_idx: Int32, + current_expert_tile_start: Int32, + current_expert_tile_end: Int32, + current_group_fc1_subphase_end: Int32, + current_group_end: Int32, + cumulative_fc1_tiles_at_group_end: Int32, + cumulative_fc2_tiles_at_group_end: Int32, + current_data_cumulative: Int32, + current_sf_cumulative: Int32, + current_token_block_cumulative: Int32, + group_start_data_cumulative: Int32, + group_start_sf_cumulative: Int32, + group_start_token_block_cumulative: Int32, + current_token_block_count: Int32, + current_expert_token_count: Int32, + ) -> None: + self.current_group_first_expert = current_group_first_expert + self.current_group_last_expert_exclusive = current_group_last_expert_exclusive + self.current_phase = current_phase + self.current_expert_idx = current_expert_idx + self.current_expert_tile_start = current_expert_tile_start + self.current_expert_tile_end = current_expert_tile_end + self.current_group_fc1_subphase_end = current_group_fc1_subphase_end + self.current_group_end = current_group_end + self.cumulative_fc1_tiles_at_group_end = cumulative_fc1_tiles_at_group_end + self.cumulative_fc2_tiles_at_group_end = cumulative_fc2_tiles_at_group_end + self.current_data_cumulative = current_data_cumulative + self.current_sf_cumulative = current_sf_cumulative + self.current_token_block_cumulative = current_token_block_cumulative + self.group_start_data_cumulative = group_start_data_cumulative + self.group_start_sf_cumulative = group_start_sf_cumulative + self.group_start_token_block_cumulative = group_start_token_block_cumulative + self.current_token_block_count = current_token_block_count + self.current_expert_token_count = current_expert_token_count + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in self._fields(): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "_Fc12TaskCursorState": + value_index = 0 + rebuilt = [] + for field in self._fields(): + field_value_count = len(extract_mlir_values(field)) + rebuilt.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) + value_index += field_value_count + if value_index != len(values): + raise ValueError( + f"_Fc12TaskCursorState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return type(self)(*rebuilt) + + def _fields(self) -> Tuple: + return ( + self.current_group_first_expert, + self.current_group_last_expert_exclusive, + self.current_phase, + self.current_expert_idx, + self.current_expert_tile_start, + self.current_expert_tile_end, + self.current_group_fc1_subphase_end, + self.current_group_end, + self.cumulative_fc1_tiles_at_group_end, + self.cumulative_fc2_tiles_at_group_end, + self.current_data_cumulative, + self.current_sf_cumulative, + self.current_token_block_cumulative, + self.group_start_data_cumulative, + self.group_start_sf_cumulative, + self.group_start_token_block_cumulative, + self.current_token_block_count, + self.current_expert_token_count, + ) + + +class Fc12TaskMappingState: + """Runtime inputs and cursor for monotonic FC12 linear-ID mapping.""" + + def __init__( + self, + expert_count, + mapping_cta_tile_shape_mnk: Tuple[int, int, int], + mapping_cluster_shape_mn: Tuple[int, int], + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + is_swap_ab: bool, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], + cursor_state: _Fc12TaskCursorState, + num_fc1_intermediate_blocks, + num_fc2_hidden_blocks, + ) -> None: + self.expert_count = expert_count + self.mapping_cta_tile_shape_mnk = mapping_cta_tile_shape_mnk + self.mapping_cluster_shape_mn = mapping_cluster_shape_mn + self.group_hint = group_hint + self.token_padding_block = token_padding_block + self.sf_padding_block = sf_padding_block + self.is_swap_ab = is_swap_ab + self.expert_token_sizes = expert_token_sizes + self.expert_token_prefix_sum = expert_token_prefix_sum + self.cursor_state = cursor_state + self.num_fc1_intermediate_blocks = num_fc1_intermediate_blocks + self.num_fc2_hidden_blocks = num_fc2_hidden_blocks + + @property + def mapping_cluster_tile_m(self) -> int: + return self.mapping_cta_tile_shape_mnk[0] * self.mapping_cluster_shape_mn[0] + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + if isinstance(self.expert_count, Int32): + values.extend(extract_mlir_values(self.expert_count)) + token_counts = self.expert_token_sizes if self.expert_token_sizes is not None else self.expert_token_prefix_sum + values.extend(extract_mlir_values(token_counts)) + values.extend(extract_mlir_values(self.cursor_state)) + if isinstance(self.num_fc1_intermediate_blocks, Int32): + values.extend(extract_mlir_values(self.num_fc1_intermediate_blocks)) + if isinstance(self.num_fc2_hidden_blocks, Int32): + values.extend(extract_mlir_values(self.num_fc2_hidden_blocks)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "Fc12TaskMappingState": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + expert_count = rebuild(self.expert_count) if isinstance(self.expert_count, Int32) else self.expert_count + if self.expert_token_sizes is not None: + expert_token_sizes = rebuild(self.expert_token_sizes) + expert_token_prefix_sum = None + else: + expert_token_sizes = None + expert_token_prefix_sum = rebuild(self.expert_token_prefix_sum) + result = type(self)( + expert_count=expert_count, + mapping_cta_tile_shape_mnk=self.mapping_cta_tile_shape_mnk, + mapping_cluster_shape_mn=self.mapping_cluster_shape_mn, + group_hint=self.group_hint, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + is_swap_ab=self.is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + cursor_state=rebuild(self.cursor_state), + num_fc1_intermediate_blocks=( + rebuild(self.num_fc1_intermediate_blocks) + if isinstance(self.num_fc1_intermediate_blocks, Int32) + else self.num_fc1_intermediate_blocks + ), + num_fc2_hidden_blocks=( + rebuild(self.num_fc2_hidden_blocks) + if isinstance(self.num_fc2_hidden_blocks, Int32) + else self.num_fc2_hidden_blocks + ), + ) + if value_index != len(values): + raise ValueError( + f"Fc12TaskMappingState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return result + + +@cute.jit +def create_fc12_task_mapping_state( + *, + expert_count, + intermediate_gateup_size, + hidden_size, + mapping_cta_tile_shape_mnk: Tuple[int, int, int], + mapping_cluster_shape_mn: Tuple[int, int], + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + is_swap_ab: bool, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], +) -> Fc12TaskMappingState: + """Create the register-resident state for one CTA's FC12 mapper.""" + cursor_state = _Fc12TaskCursorState( + current_group_first_expert=Int32(0), + current_group_last_expert_exclusive=Int32(0), + current_phase=Int32(BlockPhase.Linear1), + current_expert_idx=Int32(-1), + current_expert_tile_start=Int32(0), + current_expert_tile_end=Int32(0), + current_group_fc1_subphase_end=Int32(0), + current_group_end=Int32(0), + cumulative_fc1_tiles_at_group_end=Int32(0), + cumulative_fc2_tiles_at_group_end=Int32(0), + current_data_cumulative=Int32(0), + current_sf_cumulative=Int32(0), + current_token_block_cumulative=Int32(0), + group_start_data_cumulative=Int32(0), + group_start_sf_cumulative=Int32(0), + group_start_token_block_cumulative=Int32(0), + current_token_block_count=Int32(0), + current_expert_token_count=Int32(0), + ) + mapping_cluster_tile_n = mapping_cluster_shape_mn[1] * mapping_cta_tile_shape_mnk[1] + num_fc1_intermediate_blocks = (intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + num_fc2_hidden_blocks = (hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + return Fc12TaskMappingState( + expert_count=expert_count, + mapping_cta_tile_shape_mnk=mapping_cta_tile_shape_mnk, + mapping_cluster_shape_mn=mapping_cluster_shape_mn, + group_hint=group_hint, + token_padding_block=token_padding_block, + sf_padding_block=sf_padding_block, + is_swap_ab=is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + cursor_state=cursor_state, + num_fc1_intermediate_blocks=num_fc1_intermediate_blocks, + num_fc2_hidden_blocks=num_fc2_hidden_blocks, + ) + + +@cute.jit +def _warp_inclusive_sum(value: Int32, lane_idx: Int32) -> Int32: + inclusive = value + for step_log in cutlass.range_constexpr(5): + step = Int32(1 << step_log) + previous = Int32(cute.arch.shuffle_sync(inclusive, lane_idx - step)) + if lane_idx >= step: + inclusive = inclusive + previous + return inclusive + + +@cute.jit +def _first_matching_lane(predicate) -> Int32: + mask = Int32(cute.arch.vote_ballot_sync(predicate)) + first_lane = Int32(-1) + if mask != Int32(0): + lowbit = mask & (-mask) + first_lane = Int32(cute.arch.popc(lowbit - Int32(1))) + return first_lane + + +@cute.jit +def _load_expert_batch_metrics( + mapping_state: Fc12TaskMappingState, batch_base: Int32, active_begin: Int32, active_end: Int32, lane_idx: Int32 +) -> Tuple[Int32, Int32, Int32, Int32, Int32, Int32, Int32]: + expert_idx = batch_base + lane_idx + token_count = Int32(0) + if cutlass.const_expr(mapping_state.expert_token_sizes is not None): + if expert_idx < mapping_state.expert_count: + token_count = mapping_state.expert_token_sizes[expert_idx] + else: + prefix_end = Int32(0) + if expert_idx < mapping_state.expert_count: + prefix_end = mapping_state.expert_token_prefix_sum[expert_idx] + prefix_begin = Int32(cute.arch.shuffle_sync(prefix_end, lane_idx - Int32(1))) + if lane_idx == Int32(0): + prefix_begin = Int32(0) + if batch_base > Int32(0): + prefix_begin = mapping_state.expert_token_prefix_sum[batch_base - Int32(1)] + token_count = prefix_end - prefix_begin + if (expert_idx < active_begin) | (expert_idx >= active_end): + token_count = Int32(0) + token_blocks = (token_count + Int32(mapping_state.mapping_cluster_tile_m - 1)) // Int32( + mapping_state.mapping_cluster_tile_m + ) + data_rows = padded_expert_rows(token_count, Int32(mapping_state.token_padding_block)) + sf_rows = padded_expert_rows(token_count, Int32(mapping_state.sf_padding_block)) + fc1_tiles = token_blocks * mapping_state.num_fc1_intermediate_blocks + fc2_tiles = token_blocks * mapping_state.num_fc2_hidden_blocks + return expert_idx, token_count, token_blocks, data_rows, sf_rows, fc1_tiles, fc2_tiles + + +def make_fc12_done_tile(is_swap_ab: bool) -> SchedulerWorkTileBase: + """Build the terminal work tile for either FC12 orientation.""" + if cutlass.const_expr(is_swap_ab): + return SwapAbFc12WorkTileInfo( + expert_idx=Int32(Fc12WorkTileState.Done), + tile_m_idx=Int32(0), + tile_n_idx=Int32(0), + cumulative_data_physical_row=Int32(0), + cumulative_sf_physical_row=Int32(0), + cumulative_token_block_count=Int32(0), + valid_tokens_in_cta_tile=Int32(0), + phase_and_flags=Int32(BlockPhase.None_), + ) + return NonSwapAbFc12WorkTileInfo( + expert_idx=Int32(Fc12WorkTileState.Done), + tile_m_idx=Int32(0), + tile_n_idx=Int32(0), + cumulative_data_physical_row=Int32(0), + cumulative_sf_physical_row=Int32(0), + cumulative_token_block_count=Int32(0), + valid_tokens_in_cta_cluster_tile=Int32(0), + phase_and_flags=Int32(BlockPhase.None_), + ) + + +@cute.jit +def _switch_to_fc2(mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: + cursor = mapping_state.cursor_state + cursor.current_phase = Int32(BlockPhase.Linear2) + cursor.current_expert_idx = cursor.current_group_first_expert - Int32(1) + cursor.current_expert_tile_start = cursor.current_group_fc1_subphase_end + cursor.current_expert_tile_end = cursor.current_group_fc1_subphase_end + cursor.current_expert_token_count = Int32(0) + cursor.current_token_block_count = Int32(0) + cursor.current_data_cumulative = cursor.group_start_data_cumulative + cursor.current_sf_cumulative = cursor.group_start_sf_cumulative + cursor.current_token_block_cumulative = cursor.group_start_token_block_cumulative + mapping_state.cursor_state = cursor + return cursor + + +@cute.jit +def _sum_expert_range( + mapping_state: Fc12TaskMappingState, expert_begin: Int32, expert_end: Int32 +) -> Tuple[Int32, Int32, Int32]: + lane_idx = Int32(cute.arch.lane_idx()) + data_rows = Int32(0) + sf_rows = Int32(0) + token_blocks = Int32(0) + batch_base = (expert_begin // Int32(32)) * Int32(32) + while batch_base < expert_end: + (_, _, lane_token_blocks, lane_data_rows, lane_sf_rows, _, _) = _load_expert_batch_metrics( + mapping_state, batch_base, expert_begin, expert_end, lane_idx + ) + data_rows = data_rows + Int32(cute.arch.warp_redux_sync(lane_data_rows, "add")) + sf_rows = sf_rows + Int32(cute.arch.warp_redux_sync(lane_sf_rows, "add")) + token_blocks = token_blocks + Int32(cute.arch.warp_redux_sync(lane_token_blocks, "add")) + batch_base = batch_base + Int32(32) + return data_rows, sf_rows, token_blocks + + +@cute.jit +def _build_group_range( + mapping_state: Fc12TaskMappingState, group_first_expert: Int32, base_fc1_tiles: Int32, base_fc2_tiles: Int32 +) -> Tuple[Int32, Int32, Int32]: + lane_idx = Int32(cute.arch.lane_idx()) + group_threshold = base_fc1_tiles + Int32(mapping_state.group_hint) + group_last_expert = group_first_expert + cumulative_fc1_tiles = base_fc1_tiles + cumulative_fc2_tiles = base_fc2_tiles + batch_base = (group_first_expert // Int32(32)) * Int32(32) + + while batch_base < mapping_state.expert_count and cumulative_fc1_tiles < group_threshold: + (lane_expert_idx, _, _, _, _, lane_fc1_tiles, lane_fc2_tiles) = _load_expert_batch_metrics( + mapping_state, batch_base, group_first_expert, Int32(mapping_state.expert_count), lane_idx + ) + fc1_prefix = _warp_inclusive_sum(lane_fc1_tiles, lane_idx) + reaches_threshold = ( + (lane_expert_idx >= group_first_expert) & (lane_expert_idx < mapping_state.expert_count) + ) & (cumulative_fc1_tiles + fc1_prefix >= group_threshold) + selected_lane = _first_matching_lane(reaches_threshold) + included_fc2_tiles = lane_fc2_tiles + if selected_lane >= Int32(0): + if lane_idx > selected_lane: + included_fc2_tiles = Int32(0) + fc2_batch_tiles = Int32(cute.arch.warp_redux_sync(included_fc2_tiles, "add")) + if selected_lane >= Int32(0): + cumulative_fc1_tiles = cumulative_fc1_tiles + Int32(cute.arch.shuffle_sync(fc1_prefix, selected_lane)) + group_last_expert = batch_base + selected_lane + Int32(1) + else: + cumulative_fc1_tiles = cumulative_fc1_tiles + Int32(cute.arch.shuffle_sync(fc1_prefix, Int32(31))) + batch_base = batch_base + Int32(32) + group_last_expert = cutlass.min(batch_base, Int32(mapping_state.expert_count)) + cumulative_fc2_tiles = cumulative_fc2_tiles + fc2_batch_tiles + + return group_last_expert, cumulative_fc1_tiles, cumulative_fc2_tiles + + +@cute.jit +def _advance_group(mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: + cursor = mapping_state.cursor_state + + residual_begin = cutlass.max(cursor.current_expert_idx, cursor.current_group_first_expert) + residual_data_rows = Int32(0) + residual_sf_rows = Int32(0) + residual_token_blocks = Int32(0) + if residual_begin < cursor.current_group_last_expert_exclusive: + iket.range_push("scheduler.residual_scan") + residual_data_rows, residual_sf_rows, residual_token_blocks = _sum_expert_range( + mapping_state, residual_begin, cursor.current_group_last_expert_exclusive + ) + iket.range_pop() + cursor.current_data_cumulative = cursor.current_data_cumulative + residual_data_rows + cursor.current_sf_cumulative = cursor.current_sf_cumulative + residual_sf_rows + cursor.current_token_block_cumulative = cursor.current_token_block_cumulative + residual_token_blocks + + cursor.group_start_data_cumulative = cursor.current_data_cumulative + cursor.group_start_sf_cumulative = cursor.current_sf_cumulative + cursor.group_start_token_block_cumulative = cursor.current_token_block_cumulative + + base_fc1_tiles = cursor.cumulative_fc1_tiles_at_group_end + base_fc2_tiles = cursor.cumulative_fc2_tiles_at_group_end + cursor.current_group_first_expert = cursor.current_group_last_expert_exclusive + + iket.range_push("scheduler.group_scan") + (cursor.current_group_last_expert_exclusive, cumulative_fc1_tiles, cumulative_fc2_tiles) = _build_group_range( + mapping_state, cursor.current_group_first_expert, base_fc1_tiles, base_fc2_tiles + ) + iket.range_pop() + cursor.cumulative_fc1_tiles_at_group_end = cumulative_fc1_tiles + cursor.cumulative_fc2_tiles_at_group_end = cumulative_fc2_tiles + group_start_tile = cursor.current_group_end + cursor.current_group_fc1_subphase_end = group_start_tile + cumulative_fc1_tiles - base_fc1_tiles + cursor.current_group_end = cursor.current_group_fc1_subphase_end + cumulative_fc2_tiles - base_fc2_tiles + + cursor.current_phase = Int32(BlockPhase.Linear1) + cursor.current_expert_idx = cursor.current_group_first_expert - Int32(1) + cursor.current_expert_tile_start = group_start_tile + cursor.current_expert_tile_end = group_start_tile + cursor.current_expert_token_count = Int32(0) + cursor.current_token_block_count = Int32(0) + mapping_state.cursor_state = cursor + return cursor + + +@cute.jit +def _seek_expert_for_work_id(linear_work_id: Int32, mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: + cursor = mapping_state.cursor_state + + base_tile_end = cursor.current_expert_tile_end + base_data_cumulative = cursor.current_data_cumulative + base_sf_cumulative = cursor.current_sf_cumulative + base_token_block_cumulative = cursor.current_token_block_cumulative + if cursor.current_expert_idx >= cursor.current_group_first_expert: + current_token_count = cursor.current_expert_token_count + base_data_cumulative = base_data_cumulative + padded_expert_rows( + current_token_count, Int32(mapping_state.token_padding_block) + ) + base_sf_cumulative = base_sf_cumulative + padded_expert_rows( + current_token_count, Int32(mapping_state.sf_padding_block) + ) + base_token_block_cumulative = base_token_block_cumulative + cursor.current_token_block_count + + search_begin = cutlass.max(cursor.current_expert_idx + Int32(1), cursor.current_group_first_expert) + batch_base = (search_begin // Int32(32)) * Int32(32) + selected_expert = Int32(-1) + selected_token_count = Int32(0) + selected_token_blocks = Int32(0) + selected_tile_start = Int32(0) + selected_tile_end = Int32(0) + selected_data_cumulative = Int32(0) + selected_sf_cumulative = Int32(0) + selected_token_block_cumulative = Int32(0) + lane_idx = Int32(cute.arch.lane_idx()) + + iket.range_push("scheduler.expert_scan") + while selected_expert < Int32(0) and batch_base < cursor.current_group_last_expert_exclusive: + ( + lane_expert_idx, + lane_token_count, + lane_token_blocks, + lane_data_rows, + lane_sf_rows, + lane_fc1_tiles, + lane_fc2_tiles, + ) = _load_expert_batch_metrics( + mapping_state, batch_base, search_begin, cursor.current_group_last_expert_exclusive, lane_idx + ) + lane_phase_tiles = lane_fc1_tiles + if cursor.current_phase == Int32(BlockPhase.Linear2): + lane_phase_tiles = lane_fc2_tiles + tile_prefix = _warp_inclusive_sum(lane_phase_tiles, lane_idx) + candidate_tile_end = base_tile_end + tile_prefix + contains_work = ( + (lane_expert_idx >= search_begin) + & (lane_expert_idx < cursor.current_group_last_expert_exclusive) + & (linear_work_id < candidate_tile_end) + ) + selected_lane = _first_matching_lane(contains_work) + + included_data_rows = lane_data_rows + included_sf_rows = lane_sf_rows + included_token_blocks = lane_token_blocks + if selected_lane >= Int32(0): + if lane_idx > selected_lane: + included_data_rows = Int32(0) + included_sf_rows = Int32(0) + included_token_blocks = Int32(0) + batch_data_rows = Int32(cute.arch.warp_redux_sync(included_data_rows, "add")) + batch_sf_rows = Int32(cute.arch.warp_redux_sync(included_sf_rows, "add")) + batch_token_blocks = Int32(cute.arch.warp_redux_sync(included_token_blocks, "add")) + if selected_lane >= Int32(0): + selected_expert = batch_base + selected_lane + selected_token_count = Int32(cute.arch.shuffle_sync(lane_token_count, selected_lane)) + selected_token_blocks = Int32(cute.arch.shuffle_sync(lane_token_blocks, selected_lane)) + selected_phase_tiles = Int32(cute.arch.shuffle_sync(lane_phase_tiles, selected_lane)) + selected_tile_end = base_tile_end + Int32(cute.arch.shuffle_sync(tile_prefix, selected_lane)) + selected_tile_start = selected_tile_end - selected_phase_tiles + selected_data_rows = Int32(cute.arch.shuffle_sync(lane_data_rows, selected_lane)) + selected_sf_rows = Int32(cute.arch.shuffle_sync(lane_sf_rows, selected_lane)) + selected_data_cumulative = base_data_cumulative + batch_data_rows - selected_data_rows + selected_sf_cumulative = base_sf_cumulative + batch_sf_rows - selected_sf_rows + selected_token_block_cumulative = base_token_block_cumulative + batch_token_blocks - selected_token_blocks + else: + base_tile_end = base_tile_end + Int32(cute.arch.shuffle_sync(tile_prefix, Int32(31))) + base_data_cumulative = base_data_cumulative + batch_data_rows + base_sf_cumulative = base_sf_cumulative + batch_sf_rows + base_token_block_cumulative = base_token_block_cumulative + batch_token_blocks + batch_base = batch_base + Int32(32) + search_begin = batch_base + iket.range_pop() + + cursor.current_expert_idx = selected_expert + cursor.current_expert_token_count = selected_token_count + cursor.current_token_block_count = selected_token_blocks + cursor.current_expert_tile_start = selected_tile_start + cursor.current_expert_tile_end = selected_tile_end + cursor.current_data_cumulative = selected_data_cumulative + cursor.current_sf_cumulative = selected_sf_cumulative + cursor.current_token_block_cumulative = selected_token_block_cumulative + return cursor + + +@cute.jit +def _decode_inside_expert( + linear_work_id: Int32, cta_id_in_mapping_cluster: cute.Coord, mapping_state: Fc12TaskMappingState +) -> SchedulerWorkTileBase: + cursor = mapping_state.cursor_state + cta_tile_m = mapping_state.mapping_cta_tile_shape_mnk[0] + local_work_id = linear_work_id - cursor.current_expert_tile_start + + cluster_token_block_idx = Int32(0) + cluster_output_block_idx = Int32(0) + if cursor.current_phase == Int32(BlockPhase.Linear1): + cluster_token_block_idx = local_work_id // mapping_state.num_fc1_intermediate_blocks + cluster_output_block_idx = local_work_id - cluster_token_block_idx * mapping_state.num_fc1_intermediate_blocks + else: + cluster_token_block_idx = local_work_id // mapping_state.num_fc2_hidden_blocks + cluster_output_block_idx = local_work_id - cluster_token_block_idx * mapping_state.num_fc2_hidden_blocks + + cta_token_block_idx = ( + cluster_token_block_idx * mapping_state.mapping_cluster_shape_mn[0] + cta_id_in_mapping_cluster[0] + ) + cta_output_block_idx = ( + cluster_output_block_idx * mapping_state.mapping_cluster_shape_mn[1] + cta_id_in_mapping_cluster[1] + ) + token_start = cta_token_block_idx * Int32(cta_tile_m) + remaining_tokens = cutlass.max(cursor.current_expert_token_count - token_start, Int32(0)) + valid_tokens_in_cta_tile = cutlass.min(remaining_tokens, Int32(cta_tile_m)) + + if cutlass.const_expr(mapping_state.is_swap_ab): + return SwapAbFc12WorkTileInfo( + expert_idx=cursor.current_expert_idx, + tile_m_idx=cta_output_block_idx, + tile_n_idx=cta_token_block_idx, + cumulative_data_physical_row=cursor.current_data_cumulative, + cumulative_sf_physical_row=cursor.current_sf_cumulative, + cumulative_token_block_count=(cursor.current_token_block_cumulative), + valid_tokens_in_cta_tile=valid_tokens_in_cta_tile, + phase_and_flags=cursor.current_phase, + ) + + cluster_tile_m = mapping_state.mapping_cluster_shape_mn[0] * cta_tile_m + cluster_token_start = cluster_token_block_idx * Int32(cluster_tile_m) + remaining_cluster_tokens = cutlass.max(cursor.current_expert_token_count - cluster_token_start, Int32(0)) + valid_tokens_in_cluster_tile = cutlass.min(remaining_cluster_tokens, Int32(cluster_tile_m)) + valid_tokens_in_cta_cluster_tile = (valid_tokens_in_cta_tile << Int32(16)) | valid_tokens_in_cluster_tile + return NonSwapAbFc12WorkTileInfo( + expert_idx=cursor.current_expert_idx, + tile_m_idx=cta_token_block_idx, + tile_n_idx=cta_output_block_idx, + cumulative_data_physical_row=cursor.current_data_cumulative, + cumulative_sf_physical_row=cursor.current_sf_cumulative, + cumulative_token_block_count=(cursor.current_token_block_cumulative), + valid_tokens_in_cta_cluster_tile=(valid_tokens_in_cta_cluster_tile), + phase_and_flags=cursor.current_phase, + ) + + +@cute.jit +def map_fc12_linear_work_id( + linear_work_id: Int32, cta_id_in_mapping_cluster: cute.Coord, mapping_state: Fc12TaskMappingState +) -> Tuple[SchedulerWorkTileBase, Fc12TaskMappingState]: + """Map one monotonically increasing scalar ID to an FC12 work tile.""" + cursor = mapping_state.cursor_state + work_tile = make_fc12_done_tile(mapping_state.is_swap_ab) + + outer_group_end = cursor.current_group_end + outer_expert_end = cursor.current_group_last_expert_exclusive + while linear_work_id >= outer_group_end and outer_expert_end < mapping_state.expert_count: + mapping_state.cursor_state = _advance_group(mapping_state) + cursor = mapping_state.cursor_state + outer_group_end = cursor.current_group_end + outer_expert_end = cursor.current_group_last_expert_exclusive + cursor = mapping_state.cursor_state + + if linear_work_id < cursor.current_group_end: + if ( + cursor.current_phase == Int32(BlockPhase.Linear1) + and linear_work_id >= cursor.current_group_fc1_subphase_end + ): + mapping_state.cursor_state = _switch_to_fc2(mapping_state) + else: + mapping_state.cursor_state = mapping_state.cursor_state + cursor = mapping_state.cursor_state + + if linear_work_id >= cursor.current_expert_tile_end: + mapping_state.cursor_state = _seek_expert_for_work_id(linear_work_id, mapping_state) + else: + mapping_state.cursor_state = mapping_state.cursor_state + cursor = mapping_state.cursor_state + work_tile = _decode_inside_expert(linear_work_id, cta_id_in_mapping_cluster, mapping_state) + else: + mapping_state.cursor_state = mapping_state.cursor_state + return work_tile, mapping_state + + +class _PhaseFc12CursorState: + """Monotonic expert cursor for one phase-local FC12 work-ID stream.""" + + def __init__( + self, + expert_idx: Int32, + expert_tile_start: Int32, + expert_tile_end: Int32, + current_expert_token_count: Int32, + current_token_block_count: Int32, + data_cumulative: Int32, + sf_cumulative: Int32, + token_block_cumulative: Int32, + blocks_per_token_block: int, + ) -> None: + self.expert_idx = expert_idx + self.expert_tile_start = expert_tile_start + self.expert_tile_end = expert_tile_end + self.current_expert_token_count = current_expert_token_count + self.current_token_block_count = current_token_block_count + self.data_cumulative = data_cumulative + self.sf_cumulative = sf_cumulative + self.token_block_cumulative = token_block_cumulative + self.blocks_per_token_block = blocks_per_token_block + + def _runtime_fields(self) -> Tuple: + return ( + self.expert_idx, + self.expert_tile_start, + self.expert_tile_end, + self.current_expert_token_count, + self.current_token_block_count, + self.data_cumulative, + self.sf_cumulative, + self.token_block_cumulative, + ) + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in self._runtime_fields(): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "_PhaseFc12CursorState": + value_index = 0 + rebuilt_fields = [] + for field in self._runtime_fields(): + field_value_count = len(extract_mlir_values(field)) + rebuilt_fields.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) + value_index += field_value_count + if value_index != len(values): + raise ValueError( + f"_PhaseFc12CursorState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return type(self)(*rebuilt_fields, blocks_per_token_block=self.blocks_per_token_block) + + +class PhaseInterleavedFc12MappingState: + """Runtime inputs and independent FC1/FC2 cursors for phase-local IDs.""" + + def __init__( + self, + expert_count: int, + mapping_cta_tile_shape_mnk: Tuple[int, int, int], + mapping_cluster_shape_mn: Tuple[int, int], + token_padding_block: int, + sf_padding_block: int, + is_swap_ab: bool, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], + fc1_cursor: _PhaseFc12CursorState, + fc2_cursor: _PhaseFc12CursorState, + num_fc1_intermediate_blocks: int, + num_fc2_hidden_blocks: int, + ) -> None: + self.expert_count = expert_count + self.mapping_cta_tile_shape_mnk = mapping_cta_tile_shape_mnk + self.mapping_cluster_shape_mn = mapping_cluster_shape_mn + self.token_padding_block = token_padding_block + self.sf_padding_block = sf_padding_block + self.is_swap_ab = is_swap_ab + self.expert_token_sizes = expert_token_sizes + self.expert_token_prefix_sum = expert_token_prefix_sum + self.fc1_cursor = fc1_cursor + self.fc2_cursor = fc2_cursor + self.num_fc1_intermediate_blocks = num_fc1_intermediate_blocks + self.num_fc2_hidden_blocks = num_fc2_hidden_blocks + + @property + def mapping_cluster_tile_m(self) -> int: + return self.mapping_cta_tile_shape_mnk[0] * self.mapping_cluster_shape_mn[0] + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + token_counts = self.expert_token_sizes if self.expert_token_sizes is not None else self.expert_token_prefix_sum + values.extend(extract_mlir_values(token_counts)) + for field in (self.fc1_cursor, self.fc2_cursor): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "PhaseInterleavedFc12MappingState": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + if self.expert_token_sizes is not None: + expert_token_sizes = rebuild(self.expert_token_sizes) + expert_token_prefix_sum = None + else: + expert_token_sizes = None + expert_token_prefix_sum = rebuild(self.expert_token_prefix_sum) + result = type(self)( + expert_count=self.expert_count, + mapping_cta_tile_shape_mnk=self.mapping_cta_tile_shape_mnk, + mapping_cluster_shape_mn=self.mapping_cluster_shape_mn, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + is_swap_ab=self.is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + fc1_cursor=rebuild(self.fc1_cursor), + fc2_cursor=rebuild(self.fc2_cursor), + num_fc1_intermediate_blocks=self.num_fc1_intermediate_blocks, + num_fc2_hidden_blocks=self.num_fc2_hidden_blocks, + ) + if value_index != len(values): + raise ValueError( + f"PhaseInterleavedFc12MappingState MLIR value count mismatch: " + f"consumed {value_index}, got {len(values)}." + ) + return result + + +def _make_phase_cursor(blocks_per_token_block: int) -> _PhaseFc12CursorState: + return _PhaseFc12CursorState( + expert_idx=Int32(-1), + expert_tile_start=Int32(0), + expert_tile_end=Int32(0), + current_expert_token_count=Int32(0), + current_token_block_count=Int32(0), + data_cumulative=Int32(0), + sf_cumulative=Int32(0), + token_block_cumulative=Int32(0), + blocks_per_token_block=blocks_per_token_block, + ) + + +@cute.jit +def create_phase_interleaved_fc12_mapping_state( + *, + expert_count: int, + intermediate_gateup_size: int, + hidden_size: int, + mapping_cta_tile_shape_mnk: Tuple[int, int, int], + mapping_cluster_shape_mn: Tuple[int, int], + token_padding_block: int, + sf_padding_block: int, + is_swap_ab: bool, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], +) -> PhaseInterleavedFc12MappingState: + """Create independent monotonic mapping cursors for the FC1 and FC2 streams.""" + mapping_cluster_tile_n = mapping_cluster_shape_mn[1] * mapping_cta_tile_shape_mnk[1] + num_fc1_intermediate_blocks = (intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + num_fc2_hidden_blocks = (hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + return PhaseInterleavedFc12MappingState( + expert_count=expert_count, + mapping_cta_tile_shape_mnk=mapping_cta_tile_shape_mnk, + mapping_cluster_shape_mn=mapping_cluster_shape_mn, + token_padding_block=token_padding_block, + sf_padding_block=sf_padding_block, + is_swap_ab=is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + fc1_cursor=_make_phase_cursor(num_fc1_intermediate_blocks), + fc2_cursor=_make_phase_cursor(num_fc2_hidden_blocks), + num_fc1_intermediate_blocks=num_fc1_intermediate_blocks, + num_fc2_hidden_blocks=num_fc2_hidden_blocks, + ) + + +@cute.jit +def _advance_phase_cursor( + cursor: _PhaseFc12CursorState, mapping_state: PhaseInterleavedFc12MappingState +) -> _PhaseFc12CursorState: + previous_token_count = cursor.current_expert_token_count + cursor.data_cumulative = cursor.data_cumulative + padded_expert_rows( + previous_token_count, Int32(mapping_state.token_padding_block) + ) + cursor.sf_cumulative = cursor.sf_cumulative + padded_expert_rows( + previous_token_count, Int32(mapping_state.sf_padding_block) + ) + cursor.token_block_cumulative = cursor.token_block_cumulative + cursor.current_token_block_count + + cursor.expert_idx = cursor.expert_idx + Int32(1) + token_count = Int32(0) + if cutlass.const_expr(mapping_state.expert_token_sizes is not None): + token_count = mapping_state.expert_token_sizes[cursor.expert_idx] + else: + prefix_end = mapping_state.expert_token_prefix_sum[cursor.expert_idx] + prefix_begin = Int32(0) + if cursor.expert_idx > Int32(0): + prefix_begin = mapping_state.expert_token_prefix_sum[cursor.expert_idx - Int32(1)] + token_count = prefix_end - prefix_begin + + cursor.current_expert_token_count = token_count + cursor.current_token_block_count = (token_count + Int32(mapping_state.mapping_cluster_tile_m - 1)) // Int32( + mapping_state.mapping_cluster_tile_m + ) + cursor.expert_tile_start = cursor.expert_tile_end + cursor.expert_tile_end = cursor.expert_tile_start + cursor.current_token_block_count * Int32( + cursor.blocks_per_token_block + ) + return cursor + + +@cute.jit +def _seek_phase_cursor( + linear_work_id: Int32, cursor: _PhaseFc12CursorState, mapping_state: PhaseInterleavedFc12MappingState +) -> _PhaseFc12CursorState: + expert_tile_end = cursor.expert_tile_end + next_expert_idx = cursor.expert_idx + Int32(1) + while linear_work_id >= expert_tile_end and next_expert_idx < Int32(mapping_state.expert_count): + cursor = _advance_phase_cursor(cursor, mapping_state) + expert_tile_end = cursor.expert_tile_end + next_expert_idx = cursor.expert_idx + Int32(1) + return cursor + + +@cute.jit +def _decode_phase_work_id( + linear_work_id: Int32, + phase: Int32, + cta_id_in_mapping_cluster: cute.Coord, + cursor: _PhaseFc12CursorState, + mapping_state: PhaseInterleavedFc12MappingState, +) -> SchedulerWorkTileBase: + local_work_id = linear_work_id - cursor.expert_tile_start + cluster_token_block_idx = local_work_id // Int32(cursor.blocks_per_token_block) + cluster_output_block_idx = local_work_id - cluster_token_block_idx * Int32(cursor.blocks_per_token_block) + cta_token_block_idx = ( + cluster_token_block_idx * Int32(mapping_state.mapping_cluster_shape_mn[0]) + cta_id_in_mapping_cluster[0] + ) + cta_output_block_idx = ( + cluster_output_block_idx * Int32(mapping_state.mapping_cluster_shape_mn[1]) + cta_id_in_mapping_cluster[1] + ) + + cta_tile_m = mapping_state.mapping_cta_tile_shape_mnk[0] + token_start = cta_token_block_idx * Int32(cta_tile_m) + remaining_tokens = cutlass.max(cursor.current_expert_token_count - token_start, Int32(0)) + valid_tokens_in_cta_tile = cutlass.min(remaining_tokens, Int32(cta_tile_m)) + + if cutlass.const_expr(mapping_state.is_swap_ab): + return SwapAbFc12WorkTileInfo( + expert_idx=cursor.expert_idx, + tile_m_idx=cta_output_block_idx, + tile_n_idx=cta_token_block_idx, + cumulative_data_physical_row=cursor.data_cumulative, + cumulative_sf_physical_row=cursor.sf_cumulative, + cumulative_token_block_count=cursor.token_block_cumulative, + valid_tokens_in_cta_tile=valid_tokens_in_cta_tile, + phase_and_flags=phase, + ) + + cluster_tile_m = mapping_state.mapping_cluster_shape_mn[0] * cta_tile_m + cluster_token_start = cluster_token_block_idx * Int32(cluster_tile_m) + remaining_cluster_tokens = cutlass.max(cursor.current_expert_token_count - cluster_token_start, Int32(0)) + valid_tokens_in_cluster_tile = cutlass.min(remaining_cluster_tokens, Int32(cluster_tile_m)) + return NonSwapAbFc12WorkTileInfo( + expert_idx=cursor.expert_idx, + tile_m_idx=cta_token_block_idx, + tile_n_idx=cta_output_block_idx, + cumulative_data_physical_row=cursor.data_cumulative, + cumulative_sf_physical_row=cursor.sf_cumulative, + cumulative_token_block_count=cursor.token_block_cumulative, + valid_tokens_in_cta_cluster_tile=(valid_tokens_in_cta_tile << Int32(16)) | valid_tokens_in_cluster_tile, + phase_and_flags=phase, + ) + + +@cute.jit +def map_phase_interleaved_fc12_work_id( + linear_work_id: Int32, + phase: Int32, + cta_id_in_mapping_cluster: cute.Coord, + mapping_state: PhaseInterleavedFc12MappingState, +) -> Tuple[SchedulerWorkTileBase, Boolean, PhaseInterleavedFc12MappingState]: + """Map one phase-local ID and report whether the selected stream contains it.""" + work_tile = make_fc12_done_tile(mapping_state.is_swap_ab) + stream_has_work = Boolean(False) + fc1_cursor = mapping_state.fc1_cursor + fc2_cursor = mapping_state.fc2_cursor + + if phase == Int32(BlockPhase.Linear1): + fc1_cursor = _seek_phase_cursor(linear_work_id, fc1_cursor, mapping_state) + if linear_work_id < fc1_cursor.expert_tile_end: + work_tile = _decode_phase_work_id( + linear_work_id, phase, cta_id_in_mapping_cluster, fc1_cursor, mapping_state + ) + stream_has_work = Boolean(True) + else: + fc2_cursor = _seek_phase_cursor(linear_work_id, fc2_cursor, mapping_state) + if linear_work_id < fc2_cursor.expert_tile_end: + work_tile = _decode_phase_work_id( + linear_work_id, phase, cta_id_in_mapping_cluster, fc2_cursor, mapping_state + ) + stream_has_work = Boolean(True) + + mapping_state.fc1_cursor = fc1_cursor + mapping_state.fc2_cursor = fc2_cursor + return work_tile, stream_has_work, mapping_state + + +__all__ = [ + "BlockPhase", + "Fc12TaskMappingState", + "Fc12WorkTileState", + "NonSwapAbFc12WorkTileInfo", + "PhaseInterleavedFc12MappingState", + "SwapAbFc12WorkTileInfo", + "create_fc12_task_mapping_state", + "create_phase_interleaved_fc12_mapping_state", + "make_fc12_done_tile", + "map_fc12_linear_work_id", + "map_phase_interleaved_fc12_work_id", + "peek_ready_bit", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_scheduler.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_scheduler.py new file mode 100644 index 000000000..a0edfbc87 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_scheduler.py @@ -0,0 +1,695 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Composable grouped and phase-interleaved FC12 schedulers.""" + +import math +from typing import Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Boolean, Int32, Integer, extract_mlir_values, new_from_mlir_values + +from ...api import ImplDesc, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.smem_workspace import SmemWorkspace +from ...helpers.utils import ceil_div +from .base import SchedulerBase, SchedulerWorkTileBase, WorkIdAcquisitionMode +from .fc12_mapping import ( + BlockPhase, + NonSwapAbFc12WorkTileInfo, + SwapAbFc12WorkTileInfo, + create_fc12_task_mapping_state, + create_phase_interleaved_fc12_mapping_state, + make_fc12_done_tile, + map_fc12_linear_work_id, + map_phase_interleaved_fc12_work_id, +) +from .non_clc_mixed_cga import NonClcMixedCgaConfig, NonClcMixedCgaSchedulerWorker + + +def _make_non_clc_mixed_cga_config(impl_desc: ImplDesc) -> NonClcMixedCgaConfig: + return NonClcMixedCgaConfig( + preferred_cluster_shape=impl_desc["cluster_shape_mn"], + fallback_cluster_shape=impl_desc.get("fallback_cluster_shape_mn"), + launch_cluster_count=impl_desc.get("launch_cluster_count"), + preferred_cluster_count=impl_desc.get("preferred_cluster_count"), + fallback_cluster_count=impl_desc.get("fallback_cluster_count"), + ) + + +def _mixed_cga_impl_requirements() -> dict: + return { + "fallback_cluster_shape_mn": OptionalRequirement(Optional[tuple]), + "launch_cluster_count": OptionalRequirement(Optional[int]), + "preferred_cluster_count": OptionalRequirement(Optional[int]), + "fallback_cluster_count": OptionalRequirement(Optional[int]), + } + + +def minimum_phase_interleave_hint( + *, blocks_fc1: int, blocks_fc2: int, launch_cluster_cnt_merge_as_preferred: int +) -> int: + """Return the per-cluster FC1 prologue covering one canonical FC2 claim wave.""" + effective_wave_width = launch_cluster_cnt_merge_as_preferred + max_dependent_token_blocks = ceil_div(effective_wave_width + blocks_fc2 - 1, blocks_fc2) + required_fc1_work = max_dependent_token_blocks * blocks_fc1 + return max(1, ceil_div(required_fc1_work, launch_cluster_cnt_merge_as_preferred)) + + +def _to_fc12_mapping_cta_coord(cta_coord_in_preferred_cluster: cute.Coord, is_swap_ab: bool) -> cute.Coord: + if cutlass.const_expr(not is_swap_ab): + return cta_coord_in_preferred_cluster + return (cta_coord_in_preferred_cluster[1], cta_coord_in_preferred_cluster[0], cta_coord_in_preferred_cluster[2]) + + +class BlackwellFusedFc12Scheduler(SchedulerBase): + """Compose work-ID claim, FC12 mapping, and SMEM tile transport.""" + + pipeline_mbarriers_region = "blackwell.fc12.scheduler.pipeline_mbarriers" + work_tiles_region = "blackwell.fc12.scheduler.work_tiles" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "expert_count": StaticOrRuntimeIntegerType, + "intermediate_gateup_size": StaticOrRuntimeIntegerType, + "hidden_size": StaticOrRuntimeIntegerType, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + **super().impl_desc_require(), + "mma_tiler_mnk": tuple, + "cluster_shape_mn": tuple, + "use_2cta_instrs": bool, + "hint": Optional[int], + "token_padding_block": int, + "sf_padding_block": int, + "work_id_mode": str, + "is_swap_ab": bool, + **_mixed_cga_impl_requirements(), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + super().__init__(problem_desc, impl_desc) + + self.expert_count = problem_desc["expert_count"] + self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] + self.hidden_size = problem_desc["hidden_size"] + self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + self.cluster_shape_mn = impl_desc["cluster_shape_mn"] + self.use_2cta_instrs = impl_desc["use_2cta_instrs"] + self.hint = impl_desc["hint"] + self.group_hint = self.hint + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.work_id_mode: WorkIdAcquisitionMode = impl_desc["work_id_mode"] + self.is_swap_ab = impl_desc["is_swap_ab"] + self.non_clc_mixed_cga_config = _make_non_clc_mixed_cga_config(impl_desc) + self.launch_cluster_cnt_merge_as_preferred = self.non_clc_mixed_cga_config.launch_cluster_cnt_merge_as_preferred + self.work_tile_type = SwapAbFc12WorkTileInfo if self.is_swap_ab else NonSwapAbFc12WorkTileInfo + if self.hint is None: + self.hint = self.launch_cluster_cnt_merge_as_preferred + self.group_hint = self.hint + + self._validate_configuration() + mma_cta_count = 2 if self.use_2cta_instrs else 1 + launch_cta_tile_shape_mnk = ( + self.mma_tiler_mnk[0] // mma_cta_count, + self.mma_tiler_mnk[1], + self.mma_tiler_mnk[2], + ) + if self.is_swap_ab: + self.mapping_cta_tile_shape_mnk = ( + launch_cta_tile_shape_mnk[1], + launch_cta_tile_shape_mnk[0], + launch_cta_tile_shape_mnk[2], + ) + self.mapping_cluster_shape_mn = (self.cluster_shape_mn[1], self.cluster_shape_mn[0]) + else: + self.mapping_cta_tile_shape_mnk = launch_cta_tile_shape_mnk + self.mapping_cluster_shape_mn = self.cluster_shape_mn + + if self.work_id_mode == "cluster_launch_control": + raise NotImplementedError("cluster_launch_control is not implemented for FC12.") + self._work_id_worker = NonClcMixedCgaSchedulerWorker( + config=self.non_clc_mixed_cga_config, work_id_mode=self.work_id_mode, stream_count=1 + ) + + def _validate_configuration(self) -> None: + if len(self.mma_tiler_mnk) != 3: + raise ValueError("mma_tiler_mnk must contain three dimensions.") + if len(self.cluster_shape_mn) != 2: + raise ValueError("cluster_shape_mn must contain two dimensions.") + if any(dimension <= 0 for dimension in self.mma_tiler_mnk): + raise ValueError("mma_tiler_mnk dimensions must be positive.") + if any(dimension <= 0 for dimension in self.cluster_shape_mn): + raise ValueError("cluster_shape_mn dimensions must be positive.") + mma_cta_count = 2 if self.use_2cta_instrs else 1 + if self.mma_tiler_mnk[0] % mma_cta_count != 0: + raise ValueError("mma_tiler M must be divisible by the MMA CTA count.") + if self.group_hint is not None and self.group_hint <= 0: + raise ValueError("group_hint must be positive.") + if self.token_padding_block <= 0: + raise ValueError("token_padding_block must be positive.") + if self.sf_padding_block <= 0: + raise ValueError("sf_padding_block must be positive.") + if self.launch_cluster_cnt_merge_as_preferred <= 0: + raise ValueError("launch_cluster_cnt_merge_as_preferred must be positive.") + if self.work_id_mode not in ("grid_stride", "atomic_counter", "cluster_launch_control"): + raise ValueError( + "work_id_mode must be 'grid_stride', 'atomic_counter', or " + f"'cluster_launch_control', got {self.work_id_mode!r}." + ) + cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] + if self.work_id_mode == "atomic_counter" and cluster_size > 32: + raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per cluster.") + fallback_cluster_shape = self.non_clc_mixed_cga_config.fallback_cluster_shape + if fallback_cluster_shape is not None: + fallback_cluster_size = fallback_cluster_shape[0] * fallback_cluster_shape[1] + if self.work_id_mode == "atomic_counter" and fallback_cluster_size > 32: + raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per fallback cluster.") + for field_name in ("expert_count", "intermediate_gateup_size", "hidden_size"): + value = getattr(self, field_name) + if isinstance(value, int) and value <= 0: + raise ValueError(f"{field_name} must be positive.") + static_dimensions = ( + isinstance(self.expert_count, int), + isinstance(self.intermediate_gateup_size, int), + isinstance(self.hidden_size, int), + ) + if any(static_dimensions) and not all(static_dimensions): + raise ValueError("FC12 expert dimensions must be either all static or all runtime.") + + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: + """Register scheduler-owned SMEM transport and claim regions.""" + super().register_smem_regions(smem_workspace) + self._work_id_worker.register_smem_regions(smem_workspace) + + def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: + """Register work-ID counters and optional fixed-group fallback state.""" + self._work_id_worker.register_device_workspace(device_workspace) + + @cute.jit + def initialize_fallback_group(self) -> None: + """Register this physical fallback cluster with its fixed logical group.""" + self._work_id_worker.initialize_fallback_group() + + def get_grid_shape(self, *, max_active_clusters: Optional[int] = None, problem_desc=None) -> Tuple[int, int, int]: + """Return the persistent launch grid in GEMM-domain orientation.""" + if ( + not self.non_clc_mixed_cga_config.is_mixed + and max_active_clusters is not None + and max_active_clusters < self.launch_cluster_cnt_merge_as_preferred + ): + raise ValueError( + f"max_active_clusters ({max_active_clusters}) must be at least " + "launch_cluster_cnt_merge_as_preferred " + f"({self.launch_cluster_cnt_merge_as_preferred})." + ) + return (self.cluster_shape_mn[0], self.cluster_shape_mn[1], self.launch_cluster_cnt_merge_as_preferred) + + @cute.jit + def assign_device_members( + self, + *, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], + actual_expert_shape: Optional[Tuple], + block_idx: Tuple[Integer, Integer, Integer], + smem_workspace: SmemWorkspace, + smem_base: cute.Pointer, + device_workspace: DeviceWorkspace, + is_fallback_cluster: Optional[Boolean] = None, + ) -> None: + """Initialize all FC12 scheduler state rooted for one CTA lifetime.""" + if cutlass.const_expr((expert_token_sizes is None) == (expert_token_prefix_sum is None)): + raise ValueError("Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided.") + needs_actual_shape = not all( + isinstance(dimension, int) + for dimension in (self.expert_count, self.intermediate_gateup_size, self.hidden_size) + ) + if cutlass.const_expr(needs_actual_shape and actual_expert_shape is None): + raise ValueError("actual_expert_shape is required for runtime dimensions.") + if cutlass.const_expr(isinstance(self.expert_count, int)): + expert_count = self.expert_count + else: + expert_count = actual_expert_shape[0] + if cutlass.const_expr(isinstance(self.intermediate_gateup_size, int)): + intermediate_gateup_size = self.intermediate_gateup_size + else: + intermediate_gateup_size = actual_expert_shape[1] + if cutlass.const_expr(isinstance(self.hidden_size, int)): + hidden_size = self.hidden_size + else: + hidden_size = actual_expert_shape[2] + + self.create_scheduler_pipelines(smem_workspace, smem_base) + self._work_id_worker.assign_device_members( + is_fallback_cluster=is_fallback_cluster, + block_idx=block_idx, + smem_workspace=smem_workspace, + smem_base=smem_base, + device_workspace=device_workspace, + ) + + task_mapping_state = create_fc12_task_mapping_state( + expert_count=expert_count, + intermediate_gateup_size=intermediate_gateup_size, + hidden_size=hidden_size, + mapping_cta_tile_shape_mnk=(self.mapping_cta_tile_shape_mnk), + mapping_cluster_shape_mn=self.mapping_cluster_shape_mn, + group_hint=self.group_hint, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + is_swap_ab=self.is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + ) + + self._task_mapping_state = task_mapping_state + + @cute.jit + def gen_next_work(self) -> SchedulerWorkTileBase: + """Claim and map one work tile without first-tile prefetch.""" + work_id = self._work_id_worker.claim_next_work() + cta_id_in_mapping_cluster = _to_fc12_mapping_cta_coord( + self._work_id_worker.cta_coord_in_preferred_cluster, self.is_swap_ab + ) + work_tile, self._task_mapping_state = map_fc12_linear_work_id( + work_id, cta_id_in_mapping_cluster, self._task_mapping_state + ) + return work_tile + + def __extract_mlir_values__(self) -> list: + values = super().__extract_mlir_values__() + values.extend(extract_mlir_values(self._work_id_worker)) + values.extend(extract_mlir_values(self._task_mapping_state)) + return values + + def __new_from_mlir_values__(self, values: list) -> "BlackwellFusedFc12Scheduler": + base_value_count = len(super().__extract_mlir_values__()) + if len(values) < base_value_count: + raise ValueError( + "BlackwellFusedFc12Scheduler MLIR value count is smaller than " + f"its base state: expected at least {base_value_count}, got {len(values)}." + ) + result = super().__new_from_mlir_values__(values[:base_value_count]) + value_index = base_value_count + + def rebuild(state): + nonlocal value_index + state_value_count = len(extract_mlir_values(state)) + rebuilt_state = new_from_mlir_values(state, values[value_index : value_index + state_value_count]) + value_index += state_value_count + return rebuilt_state + + result._work_id_worker = rebuild(self._work_id_worker) + result._task_mapping_state = rebuild(self._task_mapping_state) + if value_index != len(values): + raise ValueError( + f"BlackwellFusedFc12Scheduler MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + + for field_name in ( + "expert_count", + "intermediate_gateup_size", + "hidden_size", + "mma_tiler_mnk", + "cluster_shape_mn", + "use_2cta_instrs", + "hint", + "group_hint", + "token_padding_block", + "sf_padding_block", + "work_id_mode", + "is_swap_ab", + "launch_cluster_cnt_merge_as_preferred", + "non_clc_mixed_cga_config", + "mapping_cta_tile_shape_mnk", + "mapping_cluster_shape_mn", + ): + setattr(result, field_name, getattr(self, field_name)) + return result + + +class _PhaseInterleaveControlState: + """Per-cluster phase cadence and stream exhaustion state.""" + + def __init__( + self, prologue_remaining: Int32, cycle_position: Int32, fc1_exhausted: Boolean, fc2_exhausted: Boolean + ) -> None: + self.prologue_remaining = prologue_remaining + self.cycle_position = cycle_position + self.fc1_exhausted = fc1_exhausted + self.fc2_exhausted = fc2_exhausted + + def _fields(self) -> Tuple: + return (self.prologue_remaining, self.cycle_position, self.fc1_exhausted, self.fc2_exhausted) + + def __extract_mlir_values__(self) -> list: + values = [] + for field in self._fields(): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: list) -> "_PhaseInterleaveControlState": + value_index = 0 + rebuilt_fields = [] + for field in self._fields(): + field_value_count = len(extract_mlir_values(field)) + rebuilt_fields.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) + value_index += field_value_count + if value_index != len(values): + raise ValueError( + f"_PhaseInterleaveControlState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return type(self)(*rebuilt_fields) + + +class PhaseInterleavedFc12Scheduler(SchedulerBase): + """Schedule independent FC1 and FC2 streams with per-phase atomic counters.""" + + pipeline_mbarriers_region = "fc12.phase_interleaved.scheduler.pipeline_mbarriers" + work_tiles_region = "fc12.phase_interleaved.scheduler.work_tiles" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return {"expert_count": int, "intermediate_gateup_size": int, "hidden_size": int} + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + **super().impl_desc_require(), + "mma_tiler_mnk": tuple, + "cluster_shape_mn": tuple, + "use_2cta_instrs": bool, + "hint": int, + "token_padding_block": int, + "sf_padding_block": int, + "work_id_mode": str, + "is_swap_ab": bool, + **_mixed_cga_impl_requirements(), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + super().__init__(problem_desc, impl_desc) + + self.expert_count = problem_desc["expert_count"] + self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] + self.hidden_size = problem_desc["hidden_size"] + self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + self.cluster_shape_mn = impl_desc["cluster_shape_mn"] + self.use_2cta_instrs = impl_desc["use_2cta_instrs"] + self.hint = impl_desc["hint"] + self.fc1_prologue_tiles = self.hint + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.work_id_mode: WorkIdAcquisitionMode = impl_desc["work_id_mode"] + self.is_swap_ab = impl_desc["is_swap_ab"] + self.non_clc_mixed_cga_config = _make_non_clc_mixed_cga_config(impl_desc) + self.launch_cluster_cnt_merge_as_preferred = self.non_clc_mixed_cga_config.launch_cluster_cnt_merge_as_preferred + self.work_tile_type = SwapAbFc12WorkTileInfo if self.is_swap_ab else NonSwapAbFc12WorkTileInfo + + self._validate_configuration() + mma_cta_count = 2 if self.use_2cta_instrs else 1 + launch_cta_tile_shape_mnk = ( + self.mma_tiler_mnk[0] // mma_cta_count, + self.mma_tiler_mnk[1], + self.mma_tiler_mnk[2], + ) + if self.is_swap_ab: + self.mapping_cta_tile_shape_mnk = ( + launch_cta_tile_shape_mnk[1], + launch_cta_tile_shape_mnk[0], + launch_cta_tile_shape_mnk[2], + ) + self.mapping_cluster_shape_mn = (self.cluster_shape_mn[1], self.cluster_shape_mn[0]) + else: + self.mapping_cta_tile_shape_mnk = launch_cta_tile_shape_mnk + self.mapping_cluster_shape_mn = self.cluster_shape_mn + self._work_id_worker = NonClcMixedCgaSchedulerWorker( + config=self.non_clc_mixed_cga_config, work_id_mode=self.work_id_mode, stream_count=2 + ) + + mapping_cluster_tile_n = self.mapping_cta_tile_shape_mnk[1] * self.mapping_cluster_shape_mn[1] + self.blocks_fc1 = (self.intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + self.blocks_fc2 = (self.hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + interleave_gcd = math.gcd(self.blocks_fc1, self.blocks_fc2) + self.interleave_fc2_slots = self.blocks_fc2 // interleave_gcd + self.interleave_cycle_length = (self.blocks_fc1 + self.blocks_fc2) // interleave_gcd + minimum_hint = minimum_phase_interleave_hint( + blocks_fc1=self.blocks_fc1, + blocks_fc2=self.blocks_fc2, + launch_cluster_cnt_merge_as_preferred=self.launch_cluster_cnt_merge_as_preferred, + ) + if self.fc1_prologue_tiles < minimum_hint: + raise ValueError( + f"phase_interleave hint {self.fc1_prologue_tiles} cannot cover a " + f"{self.launch_cluster_cnt_merge_as_preferred}-cluster FC2 claim wave; " + f"raise hint to at least {minimum_hint}." + ) + + def _validate_configuration(self) -> None: + if len(self.mma_tiler_mnk) != 3: + raise ValueError("mma_tiler_mnk must contain three dimensions.") + if len(self.cluster_shape_mn) != 2: + raise ValueError("cluster_shape_mn must contain two dimensions.") + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in self.mma_tiler_mnk): + raise TypeError("mma_tiler_mnk dimensions must be Python ints.") + if not all( + isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in self.cluster_shape_mn + ): + raise TypeError("cluster_shape_mn dimensions must be Python ints.") + if any(dimension <= 0 for dimension in self.mma_tiler_mnk): + raise ValueError("mma_tiler_mnk dimensions must be positive.") + if any(dimension <= 0 for dimension in self.cluster_shape_mn): + raise ValueError("cluster_shape_mn dimensions must be positive.") + mma_cta_count = 2 if self.use_2cta_instrs else 1 + if self.mma_tiler_mnk[0] % mma_cta_count != 0: + raise ValueError("mma_tiler M must be divisible by the MMA CTA count.") + if ( + isinstance(self.fc1_prologue_tiles, bool) + or not isinstance(self.fc1_prologue_tiles, int) + or self.fc1_prologue_tiles <= 0 + ): + raise ValueError("fc1_prologue_tiles must be a positive Python int resolved by the kernel frontend.") + if self.token_padding_block <= 0: + raise ValueError("token_padding_block must be positive.") + if self.sf_padding_block <= 0: + raise ValueError("sf_padding_block must be positive.") + if self.launch_cluster_cnt_merge_as_preferred <= 0: + raise ValueError("launch_cluster_cnt_merge_as_preferred must be positive.") + if self.work_id_mode != "atomic_counter": + raise ValueError("Phase-interleaved FC12 scheduling currently requires work_id_mode='atomic_counter'.") + cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] + if cluster_size > 32: + raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per cluster.") + fallback_cluster_shape = self.non_clc_mixed_cga_config.fallback_cluster_shape + if fallback_cluster_shape is not None: + fallback_cluster_size = fallback_cluster_shape[0] * fallback_cluster_shape[1] + if fallback_cluster_size > 32: + raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per fallback cluster.") + maximum_int32 = (1 << 31) - 1 + for field_name in ("expert_count", "intermediate_gateup_size", "hidden_size"): + value = getattr(self, field_name) + if isinstance(value, bool) or value <= 0: + raise ValueError(f"{field_name} must be a positive Python int.") + if value > maximum_int32: + raise ValueError(f"{field_name} must fit in a signed Int32, got {value}.") + + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: + """Register work transport and the phase-counter broadcast channel.""" + super().register_smem_regions(smem_workspace) + self._work_id_worker.register_smem_regions(smem_workspace) + + def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: + """Register independently reset FC1 and FC2 work-ID streams.""" + self._work_id_worker.register_device_workspace(device_workspace) + + @cute.jit + def initialize_fallback_group(self) -> None: + """Register this physical fallback cluster with its fixed logical group.""" + self._work_id_worker.initialize_fallback_group() + + def get_grid_shape(self, *, max_active_clusters: Optional[int] = None, problem_desc=None) -> Tuple[int, int, int]: + """Return the statically configured persistent launch grid.""" + if ( + not self.non_clc_mixed_cga_config.is_mixed + and max_active_clusters is not None + and max_active_clusters < self.launch_cluster_cnt_merge_as_preferred + ): + raise ValueError( + f"max_active_clusters ({max_active_clusters}) must be at least " + "launch_cluster_cnt_merge_as_preferred " + f"({self.launch_cluster_cnt_merge_as_preferred})." + ) + return (self.cluster_shape_mn[0], self.cluster_shape_mn[1], self.launch_cluster_cnt_merge_as_preferred) + + @cute.jit + def assign_device_members( + self, + *, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], + actual_expert_shape: Optional[Tuple], + block_idx: Tuple[Integer, Integer, Integer], + smem_workspace: SmemWorkspace, + smem_base: cute.Pointer, + device_workspace: DeviceWorkspace, + is_fallback_cluster: Optional[Boolean] = None, + ) -> None: + """Initialize phase-local mapping, cadence, and counter state.""" + if cutlass.const_expr((expert_token_sizes is None) == (expert_token_prefix_sum is None)): + raise ValueError("Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided.") + self.create_scheduler_pipelines(smem_workspace, smem_base) + self._work_id_worker.assign_device_members( + is_fallback_cluster=is_fallback_cluster, + block_idx=block_idx, + smem_workspace=smem_workspace, + smem_base=smem_base, + device_workspace=device_workspace, + ) + self._task_mapping_state = create_phase_interleaved_fc12_mapping_state( + expert_count=self.expert_count, + intermediate_gateup_size=self.intermediate_gateup_size, + hidden_size=self.hidden_size, + mapping_cta_tile_shape_mnk=self.mapping_cta_tile_shape_mnk, + mapping_cluster_shape_mn=self.mapping_cluster_shape_mn, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + is_swap_ab=self.is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + ) + self._control_state = _PhaseInterleaveControlState( + prologue_remaining=Int32(self.fc1_prologue_tiles), + cycle_position=Int32(0), + fc1_exhausted=Boolean(False), + fc2_exhausted=Boolean(False), + ) + + @cute.jit + def gen_next_work(self) -> SchedulerWorkTileBase: + """Claim until one stream yields valid work or both streams terminate.""" + work_tile = make_fc12_done_tile(self.is_swap_ab) + work_id_worker = self._work_id_worker + task_mapping_state = self._task_mapping_state + control_state = self._control_state + prologue_remaining = control_state.prologue_remaining + cycle_position = control_state.cycle_position + fc1_exhausted = control_state.fc1_exhausted + fc2_exhausted = control_state.fc2_exhausted + resolved = Boolean(False) + + while not resolved: + if fc1_exhausted and fc2_exhausted: + work_tile = make_fc12_done_tile(self.is_swap_ab) + resolved = Boolean(True) + else: + want_fc1 = Boolean(True) + if prologue_remaining <= Int32(0): + is_fc2_slot = (cycle_position * Int32(self.interleave_fc2_slots)) % Int32( + self.interleave_cycle_length + ) < Int32(self.interleave_fc2_slots) + want_fc1 = not is_fc2_slot + if want_fc1 and fc1_exhausted: + want_fc1 = Boolean(False) + if (not want_fc1) and fc2_exhausted: + want_fc1 = Boolean(True) + + atomic_counter_index = Int32(1) + if want_fc1: + atomic_counter_index = Int32(0) + linear_work_id = work_id_worker.claim_next_work(atomic_counter_index) + want_fc1 = work_id_worker.claimed_stream_index == Int32(0) + phase = Int32(BlockPhase.Linear2) + if want_fc1: + phase = Int32(BlockPhase.Linear1) + cta_id_in_mapping_cluster = _to_fc12_mapping_cta_coord( + work_id_worker.cta_coord_in_preferred_cluster, self.is_swap_ab + ) + work_tile, stream_has_work, task_mapping_state = map_phase_interleaved_fc12_work_id( + linear_work_id, phase, cta_id_in_mapping_cluster, task_mapping_state + ) + if stream_has_work: + if prologue_remaining > Int32(0): + prologue_remaining = prologue_remaining - Int32(1) + else: + cycle_position = (cycle_position + Int32(1)) % Int32(self.interleave_cycle_length) + resolved = Boolean(True) + else: + if want_fc1: + fc1_exhausted = Boolean(True) + else: + fc2_exhausted = Boolean(True) + + control_state.prologue_remaining = prologue_remaining + control_state.cycle_position = cycle_position + control_state.fc1_exhausted = fc1_exhausted + control_state.fc2_exhausted = fc2_exhausted + self._work_id_worker = work_id_worker + self._task_mapping_state = task_mapping_state + self._control_state = control_state + return work_tile + + def __extract_mlir_values__(self) -> list: + values = super().__extract_mlir_values__() + for state in (self._work_id_worker, self._task_mapping_state, self._control_state): + values.extend(extract_mlir_values(state)) + return values + + def __new_from_mlir_values__(self, values: list) -> "PhaseInterleavedFc12Scheduler": + base_value_count = len(super().__extract_mlir_values__()) + if len(values) < base_value_count: + raise ValueError( + "PhaseInterleavedFc12Scheduler MLIR value count is smaller than " + f"its base state: expected at least {base_value_count}, got {len(values)}." + ) + result = super().__new_from_mlir_values__(values[:base_value_count]) + value_index = base_value_count + + def rebuild(state): + nonlocal value_index + state_value_count = len(extract_mlir_values(state)) + rebuilt_state = new_from_mlir_values(state, values[value_index : value_index + state_value_count]) + value_index += state_value_count + return rebuilt_state + + result._work_id_worker = rebuild(self._work_id_worker) + result._task_mapping_state = rebuild(self._task_mapping_state) + result._control_state = rebuild(self._control_state) + if value_index != len(values): + raise ValueError( + f"PhaseInterleavedFc12Scheduler MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + + for field_name in ( + "expert_count", + "intermediate_gateup_size", + "hidden_size", + "mma_tiler_mnk", + "cluster_shape_mn", + "use_2cta_instrs", + "hint", + "fc1_prologue_tiles", + "token_padding_block", + "sf_padding_block", + "work_id_mode", + "is_swap_ab", + "launch_cluster_cnt_merge_as_preferred", + "non_clc_mixed_cga_config", + "mapping_cta_tile_shape_mnk", + "mapping_cluster_shape_mn", + "blocks_fc1", + "blocks_fc2", + "interleave_fc2_slots", + "interleave_cycle_length", + ): + setattr(result, field_name, getattr(self, field_name)) + return result + + +__all__ = ["BlackwellFusedFc12Scheduler", "PhaseInterleavedFc12Scheduler", "minimum_phase_interleave_hint"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/non_clc_mixed_cga.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/non_clc_mixed_cga.py new file mode 100644 index 000000000..c1df89496 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/non_clc_mixed_cga.py @@ -0,0 +1,342 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Reusable preferred/fallback cluster scheduling without hardware CLC.""" + +import dataclasses +from typing import List, Optional, Tuple + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass._mlir import ir +from cutlass.cutlass_dsl import Boolean, Int32, Int64, extract_mlir_values, new_from_mlir_values + +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.smem_workspace import SmemWorkspace +from .base import WorkIdAcquisitionMode +from .work_id_claim import ( + AtomicCounterWorkIdState, + FixedGroupMixedCgaAtomicCounterWorkIdState, + GridStrideWorkIdState, + claim_work_id, + initialize_fixed_group_mixed_cga_work_id_state, +) + + +@dataclasses.dataclass(frozen=True) +class NonClcMixedCgaConfig: + """Static launch geometry shared by non-CLC mixed-CGA schedulers.""" + + preferred_cluster_shape: Tuple[int, int] + fallback_cluster_shape: Optional[Tuple[int, int]] + launch_cluster_count: Optional[int] + preferred_cluster_count: Optional[int] + fallback_cluster_count: Optional[int] + mn_split_factors: Tuple[int, int] = dataclasses.field(init=False) + split_factor: int = dataclasses.field(init=False) + launch_cluster_cnt_merge_as_preferred: int = dataclasses.field(init=False) + total_cta_cnt: int = dataclasses.field(init=False) + is_mixed: bool = dataclasses.field(init=False) + + def __post_init__(self) -> None: + self._validate_shape(self.preferred_cluster_shape, "preferred_cluster_shape") + preferred_cluster_size = self._shape_size(self.preferred_cluster_shape) + if self.fallback_cluster_shape is not None: + self._validate_shape(self.fallback_cluster_shape, "fallback_cluster_shape") + has_no_fallback_clusters = ( + isinstance(self.fallback_cluster_count, int) + and not isinstance(self.fallback_cluster_count, bool) + and self.fallback_cluster_count == 0 + ) + if self.fallback_cluster_shape == self.preferred_cluster_shape or has_no_fallback_clusters: + object.__setattr__(self, "fallback_cluster_shape", None) + object.__setattr__(self, "preferred_cluster_count", None) + object.__setattr__(self, "fallback_cluster_count", None) + + if self.fallback_cluster_shape is None: + if not self._is_positive_int(self.launch_cluster_count): + raise ValueError("launch_cluster_count must be positive when fallback_cluster_shape is absent.") + mn_split_factors = (1, 1) + split_factor = 1 + launch_cluster_cnt_merge_as_preferred = self.launch_cluster_count + is_mixed = False + else: + if len(self.fallback_cluster_shape) != len(self.preferred_cluster_shape): + raise ValueError("preferred and fallback cluster shapes must have equal rank.") + if not self._is_positive_int(self.preferred_cluster_count): + raise ValueError("preferred_cluster_count must be positive when fallback_cluster_shape is present.") + if ( + isinstance(self.fallback_cluster_count, bool) + or not isinstance(self.fallback_cluster_count, int) + or self.fallback_cluster_count < 0 + ): + raise ValueError( + "fallback_cluster_count must be a non-negative Python int when fallback_cluster_shape is present." + ) + + mn_split_factors = ( + self.preferred_cluster_shape[0] // self.fallback_cluster_shape[0], + self.preferred_cluster_shape[1] // self.fallback_cluster_shape[1], + ) + if any( + preferred_dimension % fallback_dimension != 0 + for preferred_dimension, fallback_dimension in zip( + self.preferred_cluster_shape, self.fallback_cluster_shape + ) + ): + raise ValueError("Every preferred cluster dimension must be divisible by its fallback dimension.") + split_factor = self._shape_size(mn_split_factors) + is_mixed = self.fallback_cluster_shape != self.preferred_cluster_shape + if is_mixed: + if split_factor > 16 or split_factor & (split_factor - 1): + raise ValueError("The preferred/fallback cluster split factor must be a power of two at most 16.") + if self.fallback_cluster_count % split_factor != 0: + raise ValueError( + "fallback_cluster_count must be divisible by the preferred/fallback cluster split factor." + ) + launch_cluster_cnt_merge_as_preferred = ( + self.preferred_cluster_count + self.fallback_cluster_count // split_factor + ) + else: + launch_cluster_cnt_merge_as_preferred = self.preferred_cluster_count + self.fallback_cluster_count + + if launch_cluster_cnt_merge_as_preferred <= 0: + raise ValueError("The resolved launch must contain at least one cluster merged as preferred.") + object.__setattr__(self, "mn_split_factors", mn_split_factors) + object.__setattr__(self, "split_factor", split_factor) + object.__setattr__(self, "launch_cluster_cnt_merge_as_preferred", launch_cluster_cnt_merge_as_preferred) + object.__setattr__(self, "total_cta_cnt", launch_cluster_cnt_merge_as_preferred * preferred_cluster_size) + object.__setattr__(self, "is_mixed", is_mixed) + + @staticmethod + def _shape_size(shape: Tuple[int, int]) -> int: + result = 1 + for dimension in shape: + result *= dimension + return result + + @staticmethod + def _is_positive_int(value) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + @classmethod + def _validate_shape(cls, shape: Tuple[int, int], field_name: str) -> None: + if not isinstance(shape, tuple) or len(shape) != 2: + raise TypeError(f"{field_name} must be a two-dimensional tuple.") + if not all(cls._is_positive_int(dimension) for dimension in shape): + raise ValueError(f"{field_name} dimensions must be positive Python ints.") + + +class NonClcMixedCgaSchedulerWorker: + """Compose canonical work-ID acquisition with preferred/fallback cluster splitting.""" + + cluster_pipeline_mbarriers_region = "non_clc_mixed_cga.cluster_pipeline_mbarriers" + cluster_broadcast_region = "non_clc_mixed_cga.cluster_broadcast" + work_id_counter_region = "non_clc_mixed_cga.work_id_counters" + fallback_registration_counter_region = "non_clc_mixed_cga.fallback_registration_counter" + fallback_group_token_region = "non_clc_mixed_cga.fallback_group_tokens" + + def __init__(self, *, config: NonClcMixedCgaConfig, work_id_mode: WorkIdAcquisitionMode, stream_count: int) -> None: + if work_id_mode not in ("grid_stride", "atomic_counter"): + raise ValueError("Non-CLC mixed-CGA scheduling supports grid_stride or atomic_counter work IDs.") + if isinstance(stream_count, bool) or not isinstance(stream_count, int) or stream_count <= 0: + raise ValueError("stream_count must be a positive Python int.") + self.config = config + self.work_id_mode = work_id_mode + self.stream_count = stream_count + + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: + """Register the cluster-local claim broadcast channel.""" + if self.work_id_mode == "atomic_counter": + smem_workspace.register_mbarrier(self.cluster_pipeline_mbarriers_region, 2) + smem_workspace.register_tensor(self.cluster_broadcast_region, cutlass.Int32, (1,)) + + def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: + """Register atomic counters and optional fixed-group fallback handoff state.""" + if self.work_id_mode == "atomic_counter": + device_workspace.register( + self.work_id_counter_region, + cutlass.Int32, + (self.stream_count,), + buffer_space="local", + reset="tail_reset", + ) + if self.config.is_mixed: + device_workspace.register( + self.fallback_registration_counter_region, + cutlass.Int32, + (1,), + buffer_space="local", + reset="tail_reset", + ) + device_workspace.register( + self.fallback_group_token_region, + cutlass.Int64, + (self.config.fallback_cluster_count,), + buffer_space="local", + reset="tail_reset", + ) + + @cute.jit + def assign_device_members( + self, + *, + is_fallback_cluster: Optional[Boolean], + block_idx: Tuple, + smem_workspace: SmemWorkspace, + smem_base: cute.Pointer, + device_workspace: DeviceWorkspace, + ) -> None: + """Bind one active cluster to the configured non-CLC claim backend.""" + if cutlass.const_expr(self.config.is_mixed and is_fallback_cluster is None): + raise ValueError("is_fallback_cluster is required for a true mixed-CGA launch.") + + active_cluster_m = self.config.preferred_cluster_shape[0] + active_cluster_n = self.config.preferred_cluster_shape[1] + if cutlass.const_expr(self.config.is_mixed): + if is_fallback_cluster: + active_cluster_m = Int32(self.config.fallback_cluster_shape[0]) + active_cluster_n = Int32(self.config.fallback_cluster_shape[1]) + + cta_coord_in_active_cluster = ( + Int32(block_idx[0]) % Int32(active_cluster_m), + Int32(block_idx[1]) % Int32(active_cluster_n), + Int32(0), + ) + cta_coord_in_preferred_cluster = cta_coord_in_active_cluster + + if cutlass.const_expr(self.config.is_mixed and self.work_id_mode == "grid_stride"): + if is_fallback_cluster: + flattened_index = Int32(0) + dimension_stride = 1 + for dimension_idx in cutlass.range_constexpr(len(self.config.mn_split_factors)): + preferred_dimension = self.config.preferred_cluster_shape[dimension_idx] + fallback_dimension = self.config.fallback_cluster_shape[dimension_idx] + inner_coordinate = (Int32(block_idx[dimension_idx]) % Int32(preferred_dimension)) // Int32( + fallback_dimension + ) + flattened_index = flattened_index + inner_coordinate * Int32(dimension_stride) + dimension_stride *= self.config.mn_split_factors[dimension_idx] + cta_coord_in_preferred_cluster = self._preferred_cluster_cta_coord( + cta_coord_in_active_cluster, flattened_index + ) + self.cta_coord_in_preferred_cluster = cta_coord_in_preferred_cluster + + if cutlass.const_expr(self.work_id_mode == "atomic_counter"): + active_cluster_size = active_cluster_m * active_cluster_n + cluster_pipeline = pipeline.PipelineAsync.create( + num_stages=1, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * active_cluster_size), + barrier_storage=smem_workspace.ptr(self.cluster_pipeline_mbarriers_region, smem_base), + defer_sync=True, + ) + atomic_counter_state = AtomicCounterWorkIdState( + counter_pointer=device_workspace.ptr(self.work_id_counter_region), + counter_count=self.stream_count, + broadcast_pointer=smem_workspace.ptr(self.cluster_broadcast_region, smem_base), + is_leader_cta=(cta_coord_in_active_cluster[0] + cta_coord_in_active_cluster[1]) == Int32(0), + cluster_pipeline=cluster_pipeline, + producer_state=pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1), + consumer_state=pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1), + cluster_size=active_cluster_size, + ) + if cutlass.const_expr(self.config.is_mixed): + self._cta_coord_in_active_cluster = cta_coord_in_active_cluster + self._work_id_state = FixedGroupMixedCgaAtomicCounterWorkIdState( + atomic_counter_state=atomic_counter_state, + registration_counter_pointer=device_workspace.ptr(self.fallback_registration_counter_region), + group_token_pointer=device_workspace.ptr(self.fallback_group_token_region), + split_factor=self.config.split_factor, + fallback_cluster_count=self.config.fallback_cluster_count, + is_fallback_cluster=is_fallback_cluster, + fallback_group_idx=Int32(0), + in_group_idx=Int32(0), + previous_token=Int64(0), + next_generation=Int32(1), + claimed_counter_index=Int32(0), + ) + else: + self._work_id_state = atomic_counter_state + else: + self._work_id_state = GridStrideWorkIdState( + next_work_id=Int32(block_idx[2]), + work_id_stride=Int32(self.config.launch_cluster_cnt_merge_as_preferred), + ) + self.claimed_stream_index = Int32(0) + + @cute.jit + def initialize_fallback_group(self) -> None: + """Register one physical fallback cluster with its fixed logical group.""" + if cutlass.const_expr(self.config.is_mixed and self.work_id_mode == "atomic_counter"): + self._work_id_state = initialize_fixed_group_mixed_cga_work_id_state(self._work_id_state) + if self._work_id_state.is_fallback_cluster: + self.cta_coord_in_preferred_cluster = self._preferred_cluster_cta_coord( + self._cta_coord_in_active_cluster, self._work_id_state.in_group_idx + ) + + @cute.jit + def _preferred_cluster_cta_coord( + self, cta_coord_in_active_cluster: cute.Coord, inner_cluster_idx: Int32 + ) -> cute.Coord: + inner_cluster_m = inner_cluster_idx % Int32(self.config.mn_split_factors[0]) + inner_cluster_n = (inner_cluster_idx // Int32(self.config.mn_split_factors[0])) % Int32( + self.config.mn_split_factors[1] + ) + return ( + cta_coord_in_active_cluster[0] + inner_cluster_m * Int32(self.config.fallback_cluster_shape[0]), + cta_coord_in_active_cluster[1] + inner_cluster_n * Int32(self.config.fallback_cluster_shape[1]), + Int32(0), + ) + + @cute.jit + def claim_next_work(self, stream_index=0) -> Int32: + """Claim the next canonical work ID and update the preferred CTA coordinate.""" + claimed_work_id, self._work_id_state = claim_work_id(self._work_id_state, atomic_counter_index=stream_index) + canonical_work_id = claimed_work_id + if cutlass.const_expr(self.config.is_mixed and self.work_id_mode == "atomic_counter"): + cta_coord_in_preferred_cluster = self._cta_coord_in_active_cluster + if self._work_id_state.is_fallback_cluster: + cta_coord_in_preferred_cluster = self._preferred_cluster_cta_coord( + self._cta_coord_in_active_cluster, self._work_id_state.in_group_idx + ) + self.cta_coord_in_preferred_cluster = cta_coord_in_preferred_cluster + self.claimed_stream_index = self._work_id_state.claimed_counter_index + else: + self.claimed_stream_index = Int32(stream_index) + return canonical_work_id + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + values.extend(extract_mlir_values(self._work_id_state)) + values.extend(extract_mlir_values(self.cta_coord_in_preferred_cluster)) + values.extend(extract_mlir_values(self.claimed_stream_index)) + if self.config.is_mixed and self.work_id_mode == "atomic_counter": + values.extend(extract_mlir_values(self._cta_coord_in_active_cluster)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "NonClcMixedCgaSchedulerWorker": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + result = type(self)(config=self.config, work_id_mode=self.work_id_mode, stream_count=self.stream_count) + result._work_id_state = rebuild(self._work_id_state) + result.cta_coord_in_preferred_cluster = rebuild(self.cta_coord_in_preferred_cluster) + result.claimed_stream_index = rebuild(self.claimed_stream_index) + if self.config.is_mixed and self.work_id_mode == "atomic_counter": + result._cta_coord_in_active_cluster = rebuild(self._cta_coord_in_active_cluster) + if value_index != len(values): + raise ValueError( + f"NonClcMixedCgaSchedulerWorker MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return result + + +__all__ = ["NonClcMixedCgaConfig", "NonClcMixedCgaSchedulerWorker"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/work_id_claim.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/work_id_claim.py new file mode 100644 index 000000000..222885b64 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/work_id_claim.py @@ -0,0 +1,569 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Blackwell persistent work-ID claim backends.""" + +import dataclasses +from typing import List, Tuple + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass._mlir import ir +from cutlass.cutlass_dsl import Boolean, Int32, Int64, extract_mlir_values, new_from_mlir_values + +from ...helpers.ptx_helpers import mbarrier_arrive_expect_tx_on_peer, nanosleep, store_i32_to_peer_cluster_smem_async + + +class GridStrideWorkIdState: + """Register state for one monotonic grid-stride work-ID stream.""" + + def __init__(self, next_work_id: Int32, work_id_stride: Int32) -> None: + self.next_work_id = next_work_id + self.work_id_stride = work_id_stride + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + values.extend(extract_mlir_values(self.next_work_id)) + values.extend(extract_mlir_values(self.work_id_stride)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GridStrideWorkIdState": + next_work_id_value_count = len(extract_mlir_values(self.next_work_id)) + stride_value_count = len(extract_mlir_values(self.work_id_stride)) + expected_value_count = next_work_id_value_count + stride_value_count + if len(values) != expected_value_count: + raise ValueError( + f"GridStrideWorkIdState MLIR value count mismatch: expected {expected_value_count}, got {len(values)}." + ) + return type(self)( + next_work_id=new_from_mlir_values(self.next_work_id, values[:next_work_id_value_count]), + work_id_stride=new_from_mlir_values(self.work_id_stride, values[next_work_id_value_count:]), + ) + + +class AtomicCounterWorkIdState: + """Cluster-wide state for one of several contiguous atomic work-ID streams.""" + + def __init__( + self, + counter_pointer: cute.Pointer, + counter_count: int, + broadcast_pointer: cute.Pointer, + is_leader_cta: Boolean, + cluster_pipeline: pipeline.PipelineAsync, + producer_state, + consumer_state, + cluster_size: int | Int32, + ) -> None: + if isinstance(counter_count, bool) or not isinstance(counter_count, int) or counter_count <= 0: + raise ValueError("counter_count must be a positive Python int.") + self.counter_pointer = counter_pointer + self.counter_count = counter_count + self.broadcast_pointer = broadcast_pointer + self.is_leader_cta = is_leader_cta + self.cluster_pipeline = cluster_pipeline + self.producer_state = producer_state + self.consumer_state = consumer_state + self.cluster_size = cluster_size + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in ( + self.counter_pointer, + self.broadcast_pointer, + self.is_leader_cta, + self.producer_state, + self.consumer_state, + ): + values.extend(extract_mlir_values(field)) + if isinstance(self.cluster_size, Int32): + values.extend(extract_mlir_values(self.cluster_size)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "AtomicCounterWorkIdState": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + result = type(self)( + counter_pointer=rebuild(self.counter_pointer), + counter_count=self.counter_count, + broadcast_pointer=rebuild(self.broadcast_pointer), + is_leader_cta=rebuild(self.is_leader_cta), + cluster_pipeline=self.cluster_pipeline, + producer_state=rebuild(self.producer_state), + consumer_state=rebuild(self.consumer_state), + cluster_size=rebuild(self.cluster_size) if isinstance(self.cluster_size, Int32) else self.cluster_size, + ) + if value_index != len(values): + raise ValueError( + f"AtomicCounterWorkIdState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return result + + +class FixedGroupMixedCgaAtomicCounterWorkIdState: + """Atomic-counter state for fixed groups of physical fallback clusters.""" + + def __init__( + self, + atomic_counter_state: AtomicCounterWorkIdState, + registration_counter_pointer: cute.Pointer, + group_token_pointer: cute.Pointer, + split_factor: int, + fallback_cluster_count: int, + is_fallback_cluster: Boolean, + fallback_group_idx: Int32, + in_group_idx: Int32, + previous_token: Int64, + next_generation: Int32, + claimed_counter_index: Int32, + ) -> None: + if isinstance(split_factor, bool) or not isinstance(split_factor, int) or split_factor <= 1: + raise ValueError("split_factor must be a Python int greater than one.") + if ( + isinstance(fallback_cluster_count, bool) + or not isinstance(fallback_cluster_count, int) + or fallback_cluster_count <= 0 + ): + raise ValueError("fallback_cluster_count must be a positive Python int.") + if fallback_cluster_count % split_factor != 0: + raise ValueError("fallback_cluster_count must be divisible by split_factor.") + if atomic_counter_state.counter_count > 2: + raise ValueError("Fixed fallback groups support at most two work-ID streams.") + self.atomic_counter_state = atomic_counter_state + self.registration_counter_pointer = registration_counter_pointer + self.group_token_pointer = group_token_pointer + self.split_factor = split_factor + self.fallback_cluster_count = fallback_cluster_count + self.is_fallback_cluster = is_fallback_cluster + self.is_preferred_cluster = is_fallback_cluster == Boolean(False) + self.fallback_group_idx = fallback_group_idx + self.in_group_idx = in_group_idx + self.previous_token = previous_token + self.next_generation = next_generation + self.claimed_counter_index = claimed_counter_index + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in ( + self.atomic_counter_state, + self.registration_counter_pointer, + self.group_token_pointer, + self.is_fallback_cluster, + self.fallback_group_idx, + self.in_group_idx, + self.previous_token, + self.next_generation, + self.claimed_counter_index, + ): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "FixedGroupMixedCgaAtomicCounterWorkIdState": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + result = type(self)( + atomic_counter_state=rebuild(self.atomic_counter_state), + registration_counter_pointer=rebuild(self.registration_counter_pointer), + group_token_pointer=rebuild(self.group_token_pointer), + split_factor=self.split_factor, + fallback_cluster_count=self.fallback_cluster_count, + is_fallback_cluster=rebuild(self.is_fallback_cluster), + fallback_group_idx=rebuild(self.fallback_group_idx), + in_group_idx=rebuild(self.in_group_idx), + previous_token=rebuild(self.previous_token), + next_generation=rebuild(self.next_generation), + claimed_counter_index=rebuild(self.claimed_counter_index), + ) + if value_index != len(values): + raise ValueError( + f"FixedGroupMixedCgaAtomicCounterWorkIdState MLIR value count mismatch: " + f"consumed {value_index}, got {len(values)}." + ) + return result + + +@dataclasses.dataclass(frozen=True) +class GridWorkId: + """One CTA-specific coordinate claimed from a three-dimensional grid.""" + + grid_m: Int32 + grid_n: Int32 + grid_l: Int32 + is_valid: Boolean + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in (self.grid_m, self.grid_n, self.grid_l, self.is_valid): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GridWorkId": + if len(values) != 4: + raise ValueError(f"GridWorkId expects four MLIR values, got {len(values)}.") + fields = (self.grid_m, self.grid_n, self.grid_l, self.is_valid) + return type(self)(*(new_from_mlir_values(field, [value]) for field, value in zip(fields, values))) + + +class ClusterLaunchControlWorkIdState: + """Cluster-wide state for hardware-assisted grid-coordinate claims.""" + + def __init__( + self, + response_pending: Boolean, + grid_m: Int32, + grid_n: Int32, + grid_l: Int32, + response_is_valid: Boolean, + cta_coord_in_cluster: cute.Coord, + cluster_pipeline: pipeline.PipelineClcFetchAsync, + producer_state, + consumer_state, + is_leader_cta: Boolean, + response_pointer: cute.Pointer, + ) -> None: + self.response_pending = response_pending + self.grid_m = grid_m + self.grid_n = grid_n + self.grid_l = grid_l + self.response_is_valid = response_is_valid + self.cta_coord_in_cluster = cta_coord_in_cluster + self.cluster_pipeline = cluster_pipeline + self.producer_state = producer_state + self.consumer_state = consumer_state + self.is_leader_cta = is_leader_cta + self.response_pointer = response_pointer + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in ( + self.response_pending, + self.grid_m, + self.grid_n, + self.grid_l, + self.response_is_valid, + self.cta_coord_in_cluster, + self.producer_state, + self.consumer_state, + self.is_leader_cta, + self.response_pointer, + ): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "ClusterLaunchControlWorkIdState": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + rebuilt_field = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return rebuilt_field + + result = type(self)( + response_pending=rebuild(self.response_pending), + grid_m=rebuild(self.grid_m), + grid_n=rebuild(self.grid_n), + grid_l=rebuild(self.grid_l), + response_is_valid=rebuild(self.response_is_valid), + cta_coord_in_cluster=rebuild(self.cta_coord_in_cluster), + cluster_pipeline=self.cluster_pipeline, + producer_state=rebuild(self.producer_state), + consumer_state=rebuild(self.consumer_state), + is_leader_cta=rebuild(self.is_leader_cta), + response_pointer=rebuild(self.response_pointer), + ) + if value_index != len(values): + raise ValueError( + f"ClusterLaunchControlWorkIdState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return result + + +@cute.jit +def _claim_grid_stride_work_id(work_id_state: GridStrideWorkIdState) -> Tuple[Int32, GridStrideWorkIdState]: + """Claim the next ID from one monotonic grid-stride stream.""" + linear_work_id = work_id_state.next_work_id + work_id_state.next_work_id = linear_work_id + work_id_state.work_id_stride + return linear_work_id, work_id_state + + +@cute.jit +def _claim_atomic_counter_work_id( + work_id_state: AtomicCounterWorkIdState, atomic_counter_index=0 +) -> Tuple[Int32, AtomicCounterWorkIdState]: + """Claim from one selected counter and broadcast the ID within the cluster.""" + invalid_static_index = isinstance(atomic_counter_index, int) and ( + atomic_counter_index < 0 or atomic_counter_index >= work_id_state.counter_count + ) + if cutlass.const_expr(invalid_static_index): + raise ValueError( + f"atomic_counter_index must be in [0, {work_id_state.counter_count}), got {atomic_counter_index}." + ) + broadcast_tensor = cute.make_tensor(work_id_state.broadcast_pointer, cute.make_layout((1,))) + cluster_pipeline = work_id_state.cluster_pipeline + selected_counter_pointer = work_id_state.counter_pointer + Int32(atomic_counter_index) + + if work_id_state.is_leader_cta: + cluster_pipeline.producer_acquire(work_id_state.producer_state) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier(work_id_state.producer_state.index) + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + atomic_work_id = Int32(0) + if lane_idx == Int32(0): + atomic_work_id = cute.arch.atomic_add(selected_counter_pointer, Int32(1)) + atomic_work_id = cute.arch.shuffle_sync(atomic_work_id, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) + if lane_idx < Int32(work_id_state.cluster_size): + store_i32_to_peer_cluster_smem_async( + work_id_state.broadcast_pointer, atomic_work_id, full_barrier_pointer, lane_idx + ) + mbarrier_arrive_expect_tx_on_peer(full_barrier_pointer, Int32(4), lane_idx) + work_id_state.producer_state.advance() + + cluster_pipeline.consumer_wait(work_id_state.consumer_state) + linear_work_id = broadcast_tensor[0] + cute.arch.fence_acq_rel_cta() + cluster_pipeline.sync_object_empty.arrive(work_id_state.consumer_state.index, Int32(0)) + work_id_state.consumer_state.advance() + return linear_work_id, work_id_state + + +@cute.jit +def initialize_fixed_group_mixed_cga_work_id_state( + work_id_state: FixedGroupMixedCgaAtomicCounterWorkIdState, +) -> FixedGroupMixedCgaAtomicCounterWorkIdState: + """Register one physical fallback cluster and broadcast its fixed group coordinates.""" + atomic_counter_state = work_id_state.atomic_counter_state + if work_id_state.is_fallback_cluster: + broadcast_tensor = cute.make_tensor(atomic_counter_state.broadcast_pointer, cute.make_layout((1,))) + cluster_pipeline = atomic_counter_state.cluster_pipeline + if atomic_counter_state.is_leader_cta: + cluster_pipeline.producer_acquire(atomic_counter_state.producer_state) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier( + atomic_counter_state.producer_state.index + ) + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + fallback_ordinal = Int32(0) + if lane_idx == Int32(0): + fallback_ordinal = cute.arch.atomic_add( + work_id_state.registration_counter_pointer, Int32(1), sem="relaxed", scope="gpu" + ) + fallback_ordinal = Int32( + cute.arch.shuffle_sync(fallback_ordinal, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) + ) + if lane_idx < Int32(atomic_counter_state.cluster_size): + store_i32_to_peer_cluster_smem_async( + atomic_counter_state.broadcast_pointer, fallback_ordinal, full_barrier_pointer, lane_idx + ) + mbarrier_arrive_expect_tx_on_peer(full_barrier_pointer, Int32(4), lane_idx) + atomic_counter_state.producer_state.advance() + + cluster_pipeline.consumer_wait(atomic_counter_state.consumer_state) + fallback_ordinal = broadcast_tensor[0] + cute.arch.fence_acq_rel_cta() + cluster_pipeline.sync_object_empty.arrive(atomic_counter_state.consumer_state.index, Int32(0)) + atomic_counter_state.consumer_state.advance() + + fallback_group_idx = fallback_ordinal // Int32(work_id_state.split_factor) + work_id_state.fallback_group_idx = fallback_group_idx + work_id_state.in_group_idx = fallback_ordinal - fallback_group_idx * Int32(work_id_state.split_factor) + work_id_state.atomic_counter_state = atomic_counter_state + return work_id_state + + +@cute.jit +def _claim_fixed_group_fallback_work_id( + work_id_state: FixedGroupMixedCgaAtomicCounterWorkIdState, atomic_counter_index=0 +) -> Tuple[Int32, FixedGroupMixedCgaAtomicCounterWorkIdState]: + """Claim one canonical ID and hand it to every member of a fixed fallback group.""" + atomic_counter_state = work_id_state.atomic_counter_state + invalid_static_index = isinstance(atomic_counter_index, int) and ( + atomic_counter_index < 0 or atomic_counter_index >= atomic_counter_state.counter_count + ) + if cutlass.const_expr(invalid_static_index): + raise ValueError( + f"atomic_counter_index must be in [0, {atomic_counter_state.counter_count}), got {atomic_counter_index}." + ) + + broadcast_tensor = cute.make_tensor(atomic_counter_state.broadcast_pointer, cute.make_layout((1,))) + cluster_pipeline = atomic_counter_state.cluster_pipeline + selected_counter_pointer = atomic_counter_state.counter_pointer + Int32(atomic_counter_index) + + if atomic_counter_state.is_leader_cta: + cluster_pipeline.producer_acquire(atomic_counter_state.producer_state) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier(atomic_counter_state.producer_state.index) + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + group_base_offset = work_id_state.fallback_group_idx * Int32(work_id_state.split_factor) + group_token_pointer = work_id_state.group_token_pointer + group_base_offset + claimed_payload = Int32(0) + + if work_id_state.in_group_idx == Int32(0): + all_members_consumed = Boolean(False) + while not all_members_consumed: + observed_token = work_id_state.previous_token + if lane_idx < Int32(work_id_state.split_factor): + observed_token = cute.arch.load(group_token_pointer + lane_idx, Int64, sem="acquire", scope="gpu") + lane_is_ready = (lane_idx >= Int32(work_id_state.split_factor)) | ( + observed_token == work_id_state.previous_token + ) + ready_mask = Int32(cute.arch.vote_ballot_sync(lane_is_ready)) + all_members_consumed = ready_mask == Int32(-1) + if not all_members_consumed: + nanosleep(500) + + claimed_work_id = Int32(0) + if lane_idx == Int32(0): + claimed_work_id = cute.arch.atomic_add(selected_counter_pointer, Int32(1), sem="relaxed", scope="gpu") + claimed_work_id = Int32( + cute.arch.shuffle_sync(claimed_work_id, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) + ) + claimed_payload = claimed_work_id | (Int32(atomic_counter_index) << Int32(31)) + token = (Int64(work_id_state.next_generation) << Int64(32)) | (Int64(claimed_payload) & Int64(0xFFFFFFFF)) + if lane_idx == Int32(0): + cute.arch.store(group_token_pointer, token, sem="relaxed", scope="gpu") + work_id_state.previous_token = token + work_id_state.next_generation = work_id_state.next_generation + Int32(1) + else: + token = work_id_state.previous_token + while token == work_id_state.previous_token: + token_high = Int32(0) + token_low = Int32(0) + if lane_idx == Int32(0): + observed_token = cute.arch.load(group_token_pointer, Int64, sem="relaxed", scope="gpu") + token_high = Int32(observed_token >> Int64(32)) + token_low = Int32(observed_token & Int64(0xFFFFFFFF)) + token_high = Int32(cute.arch.shuffle_sync(token_high, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31)) + token_low = Int32(cute.arch.shuffle_sync(token_low, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31)) + token = (Int64(token_high) << Int64(32)) | (Int64(token_low) & Int64(0xFFFFFFFF)) + if token == work_id_state.previous_token: + nanosleep(500) + claimed_payload = Int32(token & Int64(0xFFFFFFFF)) + if lane_idx == Int32(0): + cute.arch.store(group_token_pointer + work_id_state.in_group_idx, token, sem="relaxed", scope="gpu") + work_id_state.previous_token = token + + if lane_idx < Int32(atomic_counter_state.cluster_size): + store_i32_to_peer_cluster_smem_async( + atomic_counter_state.broadcast_pointer, claimed_payload, full_barrier_pointer, lane_idx + ) + mbarrier_arrive_expect_tx_on_peer(full_barrier_pointer, Int32(4), lane_idx) + atomic_counter_state.producer_state.advance() + + cluster_pipeline.consumer_wait(atomic_counter_state.consumer_state) + claimed_payload = broadcast_tensor[0] + cute.arch.fence_acq_rel_cta() + cluster_pipeline.sync_object_empty.arrive(atomic_counter_state.consumer_state.index, Int32(0)) + atomic_counter_state.consumer_state.advance() + work_id_state.atomic_counter_state = atomic_counter_state + work_id_state.claimed_counter_index = (claimed_payload >> Int32(31)) & Int32(1) + linear_work_id = claimed_payload & Int32(0x7FFFFFFF) + return linear_work_id, work_id_state + + +@cute.jit +def _claim_fixed_group_mixed_cga_work_id( + work_id_state: FixedGroupMixedCgaAtomicCounterWorkIdState, atomic_counter_index=0 +) -> Tuple[Int32, FixedGroupMixedCgaAtomicCounterWorkIdState]: + """Claim directly for a preferred cluster or through its fixed fallback group.""" + linear_work_id = Int32(0) + if work_id_state.is_preferred_cluster: + linear_work_id, atomic_counter_state = _claim_atomic_counter_work_id( + work_id_state.atomic_counter_state, atomic_counter_index + ) + work_id_state.atomic_counter_state = atomic_counter_state + work_id_state.claimed_counter_index = Int32(atomic_counter_index) + else: + linear_work_id, work_id_state = _claim_fixed_group_fallback_work_id(work_id_state, atomic_counter_index) + return linear_work_id, work_id_state + + +@cute.jit +def _claim_cluster_launch_control_work_id( + work_id_state: ClusterLaunchControlWorkIdState, +) -> Tuple[GridWorkId, ClusterLaunchControlWorkIdState]: + """Claim the next canceled cluster and return this CTA's grid coordinate.""" + use_bootstrap = work_id_state.response_pending + state_before_bootstrap = work_id_state + if use_bootstrap: + work_id_state.response_pending = Boolean(False) + else: + work_id_state = state_before_bootstrap + + state_before_query = work_id_state + if not use_bootstrap: + state_before_leader_query = work_id_state + if work_id_state.is_leader_cta: + work_id_state.cluster_pipeline.producer_acquire(work_id_state.producer_state) + response_barrier = work_id_state.cluster_pipeline.producer_get_barrier(work_id_state.producer_state) + with cute.arch.elect_one(): + cute.arch.issue_clc_query(response_barrier, work_id_state.response_pointer) + else: + work_id_state = state_before_leader_query + work_id_state.producer_state.advance() + + work_id_state.cluster_pipeline.consumer_wait(work_id_state.consumer_state) + (cluster_origin_m, cluster_origin_n, grid_l, response_is_valid) = cute.arch.clc_response( + work_id_state.response_pointer + ) + cute.arch.fence_acq_rel_cta() + work_id_state.cluster_pipeline.consumer_release(work_id_state.consumer_state) + work_id_state.consumer_state.advance() + + work_id_state.grid_m = cluster_origin_m + work_id_state.cta_coord_in_cluster[0] + work_id_state.grid_n = cluster_origin_n + work_id_state.cta_coord_in_cluster[1] + work_id_state.grid_l = grid_l + work_id_state.response_is_valid = response_is_valid != Int32(0) + else: + work_id_state = state_before_query + + return ( + GridWorkId( + grid_m=work_id_state.grid_m, + grid_n=work_id_state.grid_n, + grid_l=work_id_state.grid_l, + is_valid=work_id_state.response_is_valid, + ), + work_id_state, + ) + + +@cute.jit +def claim_work_id(work_id_state, atomic_counter_index=0): + """Claim the next work ID using the backend encoded by the state type.""" + if cutlass.const_expr(isinstance(work_id_state, GridStrideWorkIdState)): + return _claim_grid_stride_work_id(work_id_state) + if cutlass.const_expr(isinstance(work_id_state, AtomicCounterWorkIdState)): + return _claim_atomic_counter_work_id(work_id_state, atomic_counter_index) + if cutlass.const_expr(isinstance(work_id_state, FixedGroupMixedCgaAtomicCounterWorkIdState)): + return _claim_fixed_group_mixed_cga_work_id(work_id_state, atomic_counter_index) + if cutlass.const_expr(isinstance(work_id_state, ClusterLaunchControlWorkIdState)): + return _claim_cluster_launch_control_work_id(work_id_state) + raise TypeError(f"Unsupported work-ID state: {type(work_id_state).__name__}.") + + +__all__ = [ + "AtomicCounterWorkIdState", + "ClusterLaunchControlWorkIdState", + "FixedGroupMixedCgaAtomicCounterWorkIdState", + "GridStrideWorkIdState", + "GridWorkId", + "claim_work_id", + "initialize_fixed_group_mixed_cga_work_id_state", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/quant_def.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/quant_def.py new file mode 100644 index 000000000..7ff5730d7 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/quant_def.py @@ -0,0 +1,249 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Canonical quantization kinds and payload formats.""" + +import dataclasses +import enum +from typing import ClassVar, Dict, Literal, Optional, Tuple, Type + +import cutlass + + +class QuantKind(str, enum.Enum): + """One admissible block-scaled quantization mode. + + A member is defined by its (weight, activation) element pair -- under swap-AB the weight is + the MMA A operand and the activation is the B operand. The selected K-throughput mode remains + explicit because one quantization kind can use multiple instruction K extents. + """ + + nvfp4 = "nvfp4" + mxfp4 = "mxfp4" + mxfp8_e4m3 = "mxfp8_e4m3" + mxfp8_e5m2 = "mxfp8_e5m2" + mxfp4_mxfp8 = "mxfp4_mxfp8" + + # Enum's __str__/__format__ for mixed-in types changed across Python 3.10/3.11/3.12, and the + # kernels fold the kind into their compiled-kernel cache key. Pin both to the member value so + # the key cannot silently become "QuantKind.nvfp4" on an interpreter upgrade. + __str__ = str.__str__ + __format__ = str.__format__ + + @property + def weight_dtype(self) -> Type[cutlass.Numeric]: + return _element_pair[self][0] + + @property + def activation_dtype(self) -> Type[cutlass.Numeric]: + """Also the dtype the FC1 epilogue must emit: FC2 consumes it as its activation.""" + return _element_pair[self][1] + + @property + def sf_vec_size(self) -> int: + return 16 if self is QuantKind.nvfp4 else 32 + + @property + def sf_dtype(self) -> Type[cutlass.Numeric]: + # Hardware allows a UE8M0 scale at vec 16 too, but nvfp4 is the only vec-16 mode we build. + return cutlass.Float8E4M3FN if self.sf_vec_size == 16 else cutlass.Float8E8M0FNU + + @property + def umma_kind(self) -> str: + """The ``tcgen05.mma.kind::`` qualifier. + + Mirrors the dispatch in ``blackwell_helpers._make_blockscaled_trivial_tiled_mma_impl``: + only an fp4 pair reaches the fp4-specific kinds, everything else -- including any mixed + pair -- falls back to mxf8f6f4. Keeping the two in sync matters because we build the tiled + MMA through that helper but emit the instruction ourselves. + """ + both_fp4 = self.weight_dtype is cutlass.Float4E2M1FN and self.activation_dtype is cutlass.Float4E2M1FN + if not both_fp4: + return "mxf8f6f4" + return "mxf4nvf4" if self.sf_vec_size == 16 else "mxf4" + + @property + def umma_scale_vec_suffix(self) -> str: + """PTX modifier after ``.block_scale``; mxf8f6f4 takes none (its scale vector is 32).""" + if self.umma_kind == "mxf8f6f4": + return "" + return ".block16" if self.sf_vec_size == 16 else ".block32" + + def instruction_k(self, mma_k_mode: Literal["1x", "2x"]) -> int: + instruction_k_1x = 32 if self.umma_kind == "mxf8f6f4" else 64 + if mma_k_mode == "1x": + return instruction_k_1x + if mma_k_mode == "2x": + return instruction_k_1x * 2 + raise ValueError(f"Invalid MMA K mode {mma_k_mode!r}; expected '1x' or '2x'.") + + @property + def weight_format_code(self) -> int: + """Instruction-descriptor ``a_format_`` under swap-AB.""" + return _instruction_descriptor_format_code(self.umma_kind, self.weight_dtype) + + @property + def activation_format_code(self) -> int: + """Instruction-descriptor ``b_format_`` under swap-AB.""" + return _instruction_descriptor_format_code(self.umma_kind, self.activation_dtype) + + @property + def scale_format_code(self) -> int: + """Instruction-descriptor ``scale_format_``: 0 = UE4M3, 1 = UE8M0.""" + return 0 if self.sf_dtype is cutlass.Float8E4M3FN else 1 + + def needs_unpack_tma(self, architecture: str) -> bool: + """Whether the narrow operand must reach SMEM in 1-byte containers (U4_UNPACK_U8 TMA). + + Blackwell mixed-width MMA consumes the narrow operand through UNPACK TMA. Rubin consumes + mixed FP4 directly from its native packed SMEM image. + """ + normalized_architecture = architecture.lower().replace("_", "") + if normalized_architecture.startswith("sm"): + normalized_architecture = normalized_architecture[2:] + if normalized_architecture not in ("100", "103", "107"): + raise ValueError(f"Unsupported TCGen05 architecture {architecture!r}.") + return normalized_architecture in ("100", "103") and self.weight_dtype.width != self.activation_dtype.width + + @property + def uses_global_scale(self) -> bool: + """Whether the caller supplies per-expert alpha / norm_const. + + Only nvfp4: an e8m0 scale is a pure power of two and already carries the whole rescale. + """ + return self is QuantKind.nvfp4 + + +_element_pair: Dict[QuantKind, Tuple[Type[cutlass.Numeric], Type[cutlass.Numeric]]] = { + QuantKind.nvfp4: (cutlass.Float4E2M1FN, cutlass.Float4E2M1FN), + QuantKind.mxfp4: (cutlass.Float4E2M1FN, cutlass.Float4E2M1FN), + QuantKind.mxfp8_e4m3: (cutlass.Float8E4M3FN, cutlass.Float8E4M3FN), + QuantKind.mxfp8_e5m2: (cutlass.Float8E5M2, cutlass.Float8E5M2), + QuantKind.mxfp4_mxfp8: (cutlass.Float4E2M1FN, cutlass.Float8E4M3FN), +} + + +# The three ``a_format_`` / ``b_format_`` bits are read against one of two disjoint enums, and the +# MMA kind picks which (CUTLASS ``UMMA::MXF4Format`` vs ``UMMA::MXF8F6F4Format`` in +# cute/arch/mma_sm100_desc.hpp). E2M1 is 1 under the first and 5 under the second. Getting it wrong +# still compiles and still runs -- it just computes garbage -- so these live in exactly one place. +_mxf4_format_code: Dict[Type[cutlass.Numeric], int] = {cutlass.Float4E2M1FN: 1} +_mxf8f6f4_format_code: Dict[Type[cutlass.Numeric], int] = { + cutlass.Float8E4M3FN: 0, + cutlass.Float8E5M2: 1, + cutlass.Float4E2M1FN: 5, +} + + +def _instruction_descriptor_format_code(umma_kind: str, dtype: Type[cutlass.Numeric]) -> int: + codes = _mxf8f6f4_format_code if umma_kind == "mxf8f6f4" else _mxf4_format_code + try: + return codes[dtype] + except KeyError as error: + raise ValueError(f"{dtype} has no instruction-descriptor format code under kind::{umma_kind}.") from error + + +@dataclasses.dataclass(frozen=True) +class CombineFormat: + """Data and scale representation of one cross-rank FC2 payload.""" + + _act_by_tag: ClassVar[Dict[str, type]] = {"e2m1": cutlass.Float4E2M1FN, "e4m3": cutlass.Float8E4M3FN} + _scale_by_tag: ClassVar[Dict[str, type]] = {"bf16": cutlass.BFloat16, "e8m0": cutlass.Float8E8M0FNU} + _rejection_reason: ClassVar[Dict[str, str]] = { + "32e5m2xe8m0": ( + "e5m2 costs the same 8.25 bits per element as e4m3 and trades a mantissa bit (6 dB of " + "SNR) for exponent range the e8m0 block scale already provides. Use 32e4m3xe8m0." + ) + } + + act_dtype: type + scale_dtype: Optional[type] + scale_block: Optional[int] + + def __post_init__(self) -> None: + allowed_act = {cutlass.BFloat16, *self._act_by_tag.values()} + allowed_scale = {None, *self._scale_by_tag.values()} + if self.act_dtype not in allowed_act: + raise ValueError(f"Unsupported combine data dtype {self.act_dtype}.") + if self.scale_dtype not in allowed_scale: + raise ValueError(f"Unsupported combine scale dtype {self.scale_dtype}.") + if self.scale_dtype is None: + if self.act_dtype is not cutlass.BFloat16 or self.scale_block is not None: + raise ValueError("The unquantized format must be bf16 without a scale block.") + return + if self.act_dtype is cutlass.BFloat16: + raise ValueError("A quantized format cannot use bf16 data.") + if self.scale_dtype is cutlass.BFloat16 and self.scale_block != 16: + raise ValueError("A bf16 amax scale requires a 16-element block.") + if self.scale_dtype is cutlass.Float8E8M0FNU and self.scale_block != 32: + raise ValueError("An e8m0 scale requires a 32-element block.") + + @property + def is_quantized(self) -> bool: + return self.scale_dtype is not None + + @property + def name(self) -> str: + if not self.is_quantized: + return "bf16" + act_tag = next(tag for tag, dtype in self._act_by_tag.items() if dtype is self.act_dtype) + scale_tag = next(tag for tag, dtype in self._scale_by_tag.items() if dtype is self.scale_dtype) + return f"{self.scale_block}{act_tag}x{scale_tag}" + + def __str__(self) -> str: + return self.name + + @classmethod + def parse(cls, text: str) -> "CombineFormat": + specs = { + "bf16": (cutlass.BFloat16, None, None), + "16e2m1xbf16": (cutlass.Float4E2M1FN, cutlass.BFloat16, 16), + "32e4m3xe8m0": (cutlass.Float8E4M3FN, cutlass.Float8E8M0FNU, 32), + } + token = text.strip().lower() + if token in cls._rejection_reason: + raise ValueError(f"Combine format {token!r} is deliberately unsupported: {cls._rejection_reason[token]}") + if token not in specs: + raise ValueError(f"Invalid combine format {text!r}; expected one of {tuple(specs)}.") + act_dtype, scale_dtype, scale_block = specs[token] + return cls(act_dtype, scale_dtype, scale_block) + + +# Every Blackwell 1x PTX hardware encoding is restated independently of the derivations above, so +# that a typo in a property fails at import rather than at the first wrong numerical result. +# Ordered as (umma_kind, scale_vec_suffix, instruction_k_1x, a_format_, b_format_, scale_format_). +_pinned_hardware_encoding: Dict[QuantKind, Tuple[str, str, int, int, int, int]] = { + QuantKind.nvfp4: ("mxf4nvf4", ".block16", 64, 1, 1, 0), + QuantKind.mxfp4: ("mxf4", ".block32", 64, 1, 1, 1), + QuantKind.mxfp8_e4m3: ("mxf8f6f4", "", 32, 0, 0, 1), + QuantKind.mxfp8_e5m2: ("mxf8f6f4", "", 32, 1, 1, 1), + QuantKind.mxfp4_mxfp8: ("mxf8f6f4", "", 32, 5, 0, 1), +} + + +def _verify_pinned_hardware_encoding() -> None: + unpinned = sorted(kind.name for kind in QuantKind if kind not in _pinned_hardware_encoding) + if unpinned: + raise AssertionError(f"QuantKind members without a pinned hardware encoding: {unpinned}.") + for kind, expected in _pinned_hardware_encoding.items(): + derived = ( + kind.umma_kind, + kind.umma_scale_vec_suffix, + kind.instruction_k("1x"), + kind.weight_format_code, + kind.activation_format_code, + kind.scale_format_code, + ) + if derived != expected: + raise AssertionError(f"QuantKind.{kind.name} derives {derived}, pinned encoding is {expected}.") + expected_instruction_k_2x = expected[2] * 2 + instruction_k_2x = kind.instruction_k("2x") + if instruction_k_2x != expected_instruction_k_2x: + raise AssertionError( + f"QuantKind.{kind.name} derives 2x instruction K {instruction_k_2x}, expected {expected_instruction_k_2x}." + ) + + +_verify_pinned_hardware_encoding() + + +__all__ = ["CombineFormat", "QuantKind"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/__init__.py new file mode 100644 index 000000000..a1362f39f --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Private MXFP8 implementation for the MegaMoE execution backend.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py new file mode 100644 index 000000000..9960c9caa --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py @@ -0,0 +1,580 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Logical MXFP8 to Rubin SM107 MegaMoE tensor staging.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from ..._contracts import Fc1WeightLayout, ValidatedForwardRequest +from ..._types import BlockScaledTensor, MoeFormat +from .._plan import PreparedResources +from .._workspace import padded_mxfp8_scale_columns +from ._config import Mxfp8KernelConfig + +_MXFP8_DATA_DTYPE = torch.float8_e4m3fn +_MXFP8_SCALE_DTYPE = torch.float8_e8m0fnu +_GATE_UP_INTERLEAVE = 32 +_WORKSPACE_GUARD_BYTE = 0xA5 + + +def _decode_moe_tensor( + tensor: torch.Tensor | BlockScaledTensor, +) -> torch.Tensor: + """Decode a public MoE tensor to float32 for host-side staging math.""" + + if isinstance(tensor, BlockScaledTensor): + return tensor.dequantize(torch.float32) + return tensor.float() + + +def _quantize_plain_mxfp8( + tensor: torch.Tensor, + *, + axis: int = 1, +) -> BlockScaledTensor: + """Stage a plain floating tensor through the backend's MXFP8 family.""" + + moved = tensor.float().movedim(axis, -1) + logical_extent = moved.shape[-1] + block_count = (logical_extent + 31) // 32 + padded_extent = block_count * 32 + if padded_extent != logical_extent: + moved = torch.nn.functional.pad( + moved, + (0, padded_extent - logical_extent), + ) + blocks = moved.reshape(*moved.shape[:-1], block_count, 32) + raw_scale = blocks.abs().amax(dim=-1) / 448.0 + safe_scale = torch.where(raw_scale > 0, raw_scale, 1.0) + scale_float = torch.where( + raw_scale > 0, + torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), + torch.zeros_like(raw_scale), + ) + scale = scale_float.to(_MXFP8_SCALE_DTYPE) + scale_for_math = scale.float() + reciprocal = torch.where( + scale_for_math > 0, + scale_for_math.reciprocal(), + 0.0, + ) + normalized = (blocks * reciprocal.unsqueeze(-1)).clamp(-448.0, 448.0) + data = normalized.to(_MXFP8_DATA_DTYPE).reshape(*moved.shape)[..., :logical_extent].movedim(-1, axis).contiguous() + return BlockScaledTensor( + data=data, + scale=scale.movedim(-1, axis).contiguous(), + format=MoeFormat.MXFP8, + logical_shape=tuple(tensor.shape), + axis=axis, + ) + + +def _as_mxfp8(tensor: torch.Tensor | BlockScaledTensor) -> BlockScaledTensor: + if isinstance(tensor, BlockScaledTensor): + if tensor.format is not MoeFormat.MXFP8: + raise NotImplementedError("MXFP8 staging cannot convert " f"{tensor.format.value!r} block-scaled input") + return tensor + return _quantize_plain_mxfp8(tensor) + + +def _typed_view( + byte_tensor: torch.Tensor, + dtype: torch.dtype, + shape: tuple[int, ...], +) -> torch.Tensor: + expected_bytes = 1 + for extent in shape: + expected_bytes *= extent + expected_bytes *= dtype.itemsize + if byte_tensor.numel() != expected_bytes: + raise ValueError(f"byte region has {byte_tensor.numel()} bytes, " f"expected {expected_bytes} for shape={shape}, dtype={dtype}") + return byte_tensor.view(dtype).reshape(shape) + + +def _typed_k_major_view( + byte_tensor: torch.Tensor, + dtype: torch.dtype, + shape: tuple[int, int], +) -> torch.Tensor: + """Return a rank-2 ``(K,N)`` view with K as the unit-stride mode.""" + + if len(shape) != 2: + raise ValueError(f"K-major view requires rank-2 shape, got {shape}") + rows, columns = shape + return _typed_view( + byte_tensor, + dtype, + (columns, rows), + ).transpose(0, 1) + + +def _as_bytes(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(torch.uint8) + + +def _validate_int32_downcast(tensor: torch.Tensor) -> None: + """Validate the only lossy public-to-kernel dtype conversion.""" + + if tensor.dtype is torch.int32: + return + if tensor.dtype is not torch.int64: + raise TypeError("topk_idx staging requires torch.int32 or torch.int64, " f"got {tensor.dtype}") + capturing = tensor.device.type == "cuda" and torch.cuda.is_current_stream_capturing() + if capturing or tensor.numel() == 0: + # The public validator checked the same tensor before capture. During + # replay callers must preserve its documented expert-id invariant. + return + limits = torch.iinfo(torch.int32) + outside_int32 = (tensor < limits.min) | (tensor > limits.max) + if bool(outside_int32.any().item()): + raise OverflowError("topk_idx contains a value outside the int32 range") + + +def _zero_workspace_prefix( + workspace: torch.Tensor, + nbytes: int, + *, + name: str, +) -> None: + if nbytes < 0 or nbytes > workspace.numel(): + raise ValueError(f"{name} zero prefix {nbytes} exceeds {workspace.numel()} bytes") + workspace[:nbytes].zero_() + + +def _zero_workspace_range( + workspace: torch.Tensor, + offset: int, + nbytes: int, + *, + name: str, +) -> None: + if offset < 0 or nbytes < 0 or offset + nbytes > workspace.numel(): + raise ValueError(f"{name} byte range [{offset}, {offset + nbytes}) exceeds " f"{workspace.numel()} bytes") + workspace.narrow(0, offset, nbytes).zero_() + + +def _interleave_gate_up_rows( + tensor: torch.Tensor, + intermediate: int, +) -> torch.Tensor: + """Convert gate-half/up-half rows to 32-row gate/up pairs.""" + + if intermediate % _GATE_UP_INTERLEAVE: + raise ValueError("MXFP8 gate/up interleave requires intermediate_size to be " f"divisible by {_GATE_UP_INTERLEAVE}, got {intermediate}") + if tensor.ndim != 3 or tensor.shape[1] != 2 * intermediate: + raise ValueError(f"expected (experts, {2 * intermediate}, K) tensor, " f"got {tuple(tensor.shape)}") + + experts, _gate_up, reduction = tensor.shape + pairs = intermediate // _GATE_UP_INTERLEAVE + gate = tensor[:, :intermediate].reshape( + experts, + pairs, + _GATE_UP_INTERLEAVE, + reduction, + ) + up = tensor[:, intermediate:].reshape( + experts, + pairs, + _GATE_UP_INTERLEAVE, + reduction, + ) + return torch.stack((gate, up), dim=2).reshape(experts, 2 * intermediate, reduction).contiguous() + + +def _to_blocked_bytes(scale_2d: torch.Tensor) -> torch.Tensor: + """Apply the kernel's 32x4x4 scale-factor atom swizzle.""" + + if scale_2d.ndim != 2: + raise ValueError(f"expected 2D scale tensor, got {scale_2d.ndim}D") + rows, columns = scale_2d.shape + if rows == 0 or columns == 0: + return scale_2d.new_empty((0,), dtype=torch.uint8) + + row_blocks = (rows + 127) // 128 + column_blocks = (columns + 3) // 4 + padded_rows = row_blocks * 128 + padded_columns = column_blocks * 4 + padded = torch.zeros( + padded_rows, + padded_columns, + dtype=torch.uint8, + device=scale_2d.device, + ) + padded[:rows, :columns].copy_(_as_bytes(scale_2d)) + blocks = padded.view(row_blocks, 128, column_blocks, 4).permute( + 0, + 2, + 1, + 3, + ) + return blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16).flatten() + + +def _stack_blocked_scales(raw_scales: torch.Tensor) -> torch.Tensor: + experts = raw_scales.shape[0] + blocked = [_to_blocked_bytes(raw_scales[e]) for e in range(experts)] + if not blocked: + return torch.empty( + (0, 0), + dtype=torch.uint8, + device=raw_scales.device, + ).view(raw_scales.dtype) + flat_size = blocked[0].numel() + output = torch.empty( + experts, + flat_size, + dtype=torch.uint8, + device=raw_scales.device, + ) + for expert, values in enumerate(blocked): + output[expert].copy_(values) + return output.view(raw_scales.dtype) + + +def _prepare_fc1( + tensor: BlockScaledTensor, + intermediate: int, + *, + already_interleaved: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build kernel-native K-major FC1 tensors.""" + + payload_nkh = _as_bytes(tensor.data).permute(0, 2, 1).contiguous() + payload_interleaved = payload_nkh if already_interleaved else _interleave_gate_up_rows(payload_nkh, intermediate) + payload = payload_interleaved.view(_MXFP8_DATA_DTYPE).permute(0, 2, 1) + + scales_nk = _as_bytes(tensor.scale).permute(0, 2, 1).contiguous() + scales_interleaved = scales_nk if already_interleaved else _interleave_gate_up_rows(scales_nk, intermediate) + scale = _stack_blocked_scales(scales_interleaved) + return payload, scale + + +def _prepare_fc2( + tensor: BlockScaledTensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Preserve logical bytes while building K-major FC2 tensors.""" + + payload_nk = _as_bytes(tensor.data).permute(0, 2, 1).contiguous() + payload = payload_nk.view(_MXFP8_DATA_DTYPE).permute(0, 2, 1) + scales_nk = _as_bytes(tensor.scale).permute(0, 2, 1).contiguous() + scale = _stack_blocked_scales(scales_nk) + return payload, scale + + +def _tensor_fingerprint(tensor: torch.Tensor) -> tuple | None: + try: + version = tensor._version + except RuntimeError: + return None + return ( + tensor.data_ptr(), + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + tensor.device, + version, + ) + + +def _block_scaled_fingerprint(tensor: BlockScaledTensor) -> tuple | None: + data = _tensor_fingerprint(tensor.data) + scale = _tensor_fingerprint(tensor.scale) + if data is None or scale is None: + return None + return data, scale + + +@dataclass(frozen=True) +class Mxfp8Weights: + fc1_weight: torch.Tensor + fc1_weight_sf: torch.Tensor + fc2_weight: torch.Tensor + fc2_weight_sf: torch.Tensor + + +@dataclass(frozen=True) +class Mxfp8LaunchInputs: + activation: torch.Tensor + activation_sf: torch.Tensor + topk_indices: torch.Tensor + topk_scores: torch.Tensor + weights: Mxfp8Weights + fc1_c: torch.Tensor | None + output_data: torch.Tensor + col_quant_data: torch.Tensor | None + col_quant_sf: torch.Tensor | None + overflow_flag: torch.Tensor + local_workspace: torch.Tensor + shared_workspace: torch.Tensor + token_count: int + + +class Mxfp8InputAdapter: + """Stateful staging adapter with mutation-aware weight transforms.""" + + def __init__(self) -> None: + self._weight_key: tuple | None = None + self._weights: Mxfp8Weights | None = None + self._weight_sources: tuple[torch.Tensor, ...] | None = None + self._weight_refresh_count = 0 + self._initialized_workspace_key: tuple[int, int] | None = None + + @property + def weight_refresh_count(self) -> int: + return self._weight_refresh_count + + def has_cached_weights(self, request: ValidatedForwardRequest) -> bool: + key = self._request_weight_key(request) + return key is not None and key == self._weight_key and self._weights is not None + + def weights_have_version_counters( + self, + request: ValidatedForwardRequest, + ) -> bool: + return self._request_weight_key(request) is not None + + @staticmethod + def _request_weight_key( + request: ValidatedForwardRequest, + ) -> tuple | None: + fc1 = _block_scaled_fingerprint(request.fc1_weight) if isinstance(request.fc1_weight, BlockScaledTensor) else _tensor_fingerprint(request.fc1_weight) + fc2 = _block_scaled_fingerprint(request.fc2_weight) if isinstance(request.fc2_weight, BlockScaledTensor) else _tensor_fingerprint(request.fc2_weight) + if fc1 is None or fc2 is None: + return None + return fc1, fc2 + + def _prepare_weights( + self, + request: ValidatedForwardRequest, + config: Mxfp8KernelConfig, + ) -> Mxfp8Weights: + key = self._request_weight_key(request) + if key is not None and key == self._weight_key and self._weights is not None: + return self._weights + + fc1_source = _as_mxfp8(request.fc1_weight) + fc2_source = _as_mxfp8(request.fc2_weight) + fc1_weight, fc1_weight_sf = _prepare_fc1( + fc1_source, + config.intermediate, + already_interleaved=(config.fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32), + ) + fc2_weight, fc2_weight_sf = _prepare_fc2(fc2_source) + weights = Mxfp8Weights( + fc1_weight=fc1_weight, + fc1_weight_sf=fc1_weight_sf, + fc2_weight=fc2_weight, + fc2_weight_sf=fc2_weight_sf, + ) + self._weight_key = key + self._weights = weights + # Retain the source storages while this entry is cached so allocator + # pointer reuse cannot produce a false cache hit. + self._weight_sources = ( + *((request.fc1_weight.data, request.fc1_weight.scale) if isinstance(request.fc1_weight, BlockScaledTensor) else (request.fc1_weight,)), + *((request.fc2_weight.data, request.fc2_weight.scale) if isinstance(request.fc2_weight, BlockScaledTensor) else (request.fc2_weight,)), + ) + self._weight_refresh_count += 1 + return weights + + def stage( + self, + request: ValidatedForwardRequest, + resources: PreparedResources, + config: Mxfp8KernelConfig, + *, + local_workspace_zero_bytes: int, + shared_workspace_zero_bytes: int, + pre_reduced_activation_offset: int | None, + pre_reduced_activation_bytes_per_token: int, + pre_reduced_activation_sf_offset: int | None, + pre_reduced_activation_sf_bytes_per_token: int, + col_quant_data_rows: int, + col_quant_sf_elements: int, + fc1_c: torch.Tensor | None = None, + ) -> Mxfp8LaunchInputs: + capacity = config.max_tokens_per_rank + token_count = request.token_count + if config.generate_c: + if fc1_c is None: + raise ValueError("generate_c=True requires an fc1_c buffer") + if ( + fc1_c.dtype is not torch.bfloat16 + or fc1_c.device != request.device + or fc1_c.ndim != 2 + or fc1_c.shape[0] <= 0 + or fc1_c.shape[1] != config.fc1_out + or not fc1_c.is_contiguous() + ): + raise ValueError("fc1_c buffer must be contiguous BF16 on the request " f"device with shape (capacity, {config.fc1_out})") + elif fc1_c is not None: + raise ValueError("generate_c=False must not receive an fc1_c buffer") + hidden_sf_columns = (config.hidden + 31) // 32 + padded_sf_columns = padded_mxfp8_scale_columns(config.hidden) + symmetric = resources.workspace.symmetric + local = resources.workspace.local + symmetric_guard = symmetric.get("symmetric_guard") + if symmetric_guard is not None: + symmetric_guard.fill_(_WORKSPACE_GUARD_BYTE) + local_guard = local.get("local_guard") + if local_guard is not None: + local_guard.fill_(_WORKSPACE_GUARD_BYTE) + + activation = _typed_view( + symmetric["activation_data"], + _MXFP8_DATA_DTYPE, + (capacity, config.hidden), + ) + activation_scale_bytes = ( + capacity * padded_sf_columns * _MXFP8_SCALE_DTYPE.itemsize + ) + activation_sf = _typed_view( + symmetric["activation_scale"][:activation_scale_bytes], + _MXFP8_SCALE_DTYPE, + (capacity, padded_sf_columns), + ) + topk_weights = _typed_view( + symmetric["topk_weights"], + torch.float32, + (capacity, config.top_k), + ) + output_data = _typed_view( + symmetric["output_data"], + torch.bfloat16, + (capacity, config.hidden), + ) + topk_indices = _typed_view( + local["topk_idx"], + torch.int32, + (capacity, config.top_k), + ) + overflow_flag = _typed_view( + local["overflow_flag"], + torch.int32, + (1,), + ) + if config.enable_col_quant: + if col_quant_data_rows <= 0 or col_quant_sf_elements <= 0: + raise ValueError("enabled column requant requires positive output capacities") + col_quant_data = _typed_k_major_view( + local["col_quant_data"], + _MXFP8_DATA_DTYPE, + (col_quant_data_rows, config.hidden), + ) + col_quant_sf = _typed_view( + local["col_quant_sf"], + torch.uint8, + (col_quant_sf_elements,), + ) + else: + if col_quant_data_rows != 0 or col_quant_sf_elements != 0: + raise ValueError("disabled column requant must not reserve output capacity") + col_quant_data = None + col_quant_sf = None + local_workspace = local["kernel_local_workspace"] + shared_workspace = symmetric["kernel_shared_workspace"] + + staged_activation = _as_mxfp8(request.activation) + _as_bytes(activation).zero_() + _as_bytes(activation[:token_count]).copy_(_as_bytes(staged_activation.data)) + _as_bytes(activation_sf).zero_() + _as_bytes(activation_sf[:token_count, :hidden_sf_columns]).copy_(_as_bytes(staged_activation.scale)) + _validate_int32_downcast(request.topk_idx) + topk_indices.fill_(-1) + topk_indices[:token_count].copy_(request.topk_idx) + topk_weights.zero_() + topk_weights[:token_count].copy_(request.topk_weights) + _as_bytes(output_data).zero_() + if col_quant_data is not None: + local["col_quant_data"].zero_() + if col_quant_sf is not None: + col_quant_sf.zero_() + overflow_flag.zero_() + workspace_key = ( + local_workspace.data_ptr(), + shared_workspace.data_ptr(), + ) + if workspace_key != self._initialized_workspace_key: + # This prefix contains both tail-reset regions and persistent + # sense-reversing NVLink barrier counters marked + # zero_on_first_allocate. Re-zeroing it on a later rank-skewed + # launch can erase a peer's signal and deadlock both kernels. + _zero_workspace_prefix( + local_workspace, + local_workspace_zero_bytes, + name="local workspace", + ) + _zero_workspace_prefix( + shared_workspace, + shared_workspace_zero_bytes, + name="shared workspace", + ) + self._initialized_workspace_key = workspace_key + if config.fc2_in_kernel_topk_reduce: + if ( + pre_reduced_activation_offset is not None + or pre_reduced_activation_bytes_per_token != 0 + or pre_reduced_activation_sf_offset is not None + or pre_reduced_activation_sf_bytes_per_token != 0 + ): + raise ValueError("in-kernel top-k reduction must not receive a " "standalone pre-reduced activation workspace") + # output_data is the in-kernel REDG accumulation base and was + # cleared above. + else: + if pre_reduced_activation_offset is None or pre_reduced_activation_bytes_per_token <= 0: + raise ValueError("standalone top-k reduction requires a pre-reduced " "activation workspace") + # The kernel writes only valid routes into this persistent combine + # plane. Clear the active token rows so dropped routes cannot reuse + # contributions from a previous launch. + _zero_workspace_range( + shared_workspace, + pre_reduced_activation_offset, + token_count * pre_reduced_activation_bytes_per_token, + name="pre-reduced activation workspace", + ) + quantized_combine = config.combine_format != "bf16" + if quantized_combine: + if pre_reduced_activation_sf_offset is None or pre_reduced_activation_sf_bytes_per_token <= 0: + raise ValueError("quantized standalone top-k reduction requires a " "pre-reduced scale workspace") + _zero_workspace_range( + shared_workspace, + pre_reduced_activation_sf_offset, + token_count * pre_reduced_activation_sf_bytes_per_token, + name="pre-reduced activation scale workspace", + ) + elif pre_reduced_activation_sf_offset is not None or pre_reduced_activation_sf_bytes_per_token != 0: + raise ValueError("BF16 standalone top-k reduction must not receive a " "pre-reduced scale workspace") + + weights = self._prepare_weights(request, config) + return Mxfp8LaunchInputs( + activation=activation, + activation_sf=activation_sf, + topk_indices=topk_indices, + topk_scores=topk_weights, + weights=weights, + fc1_c=fc1_c, + output_data=output_data, + col_quant_data=col_quant_data, + col_quant_sf=col_quant_sf, + overflow_flag=overflow_flag, + local_workspace=local_workspace, + shared_workspace=shared_workspace, + token_count=token_count, + ) + + def close(self) -> None: + self._weights = None + self._weight_key = None + self._weight_sources = None + self._initialized_workspace_key = None + + +__all__ = [ + "Mxfp8InputAdapter", + "Mxfp8LaunchInputs", + "Mxfp8Weights", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py new file mode 100644 index 000000000..2b5e65886 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py @@ -0,0 +1,287 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Single-rank and EP-subgroup MXFP8 execution backend orchestration.""" + +from __future__ import annotations + +import threading +from dataclasses import replace + +import torch +import torch.distributed as dist + +from ..._backend import BackendUnavailableError +from ..._contracts import ForwardConfig, ValidatedForwardRequest +from ..._tuning import MoeEpTuningConfig +from .._plan import ExecutionPlanOwner +from ._adapter import Mxfp8InputAdapter +from ._backward_compile import prepare_backward_kernel +from ._compile import ( + CompiledMxfp8Kernel, + PreparedMxfp8Kernel, + compile_or_get, + prepare_kernel, +) +from ._config import Mxfp8KernelConfig +from ._launch import launch_forward + + +class Mxfp8Backend: + """Own forward/backward executors and per-instance plan resources.""" + + def __init__(self, config: ForwardConfig, device: torch.device) -> None: + self.config = config + self.device = torch.device(device) + self.kernel_config = Mxfp8KernelConfig.from_operator_config(config) + self._adapter = Mxfp8InputAdapter() + self._prepared_kernel: PreparedMxfp8Kernel | None = None + self._compiled: CompiledMxfp8Kernel | None = None + self._plan: ExecutionPlanOwner | None = None + self._warmed_up = False + self._closed = False + self._completion_event: torch.cuda.Event | None = None + self._completion_recorded = False + self._device_work_may_be_pending = False + self._ep_launch_ready = config.ep_size == 1 + self._training_state = None + self._lock = threading.RLock() + + @property + def warmed_up(self) -> bool: + return self._warmed_up + + @property + def kernel_fingerprint(self) -> dict | None: + """Fingerprint of the callable compiled by the most recent launch.""" + + if self._compiled is None: + return None + return self._compiled.fingerprint + + def _ensure_prepared_kernel(self) -> PreparedMxfp8Kernel: + if self._prepared_kernel is None: + try: + self._prepared_kernel = prepare_kernel( + self.config, + self.kernel_config, + self.device, + ) + except (ImportError, OSError) as exc: + raise BackendUnavailableError( + "MoeEp MXFP8 backend requires the 'cutedsl' and 'comm' " "optional dependencies and their shared libraries" + ) from exc + return self._prepared_kernel + + def _ensure_ep_launch_ready(self, resources, stream) -> None: + if self._ep_launch_ready: + return + # First subgroup launch only: peer metadata writes begin before the + # kernel's first cross-rank device barrier. Ensure every rank's + # root-zero and staging work has completed before any rank can issue + # those writes. + stream.synchronize() + if resources.runtime.group is None: + raise RuntimeError("distributed MXFP8 launch requires a " "torch.distributed process group") + tuning_signature = self.kernel_config.tuning_signature(self._ensure_prepared_kernel().launch_cluster_count) + rank_tuning_signatures = [None] * resources.runtime.world_size + dist.all_gather_object( + rank_tuning_signatures, + tuning_signature, + group=resources.runtime.group, + ) + if any(signature != rank_tuning_signatures[0] for signature in rank_tuning_signatures[1:]): + raise RuntimeError("MoeEp tuning must match on every expert-parallel rank; " f"effective signatures by rank: {rank_tuning_signatures}") + dist.barrier(group=resources.runtime.group) + self._ep_launch_ready = True + + def forward(self, request: ValidatedForwardRequest): + with self._lock: + if self._closed: + raise RuntimeError("MoeEp MXFP8 backend is closed") + if request.device != self.device: + raise ValueError(f"MoeEp MXFP8 backend is bound to {self.device}, " f"got {request.device}") + + with torch.cuda.device(self.device): + capturing = torch.cuda.is_current_stream_capturing() + if capturing and not self._adapter.weights_have_version_counters(request): + raise NotImplementedError( + "CUDA graph capture does not support inference tensor " + "weights without version counters; eager calls remain " + "supported and repack those weights on every call" + ) + if capturing and (not self._warmed_up or not self._adapter.has_cached_weights(request)): + raise RuntimeError("MoeEp MXFP8 backend and weights must be warmed up " "before CUDA graph capture") + + stream = torch.cuda.current_stream(self.device) + if self._device_work_may_be_pending: + torch.cuda.synchronize(self.device) + self._device_work_may_be_pending = False + if self._completion_event is None: + self._completion_event = torch.cuda.Event() + elif self._completion_recorded and not capturing: + stream.wait_event(self._completion_event) + + prepared = self._ensure_prepared_kernel() + if self._plan is None: + self._plan = ExecutionPlanOwner( + self.config, + self.device, + prepared.workspace_requirements, + ) + device_work_attempted = False + try: + # Allocation zeroing, input staging, weight transforms, + # compilation, and launch can all enqueue device work. + # Record one completion event even if a later step fails so + # a retry on another stream cannot race those writes. + device_work_attempted = True + resources = self._plan.prepare(request) + inputs = self._adapter.stage( + request, + resources, + self.kernel_config, + local_workspace_zero_bytes=(prepared.local_workspace_zero_bytes), + shared_workspace_zero_bytes=(prepared.shared_workspace_zero_bytes), + pre_reduced_activation_offset=(prepared.pre_reduced_activation_offset), + pre_reduced_activation_bytes_per_token=(prepared.pre_reduced_activation_bytes_per_token), + pre_reduced_activation_sf_offset=(prepared.pre_reduced_activation_sf_offset), + pre_reduced_activation_sf_bytes_per_token=(prepared.pre_reduced_activation_sf_bytes_per_token), + col_quant_data_rows=prepared.col_quant_data_rows, + col_quant_sf_elements=prepared.col_quant_sf_elements, + fc1_c=None, + ) + self._compiled = compile_or_get( + prepared, + inputs, + resources, + ) + self._ensure_ep_launch_ready(resources, stream) + output = launch_forward( + self._compiled, + inputs, + resources, + ) + except (ImportError, OSError) as exc: + raise BackendUnavailableError( + "MoeEp MXFP8 backend requires the 'cutedsl' and 'comm' " "optional dependencies and their shared libraries" + ) from exc + finally: + if device_work_attempted and not capturing: + try: + self._completion_event.record(stream) + self._completion_recorded = True + self._device_work_may_be_pending = False + except Exception: + self._completion_recorded = False + self._device_work_may_be_pending = True + raise + + self._warmed_up = True + return output + + def prepare_training( + self, + *, + lane_count: int, + ): + """Allocate private per-lane state for stateless training calls.""" + + with self._lock: + if self._closed: + raise RuntimeError("MoeEp MXFP8 backend is closed") + if self._training_state is not None: + raise RuntimeError("MoeEp training is already prepared") + training_config = replace( + self.config, + generate_c=True, + backward_wgrad_mode="operands", + token_padding_size=128, + sf_padding_size=128, + ) + forward_kernel_config = Mxfp8KernelConfig.from_operator_config( + training_config, + tuning=training_config.tuning, + ) + backward_tuning = training_config.backward_tuning + if backward_tuning is None: + backward_tuning = MoeEpTuningConfig() + backward_kernel_config = Mxfp8KernelConfig.from_operator_config( + training_config, + tuning=backward_tuning, + ) + # Graph transport must complete its cross-rank protocol before the + # frontend applies the public trap/drop policy at graph tail. + forward_graph_kernel_config = replace( + forward_kernel_config, + drop_on_overflow=True, + # Upstream 5b89819's forward col-requant accepts token + # padding 128/256 but fixes SF atoms at 128; its dGLU + # auxiliaries require token and SF padding to match. The + # graph-only fixed-capacity intersection is therefore 128. + token_padding_block=128, + sf_padding_block=128, + ) + backward_graph_kernel_config = replace( + backward_kernel_config, + drop_on_overflow=True, + # Upstream 5b89819's forward col-requant accepts token + # padding 128/256 but fixes SF atoms at 128; its dGLU + # auxiliaries require token and SF padding to match. The + # graph-only fixed-capacity intersection is therefore 128. + token_padding_block=128, + sf_padding_block=128, + ) + forward = prepare_kernel( + training_config, + forward_graph_kernel_config, + self.device, + ) + backward = prepare_backward_kernel( + training_config, + backward_graph_kernel_config, + self.device, + ) + from ._training_resources import Mxfp8TrainingState + + state = Mxfp8TrainingState( + training_config, + self.device, + forward, + backward, + lane_count=lane_count, + ) + try: + state.prepare() + except Exception: + state.close() + raise + self._training_state = state + return state + + def close(self) -> None: + with self._lock: + if self._closed: + return + with torch.cuda.device(self.device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("MoeEp MXFP8 backend cannot be closed during " "CUDA graph capture") + if self._plan is not None or self._training_state is not None: + torch.cuda.synchronize(self.device) + self._adapter.close() + if self._training_state is not None: + self._training_state.close() + self._training_state = None + if self._plan is not None: + self._plan.close() + self._plan = None + self._prepared_kernel = None + self._compiled = None + self._completion_event = None + self._completion_recorded = False + self._device_work_may_be_pending = False + self._ep_launch_ready = self.config.ep_size == 1 + self._closed = True + + +__all__ = ["Mxfp8Backend"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py new file mode 100644 index 000000000..7636421fb --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py @@ -0,0 +1,323 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Compilation and ABI metadata for Rubin MXFP8 dGLU backward.""" + +from __future__ import annotations + +import math +import threading +from dataclasses import dataclass +from typing import Any + +import torch + +from ..._contracts import ForwardConfig +from .._plan import PreparedResources +from .._workspace import WorkspaceRequirements +from ._compile import ( + _pre_reduced_sf_workspace_metadata, + _pre_reduced_workspace_metadata, +) +from ._compile_common import _compile_kernel, _prepare_rubin_environment +from ._config import Mxfp8KernelConfig +from ._formats import combine_wire_format +from ._launch import _to_cute, _to_cute_ptr + + +@dataclass(frozen=True) +class PreparedMxfp8BackwardKernel: + config: Mxfp8KernelConfig + device: torch.device + architecture: tuple[int, int] + kernel: Any + launch_cluster_count: int + workspace_requirements: WorkspaceRequirements + pool_token_capacity: int + pre_reduced_activation_offset: int + pre_reduced_activation_bytes_per_token: int + pre_reduced_activation_sf_offset: int | None + pre_reduced_activation_sf_bytes_per_token: int + local_workspace_zero_bytes: int + shared_workspace_zero_bytes: int + dfc2_recompute: bool + dfc2_col_output: bool + enable_grad_y2_col_quant: bool + + +@dataclass(frozen=True) +class Mxfp8BackwardLaunchInputs: + grad_out: torch.Tensor + grad_out_sf: torch.Tensor + topk_idx: torch.Tensor + topk_weights: torch.Tensor + fc1_weight: torch.Tensor + fc1_weight_sf: torch.Tensor + fc2_weight: torch.Tensor + fc2_weight_sf: torch.Tensor + beta: torch.Tensor + fc1_preact: torch.Tensor + output_activation: torch.Tensor + overflow_flag: torch.Tensor + dprob: torch.Tensor + fc1_recompute: torch.Tensor + fc1_recompute_sf: torch.Tensor + fc1_col_output: torch.Tensor + fc1_col_output_sf: torch.Tensor + grad_y2: torch.Tensor + grad_y2_sf: torch.Tensor + local_workspace: torch.Tensor + shared_workspace: torch.Tensor + token_count: int + + +@dataclass(frozen=True) +class CompiledMxfp8BackwardKernel: + key: tuple + callable: Any + + +_COMPILE_LOCK = threading.RLock() +_COMPILE_CACHE: dict[tuple, CompiledMxfp8BackwardKernel] = {} + + +def prepare_backward_kernel( + forward_config: ForwardConfig, + config: Mxfp8KernelConfig, + device: torch.device, +) -> PreparedMxfp8BackwardKernel: + """Instantiate the fixed Rubin dGLU specialization.""" + + architecture, launch_cluster_count = _prepare_rubin_environment( + device, + config, + context="backward", + ) + import cutlass + + from ..cutedsl_src.kernel_src.rubin.training.mega.bwd_dglu import ( + Sm107MegaMoEMxfp8DgluKernel, + ) + from ..cutedsl_src.quant_def import CombineFormat + + group_hint = launch_cluster_count if config.group_hint is None else config.group_hint + operands_mode = forward_config.backward_wgrad_mode == "operands" + dfc2_recompute = operands_mode + dfc2_col_output = operands_mode + enable_grad_y2_col_quant = operands_mode + kernel = Sm107MegaMoEMxfp8DgluKernel.from_kwargs( + mma_tiler_mnk=config.mma_tiler_mnk, + cluster_shape_mnk=config.cluster_shape_mnk, + use_2cta_instrs=config.use_2cta_instrs, + group_hint=group_hint, + token_padding_block=config.token_padding_block, + sf_padding_block=config.sf_padding_block, + load_balance_mode=config.load_balance_mode, + static_expert_shape=( + config.num_experts, + config.intermediate, + config.hidden, + ), + force_static_sched=config.force_static_sched, + clc_bundle_size=config.clc_bundle_size, + num_sched_stages=config.num_sched_stages, + ab_dtype=cutlass.Float8E4M3FN, + sf_vec_size=config.sf_vec_size, + world_size=config.world_size, + local_rank=0, + num_topk=config.top_k, + max_tokens_per_rank=config.max_tokens_per_rank, + max_recv_size_per_rank=config.max_recv_size_per_rank, + hidden=config.hidden, + launch_cluster_count=launch_cluster_count, + drop_on_overflow=config.drop_on_overflow, + fc2_in_kernel_topk_reduce=config.fc2_in_kernel_topk_reduce, + token_back_mode=config.token_back_mode, + epi_flag_batch=config.epi_flag_batch, + flag_batch=config.flag_batch, + combine_format=CombineFormat.parse(combine_wire_format(forward_config.combine_format)), + act_func=config.act_func, + gate_up_clamp=config.gate_up_clamp, + dfc2_recompute=dfc2_recompute, + dfc2_col_output=dfc2_col_output, + enable_grad_y2_col_quant=enable_grad_y2_col_quant, + num_ctas_grad_y2_col_quant=config.col_quant_num_ctas, + ) + local_bytes, shared_bytes = kernel.get_workspace_sizes() + local_zero, shared_zero = kernel.require_zero_workspace_leading_bytes + device_workspace = kernel._mega_device_workspace + pool_capacity = int(kernel.pool_token_capacity) + fc1_preact_shape = tuple(int(extent) for extent in kernel.get_fc1_preact_shape()) + expected_preact_shape = ( + pool_capacity, + 2 * config.intermediate, + ) + if fc1_preact_shape != expected_preact_shape: + raise RuntimeError("Rubin dGLU fc1_preact shape mismatch: " f"{fc1_preact_shape} != {expected_preact_shape}") + aux_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in kernel.get_aux_output_shapes().items()} + dprob_bytes = math.prod(aux_shapes["dprob"]) * torch.float32.itemsize + aux_data_bytes = ( + max( + math.prod(aux_shapes["fc1_recompute"]), + math.prod(aux_shapes["fc1_col_output"]), + math.prod(aux_shapes["grad_y2"]), + ) + * torch.float8_e4m3fn.itemsize + ) + aux_scale_bytes = ( + max( + math.prod(aux_shapes["fc1_recompute_sf"]), + math.prod(aux_shapes["fc1_col_output_sf"]), + math.prod(aux_shapes["grad_y2_sf"]), + ) + * torch.float8_e8m0fnu.itemsize + ) + requirements = WorkspaceRequirements.for_mxfp8( + forward_config, + kernel_local_workspace_bytes=local_bytes, + kernel_shared_workspace_bytes=shared_bytes, + backward_dprob_bytes=dprob_bytes, + backward_aux_data_bytes=aux_data_bytes, + backward_aux_scale_bytes=aux_scale_bytes, + ) + pre_reduced_offset, pre_reduced_bytes_per_token = _pre_reduced_workspace_metadata( + device_workspace, + config, + shared_bytes, + ) + if pre_reduced_offset is None or pre_reduced_bytes_per_token <= 0: + raise RuntimeError("Rubin MXFP8 backward requires standalone pre-reduced activation") + pre_reduced_sf_offset, pre_reduced_sf_bytes_per_token = _pre_reduced_sf_workspace_metadata( + device_workspace, + config, + shared_bytes, + ) + return PreparedMxfp8BackwardKernel( + config=config, + device=torch.device(device), + architecture=architecture, + kernel=kernel, + launch_cluster_count=launch_cluster_count, + workspace_requirements=requirements, + pool_token_capacity=pool_capacity, + pre_reduced_activation_offset=pre_reduced_offset, + pre_reduced_activation_bytes_per_token=pre_reduced_bytes_per_token, + pre_reduced_activation_sf_offset=pre_reduced_sf_offset, + pre_reduced_activation_sf_bytes_per_token=(pre_reduced_sf_bytes_per_token), + local_workspace_zero_bytes=int(local_zero), + shared_workspace_zero_bytes=int(shared_zero), + dfc2_recompute=dfc2_recompute, + dfc2_col_output=dfc2_col_output, + enable_grad_y2_col_quant=enable_grad_y2_col_quant, + ) + + +def _layout_signature(inputs: Mxfp8BackwardLaunchInputs) -> tuple: + tensors = tuple(value for value in inputs.__dict__.values() if isinstance(value, torch.Tensor)) + return tuple((tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype) for tensor in tensors) + + +def build_backward_runtime_kwargs( + inputs: Mxfp8BackwardLaunchInputs, + resources: PreparedResources, +) -> dict[str, Any]: + import cuda.bindings.driver as cuda + + stream = resources.runtime.current_stream() + return { + "grad_out": _to_cute(inputs.grad_out), + "grad_out_sf": _to_cute(inputs.grad_out_sf), + "topk_idx": _to_cute(inputs.topk_idx), + "topk_weights": _to_cute(inputs.topk_weights, assumed_align=4), + "fc1_weight": _to_cute(inputs.fc1_weight), + "fc1_weight_sf": _to_cute(inputs.fc1_weight_sf), + "fc2_weight": _to_cute(inputs.fc2_weight), + "fc2_weight_sf": _to_cute(inputs.fc2_weight_sf), + "beta": _to_cute(inputs.beta, assumed_align=4), + "fc1_preact": _to_cute( + inputs.fc1_preact, + assumed_align=128, + dynamic_layout=False, + ), + "output_activation": _to_cute(inputs.output_activation), + "overflow_flag": _to_cute( + inputs.overflow_flag, + assumed_align=4, + dynamic_layout=False, + ), + "dprob": _to_cute(inputs.dprob, dynamic_layout=False), + "fc1_recompute": _to_cute( + inputs.fc1_recompute, + assumed_align=128, + dynamic_layout=False, + ), + "fc1_recompute_sf": _to_cute( + inputs.fc1_recompute_sf, + assumed_align=128, + dynamic_layout=False, + ), + "fc1_col_output": _to_cute( + inputs.fc1_col_output, + assumed_align=128, + dynamic_layout=False, + ), + "fc1_col_output_sf": _to_cute( + inputs.fc1_col_output_sf, + assumed_align=128, + dynamic_layout=False, + ), + "grad_y2": _to_cute( + inputs.grad_y2, + assumed_align=128, + dynamic_layout=False, + ), + "grad_y2_sf": _to_cute( + inputs.grad_y2_sf, + dynamic_layout=False, + ), + "local_workspace": _to_cute_ptr(inputs.local_workspace), + "shared_workspace": _to_cute_ptr(inputs.shared_workspace), + "peer_rank_ptr_mapper_host": (resources.workspace.peer_mapping.to_sym_buffer_host()), + "stream": cuda.CUstream(stream.cuda_stream), + } + + +def compile_backward_or_get( + prepared: PreparedMxfp8BackwardKernel, + inputs: Mxfp8BackwardLaunchInputs, + resources: PreparedResources, +) -> CompiledMxfp8BackwardKernel: + signature = _layout_signature(inputs) + key = ( + prepared.config, + prepared.device.index, + prepared.architecture, + prepared.launch_cluster_count, + prepared.dfc2_recompute, + prepared.dfc2_col_output, + prepared.enable_grad_y2_col_quant, + signature, + ) + with _COMPILE_LOCK: + cached = _COMPILE_CACHE.get(key) + if cached is not None: + return cached + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("MXFP8 backward kernel must be compiled before capture") + runtime_kwargs = build_backward_runtime_kwargs(inputs, resources) + compiled = CompiledMxfp8BackwardKernel( + key=key, + callable=_compile_kernel(prepared.kernel, runtime_kwargs), + ) + _COMPILE_CACHE[key] = compiled + return compiled + + +__all__ = [ + "CompiledMxfp8BackwardKernel", + "Mxfp8BackwardLaunchInputs", + "PreparedMxfp8BackwardKernel", + "build_backward_runtime_kwargs", + "compile_backward_or_get", + "prepare_backward_kernel", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py new file mode 100644 index 000000000..e71d27f8f --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py @@ -0,0 +1,305 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""In-process JIT cache for the vendored Rubin SM107 MXFP8 kernel.""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from typing import Any + +import torch + +from ..._contracts import ForwardConfig +from .._plan import PreparedResources +from .._workspace import WorkspaceRequirements +from ._adapter import Mxfp8LaunchInputs +from ._compile_common import _compile_kernel, _prepare_rubin_environment +from ._config import Mxfp8KernelConfig +from ._fingerprint import build_kernel_fingerprint +from ._launch import build_runtime_kwargs, layout_signature + + +@dataclass(frozen=True) +class PreparedMxfp8Kernel: + config: Mxfp8KernelConfig + device: torch.device + architecture: tuple[int, int] + kernel: Any + launch_cluster_count: int + workspace_requirements: WorkspaceRequirements + pool_token_capacity: int + col_quant_data_rows: int + col_quant_sf_elements: int + token_src_metadata_offset: int + token_src_metadata_bytes: int + col_quant_sizes_offset: int | None + col_quant_sizes_bytes: int + pre_reduced_activation_offset: int | None + pre_reduced_activation_bytes_per_token: int + pre_reduced_activation_sf_offset: int | None + pre_reduced_activation_sf_bytes_per_token: int + local_workspace_zero_bytes: int + shared_workspace_zero_bytes: int + + +@dataclass(frozen=True) +class CompiledMxfp8Kernel: + key: tuple + callable: Any + fingerprint: dict[str, Any] + + +_COMPILE_LOCK = threading.RLock() +_COMPILE_CACHE: dict[tuple, CompiledMxfp8Kernel] = {} +_TOKEN_SRC_METADATA_REGION = "nvlink.token_comm.token_src_metadata" +_COL_QUANT_SIZES_REGION = "rubin.glu_mxfp8.mega.col_quant_expert_token_sizes" +_PRE_REDUCED_ACTIVATION_REGION = "nvlink.token_comm.pre_reduced_activation" +_PRE_REDUCED_ACTIVATION_SF_REGION = "nvlink.token_comm.pre_reduced_activation_sf" + + +def _pre_reduced_workspace_metadata( + device_workspace: Any, + config: Mxfp8KernelConfig, + shared_bytes: int, +) -> tuple[int | None, int]: + """Describe the standalone combine plane, if this kernel has one.""" + if config.fc2_in_kernel_topk_reduce: + return None, 0 + + region = device_workspace.region(_PRE_REDUCED_ACTIVATION_REGION) + if region.buffer_space != "shared": + raise RuntimeError("Rubin pre_reduced_activation must reside in shared workspace") + offset = int(device_workspace.offset(_PRE_REDUCED_ACTIVATION_REGION)) + nbytes = int(device_workspace.nbytes(_PRE_REDUCED_ACTIVATION_REGION)) + wire_bits_per_element = { + "bf16": 16, + "32e4m3xe8m0": 8, + } + try: + element_bits = wire_bits_per_element[config.combine_format] + except KeyError as exc: + raise ValueError(f"unsupported combine wire format {config.combine_format!r}") from exc + wire_bits_per_token = config.top_k * config.hidden * element_bits + if wire_bits_per_token % 8: + raise RuntimeError("combine wire row is not byte aligned") + bytes_per_token = wire_bits_per_token // 8 + expected_bytes = config.max_tokens_per_rank * bytes_per_token + if nbytes != expected_bytes: + raise RuntimeError( + "Rubin pre_reduced_activation size does not match " f"combine_format={config.combine_format!r}: {nbytes} bytes, " f"expected {expected_bytes}" + ) + if offset + nbytes > shared_bytes: + raise RuntimeError("Rubin pre_reduced_activation region exceeds shared workspace") + return offset, bytes_per_token + + +def _pre_reduced_sf_workspace_metadata( + device_workspace: Any, + config: Mxfp8KernelConfig, + shared_bytes: int, +) -> tuple[int | None, int]: + """Describe the standalone quantized-combine scale plane, if present.""" + if config.fc2_in_kernel_topk_reduce or config.combine_format == "bf16": + return None, 0 + + region = device_workspace.region(_PRE_REDUCED_ACTIVATION_SF_REGION) + if region.buffer_space != "shared": + raise RuntimeError("Rubin pre_reduced_activation_sf must reside in shared workspace") + offset = int(device_workspace.offset(_PRE_REDUCED_ACTIVATION_SF_REGION)) + nbytes = int(device_workspace.nbytes(_PRE_REDUCED_ACTIVATION_SF_REGION)) + if nbytes % config.max_tokens_per_rank: + raise RuntimeError("Rubin pre_reduced_activation_sf size is not token aligned") + if offset + nbytes > shared_bytes: + raise RuntimeError("Rubin pre_reduced_activation_sf region exceeds shared workspace") + return offset, nbytes // config.max_tokens_per_rank + + +def prepare_kernel( + forward_config: ForwardConfig, + config: Mxfp8KernelConfig, + device: torch.device, +) -> PreparedMxfp8Kernel: + """Instantiate the kernel and derive exact allocation requirements.""" + + architecture, launch_cluster_count = _prepare_rubin_environment( + device, + config, + context="forward", + ) + import cutlass + + from ..cutedsl_src.kernel_src.rubin.training.mega.fwd_glu import ( + Sm107MegaMoEMxfp8GluKernel, + ) + from ..cutedsl_src.quant_def import CombineFormat + + group_hint = launch_cluster_count if config.group_hint is None else config.group_hint + kernel_kwargs = dict( + mma_tiler_mnk=config.mma_tiler_mnk, + cluster_shape_mnk=config.cluster_shape_mnk, + use_2cta_instrs=config.use_2cta_instrs, + group_hint=group_hint, + token_padding_block=config.token_padding_block, + sf_padding_block=config.sf_padding_block, + load_balance_mode=config.load_balance_mode, + static_expert_shape=( + config.num_experts, + config.fc1_out, + config.hidden, + ), + force_static_sched=config.force_static_sched, + clc_bundle_size=config.clc_bundle_size, + num_sched_stages=config.num_sched_stages, + ab_dtype=cutlass.Float8E4M3FN, + sf_vec_size=config.sf_vec_size, + world_size=config.world_size, + # Runtime rank is carried by SymmetricBufferHost. Keeping this + # descriptor rank-independent allows every EP rank to compile the + # same Rubin kernel. + local_rank=0, + num_topk=config.top_k, + max_tokens_per_rank=config.max_tokens_per_rank, + max_recv_size_per_rank=config.max_recv_size_per_rank, + hidden=config.hidden, + launch_cluster_count=launch_cluster_count, + drop_on_overflow=config.drop_on_overflow, + fc2_in_kernel_topk_reduce=config.fc2_in_kernel_topk_reduce, + token_back_mode=config.token_back_mode, + epi_flag_batch=config.epi_flag_batch, + flag_batch=config.flag_batch, + gate_up_clamp=config.gate_up_clamp, + generate_c=config.generate_c, + combine_format=CombineFormat.parse(config.combine_format), + act_func=config.act_func, + fc2_use_bulk=config.fc2_use_bulk, + fc2_tma_stages=config.fc2_tma_stages, + enable_col_quant=config.enable_col_quant, + col_quant_num_ctas=config.col_quant_num_ctas, + ) + kernel = Sm107MegaMoEMxfp8GluKernel.from_kwargs(**kernel_kwargs) + local_bytes, shared_bytes = kernel.get_workspace_sizes() + local_zero_bytes, shared_zero_bytes = kernel.require_zero_workspace_leading_bytes + for name, zero_bytes, total_bytes in ( + ("local", local_zero_bytes, local_bytes), + ("shared", shared_zero_bytes, shared_bytes), + ): + if zero_bytes < 0 or zero_bytes > total_bytes: + raise RuntimeError(f"Rubin kernel {name} zero prefix {zero_bytes} exceeds " f"workspace size {total_bytes}") + device_workspace = kernel._mega_device_workspace + metadata_region = device_workspace.region(_TOKEN_SRC_METADATA_REGION) + if metadata_region.buffer_space != "shared": + raise RuntimeError("Rubin token_src_metadata must reside in shared workspace") + token_src_metadata_offset = int(device_workspace.offset(_TOKEN_SRC_METADATA_REGION)) + token_src_metadata_bytes = int(device_workspace.nbytes(_TOKEN_SRC_METADATA_REGION)) + if token_src_metadata_offset + token_src_metadata_bytes > shared_bytes: + raise RuntimeError("Rubin token_src_metadata region exceeds shared workspace") + pool_token_capacity = int(kernel.pool_token_capacity) + if token_src_metadata_bytes != pool_token_capacity * 8: + raise RuntimeError("Rubin token_src_metadata must contain one Int64 per pool token") + col_quant_data_rows = pool_token_capacity if config.enable_col_quant else 0 + col_quant_sf_elements = int(kernel.token_comm.worst_case_sf_token_count) * (config.hidden // config.sf_vec_size) if config.enable_col_quant else 0 + if config.enable_col_quant: + col_quant_sizes_region = device_workspace.region(_COL_QUANT_SIZES_REGION) + if col_quant_sizes_region.buffer_space != "local": + raise RuntimeError("Rubin col-quant expert-size snapshot must reside in " "local workspace") + col_quant_sizes_offset = int(device_workspace.offset(_COL_QUANT_SIZES_REGION)) + col_quant_sizes_bytes = int(device_workspace.nbytes(_COL_QUANT_SIZES_REGION)) + expected_sizes_bytes = config.num_experts * torch.int32.itemsize + if col_quant_sizes_bytes != expected_sizes_bytes: + raise RuntimeError("Rubin col-quant expert-size snapshot has " f"{col_quant_sizes_bytes} bytes, expected " f"{expected_sizes_bytes}") + if col_quant_sizes_offset + col_quant_sizes_bytes > local_bytes: + raise RuntimeError("Rubin col-quant expert-size snapshot exceeds local workspace") + else: + col_quant_sizes_offset = None + col_quant_sizes_bytes = 0 + requirements = WorkspaceRequirements.for_mxfp8( + forward_config, + kernel_local_workspace_bytes=local_bytes, + kernel_shared_workspace_bytes=shared_bytes, + col_quant_data_bytes=col_quant_data_rows * config.hidden, + col_quant_sf_bytes=col_quant_sf_elements, + ) + ( + pre_reduced_activation_offset, + pre_reduced_activation_bytes_per_token, + ) = _pre_reduced_workspace_metadata( + device_workspace, + config, + shared_bytes, + ) + ( + pre_reduced_activation_sf_offset, + pre_reduced_activation_sf_bytes_per_token, + ) = _pre_reduced_sf_workspace_metadata( + device_workspace, + config, + shared_bytes, + ) + return PreparedMxfp8Kernel( + config=config, + device=torch.device(device), + architecture=architecture, + kernel=kernel, + launch_cluster_count=launch_cluster_count, + workspace_requirements=requirements, + pool_token_capacity=pool_token_capacity, + col_quant_data_rows=col_quant_data_rows, + col_quant_sf_elements=col_quant_sf_elements, + token_src_metadata_offset=token_src_metadata_offset, + token_src_metadata_bytes=token_src_metadata_bytes, + col_quant_sizes_offset=col_quant_sizes_offset, + col_quant_sizes_bytes=col_quant_sizes_bytes, + pre_reduced_activation_offset=pre_reduced_activation_offset, + pre_reduced_activation_bytes_per_token=(pre_reduced_activation_bytes_per_token), + pre_reduced_activation_sf_offset=pre_reduced_activation_sf_offset, + pre_reduced_activation_sf_bytes_per_token=(pre_reduced_activation_sf_bytes_per_token), + local_workspace_zero_bytes=int(local_zero_bytes), + shared_workspace_zero_bytes=int(shared_zero_bytes), + ) + + +def compile_or_get( + prepared: PreparedMxfp8Kernel, + inputs: Mxfp8LaunchInputs, + resources: PreparedResources, +) -> CompiledMxfp8Kernel: + signature = layout_signature(inputs) + key = ( + *prepared.config.compile_key( + prepared.device, + prepared.architecture, + prepared.launch_cluster_count, + signature, + ), + ) + with _COMPILE_LOCK: + cached = _COMPILE_CACHE.get(key) + if cached is not None: + return cached + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("MXFP8 kernel must be compiled before CUDA graph capture") + + compile_kwargs = build_runtime_kwargs( + inputs, + resources, + ) + compiled = CompiledMxfp8Kernel( + key=key, + callable=_compile_kernel(prepared.kernel, compile_kwargs), + fingerprint=build_kernel_fingerprint( + prepared, + signature, + ), + ) + _COMPILE_CACHE[key] = compiled + return compiled + + +__all__ = [ + "CompiledMxfp8Kernel", + "PreparedMxfp8Kernel", + "compile_or_get", + "prepare_kernel", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile_common.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile_common.py new file mode 100644 index 000000000..844e5d872 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile_common.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared Rubin environment setup for direction-specific MXFP8 compilers.""" + +from __future__ import annotations + +import os +from typing import Any + +import torch + +from ._config import Mxfp8KernelConfig +from ._cutedsl import require_rubin_cutedsl + + +def _prepare_rubin_environment( + device: torch.device, + config: Mxfp8KernelConfig, + *, + context: str, +) -> tuple[tuple[int, int], int]: + require_rubin_cutedsl() + torch.cuda.set_device(device) + architecture = torch.cuda.get_device_capability(device) + if architecture != (10, 7): + raise RuntimeError(f"Rubin MXFP8 {context} preparation requires compute capability " f"(10, 7), got {architecture}") + + configured_architecture = os.environ.get("CUTE_DSL_ARCH") + if configured_architecture is None: + os.environ["CUTE_DSL_ARCH"] = "sm_107a" + elif configured_architecture not in ("sm_107", "sm_107a"): + raise RuntimeError("CUTE_DSL_ARCH must target SM107 for Rubin MXFP8 " f"{context}, got {configured_architecture!r}") + + import cutlass.utils as utils + + launch_cluster_count = int(utils.HardwareInfo().get_max_active_clusters(config.cluster_size)) + if launch_cluster_count <= 0: + raise RuntimeError("hardware occupancy query returned no launchable Rubin clusters") + return architecture, launch_cluster_count + + +def _compile_kernel(kernel: Any, compile_kwargs: dict[str, Any]) -> Any: + """Import CuTeDSL only on a cache miss and compile one callable.""" + + require_rubin_cutedsl() + import cutlass.cute as cute + + return cute.compile(kernel, **compile_kwargs) + + +__all__ = ["_compile_kernel", "_prepare_rubin_environment"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py new file mode 100644 index 000000000..0cbaa8cbf --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py @@ -0,0 +1,208 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Import-light static configuration for the Rubin SM107 MXFP8 kernel.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from ..._contracts import Fc1WeightLayout, ForwardConfig +from ..._tuning import MoeEpTuningConfig +from ._formats import combine_wire_format + + +@dataclass(frozen=True) +class Mxfp8KernelConfig: + """Code-generation constants for one dense EP subgroup kernel.""" + + num_experts: int + world_size: int + local_rank: int + hidden: int + intermediate: int + top_k: int + max_tokens_per_rank: int + apply_topk_in_fc1: bool + fc1_weight_layout: Fc1WeightLayout + gate_up_clamp: float | None + generate_c: bool + max_recv_size_per_rank: int | None = None + drop_on_overflow: bool = True + enable_col_quant: bool = False + col_quant_num_ctas: int = 2368 + mma_tiler_mnk: tuple[int, int, int] = (256, 256, 128) + cluster_shape_mnk: tuple[int, int, int] = (2, 1, 1) + use_2cta_instrs: bool = True + load_balance_mode: str = "static" + force_static_sched: bool = True + clc_bundle_size: int | None = None + num_sched_stages: int | None = None + token_padding_block: int = 128 + sf_padding_block: int = 128 + sf_vec_size: int = 32 + group_hint: int | None = None + token_back_mode: str = "epi_warps" + epi_flag_batch: tuple[int, int] = (1, 1) + flag_batch: int = 1 + fc2_in_kernel_topk_reduce: bool = False + act_func: str = "swiglu" + combine_format: str = "bf16" + fc2_use_bulk: bool = False + fc2_tma_stages: int | None = None + + def __post_init__(self) -> None: + if self.max_recv_size_per_rank is not None and self.max_recv_size_per_rank <= 0: + raise ValueError("max_recv_size_per_rank must be positive") + if self.col_quant_num_ctas <= 0: + raise ValueError("col_quant_num_ctas must be positive") + + @classmethod + def from_operator_config( + cls, + config: ForwardConfig, + *, + tuning: MoeEpTuningConfig | None = None, + ) -> "Mxfp8KernelConfig": + if config.ep_size < 1: + raise ValueError("MXFP8 execution requires a positive EP size") + if config.ep_rank < 0 or config.ep_rank >= config.ep_size: + raise ValueError(f"ep_rank {config.ep_rank} is outside EP size {config.ep_size}") + if config.max_tokens_per_rank is None: + raise ValueError("MXFP8 execution requires max_tokens_per_rank") + token_padding_block = ( + config.token_padding_size + if config.backward_wgrad_mode == "operands" + else 128 if config.generate_c else config.token_padding_size + ) + # When no physical receive capacity is provided, preserve the previous + # default by accounting for worst-case per-expert padding. This is not + # equivalent to rounding the total route count once: every active + # expert owns a separately padded segment. An explicit value is already + # the physical pool size and is therefore used verbatim below. + raw_route_count = config.ep_size * config.max_tokens_per_rank * config.top_k + active_expert_count = min(config.experts_per_rank, raw_route_count) + worst_case_padded_recv_size = ( + active_expert_count + + (raw_route_count - active_expert_count) // token_padding_block + ) * token_padding_block + max_recv_size_per_rank = ( + worst_case_padded_recv_size + if config.max_recv_size_per_rank is None + else config.max_recv_size_per_rank + ) + if max_recv_size_per_rank <= 0: + raise ValueError("max_recv_size_per_rank must be positive") + if tuning is None: + tuning = config.tuning + return cls( + num_experts=config.experts_per_rank, + world_size=config.ep_size, + local_rank=config.ep_rank, + hidden=config.hidden_size, + intermediate=config.intermediate_size, + top_k=config.top_k, + max_tokens_per_rank=config.max_tokens_per_rank, + apply_topk_in_fc1=config.apply_topk_in_fc1, + fc1_weight_layout=config.fc1_weight_layout, + gate_up_clamp=config.gate_up_clamp, + generate_c=config.generate_c, + max_recv_size_per_rank=max_recv_size_per_rank, + drop_on_overflow=config.drop_on_overflow, + combine_format=combine_wire_format(config.combine_format), + enable_col_quant=(config.backward_wgrad_mode == "operands"), + token_padding_block=token_padding_block, + sf_padding_block=config.sf_padding_size, + group_hint=tuning.group_hint, + token_back_mode=tuning.token_back_mode, + epi_flag_batch=tuning.epi_flag_batch, + flag_batch=tuning.token_in_flag_batch, + fc2_in_kernel_topk_reduce=tuning.reduce_topk_in_kernel, + ) + + @property + def fc1_out(self) -> int: + return 2 * self.intermediate + + @property + def cluster_size(self) -> int: + return self.cluster_shape_mnk[0] * self.cluster_shape_mnk[1] + + def tuning_signature( + self, + launch_cluster_count: int, + ) -> tuple[str, tuple[int, int], int, int, bool]: + """Return the effective rank-independent transport/scheduler knobs.""" + + group_hint = launch_cluster_count if self.group_hint is None else self.group_hint + return ( + self.token_back_mode, + self.epi_flag_batch, + self.flag_batch, + group_hint, + self.fc2_in_kernel_topk_reduce, + ) + + def effective_config(self, launch_cluster_count: int) -> dict[str, object]: + """Return the complete JSON-safe compile-time configuration.""" + + effective_group_hint = launch_cluster_count if self.group_hint is None else self.group_hint + return { + "num_experts_per_rank": self.num_experts, + "world_size": self.world_size, + "hidden": self.hidden, + "intermediate": self.intermediate, + "top_k": self.top_k, + "max_tokens_per_rank": self.max_tokens_per_rank, + "max_recv_size_per_rank": self.max_recv_size_per_rank, + "drop_on_overflow": self.drop_on_overflow, + "apply_topk_in_fc1": self.apply_topk_in_fc1, + "fc1_weight_layout": self.fc1_weight_layout.value, + "gate_up_clamp": self.gate_up_clamp, + "generate_c": self.generate_c, + "enable_col_quant": self.enable_col_quant, + "col_quant_num_ctas": self.col_quant_num_ctas, + "combine_format": self.combine_format, + "mma_tiler_mnk": list(self.mma_tiler_mnk), + "cluster_shape_mnk": list(self.cluster_shape_mnk), + "use_2cta_instrs": self.use_2cta_instrs, + "load_balance_mode": self.load_balance_mode, + "force_static_sched": self.force_static_sched, + "clc_bundle_size": self.clc_bundle_size, + "num_sched_stages": self.num_sched_stages, + "token_padding_block": self.token_padding_block, + "sf_padding_block": self.sf_padding_block, + "sf_vec_size": self.sf_vec_size, + "effective_group_hint": effective_group_hint, + "token_back_mode": self.token_back_mode, + "epi_flag_batch": list(self.epi_flag_batch), + "token_in_flag_batch": self.flag_batch, + "fc2_in_kernel_topk_reduce": self.fc2_in_kernel_topk_reduce, + "act_func": self.act_func, + "fc2_use_bulk": self.fc2_use_bulk, + "fc2_tma_stages": self.fc2_tma_stages, + "launch_cluster_count": launch_cluster_count, + } + + def compile_key( + self, + device: torch.device, + architecture: tuple[int, int], + launch_cluster_count: int, + layout_signature: tuple, + ) -> tuple: + """Return a pointer/stream-independent in-process JIT cache key.""" + + canonical_device = torch.device(device) + return ( + self, + canonical_device.index, + architecture, + launch_cluster_count, + layout_signature, + ) + + +__all__ = ["Mxfp8KernelConfig"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_cutedsl.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_cutedsl.py new file mode 100644 index 000000000..e8fa3b48c --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_cutedsl.py @@ -0,0 +1,56 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CUTLASS DSL compatibility gate for Rubin MegaMoE kernels.""" + +from __future__ import annotations + +import importlib.metadata + +RUBIN_CUTEDSL_MIN_VERSION = (4, 8, 0) + + +def _public_cutedsl_version() -> str | None: + """Return public-wheel metadata without importing CUTLASS DSL.""" + + try: + return importlib.metadata.version("nvidia-cutlass-dsl") + except importlib.metadata.PackageNotFoundError: + return None + + +def _parse_version(version: str) -> tuple[int, int, int] | None: + """Parse the numeric release prefix and tolerate prerelease suffixes.""" + + parts = version.split("+", 1)[0].split(".") + parsed = [] + try: + for part in parts[:3]: + digits = "" + for character in part: + if not character.isdigit(): + break + digits += character + if not digits: + return None + parsed.append(int(digits)) + except (TypeError, ValueError): + return None + return tuple(parsed) if len(parsed) == 3 else None + + +def require_rubin_cutedsl() -> None: + """Reject public CUTLASS DSL wheels older than Rubin kernel support.""" + + version = _public_cutedsl_version() + parsed = None if version is None else _parse_version(version) + if parsed is not None and parsed < RUBIN_CUTEDSL_MIN_VERSION: + raise RuntimeError( + "Rubin MegaMoE MXFP8 kernels require " + "nvidia-cutlass-dsl>=4.8.0; found " + f"{version}. Other cuDNN Frontend APIs remain available with " + "the package minimum of 4.5.0" + ) + + +__all__ = ["RUBIN_CUTEDSL_MIN_VERSION", "require_rubin_cutedsl"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_fingerprint.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_fingerprint.py new file mode 100644 index 000000000..1dd1dbf31 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_fingerprint.py @@ -0,0 +1,138 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Stable, machine-readable fingerprints for compiled MXFP8 kernels.""" + +from __future__ import annotations + +import hashlib +import importlib.metadata +import json +import os +from pathlib import Path +from typing import Any + +FINGERPRINT_SCHEMA_VERSION = 1 +_KERNEL_IDENTITY_FIELDS = ( + "kernel_name", + "kernel_source", + "effective_config", + "cutlass_version", + "source_tree_sha256", + "source_git_revision", + "launch_geometry", +) + + +def canonical_json_sha256(value: object) -> str: + payload = json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def kernel_identity_sha256(fingerprint: dict[str, Any]) -> str: + """Hash fields that must match across AOT and in-process JIT paths.""" + + return canonical_json_sha256({field: fingerprint.get(field) for field in _KERNEL_IDENTITY_FIELDS}) + + +def source_tree_sha256(root: Path) -> str: + """Hash Python source content and relative paths in deterministic order.""" + + resolved = root.expanduser().resolve() + if not resolved.is_dir(): + raise RuntimeError(f"kernel source tree does not exist: {resolved}") + digest = hashlib.sha256() + sources = sorted(path for path in resolved.rglob("*.py") if path.is_file()) + if not sources: + raise RuntimeError(f"kernel source tree contains no Python files: {resolved}") + for path in sources: + relative = path.relative_to(resolved).as_posix().encode("utf-8") + digest.update(len(relative).to_bytes(8, "little")) + digest.update(relative) + payload = path.read_bytes() + digest.update(len(payload).to_bytes(8, "little")) + digest.update(payload) + return digest.hexdigest() + + +def _cutlass_version() -> str: + for distribution in ( + "nvidia-cutlass-dsl", + "nvidia-cutlass-dsl-internal", + ): + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + continue + import cutlass + + return str(getattr(cutlass, "__version__", "unknown")) + + +def _json_layout_signature(signature: tuple) -> list[object]: + return [ + ( + None + if entry is None + else { + "shape": list(entry[0]), + "stride": list(entry[1]), + "dtype": str(entry[2]), + } + ) + for entry in signature + ] + + +def build_kernel_fingerprint( + prepared: Any, + layout_signature: tuple, + *, + compiled_binary_sha256: str | None = None, +) -> dict[str, Any]: + """Describe exactly what was compiled and how the main kernel launches.""" + + kernel = prepared.kernel + source_root = Path(__file__).resolve().parents[1] / "cutedsl_src" + effective_config = prepared.config.effective_config(prepared.launch_cluster_count) + launch_geometry = { + "grid": [ + prepared.config.cluster_shape_mnk[0], + prepared.config.cluster_shape_mnk[1], + prepared.launch_cluster_count, + ], + "block": [int(kernel.threads_per_cta), 1, 1], + "cluster": list(prepared.config.cluster_shape_mnk), + "min_blocks_per_mp": int(getattr(kernel, "occupancy", 1)), + "dynamic_shared_memory_bytes": int(getattr(kernel, "smem_capacity", 0)), + } + layout = _json_layout_signature(layout_signature) + fingerprint = { + "schema_version": FINGERPRINT_SCHEMA_VERSION, + "kernel_name": str(kernel.name()), + "kernel_source": "vendored-training-mega", + "effective_config": effective_config, + "cutlass_version": _cutlass_version(), + "source_tree_sha256": source_tree_sha256(source_root), + "source_git_revision": os.environ.get("MOE_EP_SOURCE_GIT_REVISION"), + "launch_geometry": launch_geometry, + "layout_signature_sha256": canonical_json_sha256(layout), + "compiled_binary_sha256": compiled_binary_sha256, + } + fingerprint["kernel_identity_sha256"] = kernel_identity_sha256(fingerprint) + fingerprint["fingerprint_sha256"] = canonical_json_sha256(fingerprint) + return fingerprint + + +__all__ = [ + "FINGERPRINT_SCHEMA_VERSION", + "build_kernel_fingerprint", + "canonical_json_sha256", + "kernel_identity_sha256", + "source_tree_sha256", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_formats.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_formats.py new file mode 100644 index 000000000..6529ee4f3 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_formats.py @@ -0,0 +1,22 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Public MoE format names to Rubin combine-wire encodings.""" + +from __future__ import annotations + +from ..._types import MoeFormat, parse_format + +_COMBINE_WIRE_FORMATS = { + MoeFormat.BF16: "bf16", + MoeFormat.MXFP8: "32e4m3xe8m0", +} + + +def combine_wire_format(value: MoeFormat | str) -> str: + """Return the kernel encoding for one public combine format.""" + + return _COMBINE_WIRE_FORMATS[parse_format(value)] + + +__all__ = ["combine_wire_format"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_launch.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_launch.py new file mode 100644 index 000000000..a5dfbb032 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_launch.py @@ -0,0 +1,158 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CuTe tensor conversion and current-stream MXFP8 launch.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from .._plan import PreparedResources +from ._adapter import Mxfp8LaunchInputs + + +def _to_cute( + tensor: torch.Tensor, + assumed_align: int = 16, + *, + dynamic_layout: bool = True, +): + import cutlass.torch as cutlass_torch + + cute_tensor = cutlass_torch.from_dlpack( + tensor, + assumed_align=assumed_align, + enable_tvm_ffi=True, + ) + if not dynamic_layout: + return cute_tensor + return cute_tensor.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(tensor)) + + +def _to_cute_ptr(tensor: torch.Tensor, assumed_align: int = 128): + """Build the opaque byte-pointer ABI used by Rubin workspaces.""" + + import cutlass + from cutlass.cute.runtime import make_ptr + from cutlass.cute.typing import AddressSpace + + address = int(tensor.data_ptr()) + if address % assumed_align: + raise ValueError(f"Rubin workspace address {address:#x} is not " f"{assumed_align}-byte aligned") + return make_ptr( + cutlass.Uint8, + address, + AddressSpace.gmem, + assumed_align=assumed_align, + ) + + +def build_runtime_kwargs( + inputs: Mxfp8LaunchInputs, + resources: PreparedResources, +) -> dict[str, Any]: + import cuda.bindings.driver as cuda + + stream = resources.runtime.current_stream() + weights = inputs.weights + kwargs = { + "activation": _to_cute(inputs.activation), + "activation_sf": _to_cute(inputs.activation_sf), + "topk_indices": _to_cute(inputs.topk_indices), + "topk_scores": _to_cute(inputs.topk_scores, assumed_align=4), + "fc1_weight": _to_cute(weights.fc1_weight), + "fc1_weight_sf": _to_cute(weights.fc1_weight_sf), + "fc2_weight": _to_cute(weights.fc2_weight), + "fc2_weight_sf": _to_cute(weights.fc2_weight_sf), + "fc1_c": (None if inputs.fc1_c is None else _to_cute(inputs.fc1_c, dynamic_layout=False)), + "output_activation": _to_cute(inputs.output_data), + "col_quant_data": ( + None + if inputs.col_quant_data is None + else _to_cute( + inputs.col_quant_data, + assumed_align=128, + dynamic_layout=False, + ) + ), + "col_quant_sf": ( + None + if inputs.col_quant_sf is None + else _to_cute( + inputs.col_quant_sf, + dynamic_layout=False, + ) + ), + "overflow_flag": _to_cute( + inputs.overflow_flag, + assumed_align=4, + dynamic_layout=False, + ), + "local_workspace": _to_cute_ptr(inputs.local_workspace), + "shared_workspace": _to_cute_ptr(inputs.shared_workspace), + "peer_rank_ptr_mapper_host": (resources.workspace.peer_mapping.to_sym_buffer_host()), + "stream": cuda.CUstream(stream.cuda_stream), + } + return kwargs + + +def layout_signature(inputs: Mxfp8LaunchInputs) -> tuple: + tensors = ( + inputs.activation, + inputs.activation_sf, + inputs.topk_indices, + inputs.topk_scores, + inputs.weights.fc1_weight, + inputs.weights.fc1_weight_sf, + inputs.weights.fc2_weight, + inputs.weights.fc2_weight_sf, + inputs.fc1_c, + inputs.col_quant_data, + inputs.col_quant_sf, + inputs.output_data, + inputs.overflow_flag, + inputs.local_workspace, + inputs.shared_workspace, + ) + return tuple(None if tensor is None else (tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype) for tensor in tensors) + + +def _check_overflow(overflow_flag: torch.Tensor) -> None: + message = "Rubin MegaMoE receive route-pool overflow; the output is invalid for " "this routing distribution" + assert_async = getattr(torch, "_assert_async", None) + if assert_async is not None: + assert_async(overflow_flag == 0, message) + return + if torch.cuda.is_current_stream_capturing(): + raise NotImplementedError("CUDA graph capture requires torch._assert_async to surface " "Rubin MegaMoE overflow") + # Compatibility fallback for PyTorch builds without a device-side assert. + value = int(overflow_flag.item()) + if value != 0: + raise RuntimeError(f"{message} (overflow_flag={value})") + + +def launch_forward( + compiled, + inputs: Mxfp8LaunchInputs, + resources: PreparedResources, +) -> torch.Tensor: + runtime_kwargs = build_runtime_kwargs(inputs, resources) + compiled.callable(**runtime_kwargs) + _check_overflow(inputs.overflow_flag) + + output_data = torch.empty( + (inputs.token_count, inputs.output_data.shape[1]), + dtype=inputs.output_data.dtype, + device=inputs.output_data.device, + ) + output_data.copy_(inputs.output_data[: inputs.token_count]) + return output_data + + +__all__ = [ + "build_runtime_kwargs", + "launch_forward", + "layout_signature", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py new file mode 100644 index 000000000..02de65525 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py @@ -0,0 +1,404 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Ordinary/capturable stateless launch path over private lane resources.""" + +from __future__ import annotations + +import torch + +from ..._math import round_up +from ..._types import ( + BlockScaledTensor, + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpTrainingBackwardOutputs, + MoeEpTrainingForwardOutputs, + MoeEpTrainingWgradOperands, + MoeTensor, +) +from .._runtime import _runtime_debug +from .._workspace import padded_mxfp8_scale_columns +from ._adapter import ( + Mxfp8LaunchInputs, + _typed_view, +) +from ._backward_compile import ( + Mxfp8BackwardLaunchInputs, + build_backward_runtime_kwargs, + compile_backward_or_get, +) +from ._compile import compile_or_get +from ._launch import build_runtime_kwargs +from ._training_resources import ( + Mxfp8TrainingExecutionViews, + Mxfp8TrainingState, +) +from ._training_weights import ( + backward_native_to_kernel, + forward_native_to_kernel, +) +from ._training_wgrad import assemble_training_wgrad_operands + + +def _zero_pre_reduced(inputs, prepared) -> None: + capacity = prepared.config.max_tokens_per_rank + offset = prepared.pre_reduced_activation_offset + bytes_per_token = prepared.pre_reduced_activation_bytes_per_token + if offset is not None and bytes_per_token: + inputs.shared_workspace.narrow( + 0, + offset, + capacity * bytes_per_token, + ).zero_() + sf_offset = prepared.pre_reduced_activation_sf_offset + sf_bytes_per_token = prepared.pre_reduced_activation_sf_bytes_per_token + if sf_offset is not None and sf_bytes_per_token: + inputs.shared_workspace.narrow( + 0, + sf_offset, + capacity * sf_bytes_per_token, + ).zero_() + + +def _activation_views( + execution: Mxfp8TrainingExecutionViews, + *, + backward: bool, + capacity: int, + hidden: int, +) -> tuple[torch.Tensor, torch.Tensor]: + workspace = execution.backward.workspace if backward else execution.forward.workspace + return ( + _typed_view( + workspace.symmetric["activation_data"], + torch.float8_e4m3fn, + (capacity, hidden), + ), + _typed_view( + workspace.symmetric["activation_scale"], + torch.float8_e8m0fnu, + (round_up(capacity, 128), padded_mxfp8_scale_columns(hidden)), + ), + ) + + +def _stage_input( + state: Mxfp8TrainingState, + value: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + activation_data: torch.Tensor, + activation_sf: torch.Tensor, + routing_topk_idx: torch.Tensor, + routing_topk_weights: torch.Tensor, +) -> None: + """Stage into private symmetric memory, bypassing quantization for MXFP8.""" + + if not isinstance(value, BlockScaledTensor): + state.stager.stage( + value, + topk_idx, + topk_weights, + activation_data, + activation_sf, + routing_topk_idx, + routing_topk_weights, + ) + return + + token_count = int(value.logical_shape[0]) + scale_columns = int(value.scale.shape[1]) + data_in_place = value.data.data_ptr() == activation_data.data_ptr() + scale_in_place = value.scale.data_ptr() == activation_sf.data_ptr() + if data_in_place != scale_in_place: + raise ValueError( + "MXFP8 training input data and scale must either both use the " + "lane's symmetric buffers or neither use them" + ) + if not data_in_place: + activation_data.zero_() + activation_sf.zero_() + routing_topk_idx.fill_(-1) + routing_topk_weights.zero_() + if token_count == 0: + return + if not data_in_place: + activation_data[:token_count].copy_(value.data) + activation_sf[:token_count, :scale_columns].copy_(value.scale) + routing_topk_idx[:token_count].copy_(topk_idx) + routing_topk_weights[:token_count].copy_(topk_weights) + + +def _write_expert_offsets( + execution: Mxfp8TrainingExecutionViews, + padding: int, + counts: torch.Tensor, + offsets: torch.Tensor, +) -> None: + snapshot = execution.forward_expert_size_snapshot + counts.copy_(snapshot) + torch.add(counts, padding - 1, out=offsets) + torch.div(offsets, padding, rounding_mode="floor", out=offsets) + offsets.mul_(padding) + torch.cumsum(offsets, dim=0, out=offsets) + + +def launch_training_forward( + state: Mxfp8TrainingState, + execution: Mxfp8TrainingExecutionViews, + activation: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + weights: MoeEpNativeForwardWeights, + out: MoeEpTrainingForwardOutputs, +) -> torch.Tensor: + """Launch one stateless forward over caller-owned outputs.""" + + prepared = state.forward_prepared + config = prepared.config + capacity = config.max_tokens_per_rank + scratch = execution.scratch + token_count = int(activation.logical_shape[0] if isinstance(activation, BlockScaledTensor) else activation.shape[0]) + _runtime_debug( + "training-forward.begin", + lane=scratch.index, + token_count=token_count, + ) + activation_data, activation_sf = _activation_views( + execution, + backward=False, + capacity=capacity, + hidden=config.hidden, + ) + _runtime_debug("training-forward.stage.begin", lane=scratch.index) + _stage_input( + state, + activation, + topk_idx, + topk_weights, + activation_data, + activation_sf, + scratch.routing_topk_idx, + scratch.routing_topk_weights, + ) + _runtime_debug("training-forward.stage.end", lane=scratch.index) + + assert out.output is not None + assert out.fc1_a is not None + assert out.fc1_sfa is not None + assert out.valid_route_counts is not None + assert out.expert_offsets is not None + fc1_preact = out.fc1_preact + col_quant_data = out.fc1_a.transpose(0, 1) + col_quant_sf = out.fc1_sfa.view(torch.uint8).reshape(-1) + expected_elements = int(prepared.col_quant_sf_elements) + if col_quant_sf.numel() != expected_elements: + raise ValueError("out.fc1_sfa storage does not match the forward producer ABI: " f"{col_quant_sf.numel()} != {expected_elements}") + valid_route_counts = out.valid_route_counts + expert_offsets = out.expert_offsets + if out.output.data_ptr() != scratch.forward_output.data_ptr(): + raise ValueError( + "out.output must be the lane's symmetric output buffer from " + "training_symmetric_buffers()" + ) + + out.output.zero_() + scratch.forward_overflow.zero_() + col_quant_data.zero_() + # E8M0 byte 127 encodes scale 1.0. The producer only overwrites active + # expert segments, so the unused grouped-WGrad capacity must stay neutral. + col_quant_sf.fill_(127) + _runtime_debug("training-forward.reset.end", lane=scratch.index) + + workspace = execution.forward.workspace + inputs = Mxfp8LaunchInputs( + activation=activation_data, + activation_sf=activation_sf, + topk_indices=scratch.routing_topk_idx, + topk_scores=scratch.routing_topk_weights, + weights=forward_native_to_kernel(weights), + fc1_c=fc1_preact, + output_data=out.output, + col_quant_data=col_quant_data, + col_quant_sf=col_quant_sf, + overflow_flag=scratch.forward_overflow, + local_workspace=workspace.local["kernel_local_workspace"], + shared_workspace=workspace.symmetric["kernel_shared_workspace"], + token_count=token_count, + ) + _zero_pre_reduced(inputs, prepared) + _runtime_debug("training-forward.compile.begin", lane=scratch.index) + compiled = compile_or_get( + prepared, + inputs, + execution.forward, + ) + _runtime_debug("training-forward.compile.end", lane=scratch.index) + _runtime_debug("training-forward.launch.begin", lane=scratch.index) + compiled.callable(**build_runtime_kwargs(inputs, execution.forward)) + _runtime_debug("training-forward.launch.end", lane=scratch.index) + _runtime_debug("training-forward.offsets.begin", lane=scratch.index) + _write_expert_offsets( + execution, + config.token_padding_block, + valid_route_counts, + expert_offsets, + ) + _runtime_debug("training-forward.offsets.end", lane=scratch.index) + state.apply_overflow(lane=scratch.index, phase="forward") + + output = out.output[:token_count] + _runtime_debug("training-forward.end", lane=scratch.index) + return output + + +def launch_training_backward( + state: Mxfp8TrainingState, + execution: Mxfp8TrainingExecutionViews, + grad_output: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + weights: MoeEpNativeBackwardWeights, + fc1_preact: torch.Tensor, + fc1_a: torch.Tensor | None, + fc1_sfa: torch.Tensor | None, + valid_route_counts: torch.Tensor | None, + expert_offsets: torch.Tensor | None, + out: MoeEpTrainingBackwardOutputs, +) -> tuple[ + torch.Tensor, + torch.Tensor, + MoeEpTrainingWgradOperands, +]: + """Launch one stateless backward using explicit caller-owned saved state.""" + + prepared = state.backward_prepared + config = prepared.config + capacity = config.max_tokens_per_rank + scratch = execution.scratch + token_count = int(grad_output.logical_shape[0] if isinstance(grad_output, BlockScaledTensor) else grad_output.shape[0]) + _runtime_debug( + "training-backward.begin", + lane=scratch.index, + token_count=token_count, + ) + activation_data, activation_sf = _activation_views( + execution, + backward=True, + capacity=capacity, + hidden=config.hidden, + ) + _runtime_debug("training-backward.stage.begin", lane=scratch.index) + _stage_input( + state, + grad_output, + topk_idx, + topk_weights, + activation_data, + activation_sf, + scratch.routing_topk_idx, + scratch.routing_topk_weights, + ) + _runtime_debug("training-backward.stage.end", lane=scratch.index) + + assert fc1_a is not None + assert fc1_sfa is not None + assert valid_route_counts is not None + assert expert_offsets is not None + assert out.grad_activation is not None + assert out.dprob is not None + assert out.fc1_b is not None + assert out.fc1_sfb is not None + assert out.fc2_a is not None + assert out.fc2_sfa is not None + assert out.fc2_b is not None + assert out.fc2_sfb is not None + fc1_recompute = out.fc2_a.transpose(0, 1) + fc1_recompute_sf = out.fc2_sfa + fc1_col_output = out.fc1_b + fc1_col_output_sf = out.fc1_sfb + grad_y2 = out.fc2_b + grad_y2_sf = out.fc2_sfb.view(torch.uint8).reshape(-1) + if out.grad_activation.data_ptr() != scratch.backward_output.data_ptr(): + raise ValueError( + "out.grad_activation must be the lane's symmetric grad_activation " + "buffer from training_symmetric_buffers()" + ) + if out.dprob.data_ptr() != scratch.dprob.data_ptr(): + raise ValueError( + "out.dprob must be the lane's symmetric dprob buffer from " + "training_symmetric_buffers()" + ) + + out.grad_activation.zero_() + scratch.backward_overflow.zero_() + out.dprob.zero_() + fc1_recompute.zero_() + fc1_recompute_sf.view(torch.uint8).fill_(127) + fc1_col_output.zero_() + fc1_col_output_sf.view(torch.uint8).fill_(127) + grad_y2.zero_() + grad_y2_sf.fill_(127) + _runtime_debug("training-backward.reset.end", lane=scratch.index) + + workspace = execution.backward.workspace + kernel_weights = backward_native_to_kernel(weights) + inputs = Mxfp8BackwardLaunchInputs( + grad_out=activation_data, + grad_out_sf=activation_sf, + topk_idx=scratch.routing_topk_idx, + topk_weights=scratch.routing_topk_weights, + fc1_weight=kernel_weights.fc1_weight, + fc1_weight_sf=kernel_weights.fc1_weight_sf, + fc2_weight=kernel_weights.fc2_weight, + fc2_weight_sf=kernel_weights.fc2_weight_sf, + beta=state.beta, + fc1_preact=fc1_preact, + output_activation=out.grad_activation, + overflow_flag=scratch.backward_overflow, + dprob=out.dprob, + fc1_recompute=fc1_recompute, + fc1_recompute_sf=fc1_recompute_sf, + fc1_col_output=fc1_col_output, + fc1_col_output_sf=fc1_col_output_sf, + grad_y2=grad_y2, + grad_y2_sf=grad_y2_sf, + local_workspace=workspace.local["kernel_local_workspace"], + shared_workspace=workspace.symmetric["kernel_shared_workspace"], + token_count=token_count, + ) + _zero_pre_reduced(inputs, prepared) + _runtime_debug("training-backward.compile.begin", lane=scratch.index) + compiled = compile_backward_or_get( + prepared, + inputs, + execution.backward, + ) + _runtime_debug("training-backward.compile.end", lane=scratch.index) + _runtime_debug("training-backward.launch.begin", lane=scratch.index) + compiled.callable(**build_backward_runtime_kwargs(inputs, execution.backward)) + _runtime_debug("training-backward.launch.end", lane=scratch.index) + state.apply_overflow(lane=scratch.index, phase="backward") + + grad_activation = out.grad_activation[:token_count] + + dprob = out.dprob[:token_count] + + operands = assemble_training_wgrad_operands( + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + expert_offsets=expert_offsets, + valid_route_counts=valid_route_counts, + backward=out, + ) + _runtime_debug("training-backward.end", lane=scratch.index) + return grad_activation, dprob, operands + + +__all__ = [ + "launch_training_backward", + "launch_training_forward", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py new file mode 100644 index 000000000..87d5a3e83 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py @@ -0,0 +1,1026 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Private per-lane state for graph-capable stateless MXFP8 training.""" + +from __future__ import annotations + +import hashlib +import math +import threading +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Any, Mapping, Optional + +import torch +import torch.distributed as dist + +from ..._contracts import ForwardConfig +from ..._math import round_up +from ..._types import MoeEpNativeWeightLayout +from .._comm import SymmetricMemoryProvider +from .._plan import PreparedResources +from .._runtime import ( + RuntimeHandle, + RuntimeManager, + _RuntimeWatchdog, + _runtime_debug, + get_runtime_manager, +) +from .._workspace import ( + BufferRegion, + LocalMemoryProvider, + WorkspaceOwner, + WorkspaceRequirements, + WorkspaceViews, + padded_mxfp8_scale_columns, +) +from ._adapter import _typed_view +from ._backward_compile import PreparedMxfp8BackwardKernel +from ._compile import PreparedMxfp8Kernel +from ._fingerprint import canonical_json_sha256, source_tree_sha256 +from ._training_stage import Mxfp8TrainingStager + +_DATA_DTYPE = torch.float8_e4m3fn +_SCALE_DTYPE = torch.float8_e8m0fnu + +_ROUTING_SYMMETRIC = frozenset({"topk_weights"}) +_ROUTING_LOCAL = frozenset({"topk_idx"}) +_CALLER_OWNED_FORWARD_LOCAL = frozenset({"col_quant_data", "col_quant_sf"}) +_FORWARD_PRIVATE_SYMMETRIC = frozenset({"output_data", *_ROUTING_SYMMETRIC}) +_FORWARD_PRIVATE_LOCAL = frozenset({"overflow_flag", *_CALLER_OWNED_FORWARD_LOCAL, *_ROUTING_LOCAL}) +_BACKWARD_PRIVATE_SYMMETRIC = frozenset({"output_data", "backward_dprob", *_ROUTING_SYMMETRIC}) +_BACKWARD_PRIVATE_LOCAL = frozenset({"overflow_flag", "backward_aux_data", "backward_aux_scale", *_ROUTING_LOCAL}) + + +def _lane_name( + lane: int, + phase: str, + space: str, + name: str, +) -> str: + return f"lane.{lane}.{phase}.{space}.{name}" + + +def _lane_fallback_name( + lane: int, + space: str, + name: str, +) -> str: + return f"lane.{lane}.fallback.{space}.{name}" + + +def _clone_region(name: str, region: BufferRegion) -> BufferRegion: + return BufferRegion( + name=name, + nbytes=region.nbytes, + alignment=region.alignment, + ) + + +def _region_map( + requirements: WorkspaceRequirements, + space: str, +) -> dict[str, BufferRegion]: + regions = requirements.symmetric_regions if space == "symmetric" else requirements.local_regions + return {region.name: region for region in regions} + + +def _add_lane_regions( + output: list[BufferRegion], + requirements: WorkspaceRequirements, + *, + lane: int, + phase: str, + space: str, + excluded_names: frozenset[str], +) -> None: + regions = requirements.symmetric_regions if space == "symmetric" else requirements.local_regions + for region in regions: + if region.name in excluded_names: + continue + output.append( + _clone_region( + _lane_name(lane, phase, space, region.name), + region, + ) + ) + + +def build_training_workspace_requirements( + config: ForwardConfig, + forward: PreparedMxfp8Kernel, + backward: PreparedMxfp8BackwardKernel, + *, + lane_count: int, +) -> WorkspaceRequirements: + """Build one deterministic root layout for private execution lanes.""" + + if isinstance(lane_count, bool) or not isinstance(lane_count, int) or lane_count <= 0: + raise ValueError(f"lane_count must be a positive integer, got {lane_count!r}") + if not config.generate_c: + raise ValueError("training preparation requires generate_c=True") + if forward.pool_token_capacity != backward.pool_token_capacity: + raise ValueError("forward/backward pool capacities must match, got " f"{forward.pool_token_capacity} and " f"{backward.pool_token_capacity}") + + forward_requirements = forward.workspace_requirements + backward_requirements = backward.workspace_requirements + symmetric_regions: list[BufferRegion] = [] + local_regions: list[BufferRegion] = [] + + for lane in range(lane_count): + local_regions.extend( + ( + BufferRegion( + _lane_name( + lane, + "finalizer", + "local", + "global_overflow", + ), + torch.int32.itemsize, + 16, + ), + BufferRegion( + _lane_name( + lane, + "finalizer", + "local", + "overflow_ok", + ), + torch.bool.itemsize, + 16, + ), + ) + ) + _add_lane_regions( + symmetric_regions, + forward_requirements, + lane=lane, + phase="forward", + space="symmetric", + excluded_names=_FORWARD_PRIVATE_SYMMETRIC, + ) + _add_lane_regions( + local_regions, + forward_requirements, + lane=lane, + phase="forward", + space="local", + excluded_names=_FORWARD_PRIVATE_LOCAL, + ) + _add_lane_regions( + symmetric_regions, + backward_requirements, + lane=lane, + phase="backward", + space="symmetric", + excluded_names=_BACKWARD_PRIVATE_SYMMETRIC, + ) + _add_lane_regions( + local_regions, + backward_requirements, + lane=lane, + phase="backward", + space="local", + excluded_names=_BACKWARD_PRIVATE_LOCAL, + ) + + forward_symmetric = _region_map(forward_requirements, "symmetric") + forward_local = _region_map(forward_requirements, "local") + backward_symmetric = _region_map(backward_requirements, "symmetric") + backward_local = _region_map(backward_requirements, "local") + fc1_c_shape = tuple(int(extent) for extent in forward.kernel.get_aux_output_shapes()["fc1_c"]) + backward_fc1_preact_shape = tuple(int(extent) for extent in backward.kernel.get_fc1_preact_shape()) + if fc1_c_shape != backward_fc1_preact_shape: + raise ValueError("forward fc1_c and backward fc1_preact shapes differ: " f"{fc1_c_shape} != {backward_fc1_preact_shape}") + for lane in range(lane_count): + for name in sorted(_FORWARD_PRIVATE_SYMMETRIC): + if name in _ROUTING_SYMMETRIC: + continue + symmetric_regions.append( + _clone_region( + _lane_name(lane, "forward", "symmetric", name), + forward_symmetric[name], + ) + ) + for name in sorted(_BACKWARD_PRIVATE_SYMMETRIC): + if name in _ROUTING_SYMMETRIC: + continue + symmetric_regions.append( + _clone_region( + _lane_name(lane, "backward", "symmetric", name), + backward_symmetric[name], + ) + ) + for name in sorted(_FORWARD_PRIVATE_LOCAL): + if name in _ROUTING_LOCAL or name in _CALLER_OWNED_FORWARD_LOCAL: + continue + region = forward_local.get(name) + if region is not None: + local_regions.append( + _clone_region( + _lane_name(lane, "forward", "local", name), + region, + ) + ) + for name in sorted(_BACKWARD_PRIVATE_LOCAL): + if name in _ROUTING_LOCAL: + continue + local_regions.append( + _clone_region( + _lane_name(lane, "backward", "local", name), + backward_local[name], + ) + ) + local_regions.append( + BufferRegion( + _lane_fallback_name(lane, "local", "routing_topk_idx"), + int(config.max_tokens_per_rank) * config.top_k * torch.int32.itemsize, + alignment=16, + ) + ) + symmetric_regions.append( + BufferRegion( + _lane_fallback_name(lane, "symmetric", "routing_topk_weights"), + int(config.max_tokens_per_rank) * config.top_k * torch.float32.itemsize, + alignment=16, + ) + ) + return WorkspaceRequirements( + max_tokens_per_rank=int(config.max_tokens_per_rank), + symmetric_regions=tuple(symmetric_regions), + local_regions=tuple(local_regions), + ) + + +def _harmonize_symmetric_regions( + requirements: WorkspaceRequirements, + runtime: RuntimeHandle, + device: torch.device, +) -> WorkspaceRequirements: + """Make every peer-visible region size and offset identical on all ranks.""" + + if runtime.world_size <= 1: + return requirements + + regions = requirements.symmetric_regions + count = torch.tensor([len(regions)], dtype=torch.int64, device=device) + minimum_count = count.clone() + maximum_count = count.clone() + dist.all_reduce(minimum_count, op=dist.ReduceOp.MIN, group=runtime.group) + dist.all_reduce(maximum_count, op=dist.ReduceOp.MAX, group=runtime.group) + if int(minimum_count.item()) != int(maximum_count.item()): + raise RuntimeError("symmetric workspace region counts differ across EP ranks: " f"min={int(minimum_count.item())}, max={int(maximum_count.item())}") + + metadata = "\0".join(f"{region.name}:{region.alignment}" for region in regions).encode() + signature_value = int.from_bytes( + hashlib.blake2b(metadata, digest_size=8).digest(), + "little", + ) & ((1 << 63) - 1) + signature = torch.tensor( + [signature_value], + dtype=torch.int64, + device=device, + ) + minimum_signature = signature.clone() + maximum_signature = signature.clone() + dist.all_reduce( + minimum_signature, + op=dist.ReduceOp.MIN, + group=runtime.group, + ) + dist.all_reduce( + maximum_signature, + op=dist.ReduceOp.MAX, + group=runtime.group, + ) + if int(minimum_signature.item()) != int(maximum_signature.item()): + raise RuntimeError( + "symmetric workspace region names, order, or alignments differ " + "across EP ranks: " + f"local_signature={signature_value}, " + "local_regions=" + f"{tuple((region.name, region.alignment) for region in regions)}" + ) + + local_sizes = torch.tensor( + [region.nbytes for region in regions], + dtype=torch.int64, + device=device, + ) + maximum_sizes = local_sizes.clone() + dist.all_reduce(maximum_sizes, op=dist.ReduceOp.MAX, group=runtime.group) + harmonized_sizes = tuple(int(value) for value in maximum_sizes.cpu().tolist()) + changes = tuple( + f"{region.name}:{region.nbytes}->{harmonized_size}" for region, harmonized_size in zip(regions, harmonized_sizes) if region.nbytes != harmonized_size + ) + _runtime_debug( + "training-state.symmetric-layout-harmonized", + region_count=len(regions), + changed_regions=changes, + ) + if not changes: + return requirements + + return WorkspaceRequirements( + max_tokens_per_rank=requirements.max_tokens_per_rank, + symmetric_regions=tuple( + BufferRegion( + region.name, + harmonized_size, + alignment=region.alignment, + ) + for region, harmonized_size in zip(regions, harmonized_sizes) + ), + local_regions=requirements.local_regions, + ) + + +def _workspace_abi(requirements: WorkspaceRequirements) -> dict[str, object]: + def regions(values) -> list[dict[str, object]]: + return [ + { + "name": region.name, + "nbytes": int(region.nbytes), + "alignment": int(region.alignment), + } + for region in values + ] + + return { + "max_tokens_per_rank": requirements.max_tokens_per_rank, + "symmetric_regions": regions(requirements.symmetric_regions), + "local_regions": regions(requirements.local_regions), + } + + +def _prepared_kernel_abi(prepared) -> dict[str, object]: + kernel = prepared.kernel + return { + "name": str(kernel.name()), + "architecture": list(prepared.architecture), + "effective_config": prepared.config.effective_config(prepared.launch_cluster_count), + "launch": { + "cluster_count": int(prepared.launch_cluster_count), + "threads_per_cta": int(kernel.threads_per_cta), + "occupancy": int(kernel.occupancy), + "smem_capacity": int(kernel.smem_capacity), + }, + "workspace": _workspace_abi(prepared.workspace_requirements), + "pool_token_capacity": int(prepared.pool_token_capacity), + } + + +def _build_training_abi_facts( + config: ForwardConfig, + forward: PreparedMxfp8Kernel, + backward: PreparedMxfp8BackwardKernel, + requirements: WorkspaceRequirements, + *, + lane_count: int, + source_tree_digest: str | None = None, +) -> dict[str, object]: + """Return rank-independent JSON-safe facts for the stateless training ABI.""" + + if source_tree_digest is None: + source_root = Path(__file__).resolve().parents[1] / "cutedsl_src" + source_tree_digest = source_tree_sha256(source_root) + return { + "schema_version": 2, + "source_tree_sha256": source_tree_digest, + "ep": { + "size": int(config.ep_size), + "global_ranks": list(config.ep_global_ranks), + }, + "geometry": { + "num_experts": int(config.num_experts), + "experts_per_rank": int(config.experts_per_rank), + "hidden": int(config.hidden_size), + "intermediate": int(config.intermediate_size), + "top_k": int(config.top_k), + "max_tokens_per_rank": int(config.max_tokens_per_rank), + "max_recv_size_per_rank": int(forward.config.max_recv_size_per_rank), + }, + "policy": { + "drop_on_overflow": bool(config.drop_on_overflow), + "combine_format": config.combine_format, + "output_format": config.output_format, + "apply_topk_in_fc1": bool(config.apply_topk_in_fc1), + "fc1_weight_layout": config.fc1_weight_layout.value, + "gate_up_clamp": config.gate_up_clamp, + }, + "resources": { + "lane_count": int(lane_count), + "workspace": _workspace_abi(requirements), + }, + "native_weight_layouts": [layout.value for layout in MoeEpNativeWeightLayout], + "forward_kernel": _prepared_kernel_abi(forward), + "backward_kernel": _prepared_kernel_abi(backward), + } + + +def _verify_training_abi_across_ranks( + facts: dict[str, object], + runtime: RuntimeHandle, + device: torch.device, +) -> str: + """Collectively reject rank-divergent training ABI before allocation.""" + + digest = canonical_json_sha256(facts) + if runtime.world_size <= 1: + return digest + digest_value = int(digest[:16], 16) & ((1 << 63) - 1) + minimum = torch.tensor([digest_value], dtype=torch.int64, device=device) + maximum = minimum.clone() + dist.all_reduce(minimum, op=dist.ReduceOp.MIN, group=runtime.group) + dist.all_reduce(maximum, op=dist.ReduceOp.MAX, group=runtime.group) + if int(minimum.item()) == int(maximum.item()): + return digest + + rank_digests: list[Any] = [None] * runtime.world_size + dist.all_gather_object(rank_digests, digest, group=runtime.group) + raise RuntimeError( + "MoeEp training ABI differs across expert-parallel ranks before " "workspace allocation: " f"digests={rank_digests}, local_facts={facts}" + ) + + +@dataclass(frozen=True) +class Mxfp8TrainingLaneScratch: + """Private fixed-capacity transport and routing tensors for one lane.""" + + index: int + routing_topk_idx: torch.Tensor + routing_topk_weights: torch.Tensor + forward_output: torch.Tensor + backward_output: torch.Tensor + dprob: torch.Tensor + forward_overflow: torch.Tensor + backward_overflow: torch.Tensor + + +@dataclass(frozen=True) +class Mxfp8TrainingExecutionViews: + """Prepared workspaces and private scratch for one execution lane.""" + + scratch: Mxfp8TrainingLaneScratch + forward: PreparedResources + backward: PreparedResources + forward_expert_size_snapshot: torch.Tensor + + +class Mxfp8TrainingState: + """Own only private runtime and per-lane training scratch.""" + + def __init__( + self, + config: ForwardConfig, + device: torch.device, + forward: PreparedMxfp8Kernel, + backward: PreparedMxfp8BackwardKernel, + *, + lane_count: int, + runtime_manager: Optional[RuntimeManager] = None, + symmetric_provider: Optional[SymmetricMemoryProvider] = None, + local_provider: Optional[LocalMemoryProvider] = None, + ) -> None: + self.config = config + self.device = torch.device(device) + self.forward_prepared = forward + self.backward_prepared = backward + self.stager = Mxfp8TrainingStager(config.hidden_size, config.top_k) + self.beta = torch.ones( + (config.experts_per_rank,), + dtype=torch.float32, + device=self.device, + ) + self.lane_count = lane_count + self.requirements = build_training_workspace_requirements( + config, + forward, + backward, + lane_count=lane_count, + ) + self._runtime_manager = runtime_manager or get_runtime_manager() + self._symmetric_provider = symmetric_provider + self._local_provider = local_provider + self._runtime: RuntimeHandle | None = None + self._workspace: WorkspaceOwner | None = None + self._closed = False + self._lock = threading.RLock() + + def prepare(self) -> None: + with self._lock: + if self._closed: + raise RuntimeError("private training state is closed") + if self._runtime is not None and self._workspace is not None and self._workspace.allocated: + return + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("private training state must be prepared before CUDA graph capture") + _runtime_debug( + "training-state.prepare.begin", + lane_count=self.lane_count, + local_bytes=sum(region.nbytes for region in self.requirements.local_regions), + symmetric_bytes=sum(region.nbytes for region in self.requirements.symmetric_regions), + ) + _runtime_debug("training-state.runtime-acquire.begin") + runtime = self._runtime_manager.acquire(self.config, self.device) + _runtime_debug( + "training-state.runtime-acquire.end", + runtime_ref_count=self._runtime_manager.ref_count, + ) + self._runtime = runtime + try: + layout_watchdog = _RuntimeWatchdog("training-state.symmetric-layout-harmonize") + layout_watchdog.start() + _runtime_debug("training-state.symmetric-layout-harmonize.begin") + try: + self.requirements = _harmonize_symmetric_regions( + self.requirements, + runtime, + self.device, + ) + finally: + layout_watchdog.close() + _runtime_debug("training-state.symmetric-layout-harmonize.end") + if runtime.world_size > 1: + abi_watchdog = _RuntimeWatchdog("training-state.abi-handshake") + abi_watchdog.start() + _runtime_debug("training-state.abi-handshake.begin") + try: + abi_facts = _build_training_abi_facts( + self.config, + self.forward_prepared, + self.backward_prepared, + self.requirements, + lane_count=self.lane_count, + ) + abi_fingerprint = _verify_training_abi_across_ranks( + abi_facts, + runtime, + self.device, + ) + finally: + abi_watchdog.close() + _runtime_debug( + "training-state.abi-handshake.end", + fingerprint=abi_fingerprint, + ) + _runtime_debug("training-state.workspace-create.begin") + workspace = WorkspaceOwner( + self.requirements, + runtime, + symmetric_provider=self._symmetric_provider, + local_provider=self._local_provider, + ) + _runtime_debug( + "training-state.workspace-create.end", + local_bytes=workspace.local_layout.total_bytes, + symmetric_bytes=workspace.symmetric_layout.total_bytes, + ) + self._workspace = workspace + allocation_watchdog = _RuntimeWatchdog("training-state.workspace-allocate") + allocation_watchdog.start() + try: + workspace.ensure_allocated() + finally: + allocation_watchdog.close() + _runtime_debug("training-state.workspace-allocate.end") + if runtime.world_size > 1: + # Symmetric-root zeroing is asynchronous. No rank may + # enter the first device barrier until every peer has + # completed allocation and root initialization. + stream_watchdog = _RuntimeWatchdog("training-state.stream-synchronize") + stream_watchdog.start() + _runtime_debug("training-state.stream-synchronize.begin") + try: + torch.cuda.current_stream(self.device).synchronize() + finally: + stream_watchdog.close() + _runtime_debug("training-state.stream-synchronize.end") + + barrier_watchdog = _RuntimeWatchdog("training-state.rank-barrier") + barrier_watchdog.start() + _runtime_debug("training-state.rank-barrier.begin") + try: + dist.barrier(group=runtime.group) + finally: + barrier_watchdog.close() + _runtime_debug("training-state.rank-barrier.end") + except Exception: + if self._workspace is not None: + self._workspace.close() + self._workspace = None + runtime.close() + self._runtime = None + raise + _runtime_debug("training-state.prepare.end") + + def _flat_views(self, token_count: int) -> WorkspaceViews: + self.prepare() + assert self._workspace is not None + return self._workspace.views(token_count) + + @staticmethod + def _phase_workspace( + flat: WorkspaceViews, + requirements: WorkspaceRequirements, + *, + lane: int, + phase: str, + ) -> WorkspaceViews: + symmetric = {} + local = {} + for region in requirements.symmetric_regions: + if region.name in _ROUTING_SYMMETRIC: + symmetric[region.name] = flat.symmetric[_lane_fallback_name(lane, "symmetric", "routing_topk_weights")] + continue + symmetric[region.name] = flat.symmetric[_lane_name(lane, phase, "symmetric", region.name)] + for region in requirements.local_regions: + if phase == "forward" and region.name in _CALLER_OWNED_FORWARD_LOCAL: + continue + if region.name in _ROUTING_LOCAL: + local[region.name] = flat.local[_lane_fallback_name(lane, "local", "routing_topk_idx")] + continue + local[region.name] = flat.local[_lane_name(lane, phase, "local", region.name)] + return WorkspaceViews( + token_count=flat.token_count, + symmetric=MappingProxyType(symmetric), + local=MappingProxyType(local), + peer_mapping=flat.peer_mapping, + ) + + def _lane_scratch_views( + self, + flat: WorkspaceViews, + lane: int, + ) -> Mxfp8TrainingLaneScratch: + config = self.config + capacity = int(config.max_tokens_per_rank) + bwd_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in self.backward_prepared.kernel.get_aux_output_shapes().items()} + + def local_bytes(name: str) -> torch.Tensor: + return flat.local[_lane_fallback_name(lane, "local", name)] + + return Mxfp8TrainingLaneScratch( + index=lane, + routing_topk_idx=_typed_view( + local_bytes("routing_topk_idx"), + torch.int32, + (capacity, config.top_k), + ), + routing_topk_weights=_typed_view( + flat.symmetric[ + _lane_fallback_name( + lane, + "symmetric", + "routing_topk_weights", + ) + ], + torch.float32, + (capacity, config.top_k), + ), + forward_output=_typed_view( + flat.symmetric[ + _lane_name( + lane, + "forward", + "symmetric", + "output_data", + ) + ], + torch.bfloat16, + (capacity, config.hidden_size), + ), + backward_output=_typed_view( + flat.symmetric[ + _lane_name( + lane, + "backward", + "symmetric", + "output_data", + ) + ], + torch.bfloat16, + (capacity, config.hidden_size), + ), + dprob=_typed_view( + flat.symmetric[ + _lane_name( + lane, + "backward", + "symmetric", + "backward_dprob", + ) + ], + torch.float32, + bwd_shapes["dprob"], + ), + forward_overflow=_typed_view( + flat.local[ + _lane_name( + lane, + "forward", + "local", + "overflow_flag", + ) + ], + torch.int32, + (1,), + ), + backward_overflow=_typed_view( + flat.local[ + _lane_name( + lane, + "backward", + "local", + "overflow_flag", + ) + ], + torch.int32, + (1,), + ), + ) + + def public_symmetric_buffers( + self, + lane: int, + ) -> Mapping[str, torch.Tensor]: + """Return caller-visible views over one lane's symmetric I/O buffers.""" + + capacity = int(self.config.max_tokens_per_rank) + hidden = int(self.config.hidden_size) + execution = self.views(lane=lane, token_count=capacity) + + def activation_views(workspace: WorkspaceViews) -> tuple[torch.Tensor, torch.Tensor]: + return ( + _typed_view( + workspace.symmetric["activation_data"], + _DATA_DTYPE, + (capacity, hidden), + ), + _typed_view( + workspace.symmetric["activation_scale"], + _SCALE_DTYPE, + (round_up(capacity, 128), padded_mxfp8_scale_columns(hidden)), + ), + ) + + forward_data, forward_scale = activation_views(execution.forward.workspace) + backward_data, backward_scale = activation_views(execution.backward.workspace) + return MappingProxyType( + { + "forward_input": forward_data, + "forward_input_scale": forward_scale, + "backward_input": backward_data, + "backward_input_scale": backward_scale, + "output": execution.scratch.forward_output, + "grad_activation": execution.scratch.backward_output, + "dprob": execution.scratch.dprob, + } + ) + + def public_requirements( + self, + ) -> Mapping[ + str, + tuple[tuple[int, ...], tuple[int, ...], torch.dtype, int], + ]: + """Return exact caller-owned output contracts without buffer objects.""" + + config = self.config + capacity = int(config.max_tokens_per_rank) + pool_rows = int(self.forward_prepared.pool_token_capacity) + forward_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in self.forward_prepared.kernel.get_aux_output_shapes().items()} + backward_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in self.backward_prepared.kernel.get_aux_output_shapes().items()} + fc1_sfa_rows = round_up(config.hidden_size, 128) + fc1_sfa_elements = math.prod(forward_shapes["col_quant_sf"]) + if fc1_sfa_elements % fc1_sfa_rows: + raise ValueError("forward fc1_sfa producer size is not atom aligned") + fc2_sfb_rows = round_up(config.hidden_size, 128) + fc2_sfb_elements = math.prod(backward_shapes["grad_y2_sf"]) + if fc2_sfb_elements % fc2_sfb_rows: + raise ValueError("backward fc2_sfb producer size is not atom aligned") + requirements = { + "output": ( + (capacity, config.hidden_size), + (config.hidden_size, 1), + torch.bfloat16, + 16, + ), + "fc1_preact": ( + forward_shapes["fc1_c"], + (forward_shapes["fc1_c"][1], 1), + torch.bfloat16, + 128, + ), + "fc1_a": ( + (config.hidden_size, pool_rows), + (pool_rows, 1), + _DATA_DTYPE, + 128, + ), + "fc1_sfa": ( + (fc1_sfa_rows, fc1_sfa_elements // fc1_sfa_rows), + (fc1_sfa_elements // fc1_sfa_rows, 1), + _SCALE_DTYPE, + 128, + ), + "valid_route_counts": ( + (config.experts_per_rank,), + (1,), + torch.int32, + 16, + ), + "expert_offsets": ( + (config.experts_per_rank,), + (1,), + torch.int32, + 16, + ), + "grad_activation": ( + (capacity, config.hidden_size), + (config.hidden_size, 1), + torch.bfloat16, + 16, + ), + "dprob": ( + backward_shapes["dprob"], + (backward_shapes["dprob"][1], 1), + torch.float32, + 16, + ), + "fc1_b": ( + backward_shapes["fc1_col_output"], + (2 * config.intermediate_size, 1), + _DATA_DTYPE, + 128, + ), + "fc1_sfb": ( + backward_shapes["fc1_col_output_sf"], + (backward_shapes["fc1_col_output_sf"][1], 1), + _SCALE_DTYPE, + 128, + ), + "fc2_a": ( + (config.intermediate_size, pool_rows), + (1, config.intermediate_size), + _DATA_DTYPE, + 128, + ), + "fc2_sfa": ( + backward_shapes["fc1_recompute_sf"], + (backward_shapes["fc1_recompute_sf"][1], 1), + _SCALE_DTYPE, + 128, + ), + "fc2_b": ( + backward_shapes["grad_y2"], + (1, pool_rows), + _DATA_DTYPE, + 128, + ), + "fc2_sfb": ( + (fc2_sfb_rows, fc2_sfb_elements // fc2_sfb_rows), + (fc2_sfb_elements // fc2_sfb_rows, 1), + _SCALE_DTYPE, + 128, + ), + } + return MappingProxyType(requirements) + + def views( + self, + *, + lane: int, + token_count: int, + ) -> Mxfp8TrainingExecutionViews: + with self._lock: + if lane < 0 or lane >= self.lane_count: + raise ValueError(f"lane {lane} is outside [0, {self.lane_count})") + col_quant_sizes_offset = self.forward_prepared.col_quant_sizes_offset + if col_quant_sizes_offset is None: + raise RuntimeError("training preparation requires a persistent col-quant expert-size snapshot") + flat = self._flat_views(token_count) + forward_workspace = self._phase_workspace( + flat, + self.forward_prepared.workspace_requirements, + lane=lane, + phase="forward", + ) + backward_workspace = self._phase_workspace( + flat, + self.backward_prepared.workspace_requirements, + lane=lane, + phase="backward", + ) + snapshot_bytes = forward_workspace.local["kernel_local_workspace"].narrow( + 0, + col_quant_sizes_offset, + self.forward_prepared.col_quant_sizes_bytes, + ) + snapshot = _typed_view( + snapshot_bytes, + torch.int32, + (self.config.experts_per_rank,), + ) + assert self._runtime is not None + return Mxfp8TrainingExecutionViews( + scratch=self._lane_scratch_views(flat, lane), + forward=PreparedResources( + runtime=self._runtime, + workspace=forward_workspace, + ), + backward=PreparedResources( + runtime=self._runtime, + workspace=backward_workspace, + ), + forward_expert_size_snapshot=snapshot, + ) + + def close(self) -> None: + with self._lock: + if self._closed: + return + if self._workspace is not None: + self._workspace.close() + self._workspace = None + if self._runtime is not None: + self._runtime.close() + self._runtime = None + self._closed = True + + def apply_overflow( + self, + *, + lane: int, + phase: str, + ) -> torch.Tensor: + """Apply the configured policy to one phase's private overflow flag.""" + + if lane < 0 or lane >= self.lane_count: + raise ValueError(f"lane {lane} is outside [0, {self.lane_count})") + if phase not in ("forward", "backward"): + raise ValueError(f"phase must be 'forward' or 'backward', got {phase!r}") + flat = self._flat_views(0) + global_overflow = _typed_view( + flat.local[ + _lane_name( + lane, + "finalizer", + "local", + "global_overflow", + ) + ], + torch.int32, + (1,), + ) + flag = _typed_view( + flat.local[ + _lane_name( + lane, + phase, + "local", + "overflow_flag", + ) + ], + torch.int32, + (1,), + ) + global_overflow.copy_(flag) + assert self._runtime is not None + if self._runtime.world_size > 1: + dist.all_reduce( + global_overflow, + op=dist.ReduceOp.MAX, + group=self._runtime.group, + ) + if not self.config.drop_on_overflow: + assert_async = getattr(torch, "_assert_async", None) + if assert_async is None: + raise RuntimeError("drop_on_overflow=False training requires torch._assert_async") + overflow_ok = _typed_view( + flat.local[ + _lane_name( + lane, + "finalizer", + "local", + "overflow_ok", + ) + ], + torch.bool, + (1,), + ) + torch.eq(global_overflow, 0, out=overflow_ok) + assert_async( + overflow_ok, + f"Rubin MegaMoE receive route-pool overflow; the {phase} " "outputs are invalid", + ) + return global_overflow + + +__all__ = [ + "Mxfp8TrainingExecutionViews", + "Mxfp8TrainingState", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py new file mode 100644 index 000000000..74de1f3ff --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py @@ -0,0 +1,163 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Cached BF16/FP32 staging into private symmetric training scratch.""" + +from __future__ import annotations + +import threading + +import torch + +from ..._math import round_up +from ._launch import _to_cute + + +class Mxfp8TrainingStager: + """Own one compile cache; steady-state staging is allocation-free.""" + + def __init__(self, hidden: int, top_k: int) -> None: + self.hidden = int(hidden) + self.top_k = int(top_k) + self._compiled: dict[tuple, object] = {} + self._lock = threading.RLock() + + def _validate( + self, + source: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + output: torch.Tensor, + output_sf: torch.Tensor, + output_topk_idx: torch.Tensor, + output_topk_weights: torch.Tensor, + ) -> int: + if source.dtype not in (torch.bfloat16, torch.float32): + raise TypeError("training staging source must be BF16 or FP32, " f"got {source.dtype}") + if source.ndim != 2 or source.shape[1] != self.hidden: + raise ValueError(f"training staging source must have shape (T, {self.hidden})") + if not source.is_contiguous(): + raise ValueError("training staging source must be contiguous") + token_count = int(source.shape[0]) + if topk_idx.shape != (token_count, self.top_k): + raise ValueError("training staging topk_idx shape mismatch") + if topk_idx.dtype is not torch.int32 or not topk_idx.is_contiguous(): + raise TypeError("training staging topk_idx must be contiguous Int32") + if topk_weights.shape != topk_idx.shape: + raise ValueError("training staging topk_weights shape mismatch") + if topk_weights.dtype is not torch.float32 or not topk_weights.is_contiguous(): + raise TypeError("training staging topk_weights must be contiguous FP32") + if output.dtype is not torch.float8_e4m3fn or output.ndim != 2 or output.shape[1] != self.hidden or not output.is_contiguous(): + raise ValueError("training staging output must be contiguous E4M3 " f"(capacity, {self.hidden})") + if token_count > output.shape[0]: + raise ValueError(f"token count {token_count} exceeds capacity {output.shape[0]}") + logical_sf_columns = self.hidden // 32 + if ( + output_sf.dtype is not torch.float8_e8m0fnu + or output_sf.ndim != 2 + or output_sf.shape[0] != round_up(output.shape[0], 128) + or output_sf.shape[1] < logical_sf_columns + or not output_sf.is_contiguous() + ): + raise ValueError("training staging output_sf has an invalid ABI") + for name, tensor, dtype in ( + ("output_topk_idx", output_topk_idx, torch.int32), + ("output_topk_weights", output_topk_weights, torch.float32), + ): + if tensor.shape != (output.shape[0], self.top_k) or tensor.dtype is not dtype or not tensor.is_contiguous(): + raise ValueError(f"training staging {name} has an invalid ABI") + devices = { + source.device, + topk_idx.device, + topk_weights.device, + output.device, + output_sf.device, + output_topk_idx.device, + output_topk_weights.device, + } + if len(devices) != 1: + raise ValueError("all training staging tensors must share one device") + return token_count + + def stage( + self, + source: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + output: torch.Tensor, + output_sf: torch.Tensor, + output_topk_idx: torch.Tensor, + output_topk_weights: torch.Tensor, + ) -> None: + """Enqueue tail reset plus one fused quant-and-routing launch.""" + + token_count = self._validate( + source, + topk_idx, + topk_weights, + output, + output_sf, + output_topk_idx, + output_topk_weights, + ) + output_sf.zero_() + routing_in_place = topk_idx.data_ptr() == output_topk_idx.data_ptr() and topk_weights.data_ptr() == output_topk_weights.data_ptr() + routing_partially_aliased = (topk_idx.data_ptr() == output_topk_idx.data_ptr()) != (topk_weights.data_ptr() == output_topk_weights.data_ptr()) + if routing_partially_aliased: + raise ValueError("training staging routing inputs must either both alias " "their outputs or neither alias") + if not routing_in_place: + output_topk_idx.fill_(-1) + output_topk_weights.zero_() + if token_count == 0: + return + + logical_sf_columns = self.hidden // 32 + import cuda.bindings.driver as cuda + + stream = torch.cuda.current_stream(source.device) + args = ( + _to_cute(source, dynamic_layout=False), + _to_cute(topk_idx, assumed_align=4, dynamic_layout=False), + _to_cute(topk_weights, assumed_align=4, dynamic_layout=False), + _to_cute(output[:token_count], dynamic_layout=False), + _to_cute( + output_sf[:token_count, :logical_sf_columns], + assumed_align=4, + dynamic_layout=False, + ), + _to_cute( + output_topk_idx[:token_count], + assumed_align=4, + dynamic_layout=False, + ), + _to_cute( + output_topk_weights[:token_count], + assumed_align=4, + dynamic_layout=False, + ), + cuda.CUstream(stream.cuda_stream), + ) + key = ( + source.device.index, + source.dtype, + token_count, + self.hidden, + self.top_k, + tuple(output_sf.stride()), + ) + with self._lock: + compiled = self._compiled.get(key) + if compiled is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("MXFP8 training stager must be compiled before " "CUDA graph capture") + import cutlass.cute as cute + + from ._training_stage_kernel import Mxfp8TrainingStageKernel + + kernel = Mxfp8TrainingStageKernel(self.hidden, self.top_k) + compiled = cute.compile(kernel, *args) + self._compiled[key] = compiled + compiled(*args) + + +__all__ = ["Mxfp8TrainingStager"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py new file mode 100644 index 000000000..c9f5946d1 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py @@ -0,0 +1,125 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""One-launch BF16/FP32 to MXFP8 private symmetric staging.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Float32, Int32 + +from ..cutedsl_src.helpers.constants import Fp32Max, Fp8E4M3RcpLimit +from ..cutedsl_src.helpers.ptx_helpers import cvt_f32_to_fp8_to_f32 + + +class Mxfp8TrainingStageKernel: + """Quantize one token row per CTA and repack its routing metadata.""" + + _threads_per_cta = 128 + _sf_vec = 32 + + def __init__(self, hidden: int, top_k: int) -> None: + self.hidden = int(hidden) + self.top_k = int(top_k) + if self.hidden <= 0 or self.hidden % self._sf_vec: + raise ValueError("MXFP8 training stage requires hidden divisible by 32") + if self.top_k <= 0 or self.top_k > self._threads_per_cta: + raise ValueError("MXFP8 training stage requires " f"1 <= top_k <= {self._threads_per_cta}") + + @cute.jit + def __call__( + self, + source: cute.Tensor, + topk_idx: cute.Tensor, + topk_weights: cute.Tensor, + output: cute.Tensor, + output_sf: cute.Tensor, + output_topk_idx: cute.Tensor, + output_topk_weights: cute.Tensor, + stream: cuda.CUstream, + ) -> None: + self._kernel( + source, + topk_idx, + topk_weights, + output, + output_sf, + output_topk_idx, + output_topk_weights, + ).launch( + grid=[source.shape[0], 1, 1], + block=[self._threads_per_cta, 1, 1], + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.kernel + def _kernel( + self, + source: cute.Tensor, + topk_idx: cute.Tensor, + topk_weights: cute.Tensor, + output: cute.Tensor, + output_sf: cute.Tensor, + output_topk_idx: cute.Tensor, + output_topk_weights: cute.Tensor, + ) -> None: + token = cute.arch.block_idx()[0] + tid = cute.arch.thread_idx()[0] + hidden: cutlass.Constexpr[int] = self.hidden + sf_vec: cutlass.Constexpr[int] = self._sf_vec + threads: cutlass.Constexpr[int] = self._threads_per_cta + block_count: cutlass.Constexpr[int] = hidden // sf_vec + rounds: cutlass.Constexpr[int] = (block_count + threads - 1) // threads + + for block_round in cutlass.range_constexpr(rounds): + block = tid + Int32(block_round * threads) + if block < Int32(block_count): + values = cute.make_rmem_tensor((sf_vec,), Float32) + absmax = Float32(0.0) + for element in cutlass.range_constexpr(sf_vec): + value = Float32( + source[ + token, + block * Int32(sf_vec) + Int32(element), + ] + ) + values[element] = value + absmax = cute.arch.fmax( + absmax, + cute.arch.fmax(value, -value), + ) + + scale_f32 = Float32( + cvt_f32_to_fp8_to_f32( + absmax * Float32(Fp8E4M3RcpLimit), + cutlass.Float8E8M0FNU, + ) + ) + scale = scale_f32.to(cutlass.Float8E8M0FNU) + reciprocal = cute.arch.fmin( + cute.arch.rcp_approx(scale_f32), + Float32(Fp32Max), + ) + reciprocal = reciprocal * cute.arch.fmin( + scale_f32 * Float32(1.0e30), + Float32(1.0), + ) + for element in cutlass.range_constexpr(sf_vec): + output[ + token, + block * Int32(sf_vec) + Int32(element), + ] = ( + values[element] * reciprocal + ).to(cutlass.Float8E4M3FN) + output_sf[token, block] = scale + + if tid < Int32(self.top_k): + output_topk_idx[token, tid] = Int32(topk_idx[token, tid]) + output_topk_weights[token, tid] = Float32(topk_weights[token, tid]) + + +__all__ = ["Mxfp8TrainingStageKernel"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py new file mode 100644 index 000000000..c74b2552f --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py @@ -0,0 +1,437 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Stable, allocation-free layout staging for pre-quantized training weights.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from ..._contracts import Fc1WeightLayout +from ..._math import round_up +from ..._types import ( + BlockScaledTensor, + MoeEpBackwardWeightStaging, + MoeEpBackwardWeights, + MoeEpForwardWeightStaging, + MoeEpForwardWeights, + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, +) +from ..._validation import validate_training_non_aliasing +from ._adapter import Mxfp8Weights + + +def _copy_blocked_scales_plain( + target: torch.Tensor, + source: torch.Tensor, + *, + raw_rows: int, + raw_columns: int, +) -> None: + """Pack public ``(E,Kblocks,N)`` scales for a non-interleaved weight.""" + + experts = source.shape[0] + if tuple(source.shape) != (experts, raw_columns, raw_rows): + raise ValueError("plain training scale shape mismatch: " f"{tuple(source.shape)} != " f"{(experts, raw_columns, raw_rows)}") + if raw_rows % 128 or raw_columns % 4: + raise ValueError("training scale pack requires rows divisible by 128 and " "columns divisible by 4") + row_blocks = raw_rows // 128 + column_blocks = raw_columns // 4 + source_view = ( + source.view( + torch.uint8, + ) + .view( + experts, + column_blocks, + 4, + row_blocks, + 4, + 32, + ) + .permute(0, 3, 1, 5, 4, 2) + ) + target.view(torch.uint8).view( + experts, + row_blocks, + column_blocks, + 32, + 4, + 4, + ).copy_(source_view) + + +def _copy_gate_up_interleaved_last( + target: torch.Tensor, + source: torch.Tensor, + intermediate: int, +) -> None: + """Copy ``(E,K,gate||up)`` into 32-column gate/up strip order.""" + + experts, reduction, gate_up = source.shape + if target.shape != source.shape or gate_up != 2 * intermediate: + raise ValueError("forward FC1 training weight shape mismatch") + pairs = intermediate // 32 + source_view = source.view(experts, reduction, 2, pairs, 32).permute(0, 1, 3, 2, 4) + target_view = target.as_strided( + (experts, reduction, pairs, 2, 32), + ( + target.stride(0), + target.stride(1), + 64 * target.stride(2), + 32 * target.stride(2), + target.stride(2), + ), + ) + target_view.copy_(source_view) + + +def _copy_gate_up_interleaved_reduction( + target: torch.Tensor, + source: torch.Tensor, + intermediate: int, +) -> None: + """Copy ``(E,gate||up,N)`` into 32-row gate/up strip order.""" + + experts, gate_up, output = source.shape + if target.shape != source.shape or gate_up != 2 * intermediate: + raise ValueError("backward W1-transpose training weight shape mismatch") + pairs = intermediate // 32 + source_view = source.view(experts, 2, pairs, 32, output).permute(0, 2, 1, 3, 4) + target.view(experts, pairs, 2, 32, output).copy_(source_view) + + +def _copy_blocked_scales_gate_up_rows( + target: torch.Tensor, + source: torch.Tensor, + *, + intermediate: int, + reduction_blocks: int, +) -> None: + """Pack forward FC1 scales after 32-row gate/up interleave.""" + + experts = source.shape[0] + raw_rows = 2 * intermediate + if intermediate % 64 or reduction_blocks % 4: + raise ValueError("intermediate and reduction block alignment are invalid") + source_view = source.view(torch.uint8).view(experts, reduction_blocks // 4, 4, 2, raw_rows // 128, 2, 32).permute(0, 4, 1, 6, 5, 3, 2) + target.view(torch.uint8).view(experts, raw_rows // 128, reduction_blocks // 4, 32, 2, 2, 4).copy_(source_view) + + +def _copy_blocked_scales_gate_up_columns( + target: torch.Tensor, + source: torch.Tensor, + *, + intermediate: int, + output: int, +) -> None: + """Pack backward W1-transpose scales with interleaved K blocks.""" + + experts = source.shape[0] + reduction_blocks = intermediate // 32 + if output % 128 or reduction_blocks % 2: + raise ValueError("output and reduction block alignment are invalid") + source_view = source.view(torch.uint8).view(experts, 2, reduction_blocks // 2, 2, output // 128, 4, 32).permute(0, 4, 2, 6, 5, 3, 1) + target.view(torch.uint8).view(experts, output // 128, reduction_blocks // 2, 32, 4, 2, 2).copy_(source_view) + + +@dataclass(frozen=True) +class Mxfp8BackwardWeights: + """Kernel names follow the two backward FC stages.""" + + fc1_weight: torch.Tensor + fc1_weight_sf: torch.Tensor + fc2_weight: torch.Tensor + fc2_weight_sf: torch.Tensor + + +def _expect_staging_tensor( + name: str, + tensor: torch.Tensor, + *, + shape: tuple[int, ...], + stride: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> None: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.layout is not torch.strided: + raise ValueError(f"{name} must use torch.strided layout, got {tensor.layout}") + if tuple(tensor.shape) != shape or tuple(tensor.stride()) != stride: + raise ValueError(f"{name} must have shape={shape}, stride={stride}; got " f"shape={tuple(tensor.shape)}, stride={tuple(tensor.stride())}") + if tensor.dtype is not dtype: + raise ValueError(f"{name} must have dtype {dtype}, got {tensor.dtype}") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}") + if tensor.data_ptr() % 16: + raise ValueError(f"{name} must be at least 16-byte aligned") + + +def _expect_scale_staging( + name: str, + tensor: torch.Tensor, + *, + experts: int, + raw_rows: int, + raw_columns: int, + dtype: torch.dtype, + device: torch.device, +) -> None: + elements = round_up(raw_rows, 128) * round_up(raw_columns, 4) + _expect_staging_tensor( + name, + tensor, + shape=(experts, elements), + stride=(elements, 1), + dtype=dtype, + device=device, + ) + + +def forward_native_to_kernel(weights: MoeEpNativeForwardWeights) -> Mxfp8Weights: + """Create kernel views without allocating, copying, or retaining inputs.""" + + return Mxfp8Weights( + fc1_weight=weights.fc1.payload, + fc1_weight_sf=weights.fc1.scale, + fc2_weight=weights.fc2.payload, + fc2_weight_sf=weights.fc2.scale, + ) + + +def backward_native_to_kernel( + weights: MoeEpNativeBackwardWeights, +) -> Mxfp8BackwardWeights: + """Create kernel views without allocating, copying, or retaining inputs.""" + + return Mxfp8BackwardWeights( + fc1_weight=weights.w2_transpose.payload, + fc1_weight_sf=weights.w2_transpose.scale, + fc2_weight=weights.w1_transpose.payload, + fc2_weight_sf=weights.w1_transpose.scale, + ) + + +def materialize_forward( + weights: MoeEpForwardWeights, + *, + out: MoeEpForwardWeightStaging, + fc1_weight_layout: Fc1WeightLayout, +) -> MoeEpNativeForwardWeights: + """Materialize source forward weights into caller-owned native storage.""" + + if fc1_weight_layout is not Fc1WeightLayout.GATE_UP_INTERLEAVED_32: + raise ValueError("native training materialization requires weight_interleave_size=32") + if not isinstance(weights, MoeEpForwardWeights): + raise TypeError("weights must be a MoeEpForwardWeights") + if not isinstance(out, MoeEpForwardWeightStaging): + raise TypeError("out must be a MoeEpForwardWeightStaging") + fc1 = weights.fc1 + fc2 = weights.fc2 + for name, value in (("weights.fc1", fc1), ("weights.fc2", fc2)): + if not isinstance(value, BlockScaledTensor) or value.format.value != "mxfp8" or value.axis != 1: + raise TypeError(f"{name} must be an axis-1 MXFP8 BlockScaledTensor") + if not value.data.is_contiguous() or not value.scale.is_contiguous(): + raise ValueError(f"{name} data and scale must be contiguous") + experts, hidden, gate_up = fc1.data.shape + intermediate = fc2.data.shape[1] + if tuple(fc2.data.shape) != (experts, intermediate, hidden) or gate_up != 2 * intermediate: + raise ValueError("forward source weight shapes are inconsistent") + if fc2.device != fc1.device: + raise ValueError("forward source weights must share one device") + sf_dtype = torch.float8_e8m0fnu + _expect_staging_tensor( + "out.fc1_payload", + out.fc1_payload, + shape=tuple(fc1.data.shape), + stride=(hidden * gate_up, 1, hidden), + dtype=fc1.data.dtype, + device=fc1.device, + ) + _expect_staging_tensor( + "out.fc2_payload", + out.fc2_payload, + shape=tuple(fc2.data.shape), + stride=(intermediate * hidden, 1, intermediate), + dtype=fc2.data.dtype, + device=fc1.device, + ) + _expect_scale_staging( + "out.fc1_scale", + out.fc1_scale, + experts=experts, + raw_rows=gate_up, + raw_columns=hidden // 32, + dtype=sf_dtype, + device=fc1.device, + ) + _expect_scale_staging( + "out.fc2_scale", + out.fc2_scale, + experts=experts, + raw_rows=hidden, + raw_columns=intermediate // 32, + dtype=sf_dtype, + device=fc1.device, + ) + validate_training_non_aliasing( + { + "weights.fc1.data": fc1.data, + "weights.fc1.scale": fc1.scale, + "weights.fc2.data": fc2.data, + "weights.fc2.scale": fc2.scale, + "out.fc1_payload": out.fc1_payload, + "out.fc1_scale": out.fc1_scale, + "out.fc2_payload": out.fc2_payload, + "out.fc2_scale": out.fc2_scale, + } + ) + _copy_gate_up_interleaved_last(out.fc1_payload, fc1.data, intermediate) + _copy_blocked_scales_gate_up_rows( + out.fc1_scale, + fc1.scale, + intermediate=intermediate, + reduction_blocks=hidden // 32, + ) + out.fc2_payload.copy_(fc2.data) + _copy_blocked_scales_plain( + out.fc2_scale, + fc2.scale, + raw_rows=hidden, + raw_columns=intermediate // 32, + ) + return MoeEpNativeForwardWeights( + fc1=MoeEpNativeWeight( + out.fc1_payload, + out.fc1_scale, + MoeEpNativeWeightLayout.FORWARD_FC1_GATE_UP_INTERLEAVED_32_V1, + ), + fc2=MoeEpNativeWeight( + out.fc2_payload, + out.fc2_scale, + MoeEpNativeWeightLayout.FORWARD_FC2_K_MAJOR_V1, + ), + ) + + +def materialize_backward( + weights: MoeEpBackwardWeights, + *, + out: MoeEpBackwardWeightStaging, + fc1_weight_layout: Fc1WeightLayout, +) -> MoeEpNativeBackwardWeights: + """Materialize source backward weights into caller-owned native storage.""" + + if fc1_weight_layout is not Fc1WeightLayout.GATE_UP_INTERLEAVED_32: + raise ValueError("native training materialization requires weight_interleave_size=32") + if not isinstance(weights, MoeEpBackwardWeights): + raise TypeError("weights must be a MoeEpBackwardWeights") + if not isinstance(out, MoeEpBackwardWeightStaging): + raise TypeError("out must be a MoeEpBackwardWeightStaging") + w2t = weights.w2_transpose + w1t = weights.w1_transpose + for name, value in ( + ("weights.w2_transpose", w2t), + ("weights.w1_transpose", w1t), + ): + if not isinstance(value, BlockScaledTensor) or value.format.value != "mxfp8" or value.axis != 1: + raise TypeError(f"{name} must be an axis-1 MXFP8 BlockScaledTensor") + if not value.data.is_contiguous() or not value.scale.is_contiguous(): + raise ValueError(f"{name} data and scale must be contiguous") + experts, hidden, intermediate = w2t.data.shape + if tuple(w1t.data.shape) != (experts, 2 * intermediate, hidden): + raise ValueError("backward source weight shapes are inconsistent") + if w1t.device != w2t.device: + raise ValueError("backward source weights must share one device") + _expect_staging_tensor( + "out.w2_transpose_payload", + out.w2_transpose_payload, + shape=tuple(w2t.data.shape), + stride=(hidden * intermediate, intermediate, 1), + dtype=w2t.data.dtype, + device=w2t.device, + ) + _expect_staging_tensor( + "out.w1_transpose_payload", + out.w1_transpose_payload, + shape=tuple(w1t.data.shape), + stride=(2 * intermediate * hidden, hidden, 1), + dtype=w1t.data.dtype, + device=w2t.device, + ) + sf_dtype = torch.float8_e8m0fnu + _expect_scale_staging( + "out.w2_transpose_scale", + out.w2_transpose_scale, + experts=experts, + raw_rows=intermediate, + raw_columns=hidden // 32, + dtype=sf_dtype, + device=w2t.device, + ) + _expect_scale_staging( + "out.w1_transpose_scale", + out.w1_transpose_scale, + experts=experts, + raw_rows=hidden, + raw_columns=2 * intermediate // 32, + dtype=sf_dtype, + device=w2t.device, + ) + validate_training_non_aliasing( + { + "weights.w2_transpose.data": w2t.data, + "weights.w2_transpose.scale": w2t.scale, + "weights.w1_transpose.data": w1t.data, + "weights.w1_transpose.scale": w1t.scale, + "out.w2_transpose_payload": out.w2_transpose_payload, + "out.w2_transpose_scale": out.w2_transpose_scale, + "out.w1_transpose_payload": out.w1_transpose_payload, + "out.w1_transpose_scale": out.w1_transpose_scale, + } + ) + out.w2_transpose_payload.copy_(w2t.data) + _copy_blocked_scales_plain( + out.w2_transpose_scale, + w2t.scale, + raw_rows=intermediate, + raw_columns=hidden // 32, + ) + _copy_gate_up_interleaved_reduction( + out.w1_transpose_payload, + w1t.data, + intermediate, + ) + _copy_blocked_scales_gate_up_columns( + out.w1_transpose_scale, + w1t.scale, + intermediate=intermediate, + output=hidden, + ) + return MoeEpNativeBackwardWeights( + w2_transpose=MoeEpNativeWeight( + out.w2_transpose_payload, + out.w2_transpose_scale, + MoeEpNativeWeightLayout.BACKWARD_W2_TRANSPOSE_V1, + ), + w1_transpose=MoeEpNativeWeight( + out.w1_transpose_payload, + out.w1_transpose_scale, + MoeEpNativeWeightLayout.BACKWARD_W1_TRANSPOSE_GATE_UP_INTERLEAVED_32_V1, + ), + ) + + +__all__ = [ + "Mxfp8BackwardWeights", + "backward_native_to_kernel", + "forward_native_to_kernel", + "materialize_backward", + "materialize_forward", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py new file mode 100644 index 000000000..a0eea6167 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py @@ -0,0 +1,50 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Pure, zero-copy assembly of caller-owned WGrad operand views.""" + +from __future__ import annotations + +import torch + +from ..._types import ( + MoeEpTrainingBackwardOutputs, + MoeEpTrainingWgradOperands, +) + + +def assemble_training_wgrad_operands( + *, + fc1_a: torch.Tensor, + fc1_sfa: torch.Tensor, + valid_route_counts: torch.Tensor, + expert_offsets: torch.Tensor, + backward: MoeEpTrainingBackwardOutputs, +) -> MoeEpTrainingWgradOperands: + """Return non-owning views after producers wrote their final layouts.""" + + required = ( + backward.fc1_b, + backward.fc1_sfb, + backward.fc2_a, + backward.fc2_sfa, + backward.fc2_b, + backward.fc2_sfb, + ) + if any(value is None for value in required): + raise ValueError("all backward WGrad outputs are required for assembly") + return MoeEpTrainingWgradOperands( + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + fc1_b=backward.fc1_b, + fc1_sfb=backward.fc1_sfb, + fc2_a=backward.fc2_a, + fc2_sfa=backward.fc2_sfa, + fc2_b=backward.fc2_b, + fc2_sfb=backward.fc2_sfb, + expert_offsets=expert_offsets, + valid_route_counts=valid_route_counts, + ) + + +__all__ = ["assemble_training_wgrad_operands"] diff --git a/python/cudnn/moe_ep/_tuning.py b/python/cudnn/moe_ep/_tuning.py new file mode 100644 index 000000000..29f0b6ad4 --- /dev/null +++ b/python/cudnn/moe_ep/_tuning.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Public, semantic-preserving performance tuning for :class:`MoeEp`.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +AutotuneMode = Literal["inference", "training"] + +TokenBackMode = Literal[ + "epi_warps", + "standalone_warps", + "reuse_dispatch_warps", +] + +_TOKEN_BACK_MODES = frozenset( + { + "epi_warps", + "standalone_warps", + "reuse_dispatch_warps", + } +) +_EPI_FLAG_BATCHES = frozenset( + { + (4, 2), + (1, 1), + (1, 2), + (1, 4), + (2, 1), + (2, 2), + (2, 4), + (4, 4), + } +) +_TOKEN_IN_FLAG_BATCHES = frozenset({1, 2, 4, 8, 16}) +_GROUP_HINTS = frozenset({64, 128, 256, 512, 768, 1024}) + + +@dataclass(frozen=True, kw_only=True) +class MoeEpTuningConfig: + """Validated Rubin MegaMoE performance knobs. + + These fields select scheduling and transport implementations without + changing the public MoE mathematical contract. Every rank in an expert + parallel group must use the same configuration. + + ``group_hint=None`` preserves the default behavior: the backend uses the + number of hardware-resident CTA clusters. + """ + + token_back_mode: TokenBackMode = "epi_warps" + epi_flag_batch: tuple[int, int] = (1, 1) + token_in_flag_batch: int = 1 + group_hint: int | None = None + reduce_topk_in_kernel: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.token_back_mode, str) or self.token_back_mode not in _TOKEN_BACK_MODES: + raise ValueError("token_back_mode must be one of " f"{tuple(sorted(_TOKEN_BACK_MODES))}, got " f"{self.token_back_mode!r}") + if not isinstance(self.epi_flag_batch, tuple) or self.epi_flag_batch not in _EPI_FLAG_BATCHES: + raise ValueError("epi_flag_batch must be one of " f"{tuple(sorted(_EPI_FLAG_BATCHES))}, got " f"{self.epi_flag_batch!r}") + if isinstance(self.token_in_flag_batch, bool) or self.token_in_flag_batch not in _TOKEN_IN_FLAG_BATCHES: + raise ValueError("token_in_flag_batch must be one of " f"{tuple(sorted(_TOKEN_IN_FLAG_BATCHES))}, got " f"{self.token_in_flag_batch!r}") + if self.group_hint is not None and (isinstance(self.group_hint, bool) or self.group_hint not in _GROUP_HINTS): + raise ValueError("group_hint must be None or one of " f"{tuple(sorted(_GROUP_HINTS))}, got {self.group_hint!r}") + if not isinstance(self.reduce_topk_in_kernel, bool): + raise ValueError("reduce_topk_in_kernel must be a bool, got " f"{self.reduce_topk_in_kernel!r}") + if self.reduce_topk_in_kernel and self.token_back_mode != "epi_warps": + raise ValueError("reduce_topk_in_kernel requires " "token_back_mode='epi_warps'") + + +@dataclass(frozen=True) +class MoeEpAutotuneCandidateResult: + """Measured slow-rank latency for one successfully evaluated candidate.""" + + tuning: MoeEpTuningConfig + latency_ms: float + samples_ms: tuple[float, ...] + + +@dataclass(frozen=True) +class MoeEpAutotuneResult: + """Winner and measurements produced by one explicit tuning sweep.""" + + mode: AutotuneMode + winner: MoeEpTuningConfig + candidates: tuple[MoeEpAutotuneCandidateResult, ...] + + @property + def evaluated_candidates(self) -> int: + return len(self.candidates) + + +__all__ = [ + "MoeEpAutotuneCandidateResult", + "MoeEpAutotuneResult", + "MoeEpTuningConfig", +] diff --git a/python/cudnn/moe_ep/_types.py b/python/cudnn/moe_ep/_types.py new file mode 100644 index 000000000..748517eb3 --- /dev/null +++ b/python/cudnn/moe_ep/_types.py @@ -0,0 +1,349 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lightweight public tensor and format types for :mod:`cudnn.moe_ep`.""" + +from __future__ import annotations + +import operator +from dataclasses import dataclass +from enum import Enum +from typing import Tuple, Union + +import torch + +from ._math import ceil_div + + +class MoeFormat(str, Enum): + """Data formats supported by the MoE+EP interface.""" + + BF16 = "bf16" + MXFP8 = "mxfp8" + NVFP4 = "nvfp4" + + +def parse_format(value: Union[MoeFormat, str]) -> MoeFormat: + """Normalize a public format value.""" + + if isinstance(value, MoeFormat): + return value + try: + return MoeFormat(value.lower()) + except (AttributeError, ValueError) as exc: + choices = ", ".join(item.value for item in MoeFormat) + raise ValueError(f"unsupported format {value!r}; expected one of: {choices}") from exc + + +def _normalize_axis(axis: int, ndim: int) -> int: + if isinstance(axis, bool): + raise ValueError(f"axis must be an integer, got {axis!r}") + try: + axis = operator.index(axis) + except TypeError as exc: + raise ValueError(f"axis must be an integer, got {axis!r}") from exc + normalized = axis + ndim if axis < 0 else axis + if normalized < 0 or normalized >= ndim: + raise ValueError(f"axis {axis} is out of range for a {ndim}-D tensor") + return normalized + + +def _block_scaled_representation( + fmt: MoeFormat, + logical_shape: Tuple[int, ...], + axis: int, +) -> tuple[Tuple[int, ...], Tuple[int, ...], torch.dtype, torch.dtype]: + if fmt is MoeFormat.BF16: + raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") + logical_extent = logical_shape[axis] + block_size = 32 if fmt is MoeFormat.MXFP8 else 16 + payload_extent = logical_extent if fmt is MoeFormat.MXFP8 else ceil_div(logical_extent, 2) + data_shape = list(logical_shape) + data_shape[axis] = payload_extent + scale_shape = list(logical_shape) + scale_shape[axis] = ceil_div(logical_extent, block_size) + + e4m3_dtype = getattr(torch, "float8_e4m3fn", None) + if e4m3_dtype is None: + raise RuntimeError("this PyTorch build does not provide torch.float8_e4m3fn") + if fmt is MoeFormat.MXFP8: + scale_dtype = getattr(torch, "float8_e8m0fnu", None) + if scale_dtype is None: + raise RuntimeError("this PyTorch build does not provide torch.float8_e8m0fnu") + data_dtype = e4m3_dtype + else: + data_dtype = torch.uint8 + scale_dtype = e4m3_dtype + return tuple(data_shape), tuple(scale_shape), data_dtype, scale_dtype + + +@dataclass(frozen=True) +class BlockScaledTensor: + """Data-plus-scale result returned for MXFP8 and NVFP4 outputs.""" + + data: torch.Tensor + scale: torch.Tensor + format: Union[MoeFormat, str] + logical_shape: Tuple[int, ...] + axis: int = -1 + + def __post_init__(self) -> None: + if not isinstance(self.data, torch.Tensor): + raise ValueError(f"data must be a torch.Tensor, got {type(self.data).__name__}") + if not isinstance(self.scale, torch.Tensor): + raise ValueError(f"scale must be a torch.Tensor, got {type(self.scale).__name__}") + fmt = parse_format(self.format) + if fmt is MoeFormat.BF16: + raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") + if self.data.device != self.scale.device: + raise ValueError(f"data device {self.data.device} does not match " f"scale device {self.scale.device}") + try: + raw_logical_shape = tuple(self.logical_shape) + except TypeError as exc: + raise ValueError("logical_shape must be an iterable of integers") from exc + logical_shape = [] + for dim in raw_logical_shape: + if isinstance(dim, bool): + raise ValueError(f"logical_shape dimensions must be integers, got {dim!r}") + try: + dim = operator.index(dim) + except TypeError as exc: + raise ValueError(f"logical_shape dimensions must be integers, got {dim!r}") from exc + if dim < 0: + raise ValueError(f"logical_shape dimensions must be non-negative, got {dim}") + logical_shape.append(dim) + normalized_shape = tuple(logical_shape) + axis = _normalize_axis(self.axis, len(normalized_shape)) + ( + expected_data_shape, + expected_scale_shape, + expected_data_dtype, + expected_scale_dtype, + ) = _block_scaled_representation(fmt, normalized_shape, axis) + if tuple(self.data.shape) != expected_data_shape: + raise ValueError(f"{fmt.value} data shape must be {expected_data_shape}, " f"got {tuple(self.data.shape)}") + if tuple(self.scale.shape) != expected_scale_shape: + raise ValueError(f"{fmt.value} scale shape must be {expected_scale_shape}, " f"got {tuple(self.scale.shape)}") + if self.data.dtype is not expected_data_dtype: + raise ValueError(f"{fmt.value} data must have dtype {expected_data_dtype}, " f"got {self.data.dtype}") + if self.scale.dtype is not expected_scale_dtype: + raise ValueError(f"{fmt.value} scale must have dtype {expected_scale_dtype}, " f"got {self.scale.dtype}") + object.__setattr__(self, "format", fmt) + object.__setattr__(self, "logical_shape", normalized_shape) + object.__setattr__(self, "axis", axis) + + @property + def shape(self) -> Tuple[int, ...]: + return self.logical_shape + + @property + def device(self) -> torch.device: + return self.data.device + + @property + def block_size(self) -> int: + return 32 if self.format is MoeFormat.MXFP8 else 16 + + def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Decode the logical, unswizzled block-scaled representation.""" + + logical_extent = self.logical_shape[self.axis] + scale = self.scale.movedim(self.axis, -1).float() + expanded_scale = scale.repeat_interleave( + self.block_size, + dim=-1, + )[..., :logical_extent] + + if self.format is MoeFormat.MXFP8: + values = self.data.movedim(self.axis, -1).float() + else: + packed = self.data.movedim(self.axis, -1) + low = packed & 0x0F + high = packed >> 4 + codes = torch.stack((low, high), dim=-1).flatten(-2)[..., :logical_extent] + table = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=torch.float32, + device=packed.device, + ) + values = table[codes.long()] + + return (values * expanded_scale).movedim(-1, self.axis).to(dtype) + + +class MoeEpNativeWeightLayout(str, Enum): + """Versioned kernel-native MXFP8 weight layouts.""" + + FORWARD_FC1_GATE_UP_INTERLEAVED_32_V1 = "mxfp8.forward_fc1.gate_up_interleaved_32.blocked_sf.v1" + FORWARD_FC2_K_MAJOR_V1 = "mxfp8.forward_fc2.k_major.blocked_sf.v1" + BACKWARD_W2_TRANSPOSE_V1 = "mxfp8.backward_w2_transpose.contiguous.blocked_sf.v1" + BACKWARD_W1_TRANSPOSE_GATE_UP_INTERLEAVED_32_V1 = "mxfp8.backward_w1_transpose.gate_up_interleaved_32.blocked_sf.v1" + + +@dataclass(frozen=True) +class MoeEpForwardWeights: + """Logical gate-then-up MXFP8 sources accepted by the fallback packer.""" + + fc1: BlockScaledTensor + fc2: BlockScaledTensor + + +@dataclass(frozen=True) +class MoeEpBackwardWeights: + """Logical gate-then-up MXFP8 transpose sources for the fallback packer.""" + + w2_transpose: BlockScaledTensor + w1_transpose: BlockScaledTensor + + +@dataclass(frozen=True) +class MoeEpNativeWeight: + """One kernel-executable payload plus Rubin blocked/interleaved scales.""" + + payload: torch.Tensor + scale: torch.Tensor + layout_id: Union[MoeEpNativeWeightLayout, str] + + def __post_init__(self) -> None: + if not isinstance(self.payload, torch.Tensor): + raise TypeError("payload must be a torch.Tensor, " f"got {type(self.payload).__name__}") + if not isinstance(self.scale, torch.Tensor): + raise TypeError("scale must be a torch.Tensor, " f"got {type(self.scale).__name__}") + if self.payload.device != self.scale.device: + raise ValueError(f"payload device {self.payload.device} does not match " f"scale device {self.scale.device}") + try: + layout_id = MoeEpNativeWeightLayout(self.layout_id) + except (TypeError, ValueError) as exc: + choices = ", ".join(layout.value for layout in MoeEpNativeWeightLayout) + raise ValueError(f"unsupported native weight layout_id {self.layout_id!r}; " f"expected one of: {choices}") from exc + object.__setattr__(self, "layout_id", layout_id) + + @property + def device(self) -> torch.device: + return self.payload.device + + +@dataclass(frozen=True) +class MoeEpNativeForwardWeights: + """Independent kernel-native weights consumed by one forward call.""" + + fc1: MoeEpNativeWeight + fc2: MoeEpNativeWeight + + +@dataclass(frozen=True) +class MoeEpNativeBackwardWeights: + """Independent kernel-native transpose weights consumed by one backward.""" + + w2_transpose: MoeEpNativeWeight + w1_transpose: MoeEpNativeWeight + + +@dataclass(frozen=True) +class MoeEpForwardWeightStaging: + """Caller-owned destinations used by forward weight materialization.""" + + fc1_payload: torch.Tensor + fc1_scale: torch.Tensor + fc2_payload: torch.Tensor + fc2_scale: torch.Tensor + + +@dataclass(frozen=True) +class MoeEpBackwardWeightStaging: + """Caller-owned destinations used by backward weight materialization.""" + + w2_transpose_payload: torch.Tensor + w2_transpose_scale: torch.Tensor + w1_transpose_payload: torch.Tensor + w1_transpose_scale: torch.Tensor + + +@dataclass(frozen=True) +class MoeEpTrainingForwardOutputs: + """Forward destinations, including the lane's symmetric output buffer.""" + + fc1_preact: torch.Tensor + output: torch.Tensor | None = None + fc1_a: torch.Tensor | None = None + fc1_sfa: torch.Tensor | None = None + valid_route_counts: torch.Tensor | None = None + expert_offsets: torch.Tensor | None = None + + +@dataclass(frozen=True) +class MoeEpTrainingBackwardOutputs: + """Backward destinations, including the lane's symmetric input gradient.""" + + grad_activation: torch.Tensor | None = None + dprob: torch.Tensor | None = None + fc1_b: torch.Tensor | None = None + fc1_sfb: torch.Tensor | None = None + fc2_a: torch.Tensor | None = None + fc2_sfa: torch.Tensor | None = None + fc2_b: torch.Tensor | None = None + fc2_sfb: torch.Tensor | None = None + + +@dataclass(frozen=True) +class MoeEpTrainingWgradOperands: + """Non-owning views over caller-owned grouped-WGrad operand buffers.""" + + fc1_a: torch.Tensor + fc1_sfa: torch.Tensor + fc1_b: torch.Tensor + fc1_sfb: torch.Tensor + fc2_a: torch.Tensor + fc2_sfa: torch.Tensor + fc2_b: torch.Tensor + fc2_sfb: torch.Tensor + expert_offsets: torch.Tensor + valid_route_counts: torch.Tensor + + +@dataclass(frozen=True) +class MoeEpExecutionLane: + """Operator-bound index of one mutable per-stream execution lane.""" + + index: int + _operator_token: object + + +MoeTensor = Union[torch.Tensor, BlockScaledTensor] + + +__all__ = [ + "BlockScaledTensor", + "MoeEpExecutionLane", + "MoeEpBackwardWeightStaging", + "MoeEpBackwardWeights", + "MoeEpForwardWeightStaging", + "MoeEpForwardWeights", + "MoeEpNativeBackwardWeights", + "MoeEpNativeForwardWeights", + "MoeEpNativeWeight", + "MoeEpNativeWeightLayout", + "MoeEpTrainingBackwardOutputs", + "MoeEpTrainingForwardOutputs", + "MoeEpTrainingWgradOperands", + "MoeFormat", + "MoeTensor", + "parse_format", +] diff --git a/python/cudnn/moe_ep/_validation.py b/python/cudnn/moe_ep/_validation.py new file mode 100644 index 000000000..694e84af3 --- /dev/null +++ b/python/cudnn/moe_ep/_validation.py @@ -0,0 +1,621 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Pure public-contract validation for :mod:`cudnn.moe_ep`.""" + +from __future__ import annotations + +from typing import Mapping, Tuple + +import torch + +from ._contracts import Fc1WeightLayout, ForwardConfig, ValidatedForwardRequest +from ._math import round_up +from ._types import ( + BlockScaledTensor, + MoeEpBackwardWeights, + MoeEpForwardWeights, + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, + MoeEpTrainingBackwardOutputs, + MoeEpTrainingForwardOutputs, + MoeFormat, + MoeTensor, + _block_scaled_representation, +) + +_SourceWeightSpec = tuple[str, BlockScaledTensor, Tuple[int, ...]] +_NativeWeightSpec = tuple[ + str, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, + Tuple[int, ...], + Tuple[int, ...], + int, +] + + +def _require_torch_dtype(name: str) -> torch.dtype: + dtype = getattr(torch, name, None) + if dtype is None: + raise RuntimeError(f"this PyTorch build does not provide torch.{name}") + return dtype + + +def _logical_shape(tensor: MoeTensor) -> Tuple[int, ...]: + if isinstance(tensor, BlockScaledTensor): + return tensor.logical_shape + if isinstance(tensor, torch.Tensor): + return tuple(tensor.shape) + raise ValueError(f"expected torch.Tensor or BlockScaledTensor, got {type(tensor).__name__}") + + +def _tensor_device(tensor: MoeTensor) -> torch.device: + if isinstance(tensor, BlockScaledTensor): + return tensor.device + if isinstance(tensor, torch.Tensor): + return tensor.device + raise ValueError(f"expected torch.Tensor or BlockScaledTensor, got {type(tensor).__name__}") + + +def _validate_strided(name: str, tensor: torch.Tensor) -> None: + if tensor.layout is not torch.strided: + raise ValueError(f"{name} must use torch.strided layout, got {tensor.layout}") + + +def _is_training_mxfp8_scale_layout(tensor: torch.Tensor) -> bool: + """Accept compact scales or logical views into padded lane-scale storage.""" + + if tensor.is_contiguous(): + return True + return ( + tensor.ndim == 2 + and tensor.stride(1) == 1 + and tensor.stride(0) >= tensor.shape[1] + ) + + +def _validate_tensor_representation( + name: str, + tensor: MoeTensor, + expected_logical_shape: Tuple[int, ...], +) -> None: + logical_shape = _logical_shape(tensor) + if logical_shape != expected_logical_shape: + raise ValueError(f"{name} logical shape must be {expected_logical_shape}, " f"got {logical_shape}") + if isinstance(tensor, torch.Tensor): + _validate_strided(name, tensor) + if not tensor.is_floating_point(): + raise ValueError(f"{name} must be floating point, got {tensor.dtype}") + return + if not isinstance(tensor, BlockScaledTensor): + raise ValueError(f"{name} must be a torch.Tensor or BlockScaledTensor, " f"got {type(tensor).__name__}") + if tensor.axis != 1: + raise ValueError(f"{name} block-scaled axis must be 1, got {tensor.axis}") + _validate_strided(f"{name}.data", tensor.data) + _validate_strided(f"{name}.scale", tensor.scale) + + ( + expected_data_shape, + expected_scale_shape, + expected_data_dtype, + expected_scale_dtype, + ) = _block_scaled_representation( + tensor.format, + expected_logical_shape, + tensor.axis, + ) + if tuple(tensor.data.shape) != expected_data_shape: + raise ValueError(f"{name}.data shape must be {expected_data_shape}, " f"got {tuple(tensor.data.shape)}") + if tuple(tensor.scale.shape) != expected_scale_shape: + raise ValueError(f"{name}.scale shape must be {expected_scale_shape}, " f"got {tuple(tensor.scale.shape)}") + if tensor.data.dtype != expected_data_dtype: + raise ValueError(f"{name}.data must have dtype {expected_data_dtype}, " f"got {tensor.data.dtype}") + if tensor.scale.dtype != expected_scale_dtype: + raise ValueError(f"{name}.scale must have dtype {expected_scale_dtype}, " f"got {tensor.scale.dtype}") + + +def _validate_expert_ids( + config: ForwardConfig, + topk_idx: torch.Tensor, +) -> None: + valid_experts = topk_idx.reshape(-1) + valid_experts = valid_experts[valid_experts != -1] + if valid_experts.numel() > 0 and bool(((valid_experts < 0) | (valid_experts >= config.num_experts)).any().item()): + raise ValueError("topk_idx contains out-of-range expert ids") + + +def _validate_routes( + config: ForwardConfig, + token_count: int, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + validate_expert_ids: bool, +) -> None: + if not isinstance(topk_idx, torch.Tensor): + raise ValueError(f"topk_idx must be a torch.Tensor, got {type(topk_idx).__name__}") + if not isinstance(topk_weights, torch.Tensor): + raise ValueError("topk_weights must be a torch.Tensor, " f"got {type(topk_weights).__name__}") + _validate_strided("topk_idx", topk_idx) + _validate_strided("topk_weights", topk_weights) + route_shape = (token_count, config.top_k) + if tuple(topk_idx.shape) != route_shape: + raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") + if tuple(topk_weights.shape) != route_shape: + raise ValueError(f"topk_weights shape must be {route_shape}, " f"got {tuple(topk_weights.shape)}") + if topk_idx.dtype not in (torch.int32, torch.int64): + raise ValueError("topk_idx must have dtype torch.int32 or torch.int64, " f"got {topk_idx.dtype}") + if not topk_weights.is_floating_point(): + raise ValueError(f"topk_weights must be floating point, got {topk_weights.dtype}") + if config.max_tokens_per_rank is not None and token_count > config.max_tokens_per_rank: + raise ValueError(f"token count {token_count} exceeds " f"max_tokens_per_rank={config.max_tokens_per_rank}") + if validate_expert_ids: + _validate_expert_ids(config, topk_idx) + + +def validate_forward( + config: ForwardConfig, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + validate_expert_ids: bool = True, +) -> ValidatedForwardRequest: + """Validate inference-forward semantics without importing a device backend.""" + + activation_shape = _logical_shape(activation) + if len(activation_shape) != 2 or activation_shape[1] != config.hidden_size: + raise ValueError(f"activation logical shape must be (T, {config.hidden_size}), " f"got {activation_shape}") + token_count = activation_shape[0] + _validate_tensor_representation("activation", activation, activation_shape) + _validate_tensor_representation( + "fc1_weight", + fc1_weight, + ( + config.experts_per_rank, + config.hidden_size, + 2 * config.intermediate_size, + ), + ) + _validate_tensor_representation( + "fc2_weight", + fc2_weight, + ( + config.experts_per_rank, + config.intermediate_size, + config.hidden_size, + ), + ) + if config.fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32 and ( + not isinstance(fc1_weight, BlockScaledTensor) or fc1_weight.format is not MoeFormat.MXFP8 + ): + raise ValueError("weight_interleave_size=32 requires an MXFP8 BlockScaledTensor " "for fc1_weight") + _validate_routes( + config, + token_count, + topk_idx, + topk_weights, + validate_expert_ids=False, + ) + device = _tensor_device(activation) + for name, tensor in ( + ("fc1_weight", fc1_weight), + ("fc2_weight", fc2_weight), + ("topk_idx", topk_idx), + ("topk_weights", topk_weights), + ): + tensor_device = _tensor_device(tensor) + if tensor_device != device: + raise ValueError(f"{name} must be on {device}, got {tensor_device}") + if device.type == "cuda": + with torch.cuda.device(device): + capturing = torch.cuda.is_current_stream_capturing() + else: + capturing = False + if validate_expert_ids and not capturing: + _validate_expert_ids(config, topk_idx) + return ValidatedForwardRequest( + config=config, + activation=activation, + fc1_weight=fc1_weight, + fc2_weight=fc2_weight, + topk_idx=topk_idx, + topk_weights=topk_weights, + token_count=token_count, + device=device, + ) + + +def _validate_source_weight( + name: str, + tensor: BlockScaledTensor, + shape: Tuple[int, ...], +) -> None: + _validate_tensor_representation(name, tensor, shape) + if not isinstance(tensor, BlockScaledTensor): + raise TypeError(f"{name} must be an MXFP8 BlockScaledTensor") + if tensor.format is not MoeFormat.MXFP8: + raise NotImplementedError(f"{name} must use format='mxfp8', got {tensor.format.value!r}") + if not tensor.data.is_contiguous() or not tensor.scale.is_contiguous(): + raise ValueError(f"{name} data and scale must be contiguous") + + +def _validate_source_weight_pair( + expected: tuple[_SourceWeightSpec, _SourceWeightSpec], +) -> torch.device: + for name, tensor, shape in expected: + _validate_source_weight(name, tensor, shape) + device = expected[0][1].device + second_name, second_tensor, _ = expected[1] + if second_tensor.device != device: + raise ValueError(f"{second_name} must be on {device}, got {second_tensor.device}") + return device + + +def validate_forward_source_weights( + config: ForwardConfig, + weights: MoeEpForwardWeights, +) -> torch.device: + """Validate source weights used by allocation-free forward packing.""" + + if not isinstance(weights, MoeEpForwardWeights): + raise TypeError("weights must be a MoeEpForwardWeights, " f"got {type(weights).__name__}") + expected = ( + ("weights.fc1", weights.fc1, (config.experts_per_rank, config.hidden_size, 2 * config.intermediate_size)), + ("weights.fc2", weights.fc2, (config.experts_per_rank, config.intermediate_size, config.hidden_size)), + ) + return _validate_source_weight_pair(expected) + + +def validate_backward_source_weights( + config: ForwardConfig, + weights: MoeEpBackwardWeights, +) -> torch.device: + """Validate source weights used by allocation-free backward packing.""" + + if not isinstance(weights, MoeEpBackwardWeights): + raise TypeError("weights must be a MoeEpBackwardWeights, " f"got {type(weights).__name__}") + expected = ( + ("weights.w2_transpose", weights.w2_transpose, (config.experts_per_rank, config.hidden_size, config.intermediate_size)), + ("weights.w1_transpose", weights.w1_transpose, (config.experts_per_rank, 2 * config.intermediate_size, config.hidden_size)), + ) + return _validate_source_weight_pair(expected) + + +def _blocked_scale_elements(raw_rows: int, raw_columns: int) -> int: + return round_up(raw_rows, 128) * round_up(raw_columns, 4) + + +def _validate_native_weight( + name: str, + weight: MoeEpNativeWeight, + *, + layout_id: MoeEpNativeWeightLayout, + payload_shape: Tuple[int, ...], + payload_stride: Tuple[int, ...], + scale_elements: int, + scale_dtype: torch.dtype, + device: torch.device | None, +) -> torch.device: + if not isinstance(weight, MoeEpNativeWeight): + raise TypeError(f"{name} must be a MoeEpNativeWeight, got {type(weight).__name__}") + if weight.layout_id is not layout_id: + raise ValueError(f"{name}.layout_id must be {layout_id.value!r}, " f"got {weight.layout_id.value!r}") + _validate_strided(f"{name}.payload", weight.payload) + _validate_strided(f"{name}.scale", weight.scale) + if tuple(weight.payload.shape) != payload_shape: + raise ValueError(f"{name}.payload shape must be {payload_shape}, " f"got {tuple(weight.payload.shape)}") + if tuple(weight.payload.stride()) != payload_stride: + raise ValueError(f"{name}.payload stride must be {payload_stride}, " f"got {tuple(weight.payload.stride())}") + expected_payload_dtype = _require_torch_dtype("float8_e4m3fn") + if weight.payload.dtype is not expected_payload_dtype: + raise ValueError(f"{name}.payload must have dtype {expected_payload_dtype}, " f"got {weight.payload.dtype}") + expected_scale_shape = (payload_shape[0], scale_elements) + if tuple(weight.scale.shape) != expected_scale_shape: + raise ValueError(f"{name}.scale shape must be {expected_scale_shape}, " f"got {tuple(weight.scale.shape)}") + if not weight.scale.is_contiguous(): + raise ValueError(f"{name}.scale must be contiguous") + if weight.scale.dtype is not scale_dtype: + raise ValueError(f"{name}.scale must have dtype {scale_dtype}, " f"got {weight.scale.dtype}") + for field_name, tensor in ( + ("payload", weight.payload), + ("scale", weight.scale), + ): + if tensor.data_ptr() % 16: + raise ValueError(f"{name}.{field_name} must be at least 16-byte aligned") + if device is not None and weight.device != device: + raise ValueError(f"{name} must be on {device}, got {weight.device}") + return weight.device + + +def _validate_native_weight_pair( + expected: tuple[_NativeWeightSpec, _NativeWeightSpec], + *, + scale_dtype: torch.dtype, + device: torch.device | None, +) -> torch.device: + resolved = device + for name, weight, layout_id, payload_shape, payload_stride, scale_elements in expected: + resolved = _validate_native_weight( + name, + weight, + layout_id=layout_id, + payload_shape=payload_shape, + payload_stride=payload_stride, + scale_elements=scale_elements, + scale_dtype=scale_dtype, + device=resolved, + ) + assert resolved is not None + return resolved + + +def validate_native_forward_weights( + config: ForwardConfig, + weights: MoeEpNativeForwardWeights, + *, + device: torch.device | None = None, +) -> torch.device: + if not isinstance(weights, MoeEpNativeForwardWeights): + raise TypeError("weights must be a MoeEpNativeForwardWeights, " f"got {type(weights).__name__}") + if config.fc1_weight_layout is not Fc1WeightLayout.GATE_UP_INTERLEAVED_32: + raise ValueError("native training weights require weight_interleave_size=32") + experts = config.experts_per_rank + hidden = config.hidden_size + intermediate = config.intermediate_size + sf_dtype = _require_torch_dtype("float8_e8m0fnu") + expected = ( + ( + "weights.fc1", + weights.fc1, + MoeEpNativeWeightLayout.FORWARD_FC1_GATE_UP_INTERLEAVED_32_V1, + (experts, hidden, 2 * intermediate), + (hidden * 2 * intermediate, 1, hidden), + _blocked_scale_elements(2 * intermediate, hidden // 32), + ), + ( + "weights.fc2", + weights.fc2, + MoeEpNativeWeightLayout.FORWARD_FC2_K_MAJOR_V1, + (experts, intermediate, hidden), + (intermediate * hidden, 1, intermediate), + _blocked_scale_elements(hidden, intermediate // 32), + ), + ) + return _validate_native_weight_pair( + expected, + scale_dtype=sf_dtype, + device=device, + ) + + +def validate_native_backward_weights( + config: ForwardConfig, + weights: MoeEpNativeBackwardWeights, + *, + device: torch.device | None = None, +) -> torch.device: + if not isinstance(weights, MoeEpNativeBackwardWeights): + raise TypeError("weights must be a MoeEpNativeBackwardWeights, " f"got {type(weights).__name__}") + if config.fc1_weight_layout is not Fc1WeightLayout.GATE_UP_INTERLEAVED_32: + raise ValueError("native training weights require weight_interleave_size=32") + experts = config.experts_per_rank + hidden = config.hidden_size + intermediate = config.intermediate_size + sf_dtype = _require_torch_dtype("float8_e8m0fnu") + expected = ( + ( + "weights.w2_transpose", + weights.w2_transpose, + MoeEpNativeWeightLayout.BACKWARD_W2_TRANSPOSE_V1, + (experts, hidden, intermediate), + (hidden * intermediate, intermediate, 1), + _blocked_scale_elements(intermediate, hidden // 32), + ), + ( + "weights.w1_transpose", + weights.w1_transpose, + MoeEpNativeWeightLayout.BACKWARD_W1_TRANSPOSE_GATE_UP_INTERLEAVED_32_V1, + (experts, 2 * intermediate, hidden), + (2 * intermediate * hidden, hidden, 1), + _blocked_scale_elements(hidden, 2 * intermediate // 32), + ), + ) + return _validate_native_weight_pair( + expected, + scale_dtype=sf_dtype, + device=device, + ) + + +def validate_training_input( + config: ForwardConfig, + name: str, + value: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + device: torch.device, +) -> int: + logical_shape = _logical_shape(value) + if len(logical_shape) != 2 or logical_shape[1] != config.hidden_size: + raise ValueError(f"{name} logical shape must be (T, {config.hidden_size}), " f"got {logical_shape}") + _validate_tensor_representation(name, value, logical_shape) + if isinstance(value, BlockScaledTensor): + if value.format is not MoeFormat.MXFP8: + raise NotImplementedError(f"{name} only supports MXFP8 block scaling") + if not value.data.is_contiguous(): + raise ValueError(f"{name} MXFP8 data must be contiguous") + if not _is_training_mxfp8_scale_layout(value.scale): + raise ValueError( + f"{name} MXFP8 scale must be contiguous or a row-major view " + "with padded row stride" + ) + elif value.dtype not in (torch.bfloat16, torch.float32): + raise TypeError(f"{name} must be BF16, FP32, or an MXFP8 BlockScaledTensor") + elif not value.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + token_count = logical_shape[0] + if device.type == "cuda": + with torch.cuda.device(device): + capturing = torch.cuda.is_current_stream_capturing() + else: + capturing = False + _validate_routes( + config, + token_count, + topk_idx, + topk_weights, + validate_expert_ids=not capturing, + ) + if topk_idx.dtype is not torch.int32 or not topk_idx.is_contiguous(): + raise TypeError("training topk_idx must be contiguous torch.int32") + if topk_weights.dtype is not torch.float32 or not topk_weights.is_contiguous(): + raise TypeError("training topk_weights must be contiguous torch.float32") + tensors = ( + (name, value), + ("topk_idx", topk_idx), + ("topk_weights", topk_weights), + ) + for tensor_name, tensor in tensors: + tensor_device = _tensor_device(tensor) + if tensor_device != device: + raise ValueError(f"{tensor_name} must be on {device}, got {tensor_device}") + return token_count + + +def _tensor_byte_range(tensor: torch.Tensor) -> tuple[int, int]: + byte_start = tensor.data_ptr() + max_element_offset = sum((int(extent) - 1) * int(step) for extent, step in zip(tensor.shape, tensor.stride()) if int(extent) > 0) + byte_end = byte_start + (0 if tensor.numel() == 0 else (max_element_offset + 1) * tensor.element_size()) + return byte_start, byte_end + + +def _assert_no_overlap( + name: str, + byte_range: tuple[int, int], + ranges: list[tuple[int, int, str]], +) -> None: + byte_start, byte_end = byte_range + for other_start, other_end, other_name in ranges: + if byte_start < other_end and other_start < byte_end: + raise ValueError(f"{name} must not alias {other_name}") + + +def _validate_named_buffers( + tensors: Mapping[str, object], + requirements: Mapping[str, tuple[Tuple[int, ...], Tuple[int, ...], torch.dtype, int]], + *, + device: torch.device, +) -> None: + ranges: list[tuple[int, int, str]] = [] + for name, requirement in requirements.items(): + tensor = tensors[name] + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"out.{name} must be a torch.Tensor") + shape, stride, dtype, alignment = requirement + if tuple(tensor.shape) != tuple(shape): + raise ValueError(f"out.{name} shape must be {tuple(shape)}, got {tuple(tensor.shape)}") + if tuple(tensor.stride()) != tuple(stride): + raise ValueError(f"out.{name} stride must be {tuple(stride)}, " f"got {tuple(tensor.stride())}") + if tensor.dtype is not dtype: + raise ValueError(f"out.{name} dtype must be {dtype}, got {tensor.dtype}") + if tensor.device != device: + raise ValueError(f"out.{name} must be on {device}, got {tensor.device}") + if tensor.data_ptr() % alignment: + raise ValueError(f"out.{name} must be {alignment}-byte aligned") + qualified_name = f"out.{name}" + byte_start, byte_end = _tensor_byte_range(tensor) + _assert_no_overlap(qualified_name, (byte_start, byte_end), ranges) + ranges.append((byte_start, byte_end, qualified_name)) + + +def _validate_output_buffers( + output: object, + requirements: Mapping[str, tuple[Tuple[int, ...], Tuple[int, ...], torch.dtype, int]], + *, + device: torch.device, +) -> None: + _validate_named_buffers( + {name: getattr(output, name) for name in requirements}, + requirements, + device=device, + ) + + +def validate_training_non_aliasing( + tensors: Mapping[str, torch.Tensor | None], +) -> None: + """Reject overlapping caller inputs, saved state, weights, and outputs.""" + + ranges: list[tuple[int, int, str]] = [] + for name, tensor in tensors.items(): + if tensor is None or tensor.numel() == 0: + continue + byte_start, byte_end = _tensor_byte_range(tensor) + _assert_no_overlap(name, (byte_start, byte_end), ranges) + ranges.append((byte_start, byte_end, name)) + + +def validate_training_forward_outputs( + output: MoeEpTrainingForwardOutputs, + requirements: Mapping[str, tuple[Tuple[int, ...], Tuple[int, ...], torch.dtype, int]], + *, + device: torch.device, +) -> None: + if not isinstance(output, MoeEpTrainingForwardOutputs): + raise TypeError("out must be a MoeEpTrainingForwardOutputs, " f"got {type(output).__name__}") + _validate_output_buffers(output, requirements, device=device) + + +def validate_training_forward_state( + *, + fc1_preact: torch.Tensor, + fc1_a: torch.Tensor | None, + fc1_sfa: torch.Tensor | None, + valid_route_counts: torch.Tensor | None, + expert_offsets: torch.Tensor | None, + requirements: Mapping[str, tuple[Tuple[int, ...], Tuple[int, ...], torch.dtype, int]], + device: torch.device, +) -> None: + _validate_named_buffers( + { + "fc1_preact": fc1_preact, + "fc1_a": fc1_a, + "fc1_sfa": fc1_sfa, + "valid_route_counts": valid_route_counts, + "expert_offsets": expert_offsets, + }, + requirements, + device=device, + ) + + +def validate_training_backward_outputs( + output: MoeEpTrainingBackwardOutputs, + requirements: Mapping[str, tuple[Tuple[int, ...], Tuple[int, ...], torch.dtype, int]], + *, + device: torch.device, +) -> None: + if not isinstance(output, MoeEpTrainingBackwardOutputs): + raise TypeError("out must be a MoeEpTrainingBackwardOutputs, " f"got {type(output).__name__}") + _validate_output_buffers(output, requirements, device=device) + + +__all__ = [ + "validate_backward_source_weights", + "validate_forward", + "validate_forward_source_weights", + "validate_native_backward_weights", + "validate_native_forward_weights", + "validate_training_backward_outputs", + "validate_training_forward_outputs", + "validate_training_forward_state", + "validate_training_input", + "validate_training_non_aliasing", +] diff --git a/python/cudnn/moe_ep/api.py b/python/cudnn/moe_ep/api.py new file mode 100644 index 000000000..6b60e703a --- /dev/null +++ b/python/cudnn/moe_ep/api.py @@ -0,0 +1,1318 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Python API surface for fused SwiGLU MoE with expert parallelism. + +The public API performs contract validation and dispatches through a private, +lazy backend seam. Device-runtime implementation details remain outside this +module. +""" + +from __future__ import annotations + +import contextlib +import math +import threading +import warnings +from dataclasses import replace +from numbers import Real +from typing import Mapping, Optional, Sequence, Union + +import torch +import torch.distributed as dist + +from ._contracts import Fc1WeightLayout, ForwardConfig, normalize_fc1_weight_layout +from ._tuning import ( + MoeEpAutotuneCandidateResult, + MoeEpAutotuneResult, + MoeEpTuningConfig, +) +from ._types import ( + BlockScaledTensor, + MoeEpBackwardWeightStaging, + MoeEpBackwardWeights, + MoeEpExecutionLane, + MoeEpForwardWeightStaging, + MoeEpForwardWeights, + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpTrainingBackwardOutputs, + MoeEpTrainingForwardOutputs, + MoeEpTrainingWgradOperands, + MoeFormat, + MoeTensor, + parse_format as _parse_format, +) +from ._validation import ( + validate_backward_source_weights, + validate_forward, + validate_forward_source_weights, + validate_native_backward_weights, + validate_native_forward_weights, + validate_training_backward_outputs, + validate_training_forward_outputs, + validate_training_forward_state, + validate_training_input, + validate_training_non_aliasing, +) + + +def _resolve_ep_topology( + ep_group: Optional[dist.ProcessGroup], +) -> tuple[int, int, tuple[int, ...]]: + """Return dense EP rank/size plus its ordered global-rank membership.""" + + if ep_group is None: + return 1, 0, () + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("ep_group requires an initialized torch.distributed process group") + + ep_size = dist.get_world_size(ep_group) + ep_rank = dist.get_rank(ep_group) + if ep_size <= 0 or ep_rank < 0 or ep_rank >= ep_size: + raise ValueError("the current process must be a member of ep_group") + + ep_global_ranks = tuple(dist.get_global_rank(ep_group, group_rank) for group_rank in range(ep_size)) + if len(set(ep_global_ranks)) != ep_size: + raise RuntimeError("ep_group returned duplicate global ranks") + if ep_global_ranks[ep_rank] != dist.get_rank(): + raise RuntimeError("ep_group rank mapping is inconsistent with the current global rank") + return ep_size, ep_rank, ep_global_ranks + + +def _validate_training_assert_capability(config: ForwardConfig) -> None: + """Fail before allocation when graph error-mode primitives are unavailable.""" + + if config.drop_on_overflow: + return + if not callable(getattr(torch, "_assert_async", None)): + raise RuntimeError("drop_on_overflow=False training requires callable " "torch._assert_async before CUDA Graph capture") + if config.ep_size <= 1: + return + backend = dist.get_backend(config.ep_group) + if backend != dist.Backend.NCCL and str(backend).lower() != "nccl": + raise NotImplementedError("drop_on_overflow=False EP2+ training requires an NCCL " "process group for the captured scalar global overflow OR") + + +def _resolve_training_device( + device: torch.device | str | int | None, +) -> torch.device: + if device is None: + if not torch.cuda.is_available(): + raise RuntimeError("prepare_training requires an available CUDA device") + return torch.device("cuda", torch.cuda.current_device()) + if isinstance(device, bool): + raise TypeError("device must be a CUDA device, ordinal, or None") + if isinstance(device, int): + resolved = torch.device("cuda", device) + else: + resolved = torch.device(device) + if resolved.type == "cuda" and resolved.index is None: + resolved = torch.device("cuda", torch.cuda.current_device()) + if resolved.type != "cuda": + raise ValueError(f"training device must be CUDA, got {resolved}") + if resolved.index is None or resolved.index < 0 or resolved.index >= torch.cuda.device_count(): + raise ValueError(f"CUDA device {resolved} is not available") + return resolved + + +def _named_moe_tensors( + name: str, + value: MoeTensor, +) -> dict[str, torch.Tensor]: + if isinstance(value, BlockScaledTensor): + return { + f"{name}.data": value.data, + f"{name}.scale": value.scale, + } + return {name: value} + + +def pack_forward_weights( + weights: MoeEpForwardWeights, + *, + out: MoeEpForwardWeightStaging, +) -> MoeEpNativeForwardWeights: + """Standalone allocation-free forward weight materialization.""" + + from ._megamoe_backend.mxfp8._training_weights import materialize_forward + + return materialize_forward( + weights, + out=out, + fc1_weight_layout=Fc1WeightLayout.GATE_UP_INTERLEAVED_32, + ) + + +def pack_backward_weights( + weights: MoeEpBackwardWeights, + *, + out: MoeEpBackwardWeightStaging, +) -> MoeEpNativeBackwardWeights: + """Standalone allocation-free backward weight materialization.""" + + from ._megamoe_backend.mxfp8._training_weights import materialize_backward + + return materialize_backward( + weights, + out=out, + fc1_weight_layout=Fc1WeightLayout.GATE_UP_INTERLEAVED_32, + ) + + +class MoeEp: + """Fused SwiGLU MoE operator with contiguous expert parallel sharding. + + Global expert ``e`` belongs to group-relative EP rank + ``e // experts_per_rank``. The constructor captures static configuration; + calling the instance accepts runtime tensors for this rank. + + The Rubin training-Mega backend accepts plain BF16/FP32 operands + (staged to MXFP8 E4M3) or MXFP8 ``BlockScaledTensor`` operands. Final + output is BF16. ``combine_format`` may be BF16 or MXFP8; forward MXFP8 + combine quantizes each FP32 route accumulator directly before top-k + reduction. The Rubin training backend requires + ``apply_topk_in_fc1=True``. + Native NVFP4 operands and NVFP4 combine/output are not executable. + + ``__call__`` is the inference-only forward surface. Training uses + :meth:`prepare_training`, :meth:`training_forward`, and + :meth:`training_backward`. Caller-owned output bundles carry all explicit + cross-phase state; the operator retains only private runtime and lane + scratch. + + The backend is created lazily on the first supported forward call. Valid + combinations outside the current backend capability matrix fail explicitly + instead of returning uninitialized storage. Once created, a backend and its + workspaces are bound to that call's device; use a separate ``MoeEp`` + instance for another device. + """ + + def __init__( + self, + *, + num_experts: int, + hidden_size: int, + intermediate_size: int, + top_k: int, + ep_group: Optional[dist.ProcessGroup] = None, + max_tokens_per_rank: Optional[int] = None, + max_recv_size_per_rank: Optional[int] = None, + drop_on_overflow: bool = False, + output_format: Union[MoeFormat, str] = MoeFormat.BF16, + combine_format: Union[MoeFormat, str] = MoeFormat.BF16, + apply_topk_in_fc1: bool = True, + weight_interleave_size: Optional[int] = None, + gate_up_clamp: Optional[float] = None, + token_padding_size: int = 128, + sf_padding_size: int = 128, + tuning: Optional[MoeEpTuningConfig] = None, + forward_tuning: Optional[MoeEpTuningConfig] = None, + backward_tuning: Optional[MoeEpTuningConfig] = None, + ) -> None: + self._lifecycle_lock = threading.RLock() + for name, value in ( + ("num_experts", num_experts), + ("hidden_size", hidden_size), + ("intermediate_size", intermediate_size), + ("top_k", top_k), + ): + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + if top_k > num_experts: + raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})") + if max_tokens_per_rank is not None and (isinstance(max_tokens_per_rank, bool) or not isinstance(max_tokens_per_rank, int) or max_tokens_per_rank < 0): + raise ValueError("max_tokens_per_rank must be a non-negative integer or None") + if max_recv_size_per_rank is not None and ( + isinstance(max_recv_size_per_rank, bool) or not isinstance(max_recv_size_per_rank, int) or max_recv_size_per_rank <= 0 + ): + raise ValueError("max_recv_size_per_rank must be a positive integer or None") + if not isinstance(drop_on_overflow, bool): + raise ValueError("drop_on_overflow must be a bool") + if not isinstance(apply_topk_in_fc1, bool): + raise ValueError("apply_topk_in_fc1 must be a bool") + fc1_weight_layout = normalize_fc1_weight_layout(weight_interleave_size) + for name, value in ( + ("token_padding_size", token_padding_size), + ("sf_padding_size", sf_padding_size), + ): + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + if sf_padding_size % 128: + raise ValueError( + "sf_padding_size must be a positive multiple of 128, " + f"got {sf_padding_size}" + ) + for name, value in ( + ("tuning", tuning), + ("forward_tuning", forward_tuning), + ("backward_tuning", backward_tuning), + ): + if value is not None and not isinstance(value, MoeEpTuningConfig): + raise TypeError( + f"{name} must be a MoeEpTuningConfig or None, " + f"got {type(value).__name__}" + ) + if tuning is not None and forward_tuning is not None: + raise ValueError("tuning and forward_tuning are aliases; pass only one") + if gate_up_clamp is not None: + if isinstance(gate_up_clamp, bool) or not isinstance(gate_up_clamp, Real): + raise ValueError("gate_up_clamp must be a finite real number or None") + gate_up_clamp = float(gate_up_clamp) + if not math.isfinite(gate_up_clamp): + raise ValueError("gate_up_clamp must be a finite real number or None") + + if ep_group is not None and not isinstance(ep_group, dist.ProcessGroup): + raise ValueError(f"ep_group must be a torch.distributed.ProcessGroup or None, " f"got {type(ep_group).__name__}") + ep_size, ep_rank, ep_global_ranks = _resolve_ep_topology(ep_group) + if num_experts % ep_size != 0: + raise ValueError(f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})") + + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.top_k = top_k + self.ep_group = ep_group + self.ep_size = ep_size + self.ep_rank = ep_rank + self.ep_global_ranks = ep_global_ranks + self.experts_per_rank = num_experts // ep_size + self.max_tokens_per_rank = max_tokens_per_rank + self.max_recv_size_per_rank = max_recv_size_per_rank + self.drop_on_overflow = drop_on_overflow + self.output_format = _parse_format(output_format) + self.combine_format = _parse_format(combine_format) + self.apply_topk_in_fc1 = apply_topk_in_fc1 + self.weight_interleave_size = weight_interleave_size + self._fc1_weight_layout = fc1_weight_layout + self.gate_up_clamp = None if gate_up_clamp is None else abs(gate_up_clamp) + self.token_padding_size = token_padding_size + self.sf_padding_size = sf_padding_size + self.forward_tuning = ( + forward_tuning + if forward_tuning is not None + else MoeEpTuningConfig() if tuning is None else tuning + ) + self.backward_tuning = ( + backward_tuning + if backward_tuning is not None + else MoeEpTuningConfig() + ) + # Backward-compatible alias for inference and forward-only autotuning. + self.tuning = self.forward_tuning + if self.forward_tuning.reduce_topk_in_kernel and ( + self.combine_format is not MoeFormat.BF16 + or self.output_format is not MoeFormat.BF16 + or not self.apply_topk_in_fc1 + ): + raise ValueError( + "reduce_topk_in_kernel requires BF16 combine/output and " + "apply_topk_in_fc1=True" + ) + if self.backward_tuning.token_back_mode != "epi_warps": + raise ValueError("backward_tuning requires token_back_mode='epi_warps'") + if self.backward_tuning.reduce_topk_in_kernel: + raise ValueError("backward_tuning does not support reduce_topk_in_kernel=True") + + for name, fmt in ( + ("output_format", self.output_format), + ("combine_format", self.combine_format), + ): + required_multiple = 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + if hidden_size % required_multiple != 0: + raise ValueError(f"hidden_size ({hidden_size}) must be divisible by " f"{required_multiple} for {name}={fmt.value}") + + self._forward_config = ForwardConfig( + num_experts=self.num_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + top_k=self.top_k, + experts_per_rank=self.experts_per_rank, + ep_size=self.ep_size, + ep_rank=self.ep_rank, + ep_group=self.ep_group, + ep_global_ranks=self.ep_global_ranks, + max_tokens_per_rank=self.max_tokens_per_rank, + max_recv_size_per_rank=self.max_recv_size_per_rank, + drop_on_overflow=self.drop_on_overflow, + output_format=self.output_format.value, + combine_format=self.combine_format.value, + apply_topk_in_fc1=self.apply_topk_in_fc1, + fc1_weight_layout=self._fc1_weight_layout, + gate_up_clamp=self.gate_up_clamp, + generate_c=False, + token_padding_size=self.token_padding_size, + sf_padding_size=self.sf_padding_size, + tuning=self.forward_tuning, + backward_tuning=self.backward_tuning, + backward_wgrad_mode="none", + ) + self._forward_backend = None + self._forward_backend_device = None + self._validated_topk_idx = None + self._validated_topk_version = None + self._operator_token = object() + self._training_state = None + self._training_lanes: tuple[MoeEpExecutionLane, ...] = () + self._training_requirements: ( + Mapping[ + str, + tuple[tuple[int, ...], tuple[int, ...], torch.dtype, int], + ] + | None + ) = None + self._poisoned = False + self._closed = False + + @staticmethod + def _tensor_version(tensor: torch.Tensor) -> int | None: + if not isinstance(tensor, torch.Tensor): + return None + try: + return tensor._version + except RuntimeError: + return None + + def _get_backend(self, request): + """Create and cache the private backend on first supported use.""" + + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError("MoeEp is unusable after an autotune runtime failure") + from . import _backend + + if self._forward_backend is not None and request.device != self._forward_backend_device: + raise ValueError(f"MoeEp backend is bound to {self._forward_backend_device}; " f"create a separate MoeEp instance for {request.device}") + + _backend.validate_config(self._forward_config) + _backend.validate_request(request) + + if self._forward_backend is None: + self._forward_backend = _backend.create_backend( + self._forward_config, + request.device, + ) + self._forward_backend_device = request.device + return self._forward_backend + + def __call__( + self, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ) -> MoeTensor: + """Validate and dispatch one fused MoE+EP forward call. + + Expected logical shapes are ``activation=(T,H)``, + ``fc1_weight=(E_local,H,2I)``, ``fc2_weight=(E_local,I,H)``, and + ``topk_idx=topk_weights=(T,K)``. + + Training callers must use :meth:`prepare_training` followed by the + stateless training methods. + """ + + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError("MoeEp is unusable after an autotune runtime failure") + topk_version = self._tensor_version(topk_idx) + validate_expert_ids = not (self._validated_topk_idx is topk_idx and topk_version is not None and topk_version == self._validated_topk_version) + request = validate_forward( + self._forward_config, + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + validate_expert_ids=validate_expert_ids, + ) + version_after_validation = self._tensor_version(topk_idx) + if topk_version is not None and topk_version == version_after_validation: + self._validated_topk_idx = topk_idx + self._validated_topk_version = topk_version + else: + self._validated_topk_idx = None + self._validated_topk_version = None + return self._get_backend(request).forward(request) + + def autotune( + self, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + candidates: Sequence[MoeEpTuningConfig], + warmup_iters: int = 3, + timed_iters: int = 10, + max_candidates: int = 32, + ) -> MoeEpAutotuneResult: + """Collectively sweep inference configurations and apply the winner. + + Candidate compilation, allocation, and warmup are excluded from CUDA + Event timing. The measured region includes input/weight staging, the + MegaMoE launch, and the output copy performed by a normal forward. + """ + + from . import _backend + from ._autotune import ( + benchmark_candidate, + normalize_candidates, + raise_preflight_errors, + select_winner, + synchronize_candidate, + verify_candidates_across_ranks, + verify_state_across_ranks, + ) + + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError("MoeEp is unusable after an autotune runtime failure") + if self._training_state is not None: + raise RuntimeError("autotune must be called before prepare_training()") + + normalized = normalize_candidates( + self.tuning, + candidates, + warmup_iters=warmup_iters, + timed_iters=timed_iters, + max_candidates=max_candidates, + ) + verify_candidates_across_ranks(normalized, self._forward_config.ep_group) + verify_state_across_ranks( + ( + self._forward_backend is not None, + self._training_state is not None, + (None if self._forward_backend_device is None else str(self._forward_backend_device)), + ), + self._forward_config.ep_group, + ) + + candidate_requests = [] + preflight_error: BaseException | None = None + try: + for index, tuning in enumerate(normalized): + try: + config = replace(self._forward_config, tuning=tuning) + request = validate_forward( + config, + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ) + if request.device.type != "cuda": + raise ValueError(f"autotune requires CUDA inputs, got {request.device}") + with torch.cuda.device(request.device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("autotune cannot run during CUDA Graph capture") + _backend.validate_config(config) + _backend.validate_request(request) + candidate_requests.append(request) + except BaseException as exc: + raise RuntimeError(f"MoeEp autotune candidate {index} {tuning!r} " f"failed during preflight: {exc}") from exc + except BaseException as exc: + preflight_error = exc + raise_preflight_errors( + preflight_error, + phase="inference preflight", + group=self._forward_config.ep_group, + ) + assert candidate_requests + device = candidate_requests[0].device + + if self._forward_backend is not None: + try: + synchronize_candidate(device, self._forward_config.ep_group) + self._forward_backend.close() + self._forward_backend = None + self._forward_backend_device = None + if self._forward_config.ep_group is not None: + dist.barrier(group=self._forward_config.ep_group) + except BaseException as exc: + self._poisoned = True + raise RuntimeError(f"MoeEp autotune failed during active backend teardown: {exc}") from exc + + results: list[MoeEpAutotuneCandidateResult] = [] + for index, (tuning, request) in enumerate(zip(normalized, candidate_requests)): + backend = None + runtime_entered = False + phase = "backend creation" + try: + backend = _backend.create_backend(request.config, device) + phase = "compile/prime" + runtime_entered = True + with torch.cuda.device(device): + output = backend.forward(request) + del output + phase = "warmup" + for _ in range(warmup_iters): + output = backend.forward(request) + del output + phase = "pre-timing synchronize" + synchronize_candidate(device, self._forward_config.ep_group) + phase = "timing" + latency_ms, samples_ms = benchmark_candidate( + lambda: backend.forward(request), + device=device, + group=self._forward_config.ep_group, + timed_iters=timed_iters, + ) + phase = "post-timing synchronize" + synchronize_candidate(device, self._forward_config.ep_group) + results.append( + MoeEpAutotuneCandidateResult( + tuning=tuning, + latency_ms=latency_ms, + samples_ms=samples_ms, + ) + ) + phase = "teardown" + backend.close() + backend = None + if self._forward_config.ep_group is not None: + dist.barrier(group=self._forward_config.ep_group) + except BaseException as exc: + if backend is not None and self._forward_config.ep_size == 1: + with contextlib.suppress(Exception): + backend.close() + if runtime_entered: + self._poisoned = True + raise RuntimeError(f"MoeEp autotune candidate {index} {tuning!r} failed during {phase}: {exc}") from exc + + winner = select_winner(results) + winner_request = candidate_requests[normalized.index(winner.tuning)] + winner_backend = None + try: + winner_backend = _backend.create_backend( + winner_request.config, + device, + ) + with torch.cuda.device(device): + output = winner_backend.forward(winner_request) + del output + synchronize_candidate(device, self._forward_config.ep_group) + except BaseException as exc: + if winner_backend is not None and self._forward_config.ep_size == 1: + with contextlib.suppress(Exception): + winner_backend.close() + self._poisoned = True + raise RuntimeError(f"MoeEp autotune winner {winner.tuning!r} failed final validation: {exc}") from exc + + self.forward_tuning = winner.tuning + self.tuning = self.forward_tuning + self._forward_config = winner_request.config + self._forward_backend = winner_backend + self._forward_backend_device = device + self._validated_topk_idx = None + self._validated_topk_version = None + return MoeEpAutotuneResult( + mode="inference", + winner=winner.tuning, + candidates=tuple(results), + ) + + def autotune_training( + self, + activation: MoeTensor, + grad_output: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + forward_weights: MoeEpNativeForwardWeights, + backward_weights: MoeEpNativeBackwardWeights, + candidates: Sequence[MoeEpTuningConfig], + warmup_iters: int = 3, + timed_iters: int = 10, + max_candidates: int = 32, + ) -> MoeEpAutotuneResult: + """Sweep complete training forward+backward latency and apply the winner. + + This collective API uses private one-lane temporary resources. It must + run before :meth:`prepare_training` and accepts kernel-native weights + so packing allocation and source-layout conversion are not timed. + """ + + from . import _backend + from ._autotune import ( + allocate_training_outputs, + benchmark_candidate, + normalize_candidates, + raise_preflight_errors, + select_winner, + synchronize_candidate, + verify_candidates_across_ranks, + verify_state_across_ranks, + ) + from ._megamoe_backend.mxfp8._training_execute import ( + launch_training_backward, + launch_training_forward, + ) + + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError("MoeEp is unusable after an autotune runtime failure") + if self._training_state is not None: + raise RuntimeError("autotune_training must be called before prepare_training()") + + normalized = normalize_candidates( + self.tuning, + candidates, + warmup_iters=warmup_iters, + timed_iters=timed_iters, + max_candidates=max_candidates, + ) + verify_candidates_across_ranks(normalized, self._forward_config.ep_group) + verify_state_across_ranks( + ( + self._forward_backend is not None, + self._training_state is not None, + (None if self._forward_backend_device is None else str(self._forward_backend_device)), + ), + self._forward_config.ep_group, + ) + + device: torch.device | None = None + preflight_error: BaseException | None = None + candidate_configs: list[ForwardConfig] = [] + token_count = -1 + try: + device = torch.device(activation.device) + if device.type != "cuda": + raise ValueError(f"autotune_training requires CUDA inputs, got {device}") + with torch.cuda.device(device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("autotune_training cannot run during CUDA Graph capture") + for index, tuning in enumerate(normalized): + try: + config = replace(self._forward_config, tuning=tuning) + _validate_training_assert_capability(config) + _backend.validate_config(config) + activation_tokens = validate_training_input( + config, + "activation", + activation, + topk_idx, + topk_weights, + device=device, + ) + grad_tokens = validate_training_input( + config, + "grad_output", + grad_output, + topk_idx, + topk_weights, + device=device, + ) + if activation_tokens != grad_tokens: + raise ValueError("activation and grad_output must have the same token " f"count, got {activation_tokens} and {grad_tokens}") + validate_native_forward_weights(config, forward_weights, device=device) + validate_native_backward_weights(config, backward_weights, device=device) + token_count = activation_tokens + candidate_configs.append(config) + except BaseException as exc: + raise RuntimeError(f"MoeEp autotune_training candidate {index} {tuning!r} " f"failed during preflight: {exc}") from exc + except BaseException as exc: + preflight_error = exc + raise_preflight_errors( + preflight_error, + phase="training preflight", + group=self._forward_config.ep_group, + ) + assert device is not None and candidate_configs and token_count >= 0 + + if self._forward_backend is not None: + try: + synchronize_candidate(device, self._forward_config.ep_group) + self._forward_backend.close() + self._forward_backend = None + self._forward_backend_device = None + if self._forward_config.ep_group is not None: + dist.barrier(group=self._forward_config.ep_group) + except BaseException as exc: + self._poisoned = True + raise RuntimeError("MoeEp autotune_training failed during active backend " f"teardown: {exc}") from exc + + results: list[MoeEpAutotuneCandidateResult] = [] + for index, (tuning, config) in enumerate(zip(normalized, candidate_configs)): + backend = None + runtime_entered = False + phase = "backend creation" + try: + backend = _backend.create_backend(config, device) + runtime_entered = True + phase = "training preparation" + with torch.cuda.device(device): + state = backend.prepare_training(lane_count=1) + requirements = state.public_requirements() + symmetric_buffers = state.public_symmetric_buffers(0) + forward_out, backward_out = allocate_training_outputs( + requirements, + device, + symmetric_buffers, + ) + forward_names = ( + "output", + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + ) + backward_names = ( + "grad_activation", + "dprob", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + ) + validate_training_forward_outputs( + forward_out, + {name: requirements[name] for name in forward_names}, + device=device, + ) + validate_training_backward_outputs( + backward_out, + {name: requirements[name] for name in backward_names}, + device=device, + ) + validate_training_forward_state( + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + requirements={ + name: requirements[name] + for name in ( + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + ) + }, + device=device, + ) + execution = state.views(lane=0, token_count=token_count) + + def run_training_pair(): + launch_training_forward( + state, + execution, + activation, + topk_idx, + topk_weights, + weights=forward_weights, + out=forward_out, + ) + return launch_training_backward( + state, + execution, + grad_output, + topk_idx, + topk_weights, + weights=backward_weights, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, + ) + + phase = "compile/prime" + run_training_pair() + phase = "warmup" + for _ in range(warmup_iters): + run_training_pair() + phase = "pre-timing synchronize" + synchronize_candidate(device, self._forward_config.ep_group) + phase = "timing" + latency_ms, samples_ms = benchmark_candidate( + run_training_pair, + device=device, + group=self._forward_config.ep_group, + timed_iters=timed_iters, + ) + phase = "post-timing synchronize" + synchronize_candidate(device, self._forward_config.ep_group) + results.append( + MoeEpAutotuneCandidateResult( + tuning=tuning, + latency_ms=latency_ms, + samples_ms=samples_ms, + ) + ) + phase = "teardown" + backend.close() + backend = None + if self._forward_config.ep_group is not None: + dist.barrier(group=self._forward_config.ep_group) + except BaseException as exc: + if backend is not None and self._forward_config.ep_size == 1: + with contextlib.suppress(Exception): + backend.close() + if runtime_entered: + self._poisoned = True + raise RuntimeError(f"MoeEp autotune_training candidate {index} {tuning!r} " f"failed during {phase}: {exc}") from exc + + winner = select_winner(results) + winner_config = candidate_configs[normalized.index(winner.tuning)] + self.forward_tuning = winner.tuning + self.tuning = self.forward_tuning + self._forward_config = winner_config + self._forward_backend = None + self._forward_backend_device = None + self._validated_topk_idx = None + self._validated_topk_version = None + return MoeEpAutotuneResult( + mode="training", + winner=winner.tuning, + candidates=tuple(results), + ) + + def warmup( + self, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ) -> None: + """Prepare a forward plan for CUDA Graph capture. + + This runs one complete eager forward and synchronizes its CUDA device, + forcing runtime bootstrap, symmetric allocation, weight staging, JIT + compilation, and the first real kernel launch to finish before capture. + + For expert-parallel execution this method is collective by contract: + every rank in ``ep_group`` must call it concurrently with valid inputs. + It intentionally does not issue a process-group barrier; callers should + align all ranks after warmup and replay captured graphs in lockstep. + """ + + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError("MoeEp is unusable after an autotune runtime failure") + output = self( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ) + del output + device = activation.device + if device.type == "cuda": + torch.cuda.synchronize(device) + + @property + def training_lanes(self) -> tuple[MoeEpExecutionLane, ...]: + """Operator-bound lanes created by :meth:`prepare_training`.""" + + return self._training_lanes + + def training_symmetric_buffers( + self, + lane: MoeEpExecutionLane, + ) -> Mapping[str, torch.Tensor]: + """Return one lane's symmetric MXFP8 input and final-output buffers.""" + + with self._lifecycle_lock: + self._require_training_lane(lane) + assert self._training_state is not None + return self._training_state.public_symmetric_buffers(lane.index) + + def prepare_training( + self, + *, + lane_count: int = 1, + device: torch.device | str | int | None = None, + ) -> Mapping[ + str, + tuple[tuple[int, ...], tuple[int, ...], torch.dtype, int], + ]: + """Collectively prepare private training runtime and return contracts. + + ``device`` defaults to the current CUDA device. No weights are retained. + Per-lane symmetric input and final-output buffers are available through + :meth:`training_symmetric_buffers`. + """ + + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError("MoeEp is unusable after an autotune runtime failure") + if isinstance(lane_count, bool) or not isinstance(lane_count, int) or lane_count <= 0: + raise ValueError(f"lane_count must be a positive integer, got {lane_count!r}") + if self._training_state is not None: + raise RuntimeError("MoeEp training is already prepared") + if self._fc1_weight_layout is not Fc1WeightLayout.GATE_UP_INTERLEAVED_32: + raise ValueError("prepare_training requires weight_interleave_size=32") + resolved_device = _resolve_training_device(device) + _validate_training_assert_capability(self._forward_config) + from . import _backend + + _backend.validate_config(self._forward_config) + if self._forward_backend is not None and resolved_device != self._forward_backend_device: + raise ValueError(f"MoeEp backend is bound to {self._forward_backend_device}; " f"got {resolved_device}") + if self._forward_backend is None: + self._forward_backend = _backend.create_backend( + self._forward_config, + resolved_device, + ) + self._forward_backend_device = resolved_device + with torch.cuda.device(resolved_device): + state = self._forward_backend.prepare_training( + lane_count=lane_count, + ) + self._training_state = state + self._training_lanes = tuple(MoeEpExecutionLane(index, self._operator_token) for index in range(lane_count)) + self._training_requirements = state.public_requirements() + return self._training_requirements + + def _require_training_lane( + self, + lane: MoeEpExecutionLane, + ) -> None: + if self._closed: + raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError("MoeEp is unusable after an autotune runtime failure") + if self._training_state is None or self._training_requirements is None: + raise RuntimeError("prepare_training() must be called first") + if not isinstance(lane, MoeEpExecutionLane) or lane._operator_token is not self._operator_token or lane not in self._training_lanes: + raise ValueError("execution lane does not belong to this MoeEp") + + def _training_requirement_subset( + self, + names: tuple[str, ...], + ) -> dict[str, tuple[tuple[int, ...], tuple[int, ...], torch.dtype, int]]: + assert self._training_requirements is not None + return {name: self._training_requirements[name] for name in names} + + def pack_forward_weights( + self, + weights: MoeEpForwardWeights, + *, + out: MoeEpForwardWeightStaging, + ) -> MoeEpNativeForwardWeights: + """Materialize source weights into caller-owned native storage.""" + + validate_forward_source_weights(self._forward_config, weights) + from ._megamoe_backend.mxfp8._training_weights import materialize_forward + + return materialize_forward( + weights, + out=out, + fc1_weight_layout=self._forward_config.fc1_weight_layout, + ) + + def pack_backward_weights( + self, + weights: MoeEpBackwardWeights, + *, + out: MoeEpBackwardWeightStaging, + ) -> MoeEpNativeBackwardWeights: + """Materialize source transpose weights into caller-owned storage.""" + + validate_backward_source_weights(self._forward_config, weights) + from ._megamoe_backend.mxfp8._training_weights import materialize_backward + + return materialize_backward( + weights, + out=out, + fc1_weight_layout=self._forward_config.fc1_weight_layout, + ) + + def training_forward( + self, + lane: MoeEpExecutionLane, + activation: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + weights: MoeEpNativeForwardWeights, + out: MoeEpTrainingForwardOutputs, + ) -> torch.Tensor: + """Run forward into caller-owned prepared-training outputs.""" + + with self._lifecycle_lock: + self._require_training_lane(lane) + assert self._training_state is not None + assert self._training_requirements is not None + assert self._forward_backend_device is not None + token_count = validate_training_input( + self._forward_config, + "activation", + activation, + topk_idx, + topk_weights, + device=self._forward_backend_device, + ) + validate_native_forward_weights( + self._forward_config, + weights, + device=self._forward_backend_device, + ) + validate_training_forward_outputs( + out, + self._training_requirement_subset( + ( + "output", + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + ) + ), + device=self._forward_backend_device, + ) + validate_training_non_aliasing( + { + **_named_moe_tensors("activation", activation), + "topk_idx": topk_idx, + "topk_weights": topk_weights, + "weights.fc1.payload": weights.fc1.payload, + "weights.fc1.scale": weights.fc1.scale, + "weights.fc2.payload": weights.fc2.payload, + "weights.fc2.scale": weights.fc2.scale, + "out.output": out.output, + "out.fc1_preact": out.fc1_preact, + "out.fc1_a": out.fc1_a, + "out.fc1_sfa": out.fc1_sfa, + "out.valid_route_counts": out.valid_route_counts, + "out.expert_offsets": out.expert_offsets, + } + ) + from ._megamoe_backend.mxfp8._training_execute import ( + launch_training_forward, + ) + + with torch.cuda.device(self._forward_backend_device): + execution = self._training_state.views( + lane=lane.index, + token_count=token_count, + ) + return launch_training_forward( + self._training_state, + execution, + activation, + topk_idx, + topk_weights, + weights=weights, + out=out, + ) + + def training_backward( + self, + lane: MoeEpExecutionLane, + grad_output: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + weights: MoeEpNativeBackwardWeights, + fc1_preact: torch.Tensor, + fc1_a: torch.Tensor | None = None, + fc1_sfa: torch.Tensor | None = None, + valid_route_counts: torch.Tensor | None = None, + expert_offsets: torch.Tensor | None = None, + out: MoeEpTrainingBackwardOutputs | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + MoeEpTrainingWgradOperands, + ]: + """Run backward into caller-owned outputs using explicit forward state.""" + + with self._lifecycle_lock: + self._require_training_lane(lane) + assert self._training_state is not None + assert self._forward_backend_device is not None + if out is None: + raise TypeError("out must be a MoeEpTrainingBackwardOutputs") + token_count = validate_training_input( + self._forward_config, + "grad_output", + grad_output, + topk_idx, + topk_weights, + device=self._forward_backend_device, + ) + validate_native_backward_weights( + self._forward_config, + weights, + device=self._forward_backend_device, + ) + if fc1_preact is None: + raise ValueError("fc1_preact from the matching forward is required") + backward_output = out + validate_training_backward_outputs( + backward_output, + self._training_requirement_subset( + ( + "grad_activation", + "dprob", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + ) + ), + device=self._forward_backend_device, + ) + validate_training_forward_state( + fc1_preact=fc1_preact, + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + valid_route_counts=valid_route_counts, + expert_offsets=expert_offsets, + requirements=self._training_requirement_subset( + ( + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + ) + ), + device=self._forward_backend_device, + ) + validate_training_non_aliasing( + { + **_named_moe_tensors("grad_output", grad_output), + "topk_idx": topk_idx, + "topk_weights": topk_weights, + "weights.w2_transpose.payload": weights.w2_transpose.payload, + "weights.w2_transpose.scale": weights.w2_transpose.scale, + "weights.w1_transpose.payload": weights.w1_transpose.payload, + "weights.w1_transpose.scale": weights.w1_transpose.scale, + "fc1_preact": fc1_preact, + "fc1_a": fc1_a, + "fc1_sfa": fc1_sfa, + "valid_route_counts": valid_route_counts, + "expert_offsets": expert_offsets, + "out.grad_activation": backward_output.grad_activation, + "out.dprob": backward_output.dprob, + "out.fc1_b": backward_output.fc1_b, + "out.fc1_sfb": backward_output.fc1_sfb, + "out.fc2_a": backward_output.fc2_a, + "out.fc2_sfa": backward_output.fc2_sfa, + "out.fc2_b": backward_output.fc2_b, + "out.fc2_sfb": backward_output.fc2_sfb, + } + ) + from ._megamoe_backend.mxfp8._training_execute import ( + launch_training_backward, + ) + + with torch.cuda.device(self._forward_backend_device): + execution = self._training_state.views( + lane=lane.index, + token_count=token_count, + ) + return launch_training_backward( + self._training_state, + execution, + grad_output, + topk_idx, + topk_weights, + weights=weights, + fc1_preact=fc1_preact, + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + valid_route_counts=valid_route_counts, + expert_offsets=expert_offsets, + out=backward_output, + ) + + def close(self) -> None: + """Release compiled-backend instance resources; idempotent.""" + + with self._lifecycle_lock: + if self._closed: + return + if self._forward_backend is not None: + close_backend = getattr(self._forward_backend, "close", None) + if close_backend is not None: + close_backend() + self._forward_backend = None + self._forward_backend_device = None + self._validated_topk_idx = None + self._validated_topk_version = None + self._training_state = None + self._training_lanes = () + self._training_requirements = None + self._closed = True + + def __enter__(self) -> "MoeEp": + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError("MoeEp is unusable after an autotune runtime failure") + return self + + def __exit__(self, exc_type, exc_value, traceback) -> bool: + del exc_type, exc_value, traceback + self.close() + return False + + def __del__(self) -> None: + if not hasattr(self, "_closed"): + return + try: + self.close() + except Exception as exc: + # Explicit close propagates cleanup failures. During GC there is no + # safe global point to retry CUDA/NVSHMEM teardown, so report the + # failure without retaining the backend indefinitely. + with contextlib.suppress(Exception): + warnings.warn( + f"MoeEp finalizer could not release backend resources: {exc}", + ResourceWarning, + stacklevel=2, + ) + + +__all__ = [ + "BlockScaledTensor", + "MoeEp", + "MoeEpAutotuneCandidateResult", + "MoeEpAutotuneResult", + "MoeEpBackwardWeightStaging", + "MoeEpBackwardWeights", + "MoeEpExecutionLane", + "MoeEpForwardWeightStaging", + "MoeEpForwardWeights", + "MoeEpNativeBackwardWeights", + "MoeEpNativeForwardWeights", + "MoeEpTrainingBackwardOutputs", + "MoeEpTrainingForwardOutputs", + "MoeEpTrainingWgradOperands", + "MoeFormat", + "MoeTensor", + "pack_backward_weights", + "pack_forward_weights", +] diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py index 85fefa5ea..c6ec04a9d 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py @@ -761,6 +761,229 @@ def counted_compile(self): assert cache_entries == 1 +@pytest.mark.L0 +@pytest.mark.parametrize( + ("caller_owned_workspace", "expected_cache_entries"), + [(False, 2), (True, 1)], + ids=["compatibility-isolation", "caller-workspace"], +) +def test_grouped_gemm_wgrad_wrapper_explicit_dense_output_cache( + monkeypatch, + caller_owned_workspace, + expected_cache_entries, +): + from cudnn.gemm.cutedsl.grouped.wgrad import api as grouped_gemm_wgrad_api + + grouped_gemm_wgrad_api._cache_of_GroupedGemmWgradSm100Objects.clear() + compile_count = {"value": 0} + + def counted_compile(self): + compile_count["value"] += 1 + + monkeypatch.setattr(grouped_gemm_wgrad_api.GroupedGemmWgradSm100, "check_support", lambda self: True) + monkeypatch.setattr(grouped_gemm_wgrad_api.GroupedGemmWgradSm100, "compile", counted_compile) + monkeypatch.setattr(grouped_gemm_wgrad_api.GroupedGemmWgradSm100, "execute", lambda self, **kwargs: None) + monkeypatch.setattr( + grouped_gemm_wgrad_api, + "select_grouped_gemm_backend", + lambda **_: grouped_gemm_wgrad_api.GroupedGemmBackend.BLOCK_SCALED, + ) + + inputs = _make_wgrad_wrapper_cache_inputs([8, 12]) + outputs = [torch.empty((2, 32, 64), dtype=torch.bfloat16) for _ in range(2)] + workspaces = [torch.empty(512, dtype=torch.uint8) for _ in range(2)] + try: + for output, workspace in zip(outputs, workspaces): + workspace_kwargs = {"descriptor_workspace": workspace} if caller_owned_workspace else {} + cudnn.grouped_gemm_wgrad_wrapper_sm100( + **inputs, + **workspace_kwargs, + output_mode="dense", + wgrad_tensor=output, + acc_dtype=torch.float32, + wgrad_dtype=torch.bfloat16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=16, + ) + finally: + cache_entries = len(grouped_gemm_wgrad_api._cache_of_GroupedGemmWgradSm100Objects) + grouped_gemm_wgrad_api._cache_of_GroupedGemmWgradSm100Objects.clear() + + assert outputs[0].data_ptr() != outputs[1].data_ptr() + assert compile_count["value"] == expected_cache_entries + assert cache_entries == expected_cache_entries + + +@pytest.mark.L0 +def test_grouped_gemm_wgrad_wrapper_discrete_accepts_caller_workspace(monkeypatch): + from cudnn.gemm.cutedsl.grouped.wgrad import api as grouped_gemm_wgrad_api + + grouped_gemm_wgrad_api._cache_of_GroupedGemmWgradSm100Objects.clear() + compile_count = {"value": 0} + + def counted_compile(self): + compile_count["value"] += 1 + + monkeypatch.setattr( + grouped_gemm_wgrad_api.GroupedGemmWgradSm100, + "check_support", + lambda self: True, + ) + monkeypatch.setattr( + grouped_gemm_wgrad_api.GroupedGemmWgradSm100, + "compile", + counted_compile, + ) + monkeypatch.setattr( + grouped_gemm_wgrad_api.GroupedGemmWgradSm100, + "execute", + lambda self, **kwargs: None, + ) + monkeypatch.setattr( + grouped_gemm_wgrad_api, + "select_grouped_gemm_backend", + lambda **_: grouped_gemm_wgrad_api.GroupedGemmBackend.BLOCK_SCALED, + ) + + inputs = _make_wgrad_wrapper_cache_inputs([8, 12]) + outputs = [torch.empty((2, 32, 64), dtype=torch.bfloat16) for _ in range(2)] + workspaces = [torch.empty(512, dtype=torch.uint8) for _ in range(2)] + try: + for output, workspace in zip(outputs, workspaces): + cudnn.grouped_gemm_wgrad_wrapper_sm100( + **inputs, + output_mode="discrete", + wgrad_tensor=output, + descriptor_workspace=workspace, + acc_dtype=torch.float32, + wgrad_dtype=torch.bfloat16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=16, + ) + finally: + cache_entries = len(grouped_gemm_wgrad_api._cache_of_GroupedGemmWgradSm100Objects) + grouped_gemm_wgrad_api._cache_of_GroupedGemmWgradSm100Objects.clear() + + assert compile_count["value"] == 1 + assert cache_entries == 1 + + +@pytest.mark.L0 +def test_grouped_gemm_wgrad_workspace_size(): + assert cudnn.get_grouped_gemm_wgrad_workspace_size_sm100(2) == 512 + assert ( + cudnn.get_grouped_gemm_wgrad_workspace_size_sm100( + 2, + input_order="tensor_ragged", + ) + == 1024 + ) + + +@pytest.mark.L0 +def test_blockscaled_wgrad_execute_uses_caller_workspace(monkeypatch): + from cudnn.gemm.cutedsl.grouped.wgrad import _blockscaled_api + + api = object.__new__(_blockscaled_api.GroupedGemmWgradBlockScaledAPI) + api._workspace_bytes = 16 + api._workspace = torch.empty(16, dtype=torch.uint8) + api.a_desc = type("TensorDesc", (), {"device": torch.device("cpu")})() + api.weight_mode = _blockscaled_api.MoEWeightMode.DENSE + api._get_default_stream = lambda stream: stream + api._runtime_error_if = lambda condition, message: None + api._value_error_if = lambda condition, message: None + monkeypatch.setattr( + _blockscaled_api, + "from_dlpack", + lambda tensor, **kwargs: tensor, + ) + + launch_workspaces = [] + + def compiled_kernel(*args): + launch_workspaces.append(args[6]) + + api._compiled_kernel = compiled_kernel + operand = torch.empty((1, 1)) + offsets = torch.tensor([1], dtype=torch.int32) + outputs = [torch.empty((2, 32, 64), dtype=torch.bfloat16) for _ in range(2)] + workspaces = [torch.empty(16, dtype=torch.uint8) for _ in range(2)] + for output, workspace in ( + (outputs[0], workspaces[0]), + (outputs[0], workspaces[0]), + (outputs[1], workspaces[1]), + ): + api.execute( + operand, + operand, + operand, + operand, + offsets, + wgrad_tensor=output, + descriptor_workspace=workspace, + current_stream=object(), + ) + + assert launch_workspaces[0] is workspaces[0] + assert launch_workspaces[1] is workspaces[0] + assert launch_workspaces[2] is workspaces[1] + assert launch_workspaces[0].data_ptr() != launch_workspaces[2].data_ptr() + + +@pytest.mark.L0 +def test_blockscaled_discrete_wgrad_execute_uses_caller_workspace(monkeypatch): + from cudnn.gemm.cutedsl.grouped.wgrad import _blockscaled_api + + api = object.__new__(_blockscaled_api.GroupedGemmWgradBlockScaledAPI) + api._workspace_bytes = 16 + api._workspace_arg = torch.empty(16, dtype=torch.uint8) + api.a_desc = type("TensorDesc", (), {"device": torch.device("cpu")})() + api.weight_mode = _blockscaled_api.MoEWeightMode.DISCRETE + api.expert_cnt = 2 + api._get_default_stream = lambda stream: stream + api._runtime_error_if = lambda condition, message: None + api._value_error_if = lambda condition, message: None + monkeypatch.setattr( + _blockscaled_api, + "from_dlpack", + lambda tensor, **kwargs: tensor, + ) + monkeypatch.setattr( + _blockscaled_api, + "_validate_pointer_tensor", + lambda tensor, name, count: None, + ) + + launch_workspaces = [] + + def compiled_kernel(*args): + launch_workspaces.append(args[6]) + + api._compiled_kernel = compiled_kernel + operand = torch.empty((1, 1)) + offsets = torch.tensor([1, 2], dtype=torch.int32) + wgrad_ptrs = torch.empty(2, dtype=torch.int64) + workspaces = [torch.empty(16, dtype=torch.uint8) for _ in range(2)] + for workspace in (workspaces[0], workspaces[0], workspaces[1]): + api.execute( + operand, + operand, + operand, + operand, + offsets, + wgrad_ptrs=wgrad_ptrs, + descriptor_workspace=workspace, + current_stream=object(), + ) + + assert launch_workspaces[0] is workspaces[0] + assert launch_workspaces[1] is workspaces[0] + assert launch_workspaces[2] is workspaces[1] + assert launch_workspaces[0].data_ptr() != launch_workspaces[2].data_ptr() + + @pytest.mark.L0 def test_grouped_gemm_wgrad_wrapper_input_order_cache_key(monkeypatch): from cudnn.gemm.cutedsl.grouped.wgrad import api as grouped_gemm_wgrad_api diff --git a/test/python/moe_ep/moe_ep_distributed_workers.py b/test/python/moe_ep/moe_ep_distributed_workers.py new file mode 100644 index 000000000..007afd70d --- /dev/null +++ b/test/python/moe_ep/moe_ep_distributed_workers.py @@ -0,0 +1,446 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Picklable multiprocessing workers for distributed MoE EP tests.""" + +from __future__ import annotations + +from datetime import timedelta + +import torch +import torch.distributed as dist + +from moe_ep.moe_ep_test_support import ( + _allocate_stateless_training_outputs, + _allocate_training_weight_staging, + _assert_backward_matches, + _assert_grouped_wgrads_match_reference, + _assert_matches_reference, + _assert_wgrads_match_reference, + _dense_wgrads_from_grouped_kernel, + _dense_wgrads_from_operands, + _fixed_training_reference, + _fixed_training_weights, + _forward_config, + _grad_output, + _interleave_fc1_wgrad, + _output_as_float, + _reference_forward, + make_distributed_forward_inputs, + quantize_mxfp8, +) + +__all__ = [ + "_distributed_autotune_worker", + "_distributed_output_worker", + "_distributed_subgroup_output_worker", + "_run_backward_reference_case", + "_run_forward_output_case", +] + + +def _distributed_autotune_worker( + rank: int, + world_size: int, + init_file: str, +) -> None: + """Run an EP sweep and verify one rank-consistent applied winner.""" + + device = torch.device("cuda", rank) + torch.cuda.set_device(device) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + device_id=device, + timeout=timedelta(seconds=180), + ) + try: + from cudnn import MoeEp, MoeEpTuningConfig + + args = make_distributed_forward_inputs(rank, world_size, device) + config = _forward_config( + num_experts=2 * world_size, + ep_group=dist.group.WORLD, + max_tokens_per_rank=8, + ) + expected = _reference_forward(args, **config) + candidate = MoeEpTuningConfig(token_in_flag_batch=2) + op = MoeEp(**config) + try: + result = op.autotune( + *args, + candidates=[candidate], + warmup_iters=1, + timed_iters=2, + ) + actual = op(*args) + torch.cuda.synchronize(device) + winners = [None] * world_size + dist.all_gather_object(winners, result.winner) + assert all(winner == result.winner for winner in winners) + assert op.tuning == result.winner + _assert_matches_reference(actual, expected) + dist.barrier() + op.close() + op = None + dist.barrier() + finally: + if op is not None: + op.close() + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _run_forward_output_case( + *, + device: torch.device, + ep_group, + ep_rank: int, + ep_size: int, + combine_format: str = "bf16", + expected_global_ranks: tuple[int, ...] | None = None, +) -> None: + """Run inference-forward parity and dropped-route checks.""" + + from cudnn import MoeEp + + args = make_distributed_forward_inputs(ep_rank, ep_size, device) + config = _forward_config( + num_experts=2 * ep_size, + ep_group=ep_group, + max_tokens_per_rank=8, + combine_format=combine_format, + ) + expected = _reference_forward(args, **config) + op = MoeEp(**config) + try: + actual = op(*args) + actual_snapshot = _output_as_float(actual).clone() + torch.cuda.synchronize(device) + + args[3].fill_(-1) + dropped = op(*args) + dropped_snapshot = _output_as_float(dropped).clone() + torch.cuda.synchronize(device) + + dist.barrier(group=ep_group) + assertion_error = None + try: + assert op.ep_rank == ep_rank + if expected_global_ranks is not None: + assert op.ep_global_ranks == expected_global_ranks + _assert_matches_reference(actual_snapshot, expected) + assert dropped_snapshot.eq(0).all() + except BaseException as error: + assertion_error = error + dist.barrier(group=ep_group) + if assertion_error is not None: + raise assertion_error + + op.close() + op = None + dist.barrier(group=ep_group) + finally: + if op is not None: + op.close() + + +def _distributed_output_worker( + rank: int, + world_size: int, + init_file: str, + combine_format: str = "bf16", +) -> None: + device = torch.device("cuda", rank) + torch.cuda.set_device(device) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + device_id=device, + timeout=timedelta(seconds=180), + ) + try: + _run_forward_output_case( + device=device, + ep_group=dist.group.WORLD, + ep_rank=rank, + ep_size=world_size, + combine_format=combine_format, + expected_global_ranks=tuple(range(world_size)), + ) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _distributed_subgroup_output_worker( + global_rank: int, + global_world_size: int, + init_file: str, +) -> None: + """Run one of two disjoint, non-contiguous EP2 groups inside WORLD4.""" + + device = torch.device("cuda", global_rank) + torch.cuda.set_device(device) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=global_rank, + world_size=global_world_size, + device_id=device, + timeout=timedelta(seconds=180), + ) + try: + subgroup_memberships = ((0, 2), (1, 3)) + subgroups = [dist.new_group(list(members), backend="nccl") for members in subgroup_memberships] + subgroup_index = global_rank % 2 + ep_group = subgroups[subgroup_index] + ep_rank = dist.get_rank(ep_group) + ep_size = dist.get_world_size(ep_group) + actual_global_ranks = tuple(dist.get_global_rank(ep_group, group_rank) for group_rank in range(ep_size)) + + _run_forward_output_case( + device=device, + ep_group=ep_group, + ep_rank=ep_rank, + ep_size=ep_size, + expected_global_ranks=subgroup_memberships[subgroup_index], + ) + dist.barrier() + assert actual_global_ranks == subgroup_memberships[subgroup_index] + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _make_distributed_backward_inputs( + ep_rank: int, + ep_size: int, + device: torch.device, +): + """Build a minimal local/remote/drop case with one empty local expert.""" + + generator = torch.Generator(device=device).manual_seed(20260828 + ep_rank) + local_experts, token_count, hidden, intermediate = 2, 2, 128, 256 + activation = ( + torch.randn( + token_count, + hidden, + generator=generator, + device=device, + ) + / 4 + ).to(torch.bfloat16) + fc1_weight = quantize_mxfp8( + torch.randn( + local_experts, + hidden, + 2 * intermediate, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + fc2_weight = quantize_mxfp8( + torch.randn( + local_experts, + intermediate, + hidden, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + local_expert = ep_rank * local_experts + remote_expert = ((ep_rank + 1) % ep_size) * local_experts + topk_idx = torch.tensor( + [[local_expert, remote_expert], [-1, local_expert]], + dtype=torch.int32, + device=device, + ) + topk_weights = torch.tensor( + [[0.625, 0.375], [0.0, 1.0]], + dtype=torch.float32, + device=device, + ) + grad_output = _grad_output( + device, + token_count, + seed=20260901 + ep_rank, + ) + return ( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ), grad_output + + +def _run_backward_reference_case( + *, + device: torch.device, + ep_group, + ep_rank: int, + ep_size: int, + combine_format: str = "bf16", + gate_up_clamp: float | None = None, + expected_global_ranks: tuple[int, ...] | None = None, +) -> None: + """Run stateless training after the independent distributed oracle.""" + + from cudnn import MoeEp + + args, grad_output = _make_distributed_backward_inputs( + ep_rank, + ep_size, + device, + ) + num_experts = 2 * ep_size + max_recv_size_per_rank = 3 + + # Finish all collective reference work, including dense local dW, before + # constructing or launching the production operator. + expected = _fixed_training_reference( + args, + grad_output, + combine_format=combine_format, + gate_up_clamp=gate_up_clamp, + ep_group=ep_group, + num_experts=num_experts, + max_recv_size_per_rank=max_recv_size_per_rank, + drop_on_overflow=True, + ) + expected_y, expected_dx, expected_dprob, expected_wgrads = expected + expected_fc1_wgrad, expected_fc2_wgrad = expected_wgrads.dense_wgrads() + expected_dense_wgrads = ( + _interleave_fc1_wgrad(expected_fc1_wgrad), + expected_fc2_wgrad, + ) + weights = _fixed_training_weights(args) + + op = MoeEp( + num_experts=num_experts, + hidden_size=128, + intermediate_size=256, + top_k=2, + ep_group=ep_group, + max_tokens_per_rank=args[0].shape[0], + max_recv_size_per_rank=max_recv_size_per_rank, + drop_on_overflow=True, + combine_format=combine_format, + gate_up_clamp=gate_up_clamp, + weight_interleave_size=32, + ) + try: + requirements = op.prepare_training( + lane_count=1, + device=device, + ) + lane = op.training_lanes[0] + forward_staging, backward_staging = _allocate_training_weight_staging(weights) + native_forward = op.pack_forward_weights( + weights[0], + out=forward_staging, + ) + native_backward = op.pack_backward_weights( + weights[1], + out=backward_staging, + ) + forward_out, backward_out = _allocate_stateless_training_outputs( + requirements, + device, + op.training_symmetric_buffers(lane), + ) + actual_y = op.training_forward( + lane, + args[0], + args[3], + args[4], + weights=native_forward, + out=forward_out, + ) + actual_dx, actual_dprob, actual_wgrads = op.training_backward( + lane, + grad_output, + args[3], + args[4], + weights=native_backward, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, + ) + grouped_wgrads = _dense_wgrads_from_grouped_kernel(actual_wgrads) + torch.cuda.synchronize(device) + + # No rank may enter a local assertion while a peer is still inside a + # collective kernel. A second barrier keeps cleanup aligned on failure. + dist.barrier(group=ep_group) + assertion_error = None + try: + assert op.ep_rank == ep_rank + assert op.ep_size == ep_size + if expected_global_ranks is not None: + assert op.ep_global_ranks == expected_global_ranks + assert args[3][0, 0] // 2 == ep_rank + assert args[3][0, 1] // 2 == (ep_rank + 1) % ep_size + assert args[3].eq(-1).any() + assert expected_wgrads.valid_route_counts[1].eq(0) + assert actual_wgrads.valid_route_counts[1].eq(0) + _assert_matches_reference(actual_y, expected_y) + _assert_backward_matches( + (actual_dx, actual_dprob), + (expected_dx, expected_dprob), + args[3], + ) + _assert_wgrads_match_reference( + actual_wgrads, + expected_wgrads, + expected_dense=expected_dense_wgrads, + ) + expected_offsets = torch.cumsum( + torch.div( + actual_wgrads.valid_route_counts + 127, + 128, + rounding_mode="floor", + ) + * 128, + dim=0, + dtype=actual_wgrads.expert_offsets.dtype, + ) + torch.testing.assert_close( + actual_wgrads.expert_offsets, + expected_offsets, + rtol=0, + atol=0, + ) + _assert_grouped_wgrads_match_reference( + grouped_wgrads, + expected_dense_wgrads, + reference_name="the independent PyTorch MXFP8 reference", + ) + _assert_grouped_wgrads_match_reference( + grouped_wgrads, + _dense_wgrads_from_operands(actual_wgrads), + reference_name="the decoded production operand bundle", + close_kwargs={"rtol": 0.1, "atol": 0.1}, + ) + assert grouped_wgrads[0][1].eq(0).all() + assert grouped_wgrads[1][1].eq(0).all() + except BaseException as error: + assertion_error = error + dist.barrier(group=ep_group) + if assertion_error is not None: + raise assertion_error + finally: + op.close() diff --git a/test/python/moe_ep/moe_ep_reference.py b/test/python/moe_ep/moe_ep_reference.py new file mode 100644 index 000000000..b794cba7e --- /dev/null +++ b/test/python/moe_ep/moe_ep_reference.py @@ -0,0 +1,1191 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Pure PyTorch semantic reference for a SwiGLU MoE with expert parallelism. + +The implementation deliberately favors readable semantics over performance. It +supports a one-rank execution path and a variable-size ``all_to_all_single`` EP +path, plus BF16, MXFP8, and NVFP4 block-scaled public outputs. + +The quantized tensor layouts are logical (unswizzled) layouts. A production +kernel may reorder scale factors internally without changing this API contract. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Optional, Sequence, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +class MoeFormat(str, Enum): + """Public and communication formats supported by the reference.""" + + BF16 = "bf16" + MXFP8 = "mxfp8" + NVFP4 = "nvfp4" + + +def _parse_format(value: Union[MoeFormat, str]) -> MoeFormat: + if isinstance(value, MoeFormat): + return value + try: + return MoeFormat(value.lower()) + except (AttributeError, ValueError) as exc: + choices = ", ".join(item.value for item in MoeFormat) + raise ValueError(f"unsupported format {value!r}; expected one of: {choices}") from exc + + +def _require_torch_dtype(name: str) -> torch.dtype: + dtype = getattr(torch, name, None) + if dtype is None: + raise RuntimeError(f"this PyTorch build does not provide torch.{name}") + return dtype + + +def _normalize_axis(axis: int, ndim: int) -> int: + normalized = axis + ndim if axis < 0 else axis + if normalized < 0 or normalized >= ndim: + raise IndexError(f"axis {axis} is out of range for a {ndim}-D tensor") + return normalized + + +def _shape_with_axis(shape: Sequence[int], axis: int, value: int) -> Tuple[int, ...]: + result = list(shape) + result[axis] = value + return tuple(result) + + +def _ceil_div(numerator: int, denominator: int) -> int: + return (numerator + denominator - 1) // denominator + + +@dataclass(frozen=True) +class BlockScaledTensor: + """Portable data-plus-scale representation for MXFP8 or NVFP4. + + ``logical_shape`` describes the dequantized tensor. For MXFP8, ``data`` + has that shape and uses E4M3. For NVFP4, ``data`` is a uint8 tensor with + two E2M1 values per byte along ``axis`` (low nibble first). ``scale`` + replaces that axis by one scale per block. + """ + + data: torch.Tensor + scale: torch.Tensor + format: Union[MoeFormat, str] + logical_shape: Tuple[int, ...] + axis: int = -1 + + def __post_init__(self) -> None: + fmt = _parse_format(self.format) + if fmt is MoeFormat.BF16: + raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") + shape = tuple(int(dim) for dim in self.logical_shape) + if not shape or any(dim < 0 for dim in shape): + raise ValueError(f"logical_shape must contain non-negative dimensions, got {shape}") + axis = _normalize_axis(self.axis, len(shape)) + object.__setattr__(self, "format", fmt) + object.__setattr__(self, "logical_shape", shape) + object.__setattr__(self, "axis", axis) + self._validate_storage() + + @property + def block_size(self) -> int: + return 32 if self.format is MoeFormat.MXFP8 else 16 + + @property + def shape(self) -> Tuple[int, ...]: + return self.logical_shape + + @property + def device(self) -> torch.device: + return self.data.device + + def _validate_storage(self) -> None: + if self.data.device != self.scale.device: + raise ValueError("block-scaled data and scale must be on the same device") + + logical_extent = self.logical_shape[self.axis] + scale_shape = _shape_with_axis( + self.logical_shape, + self.axis, + _ceil_div(logical_extent, self.block_size), + ) + if tuple(self.scale.shape) != scale_shape: + raise ValueError(f"scale shape must be {scale_shape}, got {tuple(self.scale.shape)}") + + if self.format is MoeFormat.MXFP8: + expected_dtype = _require_torch_dtype("float8_e4m3fn") + expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") + data_shape = self.logical_shape + if self.data.dtype != expected_dtype: + raise TypeError(f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}") + else: + expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) + if self.data.dtype != torch.uint8 and self.data.dtype != fp4_dtype: + raise TypeError("nvfp4 data must be packed uint8 or torch.float4_e2m1fn_x2") + data_shape = _shape_with_axis( + self.logical_shape, + self.axis, + _ceil_div(logical_extent, 2), + ) + + if tuple(self.data.shape) != data_shape: + raise ValueError(f"data shape must be {data_shape}, got {tuple(self.data.shape)}") + if self.scale.dtype != expected_scale_dtype: + raise TypeError(f"scale must have dtype {expected_scale_dtype}, got {self.scale.dtype}") + + def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Return the logical tensor with block scales applied.""" + + logical_extent = self.logical_shape[self.axis] + scale = self.scale.movedim(self.axis, -1).float() + expanded_scale = scale.repeat_interleave(self.block_size, dim=-1)[..., :logical_extent] + + if self.format is MoeFormat.MXFP8: + values = self.data.movedim(self.axis, -1).float() + else: + packed = self.data + if packed.dtype != torch.uint8: + packed = packed.view(torch.uint8) + packed = packed.movedim(self.axis, -1) + low = packed & 0x0F + high = packed >> 4 + codes = torch.stack((low, high), dim=-1).flatten(-2)[..., :logical_extent] + table = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=torch.float32, + device=packed.device, + ) + values = table[codes.long()] + + return (values * expanded_scale).movedim(-1, self.axis).to(dtype) + + +def _nearest_e2m1_codes(values: torch.Tensor) -> torch.Tensor: + """Quantize to E2M1 nibble codes with round-to-nearest, ties-to-even.""" + + levels = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float32, + device=values.device, + ) + magnitudes = values.abs().unsqueeze(-1) + distances = (magnitudes - levels).abs() + minimum = distances.amin(dim=-1, keepdim=True) + candidates = distances == minimum + codes = torch.arange(8, dtype=torch.int64, device=values.device) + any_code = torch.where(candidates, codes, 8).amin(dim=-1) + even_code = torch.where(candidates & ((codes & 1) == 0), codes, 8).amin(dim=-1) + magnitude_code = torch.where(even_code < 8, even_code, any_code) + sign_code = torch.signbit(values).to(torch.int64) << 3 + return magnitude_code | sign_code + + +def quantize_blockwise( + tensor: torch.Tensor, + format: Union[MoeFormat, str], + *, + axis: int = -1, +) -> BlockScaledTensor: + """Quantize a floating tensor into logical MXFP8 or NVFP4 blocks. + + MXFP8 uses 32-value blocks, E4M3 payloads, and E8M0 scales rounded toward + positive infinity. NVFP4 uses 16-value blocks, packed E2M1 payloads, and + E4M3 scales rounded to nearest. + """ + + fmt = _parse_format(format) + if fmt is MoeFormat.BF16: + raise ValueError("quantize_blockwise requires mxfp8 or nvfp4") + if not tensor.is_floating_point(): + raise TypeError(f"tensor must be floating point, got {tensor.dtype}") + + axis = _normalize_axis(axis, tensor.ndim) + logical_shape = tuple(tensor.shape) + moved = tensor.float().movedim(axis, -1) + logical_extent = moved.shape[-1] + block_size = 32 if fmt is MoeFormat.MXFP8 else 16 + block_count = _ceil_div(logical_extent, block_size) + padded_extent = block_count * block_size + if padded_extent != logical_extent: + moved = F.pad(moved, (0, padded_extent - logical_extent)) + blocks = moved.reshape(*moved.shape[:-1], block_count, block_size) + + value_limit = 448.0 if fmt is MoeFormat.MXFP8 else 6.0 + scale_float = blocks.abs().amax(dim=-1) / value_limit + if fmt is MoeFormat.MXFP8: + safe_scale = torch.where(scale_float > 0, scale_float, 1.0) + scale_float = torch.where( + scale_float > 0, + torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), + torch.zeros_like(scale_float), + ) + scale_dtype = _require_torch_dtype("float8_e8m0fnu") + else: + scale_dtype = _require_torch_dtype("float8_e4m3fn") + + scale = scale_float.to(scale_dtype) + scale_for_math = scale.float() + reciprocal = torch.where(scale_for_math > 0, scale_for_math.reciprocal(), 0.0) + normalized = (blocks * reciprocal.unsqueeze(-1)).clamp(-value_limit, value_limit) + + if fmt is MoeFormat.MXFP8: + data_dtype = _require_torch_dtype("float8_e4m3fn") + data = normalized.to(data_dtype).reshape(*moved.shape)[..., :logical_extent] + else: + codes = _nearest_e2m1_codes(normalized).reshape(*moved.shape) + low = codes[..., 0::2] + high = codes[..., 1::2] + data = (low | (high << 4)).to(torch.uint8)[..., : _ceil_div(logical_extent, 2)] + + return BlockScaledTensor( + data=data.movedim(-1, axis).contiguous(), + scale=scale.movedim(-1, axis).contiguous(), + format=fmt, + logical_shape=logical_shape, + axis=axis, + ) + + +MoeTensor = Union[torch.Tensor, BlockScaledTensor] + + +@dataclass(frozen=True) +class WgradForwardStashReference: + """Logical reference for the caller-owned forward wgrad stash. + + Unlike the production object, ``fc1_a`` bundles its logical E8M0 scales + with the E4M3 payload. It represents the padded, expert-concatenated + ``x.T`` operand after input MXFP8 staging and token-axis requantization. + """ + + fc1_a: BlockScaledTensor + expert_offsets: torch.Tensor + valid_route_counts: torch.Tensor + route_metadata: torch.Tensor + + +@dataclass(frozen=True) +class WgradOperandsReference: + """Logical MXFP8 operands and dense expert-weight-gradient oracle. + + The K dimension is a concatenation of local experts. Each expert's valid + routes come first, followed by zero rows up to its 256-route boundary. + Production scale tensors use a blocked physical layout; these reference + tensors keep ordinary logical scales so their represented values are easy + to inspect. + """ + + fc1_a: BlockScaledTensor + fc1_b: BlockScaledTensor + fc2_a: BlockScaledTensor + fc2_b: BlockScaledTensor + expert_offsets: torch.Tensor + valid_route_counts: torch.Tensor + route_metadata: torch.Tensor + + def dense_wgrads(self) -> Tuple[torch.Tensor, torch.Tensor]: + """Return dense ``dW1=x.T@dC`` and ``dW2=(p*h).T@dY`` per expert.""" + + a1 = self.fc1_a.dequantize() + b1 = self.fc1_b.dequantize() + a2 = self.fc2_a.dequantize() + b2 = self.fc2_b.dequantize() + expert_count = int(self.expert_offsets.numel()) + dw1 = torch.zeros( + (expert_count, a1.shape[0], b1.shape[1]), + dtype=torch.float32, + device=a1.device, + ) + dw2 = torch.zeros( + (expert_count, a2.shape[0], b2.shape[1]), + dtype=torch.float32, + device=a2.device, + ) + begin = 0 + for expert, end_tensor in enumerate(self.expert_offsets): + end = int(end_tensor.item()) + if end > begin: + dw1[expert] = a1[:, begin:end] @ b1[begin:end] + dw2[expert] = a2[:, begin:end] @ b2[begin:end] + begin = end + return dw1, dw2 + + +@dataclass(frozen=True) +class _DispatchPlan: + """Send-side routing derived from ``topk_idx``; identical in fwd and bwd.""" + + send_expert: torch.Tensor # local expert id per sent route + send_weight: torch.Tensor # router weight per sent route + send_token_idx: torch.Tensor # source token per sent route + send_slot_idx: torch.Tensor # source top-k slot per sent route + send_counts: Tuple[int, ...] # routes sent to each rank + recv_counts: Tuple[int, ...] # routes received from each rank + + +def _tensor_device(tensor: MoeTensor) -> torch.device: + return tensor.device + + +def _decode_tensor( + tensor: MoeTensor, + *, + name: str, + expected_shape: Tuple[int, ...], + quantized_axis: int, +) -> torch.Tensor: + if isinstance(tensor, BlockScaledTensor): + if tensor.logical_shape != expected_shape: + raise ValueError(f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}") + if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): + raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") + return tensor.dequantize() + + if tuple(tensor.shape) != expected_shape: + raise ValueError(f"{name} shape must be {expected_shape}, got {tuple(tensor.shape)}") + if not tensor.is_floating_point(): + raise TypeError(f"{name} must be floating point or BlockScaledTensor, got {tensor.dtype}") + return tensor.float() + + +def _format_round_trip_axis( + tensor: torch.Tensor, + format: MoeFormat, + *, + axis: int, +) -> torch.Tensor: + if format is MoeFormat.BF16: + return tensor.to(torch.bfloat16).float() + return quantize_blockwise(tensor, format, axis=axis).dequantize() + + +def _format_round_trip(tensor: torch.Tensor, format: MoeFormat) -> torch.Tensor: + return _format_round_trip_axis(tensor, format, axis=-1) + + +def forward_combine_round_trip( + tensor: torch.Tensor, + format: MoeFormat, +) -> torch.Tensor: + """Model GLU combine conversion directly from its FP32 accumulator.""" + + return _format_round_trip(tensor, format) + + +def backward_combine_round_trip( + tensor: torch.Tensor, + format: MoeFormat, +) -> torch.Tensor: + """Model dGLU combine conversion directly from its FP32 accumulator.""" + + return _format_round_trip(tensor, format) + + +def _padded_expert_rows( + rows: torch.Tensor, + expert_rows: torch.Tensor, + valid_counts: Sequence[int], + padded_ends: Sequence[int], +) -> torch.Tensor: + """Place compact expert-grouped rows at the start of padded ranges.""" + + padded_extent = int(padded_ends[-1]) if padded_ends else 0 + padded = torch.zeros( + (padded_extent, *rows.shape[1:]), + dtype=rows.dtype, + device=rows.device, + ) + begin = 0 + for expert, (count, end) in enumerate(zip(valid_counts, padded_ends)): + positions = torch.nonzero( + expert_rows == expert, + as_tuple=False, + ).flatten() + if int(positions.numel()) != int(count): + raise ValueError(f"expert {expert} has {positions.numel()} rows, expected {count}") + if count: + padded[begin : begin + count].copy_(rows.index_select(0, positions)) + begin = int(end) + return padded + + +def _deinterleave_glu(tensor: torch.Tensor, interleave_size: int) -> torch.Tensor: + """Convert fixed-width gate/up strips to contiguous gate and up halves.""" + shape = tensor.shape + return ( + tensor.reshape( + *shape[:-1], + shape[-1] // (2 * interleave_size), + 2, + interleave_size, + ) + .transpose(-3, -2) + .reshape(shape) + ) + + +def _interleave_glu(tensor: torch.Tensor, interleave_size: int) -> torch.Tensor: + """Convert contiguous gate and up halves to fixed-width strips.""" + shape = tensor.shape + return ( + tensor.reshape( + *shape[:-1], + 2, + shape[-1] // (2 * interleave_size), + interleave_size, + ) + .transpose(-3, -2) + .reshape(shape) + ) + + +class MoeEpReference: + """Reference implementation of routed SwiGLU experts plus EP dispatch. + + Global experts are assigned contiguously: rank ``r`` owns + ``[r * experts_per_rank, (r + 1) * experts_per_rank)``. Pass an explicit + initialized process group for multi-rank execution; ``None`` means a + one-rank reference even if the default distributed group is initialized. + + ``intermediate_format`` optionally applies a post-SwiGLU, pre-FC2 format + round trip to model fused kernels that materialize their FC2 input in low + precision. ``None`` preserves the raw mathematical reference semantics. + ``backward_operand_format`` additionally models dGLU staging of grad-output + and transposed weights along their backward reduction dimensions. + """ + + def __init__( + self, + *, + num_experts: int, + hidden_size: int, + intermediate_size: int, + top_k: int, + ep_group: Optional[dist.ProcessGroup] = None, + max_tokens_per_rank: Optional[int] = None, + output_format: Union[MoeFormat, str] = MoeFormat.BF16, + combine_format: Union[MoeFormat, str] = MoeFormat.BF16, + intermediate_format: Optional[Union[MoeFormat, str]] = None, + backward_operand_format: Optional[Union[MoeFormat, str]] = None, + apply_topk_in_fc1: bool = True, + weight_interleave_size: Optional[int] = None, + gate_up_clamp: Optional[float] = None, + generate_c: bool = False, + backward_wgrad_mode: str = "none", + token_padding_size: int = 128, + ) -> None: + for name, value in ( + ("num_experts", num_experts), + ("hidden_size", hidden_size), + ("intermediate_size", intermediate_size), + ("top_k", top_k), + ): + if not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + if top_k > num_experts: + raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})") + if max_tokens_per_rank is not None and max_tokens_per_rank < 0: + raise ValueError("max_tokens_per_rank must be non-negative") + if backward_wgrad_mode not in ("none", "operands"): + raise ValueError("backward_wgrad_mode must be 'none' or 'operands'") + if backward_wgrad_mode == "operands" and not generate_c: + raise ValueError("backward_wgrad_mode='operands' requires generate_c=True") + if not isinstance(token_padding_size, int) or token_padding_size <= 0: + raise ValueError("token_padding_size must be a positive integer") + if backward_wgrad_mode == "operands" and token_padding_size != 256: + raise ValueError("backward_wgrad_mode='operands' requires " "token_padding_size=256") + if weight_interleave_size not in (None, 32): + raise ValueError("weight_interleave_size must be None or 32") + + if ep_group is None: + ep_size, ep_rank = 1, 0 + else: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("ep_group requires an initialized torch.distributed process group") + ep_size = dist.get_world_size(ep_group) + ep_rank = dist.get_rank(ep_group) + if num_experts % ep_size != 0: + raise ValueError(f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})") + + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.top_k = top_k + self.ep_group = ep_group + self.ep_size = ep_size + self.ep_rank = ep_rank + self.experts_per_rank = num_experts // ep_size + self.max_tokens_per_rank = max_tokens_per_rank + self.output_format = _parse_format(output_format) + self.combine_format = _parse_format(combine_format) + self.intermediate_format = None if intermediate_format is None else _parse_format(intermediate_format) + self.backward_operand_format = None if backward_operand_format is None else _parse_format(backward_operand_format) + self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) + self.weight_interleave_size = weight_interleave_size + self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) + self.generate_c = bool(generate_c) + self.backward_wgrad_mode = backward_wgrad_mode + self.token_padding_size = token_padding_size + + for name, fmt in (("output_format", self.output_format), ("combine_format", self.combine_format)): + required_multiple = 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + if hidden_size % required_multiple != 0: + raise ValueError(f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for {name}={fmt.value}") + + def __repr__(self) -> str: + return ( + f"{type(self).__name__}(" + f"experts={self.num_experts}, local_experts={self.experts_per_rank}, " + f"hidden={self.hidden_size}, intermediate={self.intermediate_size}, " + f"top_k={self.top_k}, ep_rank={self.ep_rank}/{self.ep_size}, " + f"output={self.output_format.value}, combine={self.combine_format.value})" + ) + + def _collective_device(self, device: torch.device) -> torch.device: + """Device the process group can run ``all_to_all_single`` on. + + Gloo only implements all-to-all for CPU tensors, so CUDA tensors are + staged through host memory; NCCL groups communicate in place. + """ + if device.type != "cpu" and dist.get_backend(self.ep_group) == "gloo": + return torch.device("cpu") + return device + + def _exchange_counts(self, send_counts: torch.Tensor) -> torch.Tensor: + if self.ep_size == 1: + return send_counts.clone() + comm_device = self._collective_device(send_counts.device) + staged = send_counts.to(comm_device) + recv_counts = torch.empty_like(staged) + dist.all_to_all_single(recv_counts, staged, group=self.ep_group) + return recv_counts.to(send_counts.device) + + def _all_to_all( + self, + send: torch.Tensor, + send_counts: Sequence[int], + recv_counts: Sequence[int], + ) -> torch.Tensor: + if self.ep_size == 1: + return send.clone() + comm_device = self._collective_device(send.device) + staged = send.contiguous().to(comm_device) + output_shape = (sum(recv_counts), *send.shape[1:]) + recv = torch.empty(output_shape, dtype=send.dtype, device=comm_device) + dist.all_to_all_single( + recv, + staged, + output_split_sizes=list(recv_counts), + input_split_sizes=list(send_counts), + group=self.ep_group, + ) + return recv.to(send.device) + + def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> _DispatchPlan: + """Route valid ``topk_idx`` entries to destination ranks, stably by rank. + + Backward reuses this so gradient re-dispatch reproduces the exact + forward route order. + """ + + device = topk_idx.device + token_count = topk_idx.shape[0] + flat_expert = topk_idx.reshape(-1).to(torch.int64) + flat_weight = topk_weights.reshape(-1).float() + valid = flat_expert != -1 + invalid_negative = flat_expert < -1 + invalid_high = flat_expert >= self.num_experts + if bool((invalid_negative | invalid_high).any().item()): + bad = flat_expert[invalid_negative | invalid_high][0].item() + raise ValueError(f"topk_idx contains out-of-range expert id {bad}") + + flat_token = torch.arange(token_count, device=device).repeat_interleave(self.top_k) + flat_slot = torch.arange(self.top_k, device=device).repeat(token_count) + expert = flat_expert[valid] + destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") + order = torch.argsort(destination, stable=True) + + send_counts_tensor = torch.bincount(destination.index_select(0, order), minlength=self.ep_size).to(torch.int64) + recv_counts_tensor = self._exchange_counts(send_counts_tensor) + return _DispatchPlan( + send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), + send_weight=flat_weight[valid].index_select(0, order), + send_token_idx=flat_token[valid].index_select(0, order), + send_slot_idx=flat_slot[valid].index_select(0, order), + send_counts=tuple(int(v) for v in send_counts_tensor.cpu().tolist()), + recv_counts=tuple(int(v) for v in recv_counts_tensor.cpu().tolist()), + ) + + def _run_local_experts( + self, + tokens: torch.Tensor, + local_expert_idx: torch.Tensor, + route_weight: torch.Tensor, + fc1_weight: torch.Tensor, + fc2_weight: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + output = torch.empty( + (tokens.shape[0], self.hidden_size), + dtype=torch.float32, + device=tokens.device, + ) + fc1_c_rows = [] if self.generate_c else None + for expert in range(self.experts_per_rank): + positions = torch.nonzero(local_expert_idx == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + expert_tokens = tokens.index_select(0, positions) + gate_up = expert_tokens @ fc1_weight[expert] + if fc1_c_rows is not None: + # Raw pre-SwiGLU accumulator: before clamp, no router weight. + fc1_c_rows.append(gate_up.to(torch.bfloat16)) + if self.weight_interleave_size is not None: + gate_up = _deinterleave_glu(gate_up, self.weight_interleave_size) + gate, up = gate_up.split(self.intermediate_size, dim=-1) + if self.gate_up_clamp is not None: + gate = gate.clamp(max=self.gate_up_clamp) + up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + intermediate = F.silu(gate) * up + weights = route_weight.index_select(0, positions).unsqueeze(-1) + if self.apply_topk_in_fc1: + intermediate = intermediate * weights + if self.intermediate_format is not None: + intermediate = _format_round_trip( + intermediate, + self.intermediate_format, + ) + expert_output = intermediate @ fc2_weight[expert] + expert_output = forward_combine_round_trip( + expert_output, + self.combine_format, + ) + if not self.apply_topk_in_fc1: + # The upstream training kernel leaves scores out of dispatch + # and applies them in standalone TopkReduce after the combine + # wire-format round trip. + expert_output = expert_output * weights + output.index_copy_(0, positions, expert_output) + fc1_c = None + if fc1_c_rows is not None: + fc1_c = torch.cat(fc1_c_rows) if fc1_c_rows else torch.empty((0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device) + return output, fc1_c + + def __call__( + self, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ) -> Union[ + MoeTensor, + Tuple[MoeTensor, torch.Tensor, torch.Tensor], + Tuple[ + MoeTensor, + torch.Tensor, + torch.Tensor, + WgradForwardStashReference, + ], + ]: + """Run dispatch, local experts, return routing, top-k reduce, and encode. + + Shapes: + activation: ``(T, H)`` + fc1_weight: ``(E_local, H, 2 * I)`` + fc2_weight: ``(E_local, I, H)`` + topk_idx/topk_weights: ``(T, K)`` + + Returns the ``(T, H)`` result, or ``(result, fc1_c, route_metadata)`` + when constructed with ``generate_c=True``. In wgrad operand mode, a + fourth :class:`WgradForwardStashReference` item is returned. ``fc1_c`` + is the BF16 + pre-SwiGLU FC1 accumulator of every route this rank's experts + processed, ``(local_routes, 2 * I)``, grouped by local expert and + ordered within each expert by (source rank, source token-major route + order); captured before the gate/up clamp, without the router weight. + ``route_metadata`` is Int32 ``(local_routes, 4)`` with columns + ``(local_expert, src_rank, src_token, src_slot)``; row ``i`` identifies + the route behind ``fc1_c`` row ``i`` for the backward gradient + re-dispatch. + """ + + if self.weight_interleave_size == 32 and (not isinstance(fc1_weight, BlockScaledTensor) or fc1_weight.format is not MoeFormat.MXFP8): + raise ValueError("weight_interleave_size=32 requires an MXFP8 BlockScaledTensor " "for fc1_weight") + if topk_idx.ndim != 2: + raise ValueError(f"topk_idx must be 2-D, got shape {tuple(topk_idx.shape)}") + token_count = topk_idx.shape[0] + route_shape = (token_count, self.top_k) + if tuple(topk_idx.shape) != route_shape: + raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") + if tuple(topk_weights.shape) != route_shape: + raise ValueError(f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}") + if topk_idx.dtype not in (torch.int32, torch.int64): + raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") + if not topk_weights.is_floating_point(): + raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") + if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: + raise ValueError(f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}") + + device = _tensor_device(activation) + inputs = { + "fc1_weight": _tensor_device(fc1_weight), + "fc2_weight": _tensor_device(fc2_weight), + "topk_idx": topk_idx.device, + "topk_weights": topk_weights.device, + } + for name, input_device in inputs.items(): + if input_device != device: + raise ValueError(f"{name} must be on {device}, got {input_device}") + + activation_float = _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + ) + # The Rubin path first stages plain activation along H, then its + # forward column requantization forms x.T scales along routed K. + wgrad_activation_float = None + if self.backward_wgrad_mode == "operands": + wgrad_activation_float = _format_round_trip( + activation_float, + MoeFormat.MXFP8, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size), + quantized_axis=1, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + ) + if self.backward_wgrad_mode == "operands": + # Plain forward operands are staged to the same public MXFP8 + # reduction-axis representation before the Rubin GEMMs. + fc1_float = _format_round_trip_axis( + fc1_float, + MoeFormat.MXFP8, + axis=1, + ) + fc2_float = _format_round_trip_axis( + fc2_float, + MoeFormat.MXFP8, + axis=1, + ) + + plan = self._dispatch_plan(topk_idx, topk_weights) + send_token_idx = plan.send_token_idx + send_slot_idx = plan.send_slot_idx + send_counts, recv_counts = plan.send_counts, plan.recv_counts + forward_activation_float = wgrad_activation_float if wgrad_activation_float is not None else activation_float + send_tokens = forward_activation_float.index_select( + 0, + send_token_idx, + ) + + recv_tokens = self._all_to_all(send_tokens, send_counts, recv_counts) + recv_wgrad_tokens = None + if wgrad_activation_float is not None: + recv_wgrad_tokens = self._all_to_all( + wgrad_activation_float.index_select(0, send_token_idx), + send_counts, + recv_counts, + ) + recv_expert = self._all_to_all(plan.send_expert, send_counts, recv_counts) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + + route_metadata = None + fc1_c_order = None + if self.generate_c: + recv_src_rank = torch.repeat_interleave( + torch.arange(self.ep_size, device=device), + torch.tensor(recv_counts, device=device), + ) + recv_token = self._all_to_all(send_token_idx, send_counts, recv_counts) + recv_slot = self._all_to_all(send_slot_idx, send_counts, recv_counts) + # Stable sort by local expert reproduces the fc1_c row order + # (grouped by expert; source order preserved within each group). + fc1_c_order = torch.argsort(recv_expert, stable=True) + route_metadata = torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1).index_select(0, fc1_c_order).to(torch.int32) + # recv rows are ordered by source rank, then that source's token-major + # route order, so the per-expert position grouping below realizes the + # documented fc1_c ordering. + recv_output, fc1_c = self._run_local_experts( + recv_tokens, + recv_expert, + recv_weight, + fc1_float, + fc2_float, + ) + + returned = self._all_to_all(recv_output, recv_counts, send_counts) + combine_plane = torch.zeros( + (token_count * self.top_k, self.hidden_size), + dtype=torch.float32, + device=device, + ) + send_flat_slot = send_token_idx * self.top_k + send_slot_idx + combine_plane.index_copy_(0, send_flat_slot, returned) + reduced = combine_plane.view(token_count, self.top_k, self.hidden_size).sum(dim=1) + + if self.output_format is MoeFormat.BF16: + output = reduced.to(torch.bfloat16) + else: + output = quantize_blockwise(reduced, self.output_format, axis=-1) + if self.generate_c: + if self.backward_wgrad_mode == "operands": + if recv_wgrad_tokens is None or fc1_c_order is None: + raise RuntimeError("wgrad forward staging was not built") + valid_counts = tuple( + int(value) + for value in torch.bincount( + recv_expert, + minlength=self.experts_per_rank, + ) + .cpu() + .tolist() + ) + padded_ends = [] + total = 0 + for count in valid_counts: + total += ( + _ceil_div( + count, + self.token_padding_size, + ) + * self.token_padding_size + ) + padded_ends.append(total) + ordered_tokens = recv_wgrad_tokens.index_select( + 0, + fc1_c_order, + ) + metadata_experts = route_metadata[:, 0].to(torch.int64) + padded_x = _padded_expert_rows( + ordered_tokens, + metadata_experts, + valid_counts, + padded_ends, + ) + wgrad_stash = WgradForwardStashReference( + fc1_a=quantize_blockwise( + padded_x.transpose(0, 1), + MoeFormat.MXFP8, + axis=1, + ), + expert_offsets=torch.tensor( + padded_ends, + dtype=torch.int32, + device=device, + ), + valid_route_counts=torch.tensor( + valid_counts, + dtype=torch.int32, + device=device, + ), + route_metadata=route_metadata, + ) + return output, fc1_c, route_metadata, wgrad_stash + return output, fc1_c, route_metadata + return output + + def backward( + self, + grad_output: torch.Tensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + fc1_c: torch.Tensor, + route_metadata: torch.Tensor, + *, + wgrad_forward_stash: Optional[WgradForwardStashReference] = None, + ) -> Union[ + Tuple[torch.Tensor, torch.Tensor], + Tuple[ + torch.Tensor, + torch.Tensor, + WgradOperandsReference, + ], + ]: + """Backward pass consuming the ``generate_c=True`` stash. + + ``fc1_c`` is the recompute source: gate/up, the clamp masks, SwiGLU, + and the FC2 input are all rebuilt from it, so no post-SwiGLU forward + intermediate needs to be saved. ``route_metadata`` alone reconstructs + the mapping between re-dispatched rows and ``fc1_c`` rows and drives + the gradient return scatter. + + Quantization round-trips (input decode, ``combine_format``, + ``output_format``) are treated as straight-through identities; + ``grad_output`` is the ``(T, H)`` gradient of the dequantized output. + + Returns ``(grad_activation, grad_topk_weights)`` in float32. In wgrad + operand mode, a third :class:`WgradOperandsReference` item models the + caller-owned grouped-GEMM operands. + """ + + if not self.generate_c: + raise RuntimeError("backward requires the operator to be constructed with generate_c=True") + if self.backward_wgrad_mode == "operands": + if not isinstance( + wgrad_forward_stash, + WgradForwardStashReference, + ): + raise TypeError("wgrad_forward_stash must be a " "WgradForwardStashReference") + if not torch.equal( + wgrad_forward_stash.route_metadata, + route_metadata, + ): + raise ValueError("wgrad_forward_stash route identity does not match " "route_metadata") + elif wgrad_forward_stash is not None: + raise ValueError("wgrad_forward_stash is only accepted in operands mode") + token_count = topk_idx.shape[0] + if tuple(grad_output.shape) != (token_count, self.hidden_size): + raise ValueError(f"grad_output shape must be {(token_count, self.hidden_size)}, got {tuple(grad_output.shape)}") + if not grad_output.is_floating_point(): + raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") + + device = _tensor_device(fc1_weight) + two_i = 2 * self.intermediate_size + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, two_i), + quantized_axis=1, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + ) + semantic_fc2_float = fc2_float + effective_backward_format = self.backward_operand_format + if effective_backward_format is None and self.backward_wgrad_mode == "operands": + effective_backward_format = MoeFormat.MXFP8 + if effective_backward_format is not None: + # The dGLU adapter requantizes both transposed weights along the + # backward GEMM reduction dimension. + fc1_float = _format_round_trip_axis( + fc1_float.transpose(1, 2), + effective_backward_format, + axis=1, + ).transpose(1, 2) + fc2_float = _format_round_trip_axis( + fc2_float.transpose(1, 2), + effective_backward_format, + axis=1, + ).transpose(1, 2) + if fc1_c.shape != (int(route_metadata.shape[0]), two_i): + raise ValueError(f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got {tuple(fc1_c.shape)}") + + # Re-dispatch router weights and output gradients along the identical + # forward routes. + plan = self._dispatch_plan(topk_idx, topk_weights) + send_counts, recv_counts = plan.send_counts, plan.recv_counts + semantic_grad_output = grad_output.float() + grad_output_float = semantic_grad_output + if effective_backward_format is not None: + grad_output_float = _format_round_trip( + grad_output_float, + effective_backward_format, + ) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + recv_grad = self._all_to_all(grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts) + recv_semantic_grad = self._all_to_all( + semantic_grad_output.index_select(0, plan.send_token_idx), + send_counts, + recv_counts, + ) + + # route_metadata rows are in fc1_c order; sorting them by + # (src_rank, src_token, src_slot) reproduces the receive order, giving + # the permutation between re-dispatched rows and fc1_c rows. + metadata = route_metadata.to(device=device, dtype=torch.int64) + local_routes = metadata.shape[0] + if local_routes > 0: + token_span = int(metadata[:, 2].max().item()) + 1 + recv_key = (metadata[:, 1] * token_span + metadata[:, 2]) * self.top_k + metadata[:, 3] + perm = torch.argsort(recv_key) # perm[j] = fc1_c row at receive position j + else: + perm = torch.empty((0,), dtype=torch.int64, device=device) + w_rows = torch.empty_like(recv_weight) + w_rows.index_copy_(0, perm, recv_weight) + dy_rows = torch.empty_like(recv_grad) + dy_rows.index_copy_(0, perm, recv_grad) + semantic_dy_rows = torch.empty_like(recv_semantic_grad) + semantic_dy_rows.index_copy_(0, perm, recv_semantic_grad) + + c_rows = fc1_c.float() + expert_rows = metadata[:, 0] + d_x_rows = torch.zeros((local_routes, self.hidden_size), dtype=torch.float32, device=device) + d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) + weighted_h_rows = torch.zeros( + (local_routes, self.intermediate_size), + dtype=torch.float32, + device=device, + ) + wgrad_dy_rows = torch.zeros( + (local_routes, self.hidden_size), + dtype=torch.float32, + device=device, + ) + dc_rows = torch.zeros( + (local_routes, two_i), + dtype=torch.float32, + device=device, + ) + for expert in range(self.experts_per_rank): + positions = torch.nonzero(expert_rows == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + c = c_rows.index_select(0, positions) + w = w_rows.index_select(0, positions).unsqueeze(-1) + d_y = dy_rows.index_select(0, positions) + semantic_d_y = semantic_dy_rows.index_select(0, positions) + + if self.weight_interleave_size is not None: + c = _deinterleave_glu(c, self.weight_interleave_size) + gate, up = c.split(self.intermediate_size, dim=-1) + if self.gate_up_clamp is not None: + g = gate.clamp(max=self.gate_up_clamp) + u = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + else: + g, u = gate, up + sig = torch.sigmoid(g) + s = g * sig + h = s * u + weighted_h_rows.index_copy_(0, positions, h * w) + wgrad_dy_rows.index_copy_(0, positions, d_y) + + if self.apply_topk_in_fc1: + d_y_pre = d_y + else: + d_y_pre = d_y * w + d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) + if self.apply_topk_in_fc1: + d_h = d_h_fc2 * w + semantic_d_h = semantic_d_y @ semantic_fc2_float[expert].transpose(0, 1) + d_w_rows[positions] = (semantic_d_h * h).sum(dim=-1) + else: + d_h = d_h_fc2 + d_w_rows[positions] = (semantic_d_y * (h @ semantic_fc2_float[expert])).sum(dim=-1) + + d_g = d_h * u * (sig * (1 + g * (1 - sig))) + d_u = d_h * s + if self.gate_up_clamp is not None: + d_gate = d_g * (gate <= self.gate_up_clamp) + d_up = d_u * ((up >= -self.gate_up_clamp) & (up <= self.gate_up_clamp)) + else: + d_gate, d_up = d_g, d_u + d_c = torch.cat((d_gate, d_up), dim=-1) + if self.weight_interleave_size is not None: + d_c = _interleave_glu(d_c, self.weight_interleave_size) + dc_rows.index_copy_(0, positions, d_c) + if self.intermediate_format is not None: + d_c = _format_round_trip(d_c, self.intermediate_format) + d_x = d_c @ fc1_float[expert].transpose(0, 1) + d_x_rows.index_copy_( + 0, + positions, + backward_combine_round_trip(d_x, self.combine_format), + ) + + # Return the route gradients to their source ranks and scatter-add. + returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) + returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) + grad_activation = torch.zeros((token_count, self.hidden_size), dtype=torch.float32, device=device) + grad_activation.index_add_(0, plan.send_token_idx, returned_dx) + grad_topk_weights = torch.zeros((token_count * self.top_k,), dtype=torch.float32, device=device) + grad_topk_weights.index_copy_(0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw) + grad_topk_weights = grad_topk_weights.view( + token_count, + self.top_k, + ) + if self.backward_wgrad_mode == "operands": + stash = wgrad_forward_stash + padded_ends = tuple(int(value) for value in stash.expert_offsets.cpu().tolist()) + valid_counts = tuple(int(value) for value in stash.valid_route_counts.cpu().tolist()) + padded_dc = _padded_expert_rows( + dc_rows, + expert_rows, + valid_counts, + padded_ends, + ) + padded_weighted_h = _padded_expert_rows( + weighted_h_rows, + expert_rows, + valid_counts, + padded_ends, + ) + padded_wgrad_dy = _padded_expert_rows( + wgrad_dy_rows, + expert_rows, + valid_counts, + padded_ends, + ) + operands = WgradOperandsReference( + fc1_a=stash.fc1_a, + fc1_b=quantize_blockwise( + padded_dc, + MoeFormat.MXFP8, + axis=0, + ), + fc2_a=quantize_blockwise( + padded_weighted_h.transpose(0, 1), + MoeFormat.MXFP8, + axis=1, + ), + fc2_b=quantize_blockwise( + padded_wgrad_dy, + MoeFormat.MXFP8, + axis=0, + ), + expert_offsets=stash.expert_offsets, + valid_route_counts=stash.valid_route_counts, + route_metadata=stash.route_metadata, + ) + return grad_activation, grad_topk_weights, operands + return grad_activation, grad_topk_weights + + +__all__ = [ + "BlockScaledTensor", + "MoeEpReference", + "MoeFormat", + "MoeTensor", + "WgradForwardStashReference", + "WgradOperandsReference", + "backward_combine_round_trip", + "forward_combine_round_trip", + "quantize_blockwise", +] diff --git a/test/python/moe_ep/moe_ep_test_support.py b/test/python/moe_ep/moe_ep_test_support.py new file mode 100644 index 000000000..f0c341871 --- /dev/null +++ b/test/python/moe_ep/moe_ep_test_support.py @@ -0,0 +1,1132 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared data, forward, and backward support for MoE EP tests.""" + +from __future__ import annotations + +# Common + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from moe_ep.moe_ep_reference import ( + BlockScaledTensor as ReferenceBlockScaledTensor, + MoeEpReference, + MoeFormat, + forward_combine_round_trip, + quantize_blockwise, +) + +__all__ = [ + "_assert_backward_matches", + "_assert_grouped_wgrads_match_reference", + "_assert_matches_reference", + "_assert_wgrads_match_reference", + "_dense_wgrads_from_operands", + "_dense_wgrads_from_grouped_kernel", + "_fixed_training_reference", + "_fixed_training_weights", + "_allocate_stateless_training_outputs", + "_allocate_training_weight_staging", + "_forward_config", + "_grad_output", + "_make_forward_case", + "_naive_reference", + "_output_as_float", + "_reference_backward", + "_reference_forward", + "_replay_cuda_graph", + "_require_distributed_sm107", + "_run_grouped_wgrad_kernel", + "_sm107_device", + "_stress_backend_reuse", + "_training_abi_prepared", + "_training_config", + "_training_prepared_pair", + "make_distributed_forward_inputs", + "make_forward_inputs", + "quantize_mxfp8", +] + + +# Data + + +def make_forward_inputs(device: torch.device): + """Build one deterministic MXFP8 forward case.""" + + generator = torch.Generator(device=device).manual_seed(20260811) + experts, tokens, hidden, intermediate = 2, 5, 128, 256 + activation = quantize_mxfp8( + torch.randn(tokens, hidden, generator=generator, device=device), + axis=1, + ) + fc1_weight = quantize_mxfp8( + torch.randn( + experts, + hidden, + 2 * intermediate, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + fc2_weight = quantize_mxfp8( + torch.randn( + experts, + intermediate, + hidden, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + topk_idx = torch.tensor( + [[0, 1], [1, 0], [0, -1], [1, 0], [0, 1]], + dtype=torch.int32, + device=device, + ) + topk_weights = torch.tensor( + [ + [0.75, 0.25], + [0.625, 0.375], + [1.0, 0.0], + [0.5, 0.5], + [0.875, 0.125], + ], + dtype=torch.bfloat16, + device=device, + ) + return activation, fc1_weight, fc2_weight, topk_idx, topk_weights + + +def make_distributed_forward_inputs( + rank: int, + world_size: int, + device: torch.device, +): + """Build rank-local inputs with one local and one remote route per token.""" + + generator = torch.Generator(device=device).manual_seed(20260811 + rank) + # Vary local shapes without exceeding the distributed tests' + # max_tokens_per_rank=8 contract at EP sizes above seven. + local_experts, tokens, hidden, intermediate = ( + 2, + rank % 7 + 2, + 128, + 256, + ) + activation = quantize_mxfp8( + torch.randn(tokens, hidden, generator=generator, device=device), + axis=1, + ) + fc1_weight = quantize_mxfp8( + torch.randn( + local_experts, + hidden, + 2 * intermediate, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + fc2_weight = quantize_mxfp8( + torch.randn( + local_experts, + intermediate, + hidden, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + remote_rank = (rank + 1) % world_size + topk_idx = torch.tensor( + [ + [ + rank * local_experts + token % local_experts, + remote_rank * local_experts + (token + 1) % local_experts, + ] + for token in range(tokens) + ], + dtype=torch.int32, + device=device, + ) + topk_weights = ( + torch.tensor( + [[0.625, 0.375]], + dtype=torch.bfloat16, + device=device, + ) + .expand(tokens, -1) + .contiguous() + ) + return activation, fc1_weight, fc2_weight, topk_idx, topk_weights + + +def quantize_mxfp8(tensor: torch.Tensor, *, axis: int = -1): + """Return a public logical MXFP8 tensor (E4M3 payload + E8M0 scales).""" + + from cudnn import BlockScaledTensor + + axis = axis % tensor.ndim + logical_shape = tuple(tensor.shape) + logical_extent = logical_shape[axis] + moved = tensor.float().movedim(axis, -1) + block_count = (logical_extent + 31) // 32 + padded_extent = block_count * 32 + if padded_extent != logical_extent: + moved = F.pad(moved, (0, padded_extent - logical_extent)) + + blocks = moved.reshape(*moved.shape[:-1], block_count, 32) + raw_scale = blocks.abs().amax(dim=-1) / 448.0 + safe_scale = torch.where(raw_scale > 0, raw_scale, 1.0) + power_of_two_scale = torch.where( + raw_scale > 0, + torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), + torch.zeros_like(raw_scale), + ) + scale = power_of_two_scale.to(torch.float8_e8m0fnu) + reciprocal = torch.where(scale.float() > 0, scale.float().reciprocal(), 0.0) + payload = (blocks * reciprocal.unsqueeze(-1)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn).reshape(*moved.shape)[..., :logical_extent] + + return BlockScaledTensor( + data=payload.movedim(-1, axis).contiguous(), + scale=scale.movedim(-1, axis).contiguous(), + format="mxfp8", + logical_shape=logical_shape, + axis=axis, + ) + + +# Training setup + + +def _training_config(**overrides): + from cudnn.moe_ep._contracts import ForwardConfig, normalize_fc1_weight_layout + from cudnn.moe_ep._tuning import MoeEpTuningConfig + + weight_interleave_size = overrides.pop("weight_interleave_size", None) + values = { + "num_experts": 2, + "hidden_size": 128, + "intermediate_size": 256, + "top_k": 2, + "experts_per_rank": 2, + "ep_size": 1, + "ep_rank": 0, + "ep_group": None, + "ep_global_ranks": (), + "max_tokens_per_rank": 4, + "max_recv_size_per_rank": 4, + "drop_on_overflow": True, + "output_format": "bf16", + "combine_format": "bf16", + "apply_topk_in_fc1": True, + "gate_up_clamp": None, + "generate_c": True, + "token_padding_size": 128, + "sf_padding_size": 128, + "tuning": MoeEpTuningConfig(), + "backward_wgrad_mode": "operands", + "fc1_weight_layout": normalize_fc1_weight_layout(weight_interleave_size), + } + values.update(overrides) + return ForwardConfig(**values) + + +def _training_prepared_pair(config, pool_rows: int = 512): + from cudnn.moe_ep._megamoe_backend._workspace import WorkspaceRequirements + + forward_shapes = { + "fc1_c": (pool_rows, 512), + "col_quant_data": (pool_rows, 128), + "col_quant_sf": (2048,), + } + backward_shapes = { + "dprob": (4, 2), + "fc1_recompute": (pool_rows, 256), + "fc1_recompute_sf": (256, 8), + "fc1_col_output": (pool_rows, 512), + "fc1_col_output_sf": (512, 8), + "grad_y2": (pool_rows, 128), + "grad_y2_sf": (2048,), + } + forward = SimpleNamespace( + pool_token_capacity=pool_rows, + workspace_requirements=WorkspaceRequirements.for_mxfp8( + config, + kernel_local_workspace_bytes=1024, + kernel_shared_workspace_bytes=2048, + col_quant_data_bytes=pool_rows * 128, + col_quant_sf_bytes=2048, + ), + kernel=SimpleNamespace(get_aux_output_shapes=lambda: forward_shapes), + col_quant_sizes_offset=0, + col_quant_sizes_bytes=8, + ) + backward = SimpleNamespace( + pool_token_capacity=pool_rows, + config=SimpleNamespace(sf_padding_block=128), + workspace_requirements=WorkspaceRequirements.for_mxfp8( + config, + kernel_local_workspace_bytes=3072, + kernel_shared_workspace_bytes=4096, + backward_dprob_bytes=4 * 2 * 4, + backward_aux_data_bytes=pool_rows * 512, + backward_aux_scale_bytes=512 * 8, + ), + kernel=SimpleNamespace( + get_aux_output_shapes=lambda: backward_shapes, + get_fc1_preact_shape=lambda: forward_shapes["fc1_c"], + ), + ) + return forward, backward + + +def _training_abi_prepared(name: str, max_recv_size: int = 4): + from cudnn.moe_ep._megamoe_backend._workspace import ( + BufferRegion, + WorkspaceRequirements, + ) + + workspace = WorkspaceRequirements( + max_tokens_per_rank=4, + symmetric_regions=(BufferRegion("symmetric", 256),), + local_regions=(BufferRegion("local", 128),), + ) + kernel_config = SimpleNamespace( + max_recv_size_per_rank=max_recv_size, + effective_config=lambda cluster_count: { + "name": name, + "max_recv_size_per_rank": max_recv_size, + "launch_cluster_count": cluster_count, + }, + ) + return SimpleNamespace( + kernel=SimpleNamespace( + name=lambda: name, + threads_per_cta=128, + occupancy=1, + smem_capacity=1024, + ), + architecture=(10, 7), + config=kernel_config, + launch_cluster_count=16, + workspace_requirements=workspace, + pool_token_capacity=512, + ) + + +# Forward + + +_DEFAULT_FORWARD_CONFIG = { + "num_experts": 2, + "hidden_size": 128, + "intermediate_size": 256, + "top_k": 2, + "max_tokens_per_rank": 5, + "apply_topk_in_fc1": True, + "combine_format": "bf16", + "output_format": "bf16", +} +_REFERENCE_CLOSE_KWARGS = {"rtol": 0.05, "atol": 0.0625} + + +def _forward_config(**overrides): + return {**_DEFAULT_FORWARD_CONFIG, **overrides} + + +def _output_as_float(output): + if isinstance(output, torch.Tensor): + return output.float() + return output.dequantize() + + +def _assert_matches_reference(actual, expected): + torch.testing.assert_close( + _output_as_float(actual), + _output_as_float(expected), + **_REFERENCE_CLOSE_KWARGS, + ) + + +def _naive_reference( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + *, + apply_topk_in_fc1, + clamp=None, + combine_format=MoeFormat.BF16, + intermediate_format=None, + apply_topk_after_combine=False, +): + token_count, top_k = topk_idx.shape + hidden_size = activation.shape[1] + intermediate_size = fc2_weight.shape[1] + combine = torch.zeros( + token_count, + top_k, + hidden_size, + dtype=torch.float32, + device=activation.device, + ) + for token in range(token_count): + for slot in range(top_k): + expert = int(topk_idx[token, slot]) + if expert == -1: + continue + gate_up = activation[token].float() @ fc1_weight[expert].float() + gate, up = gate_up.split(intermediate_size) + if clamp is not None: + gate = gate.clamp(max=clamp) + up = up.clamp(-clamp, clamp) + intermediate = F.silu(gate) * up + route_weight = topk_weights[token, slot].float() + if apply_topk_in_fc1: + intermediate = intermediate * route_weight + if intermediate_format is not None: + intermediate = quantize_blockwise( + intermediate, + intermediate_format, + ).dequantize() + result = intermediate @ fc2_weight[expert].float() + if not apply_topk_in_fc1 and not apply_topk_after_combine: + result = result * route_weight + result = forward_combine_round_trip(result, combine_format) + if not apply_topk_in_fc1 and apply_topk_after_combine: + result = result * route_weight + combine[token, slot] = result + return combine.sum(dim=1).to(torch.bfloat16) + + +def _as_reference_tensor(tensor): + if isinstance(tensor, torch.Tensor): + return tensor + return ReferenceBlockScaledTensor( + data=tensor.data, + scale=tensor.scale, + format=tensor.format.value, + logical_shape=tensor.logical_shape, + axis=tensor.axis, + ) + + +def _reference_args(args): + return ( + _as_reference_tensor(args[0]), + _as_reference_tensor(args[1]), + _as_reference_tensor(args[2]), + args[3], + args[4], + ) + + +def _reference_forward(args, **overrides): + # Rubin's fused FC1 epilogue stores the post-SwiGLU intermediate as MXFP8 + # before FC2 consumes it. Keep MoeEpReference's default raw semantics for + # its standalone tests, but model the device precision for API comparisons. + config = _forward_config(**overrides) + config.pop("tuning", None) + config.setdefault("intermediate_format", "mxfp8") + return MoeEpReference(**config)(*_reference_args(args)) + + +def _sm107_device() -> torch.device: + if not torch.cuda.is_available(): + pytest.skip("Rubin MXFP8 forward requires CUDA") + device = torch.device("cuda", 0) + if torch.cuda.get_device_capability(device) != (10, 7): + pytest.skip("Rubin MXFP8 forward requires exactly SM107 (compute capability 10.7)") + return device + + +def _require_distributed_sm107(world_size: int) -> None: + if not dist.is_available() or not dist.is_nccl_available(): + pytest.skip("multi-GPU Rubin MXFP8 forward requires NCCL") + if torch.cuda.device_count() < world_size: + pytest.skip(f"multi-GPU Rubin MXFP8 forward requires {world_size} GPUs") + if any(torch.cuda.get_device_capability(index) != (10, 7) for index in range(world_size)): + pytest.skip("multi-GPU Rubin MXFP8 forward requires exactly SM107 " "(compute capability 10.7) on every rank") + try: + import nvshmem.core # noqa: F401 + except (ImportError, OSError): + pytest.skip("multi-GPU Rubin MXFP8 forward requires NVSHMEM") + + +def _make_forward_case( + device: torch.device, + *, + experts: int, + tokens: int, + hidden: int, + intermediate: int, + top_k: int, + index_dtype: torch.dtype, + weight_dtype: torch.dtype, +): + """Build a deterministic supported case for the shape/format matrix.""" + + seed = 20260811 + experts * 1009 + tokens * 101 + hidden * 11 + intermediate + top_k + generator = torch.Generator(device=device).manual_seed(seed) + activation = quantize_mxfp8( + torch.randn(tokens, hidden, generator=generator, device=device), + axis=1, + ) + fc1_weight = quantize_mxfp8( + torch.randn( + experts, + hidden, + 2 * intermediate, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + fc2_weight = quantize_mxfp8( + torch.randn( + experts, + intermediate, + hidden, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + topk_idx = torch.arange(tokens * top_k, device=device).reshape(tokens, top_k).remainder(experts).to(index_dtype) + topk_weights = torch.arange( + 1, + tokens * top_k + 1, + dtype=torch.float32, + device=device, + ).reshape(tokens, top_k) + topk_weights /= topk_weights.sum(dim=1, keepdim=True) + return ( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights.to(weight_dtype), + ) + + +def _stress_backend_reuse( + op, + args, + original_topk_idx, + original_topk_weights, + device, + *, + check_weight_refresh, +): + backend = op._forward_backend + assert backend is not None + compiled = backend._compiled + plan_workspace = backend._plan._workspace + weight_refresh_count = backend._adapter.weight_refresh_count if check_weight_refresh else None + alternate_stream = torch.cuda.Stream(device=device) + + for iteration in range(100): + args[3].copy_(original_topk_idx) + args[4].copy_(original_topk_weights * float((iteration % 7) + 1) / 7.0) + if iteration % 10 == 0: + args[3].fill_(-1) + stream = torch.cuda.current_stream(device) if iteration % 2 == 0 else alternate_stream + with torch.cuda.stream(stream): + stressed = op(*args) + stream.synchronize() + if iteration % 10 == 0: + assert _output_as_float(stressed).eq(0).all() + else: + assert torch.isfinite(_output_as_float(stressed)).all() + assert backend._compiled is compiled + assert backend._plan._workspace is plan_workspace + if weight_refresh_count is not None: + assert backend._adapter.weight_refresh_count == weight_refresh_count + + +def _replay_cuda_graph( + op, + args, + original_topk_idx, + expected, + device, + *, + synchronize_ranks=None, +): + synchronize_ranks = synchronize_ranks or (lambda: None) + op.warmup(*args) + synchronize_ranks() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = op(*args) + synchronize_ranks() + + for replay in range(20): + if replay % 2: + args[3].fill_(-1) + else: + args[3].copy_(original_topk_idx) + synchronize_ranks() + graph.replay() + torch.cuda.synchronize(device) + if replay % 2: + assert _output_as_float(graph_output).eq(0).all() + else: + _assert_matches_reference(graph_output, expected) + + +# Backward + + +_BACKWARD_CLOSE_KWARGS = ( + {"rtol": 0.15, "atol": 0.125}, # grad_activation is BF16-rounded. + {"rtol": 0.15, "atol": 0.125}, # router-weight gradient. +) +_WGRAD_CLOSE_KWARGS = {"rtol": 0.2, "atol": 0.25} +_GROUPED_WGRAD_CLOSE_KWARGS = {"rtol": 0.1, "atol": 0.1} + + +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +def _unpack_wgrad_scale_part( + packed: torch.Tensor, + rows: int, + columns: int, +) -> torch.Tensor: + """Invert grouped-wgrad's 128x4 scale-atom swizzle.""" + + padded_rows = _round_up(rows, 128) + padded_columns = _round_up(columns, 4) + row_atoms = padded_rows // 128 + column_atoms = padded_columns // 4 + atom_count = row_atoms * column_atoms + expected = padded_rows * padded_columns + if packed.numel() != expected: + raise ValueError(f"packed scale part has {packed.numel()} bytes, expected {expected}") + blocked = ( + packed.reshape(atom_count, 32, 4, 4).transpose(1, 2).reshape(row_atoms, column_atoms, 128, 4).permute(0, 2, 1, 3).reshape(padded_rows, padded_columns) + ) + return blocked[:rows, :columns].view(torch.float8_e8m0fnu).float() + + +def _dequantize_wgrad_operand( + data: torch.Tensor, + scales: torch.Tensor, + expert_offsets: torch.Tensor, + *, + k_dim: int, +) -> torch.Tensor: + """Decode one public grouped-wgrad operand without launching a GEMM.""" + + if data.ndim != 2 or k_dim not in (0, 1): + raise ValueError("wgrad operand must be rank 2 with k_dim 0 or 1") + non_k = int(data.shape[1 - k_dim]) + padded_non_k = _round_up(non_k, 128) + flat_scales = scales.view(torch.uint8).reshape(-1) + output = torch.zeros(data.shape, dtype=torch.float32, device=data.device) + ends = [int(value) for value in expert_offsets.detach().cpu().tolist()] + k_capacity = int(data.shape[k_dim]) + previous = 0 + scale_byte_offset = 0 + for end in ends: + if end < previous or end > k_capacity: + raise ValueError("expert offsets must be nondecreasing and fit the operand " f"K capacity ({k_capacity})") + extent = end - previous + if extent % 32: + raise ValueError("each padded expert K extent must be divisible by 32") + if extent == 0: + continue + scale_columns = _round_up(extent // 32, 4) + scale_byte_count = padded_non_k * scale_columns + if scale_byte_offset + scale_byte_count > flat_scales.numel(): + raise ValueError("expert offsets exceed the scale tensor") + part = flat_scales.narrow( + 0, + scale_byte_offset, + scale_byte_count, + ) + logical_scale = _unpack_wgrad_scale_part( + part, + non_k, + extent // 32, + ) + if k_dim == 1: + expanded_scale = logical_scale.repeat_interleave(32, dim=1) + output[:, previous:end] = data[:, previous:end].float() * expanded_scale + else: + expanded_scale = logical_scale.repeat_interleave( + 32, + dim=1, + ).transpose(0, 1) + output[previous:end, :] = data[previous:end, :].float() * expanded_scale + previous = end + scale_byte_offset += scale_byte_count + + if previous < k_capacity: + capacity_tail = data.narrow(k_dim, previous, k_capacity - previous) + if bool(capacity_tail.float().ne(0).any().item()): + raise ValueError("unused WGrad operand capacity tail must contain zero data") + scale_tail = flat_scales[scale_byte_offset:] + if scale_tail.numel() and bool(scale_tail.ne(127).any().item()): + raise ValueError("unused WGrad operand capacity tail must contain neutral E8M0 scales") + return output + + +def _dense_wgrads_from_operands(operands): + """Reference grouped matmuls over the producer-native operand ABI.""" + + fc1_a = _dequantize_wgrad_operand( + operands.fc1_a, + operands.fc1_sfa, + operands.expert_offsets, + k_dim=1, + ) + fc1_b = _dequantize_wgrad_operand( + operands.fc1_b, + operands.fc1_sfb, + operands.expert_offsets, + k_dim=0, + ) + fc2_a = _dequantize_wgrad_operand( + operands.fc2_a, + operands.fc2_sfa, + operands.expert_offsets, + k_dim=1, + ) + fc2_b = _dequantize_wgrad_operand( + operands.fc2_b, + operands.fc2_sfb, + operands.expert_offsets, + k_dim=0, + ) + fc1_parts = [] + fc2_parts = [] + ends = [int(value) for value in operands.expert_offsets.detach().cpu().tolist()] + valid_counts = [int(value) for value in operands.valid_route_counts.detach().cpu().tolist()] + if len(ends) != len(valid_counts): + raise ValueError("expert offsets and valid route counts must have equal size") + + previous = 0 + for expert, (end, valid_count) in enumerate(zip(ends, valid_counts)): + extent = end - previous + if valid_count < 0 or valid_count > extent: + raise ValueError(f"expert {expert} valid route count {valid_count} exceeds " f"its padded extent {extent}") + valid_end = previous + valid_count + for name, tensor, k_dim in ( + ("fc1_a", fc1_a, 1), + ("fc1_b", fc1_b, 0), + ("fc2_a", fc2_a, 1), + ("fc2_b", fc2_b, 0), + ): + padding = tensor.narrow(k_dim, valid_end, end - valid_end) + if bool(padding.ne(0).any().item()): + raise ValueError(f"{name} expert {expert} padded rows must decode to zero") + fc1_parts.append(fc1_a[:, previous:valid_end] @ fc1_b[previous:valid_end, :]) + fc2_parts.append(fc2_a[:, previous:valid_end] @ fc2_b[previous:valid_end, :]) + previous = end + return torch.stack(fc1_parts), torch.stack(fc2_parts) + + +def _run_grouped_wgrad_kernel( + operands, + prefix: str, + *, + wgrad_tensor=None, + accumulate_on_output: bool = False, + current_stream=None, +): + """Run one fixed-capacity operand bundle through production WGrad.""" + + import cudnn + + if prefix not in ("fc1", "fc2"): + raise ValueError(f"prefix must be 'fc1' or 'fc2', got {prefix!r}") + # Graph callers provide one persistent output per training lane. This is + # currently also the isolation key for a temporary production-WGrad + # workaround: an EP2 graph with two same-signature calls produced correct + # operands but corrupted the second WGrad when both calls shared the + # cached API object's mutable TMA descriptor workspace. Distinct fixed + # outputs make the calls use distinct workspaces. The production fix + # should instead share the compiled kernel while owning descriptor + # workspace per graph call site, after which output identity must no + # longer participate in the compile cache key. + return cudnn.grouped_gemm_wgrad_wrapper_sm100( + a_tensor=getattr(operands, f"{prefix}_a"), + b_tensor=getattr(operands, f"{prefix}_b"), + sfa_tensor=getattr(operands, f"{prefix}_sfa"), + sfb_tensor=getattr(operands, f"{prefix}_sfb"), + offsets_tensor=operands.expert_offsets, + output_mode="dense", + wgrad_tensor=wgrad_tensor, + wgrad_dtype=torch.bfloat16, + acc_dtype=torch.float32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=32, + accumulate_on_output=accumulate_on_output, + input_order="tensor2d", + current_stream=current_stream, + )["wgrad_tensor"] + + +def _dense_wgrads_from_grouped_kernel( + operands, + *, + wgrad_tensors=None, + accumulate_on_output: bool = False, + current_stream=None, +): + """Run both fixed-capacity operand bundles through production WGrad.""" + + if wgrad_tensors is None: + wgrad_tensors = (None, None) + if len(wgrad_tensors) != 2: + raise ValueError("wgrad_tensors must contain FC1 and FC2 outputs") + return tuple( + _run_grouped_wgrad_kernel( + operands, + prefix, + wgrad_tensor=output, + accumulate_on_output=accumulate_on_output, + current_stream=current_stream, + ) + for prefix, output in zip(("fc1", "fc2"), wgrad_tensors) + ) + + +def _assert_grouped_wgrads_match_reference( + actual, + expected, + *, + reference_name: str, + close_kwargs=None, +) -> None: + """Compare grouped-kernel FC1/FC2 outputs and report useful error maxima.""" + + if close_kwargs is None: + close_kwargs = _GROUPED_WGRAD_CLOSE_KWARGS + for name, actual_dw, expected_dw in zip( + ("grad_fc1_weight", "grad_fc2_weight"), + actual, + expected, + ): + actual_fp32 = actual_dw.float() + expected_fp32 = expected_dw.float() + absolute_error = (actual_fp32 - expected_fp32).abs() + max_absolute_error = absolute_error.max().item() + max_relative_error = (absolute_error / expected_fp32.abs().clamp_min(1.0e-6)).max().item() + torch.testing.assert_close( + actual_fp32, + expected_fp32, + msg=lambda default, name=name: ( + f"{name} does not match {reference_name}; " f"max_abs_error={max_absolute_error:.6g}, " f"max_rel_error={max_relative_error:.6g}\n{default}" + ), + **close_kwargs, + ) + + +def _reference_backward(config) -> MoeEpReference: + options = dict(config) + for production_only in ( + "drop_on_overflow", + "ep_global_ranks", + "ep_rank", + "ep_size", + "experts_per_rank", + "max_recv_size_per_rank", + "sf_padding_size", + "tuning", + ): + options.pop(production_only, None) + options["intermediate_format"] = "mxfp8" + options["backward_operand_format"] = "mxfp8" + return MoeEpReference(**options) + + +def _fixed_training_weights(args): + """Build independent source packs for allocation-free native packing.""" + + from cudnn.moe_ep import MoeEpBackwardWeights, MoeEpForwardWeights + from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( + _quantize_plain_mxfp8, + ) + + fc1_weight = args[1] + fc2_weight = args[2] + dense_fc1 = fc1_weight if isinstance(fc1_weight, torch.Tensor) else fc1_weight.dequantize() + dense_fc2 = fc2_weight if isinstance(fc2_weight, torch.Tensor) else fc2_weight.dequantize() + forward = MoeEpForwardWeights( + fc1=(_quantize_plain_mxfp8(dense_fc1, axis=1) if isinstance(fc1_weight, torch.Tensor) else fc1_weight), + fc2=(_quantize_plain_mxfp8(dense_fc2, axis=1) if isinstance(fc2_weight, torch.Tensor) else fc2_weight), + ) + backward = MoeEpBackwardWeights( + w2_transpose=_quantize_plain_mxfp8( + dense_fc2.transpose(1, 2).contiguous(), + axis=1, + ), + w1_transpose=_quantize_plain_mxfp8( + dense_fc1.transpose(1, 2).contiguous(), + axis=1, + ), + ) + return forward, backward + + +def _allocate_training_weight_staging(weights): + """Allocate caller-owned native pack destinations for one source pair.""" + + from cudnn.moe_ep import ( + MoeEpBackwardWeightStaging, + MoeEpForwardWeightStaging, + ) + + forward, backward = weights + fc1 = forward.fc1 + fc2 = forward.fc2 + experts, hidden, gate_up = fc1.data.shape + intermediate = fc2.data.shape[1] + + def scale(elements): + return torch.empty( + (experts, elements), + dtype=torch.float8_e8m0fnu, + device=fc1.device, + ) + + def blocked_elements(rows, columns): + return ((rows + 127) // 128 * 128) * ((columns + 3) // 4 * 4) + + forward_out = MoeEpForwardWeightStaging( + fc1_payload=torch.empty_strided( + fc1.data.shape, + (hidden * gate_up, 1, hidden), + dtype=fc1.data.dtype, + device=fc1.device, + ), + fc1_scale=scale(blocked_elements(gate_up, hidden // 32)), + fc2_payload=torch.empty_strided( + fc2.data.shape, + (intermediate * hidden, 1, intermediate), + dtype=fc2.data.dtype, + device=fc2.device, + ), + fc2_scale=scale(blocked_elements(hidden, intermediate // 32)), + ) + w2t = backward.w2_transpose + w1t = backward.w1_transpose + backward_out = MoeEpBackwardWeightStaging( + w2_transpose_payload=torch.empty_like(w2t.data), + w2_transpose_scale=scale(blocked_elements(intermediate, hidden // 32)), + w1_transpose_payload=torch.empty_like(w1t.data), + w1_transpose_scale=scale(blocked_elements(hidden, gate_up // 32)), + ) + return forward_out, backward_out + + +def _allocate_stateless_training_outputs(requirements, device, symmetric_buffers): + """Bind symmetric final outputs and allocate the remaining contracts.""" + + from cudnn.moe_ep import ( + MoeEpTrainingBackwardOutputs, + MoeEpTrainingForwardOutputs, + ) + + def allocate(name): + shape, stride, dtype, _alignment = requirements[name] + return torch.empty_strided( + shape, + stride, + dtype=dtype, + device=device, + ) + + forward = MoeEpTrainingForwardOutputs( + output=symmetric_buffers["output"], + fc1_preact=allocate("fc1_preact"), + fc1_a=allocate("fc1_a"), + fc1_sfa=allocate("fc1_sfa"), + valid_route_counts=allocate("valid_route_counts"), + expert_offsets=allocate("expert_offsets"), + ) + backward = MoeEpTrainingBackwardOutputs( + grad_activation=symmetric_buffers["grad_activation"], + dprob=symmetric_buffers["dprob"], + fc1_b=allocate("fc1_b"), + fc1_sfb=allocate("fc1_sfb"), + fc2_a=allocate("fc2_a"), + fc2_sfa=allocate("fc2_sfa"), + fc2_b=allocate("fc2_b"), + fc2_sfb=allocate("fc2_sfb"), + ) + return forward, backward + + +def _fixed_training_reference( + args, + grad_output, + *, + combine_format, + gate_up_clamp, + ep_group=None, + num_experts=None, + **config_overrides, +): + """Run the standalone oracle for EP1 or a distributed EP group.""" + + ep_size = 1 if ep_group is None else dist.get_world_size(ep_group) + local_experts = int(args[1].shape[0]) + if num_experts is None: + num_experts = local_experts * ep_size + reference_config = _forward_config(**config_overrides) + reference_config.update( + num_experts=num_experts, + hidden_size=int(args[0].shape[1]), + intermediate_size=int(args[2].shape[1]), + top_k=int(args[3].shape[1]), + max_tokens_per_rank=config_overrides.get( + "max_tokens_per_rank", + int(args[0].shape[0]), + ), + ep_group=ep_group, + combine_format=combine_format, + gate_up_clamp=gate_up_clamp, + generate_c=True, + backward_wgrad_mode="operands", + # The standalone operand oracle's legacy ABI uses 256-row + # segments. The stateless producer ABI uses 128-row segments; + # their represented dense gradients are compared below. + token_padding_size=256, + ) + reference = _reference_backward(reference_config) + reference_args = _reference_args(args) + output, fc1_c, route_metadata, forward_stash = reference(*reference_args) + grad_activation, grad_topk_weights, wgrad_operands = reference.backward( + grad_output, + *reference_args[1:], + fc1_c, + route_metadata, + wgrad_forward_stash=forward_stash, + ) + return ( + output, + grad_activation, + grad_topk_weights, + wgrad_operands, + ) + + +def _grad_output( + device: torch.device, + token_count: int, + *, + seed: int, +) -> torch.Tensor: + generator = torch.Generator(device=device).manual_seed(seed) + return ( + torch.randn( + token_count, + 128, + generator=generator, + dtype=torch.float32, + device=device, + ) + / 8 + ) + + +def _assert_backward_matches(actual, expected, topk_idx) -> None: + assert len(actual) == len(expected) == 2 + for name, gradient, reference, close_kwargs in zip( + ("grad_activation", "grad_topk_weights"), + actual, + expected, + _BACKWARD_CLOSE_KWARGS, + ): + assert gradient.shape == reference.shape + assert gradient.dtype == torch.float32 + assert torch.isfinite(gradient).all() + torch.testing.assert_close( + gradient, + reference, + msg=lambda default, name=name: (f"{name} does not match the backward reference\n{default}"), + **close_kwargs, + ) + + dropped = topk_idx == -1 + assert actual[1][dropped].eq(0).all() + + +def _interleave_fc1_wgrad( + tensor: torch.Tensor, + interleave_size: int = 32, +) -> torch.Tensor: + """Convert logical gate-then-up columns to producer-native strip order.""" + + out_features = tensor.shape[-1] + return ( + tensor.view( + *tensor.shape[:-1], + 2, + out_features // (2 * interleave_size), + interleave_size, + ) + .transpose(-3, -2) + .reshape(tensor.shape) + ) + + +def _assert_wgrads_match_reference( + actual, + expected, + *, + expected_dense=None, + weight_interleave_size=None, +) -> None: + """Compare fixed-capacity production operands with standalone dense dW.""" + + torch.testing.assert_close( + actual.valid_route_counts, + expected.valid_route_counts, + rtol=0, + atol=0, + msg="valid route counts differ from the independent reference", + ) + actual_dense = _dense_wgrads_from_operands(actual) + if expected_dense is None: + expected_dense = expected.dense_wgrads() + if weight_interleave_size is not None: + expected_fc1, expected_fc2 = expected_dense + expected_dense = ( + _interleave_fc1_wgrad(expected_fc1, weight_interleave_size), + expected_fc2, + ) + for name, actual_dw, expected_dw in zip( + ("grad_fc1_weight", "grad_fc2_weight"), + actual_dense, + expected_dense, + ): + torch.testing.assert_close( + actual_dw, + expected_dw, + msg=lambda default, name=name: (f"{name} does not match the independent reference\n{default}"), + **_WGRAD_CLOSE_KWARGS, + ) diff --git a/test/python/moe_ep/probe_moe_ep_training_graph.py b/test/python/moe_ep/probe_moe_ep_training_graph.py new file mode 100644 index 000000000..801ee8fb9 --- /dev/null +++ b/test/python/moe_ep/probe_moe_ep_training_graph.py @@ -0,0 +1,253 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Stateless SM107 multi-rank CUDA Graph training probe.""" + +from __future__ import annotations + +import argparse +import os +from datetime import timedelta + +import torch +import torch.distributed as dist + +from cudnn import MoeEp +from moe_ep.moe_ep_test_support import ( + _allocate_stateless_training_outputs, + _allocate_training_weight_staging, + _fixed_training_weights, + _grad_output, + make_distributed_forward_inputs, +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--diagnostic-replays", type=int, default=2) + parser.add_argument("--burst-replays", type=int, default=100) + parser.add_argument("--multistream-replays", type=int, default=10) + parser.add_argument("--max-recv-size-per-rank", type=int, default=1) + parser.add_argument("--cycles", type=int, default=2) + parser.add_argument("--timeout-seconds", type=int, default=600) + parser.add_argument("--skip-multistream", action="store_true") + parser.add_argument("--expect-overflow-assert", action="store_true") + return parser.parse_args() + + +def _positive(name: str, value: int) -> None: + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + + +def _capture_training_graph( + op: MoeEp, + lane, + args, + grad_output, + native_forward, + native_backward, + forward_out, + backward_out, +) -> torch.cuda.CUDAGraph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op.training_forward( + lane, + args[0], + args[3], + args[4], + weights=native_forward, + out=forward_out, + ) + op.training_backward( + lane, + grad_output, + args[3], + args[4], + weights=native_backward, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, + ) + return graph + + +def _prepare_case( + *, + device: torch.device, + rank: int, + world_size: int, + lane_count: int, + max_recv_size_per_rank: int, + drop_on_overflow: bool, +): + args = make_distributed_forward_inputs(rank, world_size, device) + args = (*args[:4], args[4].float().contiguous()) + grad_output = _grad_output(device, args[0].shape[0], seed=7000 + rank) + source_weights = _fixed_training_weights(args) + op = MoeEp( + num_experts=2 * world_size, + hidden_size=128, + intermediate_size=256, + top_k=2, + ep_group=dist.group.WORLD, + # This is a collective ABI capacity, not the rank-local token count. + # make_distributed_forward_inputs intentionally varies local shapes. + max_tokens_per_rank=8, + max_recv_size_per_rank=max_recv_size_per_rank, + drop_on_overflow=drop_on_overflow, + combine_format="bf16", + weight_interleave_size=32, + ) + requirements = op.prepare_training(lane_count=lane_count, device=device) + forward_staging, backward_staging = _allocate_training_weight_staging(source_weights) + native_forward = op.pack_forward_weights( + source_weights[0], + out=forward_staging, + ) + native_backward = op.pack_backward_weights( + source_weights[1], + out=backward_staging, + ) + output_pairs = tuple( + _allocate_stateless_training_outputs( + requirements, + device, + op.training_symmetric_buffers(lane), + ) + for lane in op.training_lanes + ) + return ( + op, + args, + grad_output, + native_forward, + native_backward, + output_pairs, + ) + + +def _run_cycle(args: argparse.Namespace, *, device: torch.device, rank: int, world_size: int) -> None: + lane_count = 1 if args.skip_multistream else 2 + case = _prepare_case( + device=device, + rank=rank, + world_size=world_size, + lane_count=lane_count, + max_recv_size_per_rank=args.max_recv_size_per_rank, + drop_on_overflow=not args.expect_overflow_assert, + ) + op, inputs, grad_output, native_forward, native_backward, output_pairs = case + try: + # Warm each lane and every kernel specialization before capture. + for lane, (forward_out, backward_out) in zip( + op.training_lanes, + output_pairs, + ): + op.training_forward( + lane, + inputs[0], + inputs[3], + inputs[4], + weights=native_forward, + out=forward_out, + ) + op.training_backward( + lane, + grad_output, + inputs[3], + inputs[4], + weights=native_backward, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, + ) + torch.cuda.synchronize(device) + dist.barrier(group=dist.group.WORLD, device_ids=[device.index]) + + graphs = tuple( + _capture_training_graph( + op, + lane, + inputs, + grad_output, + native_forward, + native_backward, + forward_out, + backward_out, + ) + for lane, (forward_out, backward_out) in zip( + op.training_lanes, + output_pairs, + ) + ) + dist.barrier(group=dist.group.WORLD, device_ids=[device.index]) + + replay_count = args.diagnostic_replays + args.burst_replays + if lane_count > 1: + replay_count += args.multistream_replays + caught = None + try: + for replay in range(replay_count): + graphs[replay % len(graphs)].replay() + torch.cuda.synchronize(device) + except Exception as error: + caught = error + + if args.expect_overflow_assert: + if caught is None: + raise AssertionError("expected the captured overflow assertion") + elif caught is not None: + raise caught + dist.barrier(group=dist.group.WORLD, device_ids=[device.index]) + finally: + op.close() + + +def main() -> None: + args = _parse_args() + for name in ( + "diagnostic_replays", + "burst_replays", + "multistream_replays", + "max_recv_size_per_rank", + "cycles", + "timeout_seconds", + ): + _positive(name, getattr(args, name)) + + local_rank = int(os.environ["LOCAL_RANK"]) + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + dist.init_process_group( + backend="nccl", + timeout=timedelta(seconds=args.timeout_seconds), + device_id=device, + ) + try: + if torch.cuda.get_device_capability(device) != (10, 7): + raise RuntimeError("stateless training graph probe requires SM107") + for _ in range(args.cycles): + _run_cycle(args, device=device, rank=rank, world_size=world_size) + if rank == 0: + print( + "stateless MoeEP training graph probe passed: " f"world_size={world_size}, cycles={args.cycles}", + flush=True, + ) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/test/python/moe_ep/test_moe_ep_autotune.py b/test/python/moe_ep/test_moe_ep_autotune.py new file mode 100644 index 000000000..c80bfd9e6 --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_autotune.py @@ -0,0 +1,616 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts and smoke coverage for the explicit MoeEP sweep autotuner.""" + +from __future__ import annotations + +import contextlib +from types import SimpleNamespace + +import pytest +import torch + +from moe_ep.moe_ep_test_support import ( + _allocate_stateless_training_outputs, + _allocate_training_weight_staging, + _assert_backward_matches, + _assert_matches_reference, + _fixed_training_reference, + _fixed_training_weights, + _forward_config, + _grad_output, + _reference_forward, + _replay_cuda_graph, + _sm107_device, + make_forward_inputs, +) + + +def _validated_request(config, *args, **kwargs): + del args, kwargs + return SimpleNamespace(config=config, device=torch.device("cuda", 0)) + + +def _patch_common_inference_dependencies( + patch, + *, + api_module, + backend_module, + create_backend, +) -> None: + patch.setattr(api_module, "validate_forward", _validated_request) + patch.setattr(backend_module, "validate_config", lambda config: None) + patch.setattr(backend_module, "validate_request", lambda request: None) + patch.setattr(backend_module, "create_backend", create_backend) + patch.setattr(torch.cuda, "device", lambda device: contextlib.nullcontext()) + patch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + +@pytest.mark.L0 +def test_autotune_core_contracts(monkeypatch): + import cudnn.moe_ep._autotune as autotune_module + from cudnn import ( + MoeEpAutotuneCandidateResult, + MoeEpAutotuneResult, + MoeEpTuningConfig, + ) + + baseline = MoeEpTuningConfig() + candidate = MoeEpTuningConfig(token_in_flag_batch=2) + normalize = autotune_module.normalize_candidates + assert normalize( + baseline, + [candidate, baseline, candidate], + warmup_iters=0, + timed_iters=1, + max_candidates=2, + ) == (baseline, candidate) + + invalid_limits = ( + ({"warmup_iters": -1}, "warmup_iters"), + ({"timed_iters": 0}, "timed_iters"), + ({"max_candidates": 33}, "max_candidates"), + ) + for overrides, message in invalid_limits: + arguments = { + "warmup_iters": 0, + "timed_iters": 1, + "max_candidates": 32, + **overrides, + } + with pytest.raises(ValueError, match=message): + normalize(baseline, [baseline], **arguments) + + with pytest.raises(ValueError, match="does not sweep reduce_topk_in_kernel"): + normalize( + baseline, + [MoeEpTuningConfig(reduce_topk_in_kernel=True)], + warmup_iters=0, + timed_iters=1, + max_candidates=32, + ) + with pytest.raises(ValueError, match="exceeding max_candidates=1"): + normalize( + baseline, + [candidate], + warmup_iters=0, + timed_iters=1, + max_candidates=1, + ) + + first = MoeEpAutotuneCandidateResult(baseline, 1.0, (1.0,)) + second = MoeEpAutotuneCandidateResult(candidate, 1.0, (1.0,)) + assert autotune_module.select_winner((first, second)) is first + result = MoeEpAutotuneResult("inference", first.tuning, (first, second)) + assert result.evaluated_candidates == 2 + + requirements = { + name: ((2, 3), (3, 1), torch.float32, 1) + for name in ( + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + ) + } + symmetric_buffers = { + "output": torch.empty(2, 3), + "grad_activation": torch.empty(2, 3), + "dprob": torch.empty(2, 3), + } + forward_out, backward_out = autotune_module.allocate_training_outputs( + requirements, + torch.device("cpu"), + symmetric_buffers, + ) + assert forward_out.output is symmetric_buffers["output"] + assert backward_out.grad_activation is symmetric_buffers["grad_activation"] + assert backward_out.dprob is symmetric_buffers["dprob"] + + with monkeypatch.context() as patch: + remote = (candidate,) + patch.setattr(autotune_module.dist, "get_world_size", lambda group: 2) + + def gather(output, value, *, group): + del group + output[:] = [value, remote] + + patch.setattr(autotune_module.dist, "all_gather_object", gather) + with pytest.raises(RuntimeError, match="must match on every EP rank"): + autotune_module.verify_candidates_across_ranks((baseline,), object()) + + with monkeypatch.context() as patch: + local_samples = iter((1.0, 7.0, 3.0)) + + class Event: + def record(self, stream): + del stream + + def synchronize(self): + pass + + def elapsed_time(self, end): + del end + return next(local_samples) + + patch.setattr(torch.cuda, "current_stream", lambda device: object()) + patch.setattr(torch.cuda, "Event", lambda enable_timing: Event()) + patch.setattr( + autotune_module.dist, + "all_reduce", + lambda values, **kwargs: values.copy_(torch.tensor([5.0, 8.0, 4.0], dtype=values.dtype)), + ) + latency, samples = autotune_module.benchmark_candidate( + lambda: None, + device=torch.device("cpu"), + group=object(), + timed_iters=3, + ) + assert samples == (5.0, 8.0, 4.0) + assert latency == 5.0 + + +@pytest.mark.L0 +def test_autotune_api_transactions(monkeypatch): + import cudnn.moe_ep._autotune as autotune_module + import cudnn.moe_ep._backend as backend_module + import cudnn.moe_ep._megamoe_backend.mxfp8._training_execute as execute_module + import cudnn.moe_ep.api as api_module + from cudnn import MoeEp, MoeEpTuningConfig + + baseline = MoeEpTuningConfig() + candidate = MoeEpTuningConfig(token_in_flag_batch=2) + + # Validation failures happen before teardown and preserve active state. + op = MoeEp(**_forward_config()) + active_backend = object() + op._forward_backend = active_backend + with pytest.raises(ValueError, match="does not sweep reduce_topk_in_kernel"): + op.autotune( + None, + None, + None, + None, + None, + candidates=[MoeEpTuningConfig(reduce_topk_in_kernel=True)], + warmup_iters=0, + timed_iters=1, + ) + assert op.tuning == baseline + assert op._forward_backend is active_backend + op._forward_backend = None + op.close() + + # A runtime failure is fail-fast and permanently poisons the instance. + with monkeypatch.context() as patch: + calls = [] + + class FailingBackend: + def forward(self, request): + calls.append(request.config.tuning) + raise RuntimeError("launch failed") + + def close(self): + pass + + _patch_common_inference_dependencies( + patch, + api_module=api_module, + backend_module=backend_module, + create_backend=lambda config, device: FailingBackend(), + ) + patch.setattr( + autotune_module, + "verify_state_across_ranks", + lambda state, group: None, + ) + op = MoeEp(**_forward_config()) + with pytest.raises(RuntimeError, match="candidate 0.*compile/prime"): + op.autotune( + None, + None, + None, + None, + None, + candidates=[candidate], + warmup_iters=0, + timed_iters=1, + ) + assert calls == [baseline] + with pytest.raises(RuntimeError, match="unusable"): + op.autotune( + None, + None, + None, + None, + None, + candidates=[candidate], + warmup_iters=0, + timed_iters=1, + ) + op._forward_backend = None + op.close() + + # Inference commits only the measured winner and retains its rebuilt backend. + with monkeypatch.context() as patch: + active_backends = [] + + class InferenceBackend: + def __init__(self, config): + self.config = config + self.closed = False + + def forward(self, request): + assert request.config == self.config + return object() + + def close(self): + self.closed = True + + def create_inference_backend(config, device): + del device + active_backends.append(InferenceBackend(config)) + return active_backends[-1] + + def benchmark_inference(run, *, device, group, timed_iters): + del device, group, timed_iters + run() + tuning = active_backends[-1].config.tuning + latency = 1.0 if tuning == candidate else 2.0 + return latency, (latency,) + + _patch_common_inference_dependencies( + patch, + api_module=api_module, + backend_module=backend_module, + create_backend=create_inference_backend, + ) + patch.setattr( + autotune_module, + "benchmark_candidate", + benchmark_inference, + ) + patch.setattr( + autotune_module, + "synchronize_candidate", + lambda device, group: None, + ) + op = MoeEp(**_forward_config()) + result = op.autotune( + None, + None, + None, + None, + None, + candidates=[candidate], + warmup_iters=0, + timed_iters=1, + ) + assert result.winner == candidate + assert result.evaluated_candidates == 2 + assert op.tuning == op._forward_config.tuning == candidate + assert op._forward_backend is active_backends[-1] + assert not active_backends[-1].closed + op.close() + + # Training times forward/backward pairs and leaves preparation to the caller. + with monkeypatch.context() as patch: + launches = [] + active_backends = [] + requirement_names = ( + "output", + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + "grad_activation", + "dprob", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + ) + symmetric_buffers = { + "output": object(), + "grad_activation": object(), + "dprob": object(), + } + + class TrainingState: + def public_requirements(self): + return {name: None for name in requirement_names} + + def public_symmetric_buffers(self, lane): + assert lane == 0 + return symmetric_buffers + + def views(self, *, lane, token_count): + return lane, token_count + + class TrainingBackend: + def __init__(self, config): + self.config = config + + def prepare_training(self, *, lane_count): + assert lane_count == 1 + return TrainingState() + + def close(self): + pass + + forward_outputs = SimpleNamespace( + fc1_preact=object(), + output=object(), + fc1_a=object(), + fc1_sfa=object(), + valid_route_counts=object(), + expert_offsets=object(), + ) + backward_outputs = SimpleNamespace() + + patch.setattr( + api_module, + "_validate_training_assert_capability", + lambda config: None, + ) + for validation_name in ( + "validate_native_forward_weights", + "validate_native_backward_weights", + "validate_training_forward_outputs", + "validate_training_backward_outputs", + "validate_training_forward_state", + ): + patch.setattr( + api_module, + validation_name, + lambda *args, **kwargs: None, + ) + patch.setattr( + api_module, + "validate_training_input", + lambda *args, **kwargs: 2, + ) + patch.setattr(backend_module, "validate_config", lambda config: None) + + def create_training_backend(config, device): + del device + active_backends.append(TrainingBackend(config)) + return active_backends[-1] + + patch.setattr(backend_module, "create_backend", create_training_backend) + patch.setattr( + autotune_module, + "allocate_training_outputs", + lambda requirements, device, symmetric: ( + (forward_outputs, backward_outputs) + if symmetric is symmetric_buffers + else pytest.fail("autotune used the wrong symmetric buffers") + ), + ) + patch.setattr( + autotune_module, + "synchronize_candidate", + lambda device, group: None, + ) + + def benchmark_training(run, *, device, group, timed_iters): + del device, group, timed_iters + run() + tuning = active_backends[-1].config.tuning + latency = 1.0 if tuning == candidate else 2.0 + return latency, (latency,) + + patch.setattr( + autotune_module, + "benchmark_candidate", + benchmark_training, + ) + patch.setattr( + torch.cuda, + "device", + lambda device: contextlib.nullcontext(), + ) + patch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + patch.setattr( + execute_module, + "launch_training_forward", + lambda *args, **kwargs: launches.append("forward"), + ) + patch.setattr( + execute_module, + "launch_training_backward", + lambda *args, **kwargs: launches.append("backward"), + ) + + op = MoeEp(**_forward_config(), weight_interleave_size=32) + value = SimpleNamespace(device=torch.device("cuda", 0)) + result = op.autotune_training( + value, + value, + None, + None, + forward_weights=object(), + backward_weights=object(), + candidates=[candidate], + warmup_iters=0, + timed_iters=1, + ) + assert result.mode == "training" + assert result.winner == candidate + assert launches == ["forward", "backward"] * 4 + assert op._training_state is None + assert op._forward_backend is None + op.close() + + +def _print_candidate_timings(label, result) -> None: + print(f"\n{label} autotune timings:", flush=True) + for index, measurement in enumerate(result.candidates): + samples = ", ".join(f"{sample:.4f}" for sample in measurement.samples_ms) + print( + f" [{index}] median={measurement.latency_ms:.4f} ms " f"samples=[{samples}] tuning={measurement.tuning}", + flush=True, + ) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_autotune_sm107_inference_training_and_graph(): + from cudnn import MoeEp, MoeEpTuningConfig + + device = _sm107_device() + candidates = [ + MoeEpTuningConfig(), + MoeEpTuningConfig(token_in_flag_batch=2), + MoeEpTuningConfig(token_in_flag_batch=4), + MoeEpTuningConfig(epi_flag_batch=(2, 1)), + MoeEpTuningConfig(group_hint=64), + ] + + inference_args = make_forward_inputs(device) + inference_expected = _reference_forward(inference_args) + original_topk_idx = inference_args[3].clone() + with MoeEp(**_forward_config()) as op: + result = op.autotune( + *inference_args, + candidates=candidates, + warmup_iters=1, + timed_iters=2, + ) + _print_candidate_timings("inference", result) + actual = op(*inference_args) + torch.cuda.synchronize(device) + assert result.evaluated_candidates == len(candidates) == 5 + assert result.winner in candidates + assert op.tuning == result.winner + assert op._forward_backend is not None + _assert_matches_reference(actual, inference_expected) + _replay_cuda_graph( + op, + inference_args, + original_topk_idx, + inference_expected, + device, + ) + + base_args = make_forward_inputs(device) + training_args = ( + base_args[0].dequantize(torch.bfloat16), + base_args[1], + base_args[2], + base_args[3], + base_args[4].float().contiguous(), + ) + grad_output = _grad_output( + device, + training_args[0].shape[0], + seed=20260903, + ) + training_expected = _fixed_training_reference( + training_args, + grad_output, + combine_format="bf16", + gate_up_clamp=None, + ) + source_weights = _fixed_training_weights(training_args) + with MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=training_args[0].shape[0], + max_recv_size_per_rank=2 * 128, + drop_on_overflow=True, + combine_format="bf16", + weight_interleave_size=32, + ) as op: + forward_staging, backward_staging = _allocate_training_weight_staging(source_weights) + native_forward = op.pack_forward_weights( + source_weights[0], + out=forward_staging, + ) + native_backward = op.pack_backward_weights( + source_weights[1], + out=backward_staging, + ) + result = op.autotune_training( + training_args[0], + grad_output, + training_args[3], + training_args[4], + forward_weights=native_forward, + backward_weights=native_backward, + candidates=candidates, + warmup_iters=1, + timed_iters=2, + ) + _print_candidate_timings("training", result) + requirements = op.prepare_training(lane_count=1, device=device) + lane = op.training_lanes[0] + forward_out, backward_out = _allocate_stateless_training_outputs( + requirements, + device, + op.training_symmetric_buffers(lane), + ) + actual_y = op.training_forward( + lane, + training_args[0], + training_args[3], + training_args[4], + weights=native_forward, + out=forward_out, + ) + actual_dx, actual_dprob, _ = op.training_backward( + lane, + grad_output, + training_args[3], + training_args[4], + weights=native_backward, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, + ) + torch.cuda.synchronize(device) + assert result.evaluated_candidates == len(candidates) == 5 + assert result.winner in candidates + assert op.tuning == result.winner + _assert_matches_reference(actual_y, training_expected[0]) + _assert_backward_matches( + (actual_dx, actual_dprob), + (training_expected[1], training_expected[2]), + training_args[3], + ) diff --git a/test/python/moe_ep/test_moe_ep_backward.py b/test/python/moe_ep/test_moe_ep_backward.py new file mode 100644 index 000000000..9fe58194b --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_backward.py @@ -0,0 +1,1195 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Stateless MoE EP training contracts.""" + +from __future__ import annotations + +from dataclasses import fields + +import cudnn +import pytest +import torch + +from cudnn.moe_ep import ( + BlockScaledTensor, + MoeEp, + MoeEpBackwardWeightStaging, + MoeEpBackwardWeights, + MoeEpExecutionLane, + MoeEpForwardWeights, + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, + MoeEpTrainingBackwardOutputs, + MoeEpTrainingForwardOutputs, + MoeEpTrainingWgradOperands, + pack_backward_weights, + pack_forward_weights, +) +from cudnn.moe_ep._megamoe_backend.mxfp8._training_resources import ( + Mxfp8TrainingState, + _build_training_abi_facts, + _harmonize_symmetric_regions, +) +from cudnn.moe_ep._megamoe_backend.mxfp8._training_execute import _stage_input +from cudnn.moe_ep._megamoe_backend.mxfp8._training_weights import ( + backward_native_to_kernel, + forward_native_to_kernel, +) +from cudnn.moe_ep._megamoe_backend.mxfp8._training_wgrad import ( + assemble_training_wgrad_operands, +) +from cudnn.moe_ep._megamoe_backend._workspace import ( + BufferRegion, + WorkspaceRequirements, + WorkspaceViews, +) +from cudnn.moe_ep._megamoe_backend.mxfp8._fingerprint import canonical_json_sha256 +from cudnn.moe_ep._validation import ( + validate_native_backward_weights, + validate_native_forward_weights, + validate_training_backward_outputs, + validate_training_forward_outputs, + validate_training_forward_state, + validate_training_input, + validate_training_non_aliasing, +) +from cudnn.moe_ep.api import _resolve_training_device +from moe_ep.moe_ep_test_support import ( + _allocate_stateless_training_outputs, + _allocate_training_weight_staging, + _assert_backward_matches, + _assert_grouped_wgrads_match_reference, + _assert_matches_reference, + _assert_wgrads_match_reference, + _dense_wgrads_from_grouped_kernel, + _fixed_training_reference, + _fixed_training_weights, + _grad_output, + _interleave_fc1_wgrad, + _sm107_device, + _training_abi_prepared, + _training_config, + _training_prepared_pair, + make_forward_inputs, + quantize_mxfp8, +) + + +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +def _blocked_scale_elements(rows: int, columns: int) -> int: + return _round_up(rows, 128) * _round_up(columns, 4) + + +def _native_forward( + config, + *, + device: torch.device = torch.device("cpu"), +) -> MoeEpNativeForwardWeights: + e = config.experts_per_rank + h = config.hidden_size + i = config.intermediate_size + fc1 = torch.empty_strided( + (e, h, 2 * i), + (h * 2 * i, 1, h), + dtype=torch.float8_e4m3fn, + device=device, + ) + fc2 = torch.empty_strided( + (e, i, h), + (i * h, 1, i), + dtype=torch.float8_e4m3fn, + device=device, + ) + return MoeEpNativeForwardWeights( + fc1=MoeEpNativeWeight( + fc1, + torch.empty( + (e, _blocked_scale_elements(2 * i, h // 32)), + dtype=torch.float8_e8m0fnu, + device=device, + ), + MoeEpNativeWeightLayout.FORWARD_FC1_GATE_UP_INTERLEAVED_32_V1, + ), + fc2=MoeEpNativeWeight( + fc2, + torch.empty( + (e, _blocked_scale_elements(h, i // 32)), + dtype=torch.float8_e8m0fnu, + device=device, + ), + MoeEpNativeWeightLayout.FORWARD_FC2_K_MAJOR_V1, + ), + ) + + +def _native_backward( + config, + *, + device: torch.device = torch.device("cpu"), +) -> MoeEpNativeBackwardWeights: + e = config.experts_per_rank + h = config.hidden_size + i = config.intermediate_size + return MoeEpNativeBackwardWeights( + w2_transpose=MoeEpNativeWeight( + torch.empty( + (e, h, i), + dtype=torch.float8_e4m3fn, + device=device, + ), + torch.empty( + (e, _blocked_scale_elements(i, h // 32)), + dtype=torch.float8_e8m0fnu, + device=device, + ), + MoeEpNativeWeightLayout.BACKWARD_W2_TRANSPOSE_V1, + ), + w1_transpose=MoeEpNativeWeight( + torch.empty( + (e, 2 * i, h), + dtype=torch.float8_e4m3fn, + device=device, + ), + torch.empty( + (e, _blocked_scale_elements(h, 2 * i // 32)), + dtype=torch.float8_e8m0fnu, + device=device, + ), + MoeEpNativeWeightLayout.BACKWARD_W1_TRANSPOSE_GATE_UP_INTERLEAVED_32_V1, + ), + ) + + +def _source_weights(config): + e = config.experts_per_rank + h = config.hidden_size + i = config.intermediate_size + + def block_scaled(shape): + scale_shape = list(shape) + scale_shape[1] //= 32 + return BlockScaledTensor( + data=torch.empty(shape, dtype=torch.float8_e4m3fn), + scale=torch.empty(scale_shape, dtype=torch.float8_e8m0fnu), + format="mxfp8", + logical_shape=shape, + axis=1, + ) + + return ( + MoeEpForwardWeights( + block_scaled((e, h, 2 * i)), + block_scaled((e, i, h)), + ), + MoeEpBackwardWeights( + block_scaled((e, h, i)), + block_scaled((e, 2 * i, h)), + ), + ) + + +@pytest.mark.L0 +def test_only_stateless_training_types_are_public(): + removed = ( + "MoeEpTrainingResources", + "MoeEpTrainingSlot", + "MoeEpTrainingWeights", + ) + for name in removed: + assert not hasattr(cudnn, name) + for name in ( + "MoeEpForwardWeights", + "MoeEpBackwardWeights", + "MoeEpNativeWeight", + "MoeEpTrainingForwardOutputs", + "MoeEpTrainingBackwardOutputs", + "pack_forward_weights", + "pack_backward_weights", + ): + assert hasattr(cudnn, name) + + +@pytest.mark.L0 +def test_training_device_prefers_explicit_then_current(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 2) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 4) + + assert _resolve_training_device(None) == torch.device("cuda:2") + assert _resolve_training_device("cuda") == torch.device("cuda:2") + assert _resolve_training_device(1) == torch.device("cuda:1") + with pytest.raises(ValueError, match="must be CUDA"): + _resolve_training_device("cpu") + + +@pytest.mark.L0 +def test_training_input_rejects_noncontiguous_plain_tensor(): + config = _training_config(weight_interleave_size=32) + activation = torch.empty((config.hidden_size, 2), dtype=torch.bfloat16).t() + topk_idx = torch.tensor([[0, 1], [1, 0]], dtype=torch.int32) + topk_weights = torch.ones((2, 2), dtype=torch.float32) + + with pytest.raises(ValueError, match="activation must be contiguous"): + validate_training_input( + config, + "activation", + activation, + topk_idx, + topk_weights, + device=torch.device("cpu"), + ) + + +@pytest.mark.L0 +def test_training_input_accepts_logical_view_of_padded_lane_scale(): + config = _training_config( + weight_interleave_size=32, + max_tokens_per_rank=5, + ) + token_count = 5 + logical_scale_columns = config.hidden_size // 32 + lane_scale = torch.empty( + (128, 16), + dtype=torch.float8_e8m0fnu, + ) + scale = lane_scale[:token_count, :logical_scale_columns] + activation = BlockScaledTensor( + data=torch.empty( + (token_count, config.hidden_size), + dtype=torch.float8_e4m3fn, + ), + scale=scale, + format="mxfp8", + logical_shape=(token_count, config.hidden_size), + axis=1, + ) + topk_idx = torch.zeros((token_count, config.top_k), dtype=torch.int32) + topk_weights = torch.ones((token_count, config.top_k), dtype=torch.float32) + + assert scale.shape == (5, 4) + assert scale.stride() == (16, 1) + assert not scale.is_contiguous() + assert ( + validate_training_input( + config, + "activation", + activation, + topk_idx, + topk_weights, + device=torch.device("cpu"), + ) + == token_count + ) + + +@pytest.mark.L0 +def test_training_bundle_fields_match_public_contracts(): + dummy = object() + assert [field.name for field in fields(MoeEpForwardWeights)] == ["fc1", "fc2"] + assert [field.name for field in fields(MoeEpBackwardWeights)] == [ + "w2_transpose", + "w1_transpose", + ] + assert MoeEpForwardWeights(dummy, dummy).fc1 is dummy + assert MoeEpBackwardWeights(dummy, dummy).w2_transpose is dummy + assert [field.name for field in fields(MoeEpTrainingForwardOutputs)] == [ + "fc1_preact", + "output", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + ] + assert [field.name for field in fields(MoeEpTrainingBackwardOutputs)] == [ + "grad_activation", + "dprob", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + ] + + +@pytest.mark.L0 +def test_native_weight_validation_and_kernel_views_are_zero_copy(): + config = _training_config(weight_interleave_size=32) + forward = _native_forward(config) + backward = _native_backward(config) + + assert validate_native_forward_weights(config, forward) == torch.device("cpu") + assert validate_native_backward_weights(config, backward) == torch.device("cpu") + + forward_kernel = forward_native_to_kernel(forward) + backward_kernel = backward_native_to_kernel(backward) + assert forward_kernel.fc1_weight.data_ptr() == forward.fc1.payload.data_ptr() + assert forward_kernel.fc1_weight_sf.data_ptr() == forward.fc1.scale.data_ptr() + assert forward_kernel.fc2_weight.data_ptr() == forward.fc2.payload.data_ptr() + assert backward_kernel.fc1_weight.data_ptr() == backward.w2_transpose.payload.data_ptr() + assert backward_kernel.fc2_weight_sf.data_ptr() == backward.w1_transpose.scale.data_ptr() + + +@pytest.mark.L0 +def test_standalone_weight_packers_write_only_caller_staging(): + config = _training_config(weight_interleave_size=32) + source = _source_weights(config) + forward_out, backward_out = _allocate_training_weight_staging(source) + + native_forward = pack_forward_weights(source[0], out=forward_out) + native_backward = pack_backward_weights(source[1], out=backward_out) + + assert native_forward.fc1.payload is forward_out.fc1_payload + assert native_forward.fc2.scale is forward_out.fc2_scale + assert native_backward.w2_transpose.payload is backward_out.w2_transpose_payload + assert native_backward.w1_transpose.scale is backward_out.w1_transpose_scale + validate_native_forward_weights(config, native_forward) + validate_native_backward_weights(config, native_backward) + + +@pytest.mark.L0 +def test_weight_packing_rejects_source_staging_alias(): + config = _training_config(weight_interleave_size=32) + source = _source_weights(config) + _, backward_out = _allocate_training_weight_staging(source) + aliased_out = MoeEpBackwardWeightStaging( + w2_transpose_payload=source[1].w2_transpose.data, + w2_transpose_scale=backward_out.w2_transpose_scale, + w1_transpose_payload=backward_out.w1_transpose_payload, + w1_transpose_scale=backward_out.w1_transpose_scale, + ) + + with pytest.raises(ValueError, match="must not alias"): + pack_backward_weights(source[1], out=aliased_out) + + +@pytest.mark.L0 +def test_mxfp8_training_input_bypasses_quantization_stager(): + class RejectingStager: + def stage(self, *args, **kwargs): + raise AssertionError("MXFP8 input must bypass the quantization stager") + + token_count = 2 + hidden = 32 + top_k = 2 + value = BlockScaledTensor( + data=torch.ones((token_count, hidden), dtype=torch.float8_e4m3fn), + scale=torch.ones((token_count, hidden // 32), dtype=torch.float8_e8m0fnu), + format="mxfp8", + logical_shape=(token_count, hidden), + axis=1, + ) + topk_idx = torch.tensor([[0, 1], [1, -1]], dtype=torch.int32) + topk_weights = torch.tensor([[0.75, 0.25], [1.0, 0.0]], dtype=torch.float32) + activation_data = torch.empty((4, hidden), dtype=torch.float8_e4m3fn) + activation_sf = torch.empty((4, hidden // 32), dtype=torch.float8_e8m0fnu) + routing_idx = torch.empty((4, top_k), dtype=torch.int32) + routing_weights = torch.empty((4, top_k), dtype=torch.float32) + + _stage_input( + type("Owner", (), {"stager": RejectingStager()})(), + value, + topk_idx, + topk_weights, + activation_data, + activation_sf, + routing_idx, + routing_weights, + ) + + torch.testing.assert_close(activation_data[:token_count], value.data) + torch.testing.assert_close(activation_sf[:token_count], value.scale) + torch.testing.assert_close(routing_idx[:token_count], topk_idx) + torch.testing.assert_close(routing_weights[:token_count], topk_weights) + assert activation_data[token_count:].eq(0).all() + assert activation_sf[token_count:].view(torch.uint8).eq(0).all() + assert routing_idx[token_count:].eq(-1).all() + assert routing_weights[token_count:].eq(0).all() + + +@pytest.mark.L0 +def test_mxfp8_training_input_already_in_symmetric_staging_is_not_copied(): + token_count = 2 + capacity = 4 + hidden = 32 + top_k = 2 + activation_data = torch.ones((capacity, hidden), dtype=torch.float8_e4m3fn) + activation_sf = torch.ones((capacity, hidden // 32), dtype=torch.float8_e8m0fnu) + value = BlockScaledTensor( + data=activation_data[:token_count], + scale=activation_sf[:token_count], + format="mxfp8", + logical_shape=(token_count, hidden), + axis=1, + ) + topk_idx = torch.tensor([[0, 1], [1, -1]], dtype=torch.int32) + topk_weights = torch.tensor([[0.75, 0.25], [1.0, 0.0]], dtype=torch.float32) + routing_idx = torch.empty((capacity, top_k), dtype=torch.int32) + routing_weights = torch.empty((capacity, top_k), dtype=torch.float32) + tail_data = activation_data[token_count:].clone() + tail_scale = activation_sf[token_count:].clone() + + _stage_input( + object(), + value, + topk_idx, + topk_weights, + activation_data, + activation_sf, + routing_idx, + routing_weights, + ) + + torch.testing.assert_close(activation_data[token_count:], tail_data) + torch.testing.assert_close(activation_sf[token_count:], tail_scale) + torch.testing.assert_close(routing_idx[:token_count], topk_idx) + torch.testing.assert_close(routing_weights[:token_count], topk_weights) + + +@pytest.mark.L0 +def test_native_execution_rejects_compact_or_wrong_layout_scales(): + config = _training_config(weight_interleave_size=32) + native = _native_forward(config) + bad = MoeEpNativeForwardWeights( + fc1=MoeEpNativeWeight( + native.fc1.payload, + torch.empty( + (config.experts_per_rank, config.hidden_size // 32, 2 * config.intermediate_size), + dtype=torch.float8_e8m0fnu, + ), + native.fc1.layout_id, + ), + fc2=native.fc2, + ) + with pytest.raises(ValueError, match=r"weights\.fc1\.scale shape"): + validate_native_forward_weights(config, bad) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("phase", "missing"), + ( + ("forward", "fc1_preact"), + ("forward", "output"), + ("forward", "fc1_a"), + ("forward", "fc1_sfa"), + ("forward", "valid_route_counts"), + ("forward", "expert_offsets"), + ("backward", "grad_activation"), + ("backward", "dprob"), + ("backward", "fc1_b"), + ("backward", "fc1_sfb"), + ("backward", "fc2_a"), + ("backward", "fc2_sfa"), + ("backward", "fc2_b"), + ("backward", "fc2_sfb"), + ), +) +def test_training_output_requirements_reject_missing_fields(phase, missing): + requirement = ((1,), (1,), torch.float32, 1) + if phase == "forward": + names = ("output", "fc1_preact", "fc1_a", "fc1_sfa", "valid_route_counts", "expert_offsets") + values = {name: torch.empty(1) for name in names} + values[missing] = None + output = MoeEpTrainingForwardOutputs(**values) + validate = validate_training_forward_outputs + else: + names = ("grad_activation", "dprob", "fc1_b", "fc1_sfb", "fc2_a", "fc2_sfa", "fc2_b", "fc2_sfb") + values = {name: torch.empty(1) for name in names} + values[missing] = None + output = MoeEpTrainingBackwardOutputs(**values) + validate = validate_training_backward_outputs + requirements = {name: requirement for name in names} + with pytest.raises(TypeError, match=rf"out\.{missing} must be a torch.Tensor"): + validate(output, requirements, device=torch.device("cpu")) + + +@pytest.mark.L0 +def test_training_output_types_remain_optional_before_validation(): + with pytest.raises(TypeError, match="fc1_preact"): + MoeEpTrainingForwardOutputs() + forward = MoeEpTrainingForwardOutputs(fc1_preact=torch.empty(1)) + backward = MoeEpTrainingBackwardOutputs() + assert forward.output is None + assert backward.grad_activation is None + + +@pytest.mark.L0 +def test_training_forward_state_validation_uses_output_contract_names(): + requirement = ((1,), (1,), torch.float32, 1) + requirements = { + name: requirement + for name in ( + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + ) + } + with pytest.raises( + TypeError, + match=r"out\.fc1_a must be a torch.Tensor", + ): + validate_training_forward_state( + fc1_preact=torch.empty(1), + fc1_a=None, + fc1_sfa=torch.empty(1), + valid_route_counts=torch.empty(1), + expert_offsets=torch.empty(1), + requirements=requirements, + device=torch.device("cpu"), + ) + + +@pytest.mark.L0 +def test_training_backward_rejects_missing_output_bundle_after_prepare(): + op = MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=4, + max_recv_size_per_rank=4, + weight_interleave_size=32, + ) + lane = MoeEpExecutionLane(0, op._operator_token) + op._training_state = object() + op._training_requirements = {} + op._training_lanes = (lane,) + op._forward_backend_device = torch.device("cpu") + with pytest.raises(TypeError, match="out must be a MoeEpTrainingBackwardOutputs"): + op.training_backward( + lane, + torch.empty((0, 128), dtype=torch.bfloat16), + torch.empty((0, 2), dtype=torch.int32), + torch.empty((0, 2), dtype=torch.float32), + weights=None, + fc1_preact=torch.empty((0, 512), dtype=torch.bfloat16), + out=None, + ) + + +@pytest.mark.L0 +def test_wgrad_assembly_returns_only_caller_owned_views(): + buffers = { + name: torch.empty(1) + for name in ( + "fc1_a", + "fc1_sfa", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + "valid_route_counts", + "expert_offsets", + ) + } + backward = MoeEpTrainingBackwardOutputs( + fc1_b=buffers["fc1_b"], + fc1_sfb=buffers["fc1_sfb"], + fc2_a=buffers["fc2_a"], + fc2_sfa=buffers["fc2_sfa"], + fc2_b=buffers["fc2_b"], + fc2_sfb=buffers["fc2_sfb"], + ) + operands = assemble_training_wgrad_operands( + fc1_a=buffers["fc1_a"], + fc1_sfa=buffers["fc1_sfa"], + valid_route_counts=buffers["valid_route_counts"], + expert_offsets=buffers["expert_offsets"], + backward=backward, + ) + assert isinstance(operands, MoeEpTrainingWgradOperands) + for name in buffers: + assert getattr(operands, name).data_ptr() == buffers[name].data_ptr() + + +@pytest.mark.L0 +def test_private_training_state_has_no_bound_weights_or_wgrad_exporter(): + config = _training_config(weight_interleave_size=32) + forward, backward = _training_prepared_pair(config) + state = Mxfp8TrainingState( + config, + torch.device("cpu"), + forward, + backward, + lane_count=2, + ) + assert not hasattr(state, "weight_bindings") + assert not hasattr(state, "wgrad_exporter") + assert not hasattr(state, "slot_count") + assert state.lane_count == 2 + assert all( + "fc1_preact" not in region.name + for region in ( + *state.requirements.symmetric_regions, + *state.requirements.local_regions, + ) + ) + + requirements = state.public_requirements() + assert requirements["fc1_a"] == ( + (config.hidden_size, forward.pool_token_capacity), + (forward.pool_token_capacity, 1), + torch.float8_e4m3fn, + 128, + ) + assert requirements["fc1_b"][1] == (2 * config.intermediate_size, 1) + assert requirements["fc2_a"][1] == (1, config.intermediate_size) + assert requirements["fc2_b"][1] == (1, forward.pool_token_capacity) + assert requirements["grad_activation"][2] is torch.bfloat16 + + +@pytest.mark.L0 +def test_private_training_workspace_keeps_only_live_lane_scratch(): + config = _training_config(weight_interleave_size=32) + forward, backward = _training_prepared_pair(config) + state = Mxfp8TrainingState( + config, + torch.device("cpu"), + forward, + backward, + lane_count=2, + ) + flat = WorkspaceViews( + token_count=0, + symmetric={region.name: torch.empty(region.nbytes, dtype=torch.uint8) for region in state.requirements.symmetric_regions}, + local={region.name: torch.empty(region.nbytes, dtype=torch.uint8) for region in state.requirements.local_regions}, + peer_mapping=object(), + ) + + first = state._lane_scratch_views(flat, 0) + second = state._lane_scratch_views(flat, 1) + forward_workspace = state._phase_workspace( + flat, + forward.workspace_requirements, + lane=0, + phase="forward", + ) + backward_workspace = state._phase_workspace( + flat, + backward.workspace_requirements, + lane=0, + phase="backward", + ) + assert "col_quant_data" not in forward_workspace.local + assert "col_quant_sf" not in forward_workspace.local + assert "kernel_local_workspace" in forward_workspace.local + assert "backward_aux_data" in backward_workspace.local + assert "backward_aux_scale" in backward_workspace.local + for field in fields(first): + first_value = getattr(first, field.name) + second_value = getattr(second, field.name) + if isinstance(first_value, torch.Tensor): + assert first_value.data_ptr() != second_value.data_ptr() + names = {region.name for region in (*state.requirements.symmetric_regions, *state.requirements.local_regions)} + removed = ( + "valid_route_counts", + "expert_offsets", + "fc1_recompute", + "fc1_recompute_sf", + "fc1_col_output", + "fc1_col_output_sf", + "grad_y2", + "grad_y2_sf", + "col_quant_data", + "col_quant_sf", + ) + assert not any(any(name.endswith(removed_name) for removed_name in removed) for name in names) + for lane in range(2): + assert f"lane.{lane}.fallback.local.routing_topk_idx" in names + assert f"lane.{lane}.fallback.symmetric.routing_topk_weights" in names + assert f"lane.{lane}.backward.local.backward_aux_data" in names + assert f"lane.{lane}.backward.local.backward_aux_scale" in names + assert f"lane.{lane}.forward.symmetric.output_data" in names + assert f"lane.{lane}.backward.symmetric.output_data" in names + assert f"lane.{lane}.backward.symmetric.backward_dprob" in names + + +@pytest.mark.L0 +def test_training_views_require_col_quant_snapshot(): + config = _training_config(weight_interleave_size=32) + forward, backward = _training_prepared_pair(config) + forward.col_quant_sizes_offset = None + state = Mxfp8TrainingState( + config, + torch.device("cpu"), + forward, + backward, + lane_count=1, + ) + with pytest.raises(RuntimeError, match="persistent col-quant expert-size snapshot"): + state.views(lane=0, token_count=0) + + +@pytest.mark.L0 +def test_training_abi_fingerprint_covers_lanes_and_native_layouts(): + config = _training_config( + ep_size=2, + ep_global_ranks=(0, 1), + weight_interleave_size=32, + ) + forward = _training_abi_prepared("forward") + backward = _training_abi_prepared("backward") + requirements = WorkspaceRequirements( + max_tokens_per_rank=4, + symmetric_regions=(BufferRegion("symmetric", 256),), + local_regions=(BufferRegion("local", 128),), + ) + first = _build_training_abi_facts( + config, + forward, + backward, + requirements, + lane_count=1, + source_tree_digest="source", + ) + repeated = _build_training_abi_facts( + config, + forward, + backward, + requirements, + lane_count=1, + source_tree_digest="source", + ) + changed_lanes = _build_training_abi_facts( + config, + forward, + backward, + requirements, + lane_count=2, + source_tree_digest="source", + ) + + assert first["schema_version"] == 2 + assert first["native_weight_layouts"] == [layout.value for layout in MoeEpNativeWeightLayout] + assert canonical_json_sha256(first) == canonical_json_sha256(repeated) + assert canonical_json_sha256(first) != canonical_json_sha256(changed_lanes) + + +@pytest.mark.L0 +def test_training_workspace_harmonizes_symmetric_regions(monkeypatch): + requirements = WorkspaceRequirements( + max_tokens_per_rank=1, + symmetric_regions=( + BufferRegion("first", 1), + BufferRegion("second", 257), + ), + local_regions=(BufferRegion("local", 1),), + ) + runtime = type("Runtime", (), {"world_size": 2, "group": object()})() + + def all_reduce(tensor, *, op, group): + assert group is runtime.group + if tensor.numel() == 2 and op == torch.distributed.ReduceOp.MAX: + tensor.copy_(torch.tensor([257, 257], dtype=torch.int64)) + + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + harmonized = _harmonize_symmetric_regions( + requirements, + runtime, + torch.device("cpu"), + ) + + assert tuple(region.nbytes for region in harmonized.symmetric_regions) == ( + 257, + 257, + ) + assert harmonized.local_regions == requirements.local_regions + + +@pytest.mark.L0 +def test_training_contract_rejects_cross_bundle_aliases(): + storage = torch.empty(16) + with pytest.raises(ValueError, match="out must not alias saved"): + validate_training_non_aliasing( + { + "saved": storage[:8], + "out": storage[4:12], + } + ) + + +@pytest.mark.L0 +def test_training_methods_require_prepare_and_do_not_expose_cleanup(): + op = MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=4, + max_recv_size_per_rank=4, + weight_interleave_size=32, + ) + assert hasattr(op, "prepare_training") + assert hasattr(op, "training_forward") + assert hasattr(op, "training_backward") + assert not hasattr(op, "prepare_training_resources") + assert not hasattr(op, "refresh_weights") + assert not hasattr(op, "finalize_overflow") + with pytest.raises(RuntimeError, match="prepare_training"): + op.training_forward( + object(), + torch.empty((0, 128), dtype=torch.bfloat16), + torch.empty((0, 2), dtype=torch.int32), + torch.empty((0, 2), dtype=torch.float32), + weights=_native_forward(_training_config(weight_interleave_size=32)), + out=MoeEpTrainingForwardOutputs( + fc1_preact=torch.empty((0, 512), dtype=torch.bfloat16), + ), + ) + + conventional = MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=4, + max_recv_size_per_rank=4, + ) + with pytest.raises(ValueError, match="weight_interleave_size=32"): + conventional.prepare_training() + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_stateless_training_padded_capacity_eager_and_cuda_graph_match_reference(): + device = _sm107_device() + base_args = make_forward_inputs(device) + args = ( + base_args[0].dequantize(torch.bfloat16), + base_args[1], + base_args[2], + base_args[3], + base_args[4].float().contiguous(), + ) + grad_output = _grad_output(device, args[0].shape[0], seed=20260902) + expected = _fixed_training_reference( + args, + grad_output, + combine_format="bf16", + gate_up_clamp=None, + ) + source_weights = _fixed_training_weights(args) + capacity = args[0].shape[0] + 1 + assert capacity % 128 != 0 + + with MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=capacity, + max_recv_size_per_rank=2 * 128, + drop_on_overflow=True, + combine_format="bf16", + weight_interleave_size=32, + ) as op: + requirements = op.prepare_training(lane_count=1, device=device) + forward_staging, backward_staging = _allocate_training_weight_staging(source_weights) + native_forward = op.pack_forward_weights( + source_weights[0], + out=forward_staging, + ) + native_backward = op.pack_backward_weights( + source_weights[1], + out=backward_staging, + ) + lane = op.training_lanes[0] + symmetric = op.training_symmetric_buffers(lane) + assert symmetric["forward_input_scale"].shape[0] == 128 + assert symmetric["backward_input_scale"].shape[0] == 128 + forward_out, backward_out = _allocate_stateless_training_outputs( + requirements, + device, + symmetric, + ) + + def run(): + y = op.training_forward( + lane, + args[0], + args[3], + args[4], + weights=native_forward, + out=forward_out, + ) + dx, dprob, operands = op.training_backward( + lane, + grad_output, + args[3], + args[4], + weights=native_backward, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, + ) + assert operands is not None + return y, dx, dprob, operands + + def assert_matches(actual): + y, dx, dprob, operands = actual + _assert_matches_reference(y, expected[0]) + _assert_backward_matches( + (dx, dprob), + (expected[1], expected[2]), + args[3], + ) + _assert_wgrads_match_reference( + operands, + expected[3], + weight_interleave_size=32, + ) + + eager = run() + grouped_wgrads = _dense_wgrads_from_grouped_kernel(eager[3]) + torch.cuda.synchronize(device) + assert_matches(eager) + expected_fc1_wgrad, expected_fc2_wgrad = expected[3].dense_wgrads() + _assert_grouped_wgrads_match_reference( + grouped_wgrads, + ( + _interleave_fc1_wgrad(expected_fc1_wgrad), + expected_fc2_wgrad, + ), + reference_name="the independent PyTorch MXFP8 reference", + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run() + pointers = tuple(tensor.data_ptr() for bundle in (forward_out, backward_out) for tensor in vars(bundle).values() if tensor is not None) + for _ in range(2): + graph.replay() + torch.cuda.synchronize(device) + assert pointers == tuple(tensor.data_ptr() for bundle in (forward_out, backward_out) for tensor in vars(bundle).values() if tensor is not None) + assert_matches(captured) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_native_io_mxfp8_cuda_graph_replay(): + """Exercise native MXFP8 I/O and weight contracts with graph replay.""" + + device = _sm107_device() + base_args = make_forward_inputs(device) + activation = base_args[0] + topk_idx = base_args[3] + topk_weights = base_args[4].float().contiguous() + args = ( + activation, + base_args[1], + base_args[2], + topk_idx, + topk_weights, + ) + grad_output_plain = _grad_output( + device, + activation.shape[0], + seed=20260903, + ) + grad_output = quantize_mxfp8(grad_output_plain, axis=1) + expected = _fixed_training_reference( + args, + grad_output.dequantize(torch.float32), + combine_format="bf16", + gate_up_clamp=None, + ) + + op = MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=activation.shape[0], + max_recv_size_per_rank=2 * 128, + drop_on_overflow=True, + output_format="bf16", + combine_format="bf16", + weight_interleave_size=32, + ) + try: + requirements = op.prepare_training(lane_count=1, device=device) + lane = op.training_lanes[0] + symmetric = op.training_symmetric_buffers(lane) + forward_out, backward_out = _allocate_stateless_training_outputs( + requirements, + device, + symmetric, + ) + + def stage_symmetric(source, prefix): + token_count = source.logical_shape[0] + data = symmetric[prefix][:token_count] + scale = symmetric[f"{prefix}_scale"][ + :token_count, + : source.scale.shape[1], + ] + data.copy_(source.data) + scale.copy_(source.scale) + return BlockScaledTensor( + data=data, + scale=scale, + format="mxfp8", + logical_shape=source.logical_shape, + axis=1, + ) + + activation = stage_symmetric(activation, "forward_input") + grad_output = stage_symmetric(grad_output, "backward_input") + + # The production call below receives only native packs. The existing + # fallback packer is used once here as a test oracle to create known-good + # native contents, then copied into independent caller-owned tensors. + source_weights = _fixed_training_weights(args) + forward_staging, backward_staging = _allocate_training_weight_staging(source_weights) + packed_forward = op.pack_forward_weights( + source_weights[0], + out=forward_staging, + ) + packed_backward = op.pack_backward_weights( + source_weights[1], + out=backward_staging, + ) + + def clone_native_tensor(tensor): + return torch.empty_strided( + tensor.shape, + tensor.stride(), + dtype=tensor.dtype, + device=tensor.device, + ).copy_(tensor) + + native_forward = MoeEpNativeForwardWeights( + fc1=MoeEpNativeWeight( + clone_native_tensor(packed_forward.fc1.payload), + clone_native_tensor(packed_forward.fc1.scale), + MoeEpNativeWeightLayout.FORWARD_FC1_GATE_UP_INTERLEAVED_32_V1, + ), + fc2=MoeEpNativeWeight( + clone_native_tensor(packed_forward.fc2.payload), + clone_native_tensor(packed_forward.fc2.scale), + MoeEpNativeWeightLayout.FORWARD_FC2_K_MAJOR_V1, + ), + ) + native_backward = MoeEpNativeBackwardWeights( + w2_transpose=MoeEpNativeWeight( + clone_native_tensor(packed_backward.w2_transpose.payload), + clone_native_tensor(packed_backward.w2_transpose.scale), + MoeEpNativeWeightLayout.BACKWARD_W2_TRANSPOSE_V1, + ), + w1_transpose=MoeEpNativeWeight( + clone_native_tensor(packed_backward.w1_transpose.payload), + clone_native_tensor(packed_backward.w1_transpose.scale), + MoeEpNativeWeightLayout.BACKWARD_W1_TRANSPOSE_GATE_UP_INTERLEAVED_32_V1, + ), + ) + assert native_forward.fc1.payload.data_ptr() != packed_forward.fc1.payload.data_ptr() + assert native_backward.w1_transpose.scale.data_ptr() != packed_backward.w1_transpose.scale.data_ptr() + + def run(): + output = op.training_forward( + lane, + activation, + topk_idx, + topk_weights, + weights=native_forward, + out=forward_out, + ) + grad_activation, dprob, operands = op.training_backward( + lane, + grad_output, + topk_idx, + topk_weights, + weights=native_backward, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, + ) + return output, grad_activation, dprob, operands + + def assert_matches(result): + output, grad_activation, dprob, operands = result + assert output.data_ptr() == forward_out.output.data_ptr() + assert grad_activation.data_ptr() == backward_out.grad_activation.data_ptr() + assert dprob.data_ptr() == backward_out.dprob.data_ptr() + _assert_matches_reference(output, expected[0]) + _assert_backward_matches( + (grad_activation, dprob), + (expected[1], expected[2]), + topk_idx, + ) + _assert_wgrads_match_reference( + operands, + expected[3], + weight_interleave_size=32, + ) + + eager = run() + grouped_wgrads = _dense_wgrads_from_grouped_kernel(eager[3]) + torch.cuda.synchronize(device) + assert_matches(eager) + expected_fc1_wgrad, expected_fc2_wgrad = expected[3].dense_wgrads() + _assert_grouped_wgrads_match_reference( + grouped_wgrads, + ( + _interleave_fc1_wgrad(expected_fc1_wgrad), + expected_fc2_wgrad, + ), + reference_name="the independent PyTorch MXFP8 reference", + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run() + output_pointers = ( + forward_out.output.data_ptr(), + backward_out.grad_activation.data_ptr(), + backward_out.dprob.data_ptr(), + ) + native_weight_pointers = ( + native_forward.fc1.payload.data_ptr(), + native_forward.fc1.scale.data_ptr(), + native_forward.fc2.payload.data_ptr(), + native_forward.fc2.scale.data_ptr(), + native_backward.w2_transpose.payload.data_ptr(), + native_backward.w2_transpose.scale.data_ptr(), + native_backward.w1_transpose.payload.data_ptr(), + native_backward.w1_transpose.scale.data_ptr(), + ) + for _ in range(2): + graph.replay() + torch.cuda.synchronize(device) + assert output_pointers == ( + forward_out.output.data_ptr(), + backward_out.grad_activation.data_ptr(), + backward_out.dprob.data_ptr(), + ) + assert native_weight_pointers == ( + native_forward.fc1.payload.data_ptr(), + native_forward.fc1.scale.data_ptr(), + native_forward.fc2.payload.data_ptr(), + native_forward.fc2.scale.data_ptr(), + native_backward.w2_transpose.payload.data_ptr(), + native_backward.w2_transpose.scale.data_ptr(), + native_backward.w1_transpose.payload.data_ptr(), + native_backward.w1_transpose.scale.data_ptr(), + ) + assert_matches(captured) + finally: + op.close() diff --git a/test/python/moe_ep/test_moe_ep_cutedsl.py b/test/python/moe_ep/test_moe_ep_cutedsl.py new file mode 100644 index 000000000..48800758a --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_cutedsl.py @@ -0,0 +1,119 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CUTLASS DSL version-gate tests for Rubin MegaMoE.""" + +import pytest +import torch + + +@pytest.mark.L0 +def test_rubin_cutedsl_gate_rejects_public_wheels_below_4_8( + monkeypatch, +): + from cudnn.moe_ep._megamoe_backend.mxfp8 import _cutedsl + + monkeypatch.setattr(_cutedsl, "_public_cutedsl_version", lambda: "4.7.0") + + with pytest.raises( + RuntimeError, + match=r"nvidia-cutlass-dsl>=4\.8\.0", + ): + _cutedsl.require_rubin_cutedsl() + + +@pytest.mark.L0 +@pytest.mark.parametrize("version", ["4.8.0", "4.8.0rc1", "4.9.0"]) +def test_rubin_cutedsl_gate_accepts_4_8_or_newer(monkeypatch, version): + from cudnn.moe_ep._megamoe_backend.mxfp8 import _cutedsl + + monkeypatch.setattr(_cutedsl, "_public_cutedsl_version", lambda: version) + + _cutedsl.require_rubin_cutedsl() + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "module_name, function_name", + [ + ( + "cudnn.moe_ep._megamoe_backend.mxfp8._compile", + "prepare_kernel", + ), + ( + "cudnn.moe_ep._megamoe_backend.mxfp8._backward_compile", + "prepare_backward_kernel", + ), + ], +) +def test_rubin_prepare_gates_before_cuda_initialization( + monkeypatch, + module_name, + function_name, +): + module = __import__(module_name, fromlist=[function_name]) + from cudnn.moe_ep._megamoe_backend.mxfp8 import _compile_common + + class GateReached(RuntimeError): + pass + + def reject(): + raise GateReached + + monkeypatch.setattr(_compile_common, "require_rubin_cutedsl", reject) + + with pytest.raises(GateReached): + getattr(module, function_name)(None, None, None) + + +@pytest.mark.L0 +@pytest.mark.parametrize("context", ["forward", "backward"]) +def test_rubin_environment_errors_include_compile_context( + monkeypatch, + context, +): + from cudnn.moe_ep._megamoe_backend.mxfp8 import _compile_common + + monkeypatch.setattr(_compile_common, "require_rubin_cutedsl", lambda: None) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + monkeypatch.setattr( + torch.cuda, + "get_device_capability", + lambda device: (10, 0), + ) + + with pytest.raises( + RuntimeError, + match=rf"Rubin MXFP8 {context} preparation", + ): + _compile_common._prepare_rubin_environment( + torch.device("cuda", 0), + 2, + context=context, + ) + + +@pytest.mark.L0 +def test_rubin_environment_rejects_incompatible_arch_override( + monkeypatch, +): + from cudnn.moe_ep._megamoe_backend.mxfp8 import _compile_common + + monkeypatch.setattr(_compile_common, "require_rubin_cutedsl", lambda: None) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + monkeypatch.setattr( + torch.cuda, + "get_device_capability", + lambda device: (10, 7), + ) + monkeypatch.setenv("CUTE_DSL_ARCH", "sm_100a") + + with pytest.raises( + RuntimeError, + match=r"CUTE_DSL_ARCH.*forward.*sm_100a", + ): + _compile_common._prepare_rubin_environment( + torch.device("cuda", 0), + 2, + context="forward", + ) diff --git a/test/python/moe_ep/test_moe_ep_forward.py b/test/python/moe_ep/test_moe_ep_forward.py new file mode 100644 index 000000000..1e0b20d52 --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_forward.py @@ -0,0 +1,1698 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Core MoE EP forward contract, parity, runtime, and distributed tests.""" + +from __future__ import annotations + +import os +import sys +from dataclasses import replace +from types import ModuleType, SimpleNamespace + +import numpy as np +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from moe_ep.moe_ep_distributed_workers import ( + _distributed_autotune_worker, + _distributed_output_worker, + _distributed_subgroup_output_worker, +) +from moe_ep.moe_ep_test_support import ( + _assert_matches_reference, + _forward_config, + _make_forward_case, + _naive_reference, + _output_as_float, + _reference_forward, + _replay_cuda_graph, + _require_distributed_sm107, + _sm107_device, + _stress_backend_reuse, + make_forward_inputs, + quantize_mxfp8, +) +from moe_ep.moe_ep_reference import ( + MoeEpReference, + MoeFormat, + forward_combine_round_trip, + quantize_blockwise, +) + + +def _public_nvfp4(data, scale, logical_shape): + from cudnn import BlockScaledTensor + + return BlockScaledTensor( + data=data, + scale=scale, + format="nvfp4", + logical_shape=logical_shape, + axis=1, + ) + + +def _request(activation, fc1, fc2): + return SimpleNamespace( + activation=activation, + fc1_weight=fc1, + fc2_weight=fc2, + ) + + +# Public API, capability, layout, and workspace contracts. + + +@pytest.mark.L0 +def test_moe_ep_finalizer_warns_without_retaining_failed_backend(): + import cudnn.moe_ep.api as api_module + from cudnn import MoeEp + + class Backend: + close_calls = 0 + + def close(self): + self.close_calls += 1 + raise RuntimeError("cleanup failed") + + operator = MoeEp(**_forward_config()) + backend = Backend() + operator._forward_backend = backend + + with pytest.warns(ResourceWarning, match="cleanup failed"): + operator.__del__() + + assert backend.close_calls == 1 + assert not hasattr(api_module, "_FAILED_FINALIZER_BACKENDS") + operator._forward_backend = None + operator._closed = True + + +@pytest.mark.L0 +def test_moe_ep_tuning_public_contract_mapping_and_cache_key(): + from cudnn import MoeEp, MoeEpTuningConfig + from cudnn.moe_ep import ( + MoeEpTuningConfig as PackageMoeEpTuningConfig, + ) + from cudnn.moe_ep._megamoe_backend.mxfp8._config import ( + Mxfp8KernelConfig, + ) + + assert PackageMoeEpTuningConfig is MoeEpTuningConfig + forward_tuning = MoeEpTuningConfig( + token_back_mode="standalone_warps", + epi_flag_batch=(4, 2), + token_in_flag_batch=4, + group_hint=768, + ) + backward_tuning = MoeEpTuningConfig( + token_back_mode="epi_warps", + epi_flag_batch=(2, 2), + token_in_flag_batch=8, + group_hint=512, + ) + with MoeEp( + **_forward_config(), + forward_tuning=forward_tuning, + backward_tuning=backward_tuning, + ) as op: + assert op.tuning is forward_tuning + assert op.forward_tuning is forward_tuning + assert op.backward_tuning is backward_tuning + forward_kernel_config = Mxfp8KernelConfig.from_operator_config( + op._forward_config, + tuning=op.forward_tuning, + ) + backward_kernel_config = Mxfp8KernelConfig.from_operator_config( + op._forward_config, + tuning=op.backward_tuning, + ) + + assert forward_kernel_config.tuning_signature(123) == ( + "standalone_warps", + (4, 2), + 4, + 768, + False, + ) + assert backward_kernel_config.tuning_signature(123) == ( + "epi_warps", + (2, 2), + 8, + 512, + False, + ) + + with MoeEp(**_forward_config()) as default_op: + default_config = Mxfp8KernelConfig.from_operator_config( + default_op._forward_config + ) + key_args = ( + torch.device("cuda", 0), + (10, 7), + 123, + (), + ) + assert forward_kernel_config.compile_key( + *key_args + ) != default_config.compile_key(*key_args) + assert backward_kernel_config.compile_key( + *key_args + ) != forward_kernel_config.compile_key(*key_args) + + +@pytest.mark.L0 +def test_internal_column_requant_config_is_disabled_by_default_and_cache_distinct(): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend.mxfp8._config import ( + Mxfp8KernelConfig, + ) + + with MoeEp(**_forward_config()) as op: + default_forward = op._forward_config + default_config = Mxfp8KernelConfig.from_operator_config(default_forward) + enabled_config = replace( + default_config, + enable_col_quant=True, + col_quant_num_ctas=512, + ) + + assert default_config.enable_col_quant is False + raw_route_count = ( + default_forward.ep_size + * default_forward.max_tokens_per_rank + * default_forward.top_k + ) + active_expert_count = min(default_forward.experts_per_rank, raw_route_count) + expected_padded_capacity = ( + active_expert_count + (raw_route_count - active_expert_count) // 128 + ) * 128 + assert default_config.max_recv_size_per_rank == expected_padded_capacity + assert enabled_config.enable_col_quant is True + assert enabled_config.col_quant_num_ctas == 512 + with pytest.raises(ValueError, match="max_recv_size_per_rank"): + replace(default_config, max_recv_size_per_rank=0) + with pytest.raises(ValueError, match="col_quant_num_ctas"): + replace(default_config, col_quant_num_ctas=0) + key_args = (torch.device("cuda", 0), (10, 7), 123, ()) + assert default_config.compile_key(*key_args) != enabled_config.compile_key(*key_args) + + +@pytest.mark.L0 +def test_bounded_receive_capacity_propagates_to_kernel_config(): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend.mxfp8._config import ( + Mxfp8KernelConfig, + ) + + with MoeEp( + **_forward_config(), + max_recv_size_per_rank=7, + drop_on_overflow=False, + ) as op: + config = Mxfp8KernelConfig.from_operator_config(op._forward_config) + + assert config.max_recv_size_per_rank == 7 + assert config.drop_on_overflow is False + + with pytest.raises(ValueError, match="max_recv_size_per_rank"): + MoeEp(**_forward_config(), max_recv_size_per_rank=0) + with pytest.raises(ValueError, match="drop_on_overflow"): + MoeEp(**_forward_config(), drop_on_overflow=1) + + +@pytest.mark.L0 +def test_receive_capacity_is_the_physical_pool_size(): + from cudnn.moe_ep._megamoe_backend.cutedsl_src.communication.nvlink_domain.token_comm_deterministic import ( + _compute_receive_capacity, + ) + + capacity = _compute_receive_capacity( + world_size=4, + max_tokens_per_rank=64, + topk=2, + max_recv_size_per_rank=256, + ) + + assert capacity.raw_route_count == 512 + assert capacity.padded_route_count == 256 + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("public_format", "wire_format"), + [ + ("bf16", "bf16"), + ("mxfp8", "32e4m3xe8m0"), + ], +) +def test_combine_format_maps_to_contract_wire(public_format, wire_format): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend.mxfp8._config import ( + Mxfp8KernelConfig, + ) + + with MoeEp(**_forward_config(combine_format=public_format)) as op: + kernel_config = Mxfp8KernelConfig.from_operator_config(op._forward_config) + + assert kernel_config.combine_format == wire_format + + +@pytest.mark.L0 +def test_distributed_topk_can_exceed_local_expert_count(): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend._capability import validate_config + + with MoeEp( + **_forward_config( + num_experts=4, + top_k=3, + ) + ) as op: + distributed = replace( + op._forward_config, + experts_per_rank=2, + ep_size=2, + ep_global_ranks=(0, 1), + ) + + validate_config(distributed) + + +@pytest.mark.L0 +def test_overflow_check_prefers_device_assert(monkeypatch): + from cudnn.moe_ep._megamoe_backend.mxfp8._launch import ( + _check_overflow, + ) + + calls = [] + + def record_assert(condition, message): + calls.append((condition.clone(), message)) + + monkeypatch.setattr(torch, "_assert_async", record_assert) + flag = torch.zeros(1, dtype=torch.int32) + _check_overflow(flag) + + assert len(calls) == 1 + assert bool(calls[0][0]) + assert "route-pool overflow" in calls[0][1] + + +@pytest.mark.L0 +def test_overflow_check_fallback_rejects_nonzero_flag(monkeypatch): + from cudnn.moe_ep._megamoe_backend.mxfp8._launch import ( + _check_overflow, + ) + + monkeypatch.setattr(torch, "_assert_async", None) + monkeypatch.setattr( + torch.cuda, + "is_current_stream_capturing", + lambda: False, + ) + with pytest.raises(RuntimeError, match="receive route-pool overflow"): + _check_overflow(torch.ones(1, dtype=torch.int32)) + + +@pytest.mark.L0 +def test_in_kernel_topk_reduce_omits_standalone_combine_workspace(): + from cudnn.moe_ep._megamoe_backend.mxfp8._compile import ( + _pre_reduced_workspace_metadata, + ) + + class NoStandaloneWorkspace: + def region(self, _name): + raise AssertionError("in-kernel reduction must not query region") + + def offset(self, _name): + raise AssertionError("in-kernel reduction must not query offset") + + def nbytes(self, _name): + raise AssertionError("in-kernel reduction must not query size") + + config = SimpleNamespace( + fc2_in_kernel_topk_reduce=True, + top_k=6, + hidden=7168, + max_tokens_per_rank=4096, + combine_format="bf16", + ) + + assert _pre_reduced_workspace_metadata( + NoStandaloneWorkspace(), + config, + shared_bytes=0, + ) == (None, 0) + + bytes_per_token = config.top_k * config.hidden * 2 + total_bytes = config.max_tokens_per_rank * bytes_per_token + + class StandaloneWorkspace: + def region(self, _name): + return SimpleNamespace(buffer_space="shared") + + def offset(self, _name): + return 256 + + def nbytes(self, _name): + return total_bytes + + config.fc2_in_kernel_topk_reduce = False + assert _pre_reduced_workspace_metadata( + StandaloneWorkspace(), + config, + shared_bytes=256 + total_bytes, + ) == (256, bytes_per_token) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("combine_format", "bits_per_element"), + [ + ("bf16", 16), + ("32e4m3xe8m0", 8), + ], +) +def test_standalone_combine_workspace_tracks_wire_width( + combine_format, + bits_per_element, +): + from cudnn.moe_ep._megamoe_backend.mxfp8._compile import ( + _pre_reduced_workspace_metadata, + ) + + config = SimpleNamespace( + fc2_in_kernel_topk_reduce=False, + top_k=2, + hidden=128, + max_tokens_per_rank=5, + combine_format=combine_format, + ) + bytes_per_token = config.top_k * config.hidden * bits_per_element // 8 + total_bytes = config.max_tokens_per_rank * bytes_per_token + + class StandaloneWorkspace: + def region(self, _name): + return SimpleNamespace(buffer_space="shared") + + def offset(self, _name): + return 128 + + def nbytes(self, _name): + return total_bytes + + assert _pre_reduced_workspace_metadata( + StandaloneWorkspace(), + config, + shared_bytes=128 + total_bytes, + ) == (128, bytes_per_token) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("combine_format", "expected"), + [ + ("bf16", (None, 0)), + ("32e4m3xe8m0", (128, 64)), + ], +) +def test_standalone_combine_scale_workspace_metadata( + combine_format, + expected, +): + from cudnn.moe_ep._megamoe_backend.mxfp8._compile import ( + _pre_reduced_sf_workspace_metadata, + ) + + config = SimpleNamespace( + fc2_in_kernel_topk_reduce=False, + max_tokens_per_rank=5, + combine_format=combine_format, + ) + + class StandaloneWorkspace: + def region(self, _name): + return SimpleNamespace(buffer_space="shared") + + def offset(self, _name): + return 128 + + def nbytes(self, _name): + return config.max_tokens_per_rank * 64 + + workspace = StandaloneWorkspace() + if combine_format == "bf16": + + class NoScaleWorkspace: + def region(self, _name): + raise AssertionError("BF16 combine must not query scale region") + + workspace = NoScaleWorkspace() + + assert ( + _pre_reduced_sf_workspace_metadata( + workspace, + config, + shared_bytes=128 + config.max_tokens_per_rank * 64, + ) + == expected + ) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "kwargs", + [ + {"token_padding_size": True}, + {"token_padding_size": 0}, + {"token_padding_size": -1}, + {"token_padding_size": 64.0}, + {"sf_padding_size": True}, + {"sf_padding_size": 0}, + {"sf_padding_size": 64}, + {"sf_padding_size": 128.0}, + ], +) +def test_moe_ep_rejects_invalid_padding(kwargs): + from cudnn import MoeEp + + with pytest.raises(ValueError): + MoeEp(**_forward_config(), **kwargs) + + +@pytest.mark.L0 +def test_moe_ep_rejects_unsupported_weight_interleave_size(): + from cudnn import MoeEp + + with pytest.raises(ValueError, match="weight_interleave_size must be None or 32"): + MoeEp(**_forward_config(), weight_interleave_size=16) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("weight_interleave_size", "expected_layout"), + [ + (None, "gate_then_up"), + (32, "gate_up_interleaved_32"), + ], +) +def test_moe_ep_normalizes_fc1_weight_layout( + weight_interleave_size, + expected_layout, +): + from cudnn import MoeEp + + with MoeEp( + **_forward_config(), + weight_interleave_size=weight_interleave_size, + ) as op: + assert op._forward_config.fc1_weight_layout.value == expected_layout + + +@pytest.mark.L0 +def test_moe_ep_rejects_interleaved_plain_fc1_weight(): + from cudnn import MoeEp + + config = _forward_config() + activation = torch.zeros((1, config["hidden_size"])) + fc1_weight = torch.zeros((config["num_experts"], config["hidden_size"], 2 * config["intermediate_size"])) + fc2_weight = torch.zeros((config["num_experts"], config["intermediate_size"], config["hidden_size"])) + topk_idx = torch.zeros((1, config["top_k"]), dtype=torch.int32) + topk_weights = torch.ones((1, config["top_k"])) + with MoeEp(**config, weight_interleave_size=32) as op: + with pytest.raises(ValueError, match="requires an MXFP8"): + op(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + + +@pytest.mark.L0 +@pytest.mark.parametrize("name", ("tuning", "forward_tuning", "backward_tuning")) +def test_moe_ep_rejects_untyped_tuning(name): + from cudnn import MoeEp + + with pytest.raises(TypeError, match="MoeEpTuningConfig"): + MoeEp(**_forward_config(), **{name: {"group_hint": 768}}) + + +@pytest.mark.L0 +def test_moe_ep_rejects_conflicting_forward_tuning_aliases(): + from cudnn import MoeEp, MoeEpTuningConfig + + with pytest.raises(ValueError, match="aliases"): + MoeEp( + **_forward_config(), + tuning=MoeEpTuningConfig(), + forward_tuning=MoeEpTuningConfig(), + ) + + +@pytest.mark.L0 +def test_moe_ep_backward_tuning_default_is_independent(): + from cudnn import MoeEp, MoeEpTuningConfig + + forward_tuning = MoeEpTuningConfig( + token_back_mode="standalone_warps", + epi_flag_batch=(4, 2), + token_in_flag_batch=8, + group_hint=768, + ) + with MoeEp(**_forward_config(), forward_tuning=forward_tuning) as op: + assert op.forward_tuning is forward_tuning + assert op.backward_tuning == MoeEpTuningConfig() + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "kwargs", + [ + {"combine_format": "mxfp8"}, + {"output_format": "mxfp8"}, + {"apply_topk_in_fc1": False}, + ], +) +def test_moe_ep_rejects_incompatible_in_kernel_topk_reduce(kwargs): + from cudnn import MoeEp, MoeEpTuningConfig + + config = _forward_config() + config.update(kwargs) + with pytest.raises(ValueError, match="reduce_topk_in_kernel requires"): + MoeEp( + **config, + tuning=MoeEpTuningConfig(reduce_topk_in_kernel=True), + ) + + +@pytest.mark.L0 +def test_distributed_launch_rejects_mismatched_tuning_before_barrier( + monkeypatch, +): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend.mxfp8._backend import ( + Mxfp8Backend, + ) + + with MoeEp(**_forward_config()) as op: + backend = Mxfp8Backend( + op._forward_config, + torch.device("cuda", 0), + ) + backend._ep_launch_ready = False + + stream = SimpleNamespace(synchronize=lambda: None) + resources = SimpleNamespace(runtime=SimpleNamespace(group=object(), world_size=2)) + prepared = SimpleNamespace(launch_cluster_count=123) + monkeypatch.setattr( + backend, + "_ensure_prepared_kernel", + lambda: prepared, + ) + + def mismatched_all_gather(output, signature, *, group): + assert group is resources.runtime.group + output[:] = [ + signature, + ("standalone_warps", (1, 1), 1, 123), + ] + + barrier_called = False + + def unexpected_barrier(*, group): + nonlocal barrier_called + barrier_called = True + + monkeypatch.setattr(dist, "all_gather_object", mismatched_all_gather) + monkeypatch.setattr(dist, "barrier", unexpected_barrier) + + with pytest.raises(RuntimeError, match="MoeEp tuning must match"): + backend._ensure_ep_launch_ready(resources, stream) + assert not barrier_called + assert not backend._ep_launch_ready + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "kwargs", + [ + {"combine_format": "nvfp4"}, + {"output_format": "mxfp8"}, + {"apply_topk_in_fc1": False}, + ], +) +def test_training_megamoe_rejects_unsupported_config_before_backend(kwargs): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend._capability import validate_config + + with MoeEp(**_forward_config(**kwargs)) as op: + with pytest.raises(NotImplementedError, match="training MegaMoE"): + validate_config(op._forward_config) + + +@pytest.mark.L0 +def test_training_megamoe_rejects_nvfp4_operand_before_cuda_query(monkeypatch): + from cudnn.moe_ep._megamoe_backend._capability import validate_request + + operand = _public_nvfp4( + torch.zeros(2, 64, dtype=torch.uint8), + torch.ones(2, 8).to(torch.float8_e4m3fn), + (2, 128), + ) + request = _request(operand, operand, operand) + request.device = torch.device("cuda", 0) + + monkeypatch.setattr( + torch.cuda, + "get_device_capability", + lambda _device: pytest.fail("CUDA capability queried too early"), + ) + with pytest.raises(NotImplementedError, match="only MXFP8"): + validate_request(request) + + +# Single-rank and distributed forward numerical parity. + + +@pytest.mark.L0 +def test_bf16_forward_matches_reference_and_returns_fresh_outputs(): + from cudnn import MoeEp + + device = _sm107_device() + args = make_forward_inputs(device) + activation, fc1_weight, fc2_weight = args[:3] + expected = _reference_forward(args) + + assert activation.logical_shape == (5, 128) + assert fc1_weight.logical_shape == (2, 128, 512) + assert fc2_weight.logical_shape == (2, 256, 128) + + with MoeEp(**_forward_config()) as op: + first = op(*args) + snapshot = first.clone() + second = op(*args) + args[3].fill_(-1) + dropped = op(*args) + torch.cuda.synchronize(device) + + assert isinstance(first, torch.Tensor) + assert isinstance(second, torch.Tensor) + assert first.shape == second.shape == (5, 128) + assert first.dtype == second.dtype == torch.bfloat16 + assert first.device == second.device == device + assert first is not second + assert first.data_ptr() != second.data_ptr() + torch.testing.assert_close(first, snapshot, rtol=0, atol=0) + torch.testing.assert_close(first, second, rtol=0, atol=0) + _assert_matches_reference(first, expected) + assert dropped.eq(0).all() + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_mxfp8_combine_matches_direct_fp32_training_reference(): + from cudnn import MoeEp + + device = _sm107_device() + args = make_forward_inputs(device) + config = _forward_config( + combine_format="mxfp8", + ) + expected = _reference_forward(args, **config) + + with MoeEp(**config) as op: + actual = op(*args) + torch.cuda.synchronize(device) + + _assert_matches_reference(actual, expected) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +@pytest.mark.parametrize( + ("plain_mask", "plain_dtype"), + [ + pytest.param( + (True, False, False), + torch.bfloat16, + id="activation-bf16", + ), + pytest.param( + (False, True, False), + torch.float16, + id="fc1-fp16", + ), + pytest.param( + (False, False, True), + torch.float32, + id="fc2-fp32", + ), + ], +) +def test_plain_and_mixed_inputs_match_staged_reference( + plain_mask, + plain_dtype, +): + from cudnn import MoeEp + + device = _sm107_device() + args = list(make_forward_inputs(device)) + for index, make_plain in enumerate(plain_mask): + if make_plain: + args[index] = args[index].dequantize(dtype=plain_dtype) + args = tuple(args) + expected = _reference_forward(args) + + with MoeEp(**_forward_config()) as op: + actual = op(*args) + torch.cuda.synchronize(device) + + _assert_matches_reference(actual, expected) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_one_operator_switches_mxfp8_and_plain_weight_families(): + from cudnn import MoeEp + + device = _sm107_device() + quantized_args = make_forward_inputs(device) + plain_args = ( + quantized_args[0].dequantize(dtype=torch.bfloat16), + quantized_args[1].dequantize(dtype=torch.bfloat16), + quantized_args[2].dequantize(dtype=torch.bfloat16), + *quantized_args[3:], + ) + expected_quantized = _reference_forward(quantized_args) + expected_plain = _reference_forward(plain_args) + + with MoeEp(**_forward_config()) as op: + quantized = op(*quantized_args) + backend = op._forward_backend + refresh_before = backend._adapter.weight_refresh_count + plain = op(*plain_args) + refresh_after = backend._adapter.weight_refresh_count + torch.cuda.synchronize(device) + + assert op._forward_backend is None + assert refresh_after == refresh_before + 1 + _assert_matches_reference(quantized, expected_quantized) + _assert_matches_reference(plain, expected_plain) + + +@pytest.mark.L0 +def test_nondefault_moe_ep_tuning_matches_reference_and_reuses_plan(): + from cudnn import MoeEp, MoeEpTuningConfig + + device = _sm107_device() + args = make_forward_inputs(device) + expected = _reference_forward(args) + tuning = MoeEpTuningConfig( + token_back_mode="standalone_warps", + epi_flag_batch=(4, 2), + token_in_flag_batch=4, + group_hint=64, + ) + + with MoeEp(**_forward_config(), tuning=tuning) as op: + first = op(*args) + backend = op._forward_backend + assert backend is not None + compiled = backend._compiled + workspace = backend._plan._workspace + second = op(*args) + torch.cuda.synchronize(device) + + assert backend._compiled is compiled + assert backend._plan._workspace is workspace + assert backend.kernel_config.tuning_signature(backend._prepared_kernel.launch_cluster_count) == ("standalone_warps", (4, 2), 4, 64, False) + + _assert_matches_reference(first, expected) + _assert_matches_reference(second, expected) + + +@pytest.mark.L0 +def test_gate_up_clamp_matches_moe_ep_reference(): + from cudnn import MoeEp + + device = _sm107_device() + args = make_forward_inputs(device) + clamp = 0.5 + expected = _reference_forward(args, gate_up_clamp=clamp) + unclamped = _reference_forward(args) + + with MoeEp(**_forward_config(gate_up_clamp=clamp)) as op: + actual = op(*args) + torch.cuda.synchronize(device) + + assert not torch.equal(expected, unclamped) + _assert_matches_reference(actual, expected) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +@pytest.mark.parametrize( + ( + "experts", + "tokens", + "hidden", + "intermediate", + "top_k", + "index_dtype", + "weight_dtype", + ), + [ + pytest.param( + 2, + 3, + 128, + 256, + 1, + torch.int32, + torch.bfloat16, + id="topk1-h128-i256-int32-bf16", + ), + pytest.param( + 2, + 5, + 128, + 256, + 2, + torch.int64, + torch.float32, + id="topk2-h128-i256-int64-fp32", + ), + pytest.param( + 4, + 3, + 256, + 256, + 4, + torch.int32, + torch.float16, + id="topk4-h256-i256-int32-fp16", + ), + pytest.param( + 32, + 1, + 128, + 256, + 32, + torch.int64, + torch.float32, + id="topk32-boundary-int64-fp32", + ), + ], +) +def test_supported_topk_shape_and_routing_format_matrix( + experts, + tokens, + hidden, + intermediate, + top_k, + index_dtype, + weight_dtype, +): + from cudnn import MoeEp + + device = _sm107_device() + args = _make_forward_case( + device, + experts=experts, + tokens=tokens, + hidden=hidden, + intermediate=intermediate, + top_k=top_k, + index_dtype=index_dtype, + weight_dtype=weight_dtype, + ) + config = _forward_config( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, + max_tokens_per_rank=tokens, + ) + expected = _reference_forward(args, **config) + + with MoeEp(**config) as op: + actual = op(*args) + torch.cuda.synchronize(device) + + assert actual.shape == (tokens, hidden) + assert actual.dtype == torch.bfloat16 + _assert_matches_reference(actual, expected) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_single_gpu_stress_and_cuda_graph_replay(): + from cudnn import MoeEp + + device = _sm107_device() + args = make_forward_inputs(device) + original_topk_idx = args[3].clone() + original_topk_weights = args[4].clone() + config = _forward_config() + expected = _reference_forward(args, **config) + + with MoeEp(**config) as op: + op.warmup(*args) + _stress_backend_reuse( + op, + args, + original_topk_idx, + original_topk_weights, + device, + check_weight_refresh=True, + ) + + args[3].copy_(original_topk_idx) + args[4].copy_(original_topk_weights) + eager = op(*args) + torch.cuda.synchronize(device) + _assert_matches_reference(eager, expected) + _replay_cuda_graph(op, args, original_topk_idx, expected, device) + + +@pytest.mark.L0 +def test_forward_mxfp8_combine_is_direct_fp32(): + generator = torch.Generator().manual_seed(20260819) + accumulator = torch.randn(4, 128, generator=generator) * 3.25 + + forward = forward_combine_round_trip(accumulator, MoeFormat.MXFP8) + direct_fp32 = quantize_blockwise( + accumulator, + MoeFormat.MXFP8, + ).dequantize() + bf16_preround = quantize_blockwise( + accumulator.to(torch.bfloat16).float(), + MoeFormat.MXFP8, + ).dequantize() + + torch.testing.assert_close( + forward, + direct_fp32, + rtol=0, + atol=0, + ) + assert not torch.equal(forward, bf16_preround) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +@pytest.mark.parametrize( + ("world_size", "combine_format"), + [ + pytest.param(2, "bf16", id="ep2-bf16"), + pytest.param(2, "mxfp8", id="ep2-mxfp8"), + pytest.param(3, "mxfp8", id="ep3-mxfp8"), + pytest.param(4, "bf16", id="ep4-bf16"), + ], +) +def test_mxfp8_forward_multi_gpu_matches_reference( + world_size, + combine_format, + tmp_path, +): + _require_distributed_sm107(world_size) + os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") + init_file = tmp_path / f"{combine_format}_combine_ep{world_size}.init" + mp.spawn( + _distributed_output_worker, + args=(world_size, str(init_file), combine_format), + nprocs=world_size, + join=True, + ) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_mxfp8_forward_ep2_autotune_is_rank_consistent(tmp_path): + world_size = 2 + _require_distributed_sm107(world_size) + os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") + init_file = tmp_path / "autotune_ep2.init" + mp.spawn( + _distributed_autotune_worker, + args=(world_size, str(init_file)), + nprocs=world_size, + join=True, + ) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_mxfp8_forward_noncontiguous_ep_subgroups(tmp_path): + global_world_size = 4 + _require_distributed_sm107(global_world_size) + os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") + init_file = tmp_path / "two_noncontiguous_ep2.init" + mp.spawn( + _distributed_subgroup_output_worker, + args=(global_world_size, str(init_file)), + nprocs=global_world_size, + join=True, + ) + + +# Input staging and workspace layout. + + +@pytest.mark.L0 +def test_plain_tensor_staging_matches_logical_mxfp8_quantization(): + if not torch.cuda.is_available(): + pytest.skip("MXFP8 staging test requires CUDA") + from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( + _quantize_plain_mxfp8, + ) + + device = torch.device("cuda", 0) + plain = torch.randn(2, 128, 3, device=device).to(torch.bfloat16) + actual = _quantize_plain_mxfp8(plain, axis=1) + expected = quantize_mxfp8(plain, axis=1) + + torch.testing.assert_close( + actual.data.view(torch.uint8), + expected.data.view(torch.uint8), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + actual.scale.view(torch.uint8), + expected.scale.view(torch.uint8), + rtol=0, + atol=0, + ) + + +@pytest.mark.L0 +def test_intermediate_requires_full_mma_n_tile(): + from cudnn import MoeEp + + args = _make_forward_case( + torch.device("cpu"), + experts=2, + tokens=3, + hidden=128, + intermediate=128, + top_k=2, + index_dtype=torch.int32, + weight_dtype=torch.bfloat16, + ) + with MoeEp(**_forward_config(intermediate_size=128, max_tokens_per_rank=3)) as op: + with pytest.raises( + NotImplementedError, + match=r"intermediate_size .*divisible by 256", + ): + op(*args) + assert op._forward_backend is None + + +@pytest.mark.L0 +def test_inference_activation_scale_uses_unpadded_prefix(monkeypatch): + import cudnn.moe_ep._megamoe_backend.mxfp8._adapter as adapter_module + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend._workspace import ( + WorkspaceRequirements, + padded_mxfp8_scale_columns, + ) + from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( + Mxfp8InputAdapter, + Mxfp8Weights, + ) + + assert padded_mxfp8_scale_columns(128) == 16 + assert padded_mxfp8_scale_columns(512) == 16 + assert padded_mxfp8_scale_columns(640) == 32 + + with MoeEp(**_forward_config()) as op: + requirements = WorkspaceRequirements.for_mxfp8( + op._forward_config, + kernel_local_workspace_bytes=128, + kernel_shared_workspace_bytes=128, + ) + activation_scale = next( + region + for region in requirements.symmetric_regions + if region.name == "activation_scale" + ) + assert activation_scale.nbytes == 128 * 16 + + capacity = 5 + hidden = 128 + top_k = 2 + symmetric = { + "activation_data": torch.empty(capacity * hidden, dtype=torch.uint8), + "activation_scale": torch.empty(activation_scale.nbytes, dtype=torch.uint8), + "topk_weights": torch.empty(capacity * top_k * 4, dtype=torch.uint8), + "output_data": torch.empty(capacity * hidden * 2, dtype=torch.uint8), + "kernel_shared_workspace": torch.empty(128, dtype=torch.uint8), + } + local = { + "topk_idx": torch.empty(capacity * top_k * 4, dtype=torch.uint8), + "overflow_flag": torch.empty(4, dtype=torch.uint8), + "kernel_local_workspace": torch.empty(128, dtype=torch.uint8), + } + request = SimpleNamespace( + token_count=1, + activation=object(), + topk_idx=torch.zeros((1, top_k), dtype=torch.int32), + topk_weights=torch.ones((1, top_k), dtype=torch.float32), + ) + staged_activation = SimpleNamespace( + data=torch.zeros((1, hidden), dtype=torch.float8_e4m3fn), + scale=torch.zeros((1, hidden // 32), dtype=torch.float8_e8m0fnu), + ) + config = SimpleNamespace( + max_tokens_per_rank=capacity, + hidden=hidden, + top_k=top_k, + generate_c=False, + enable_col_quant=False, + fc2_in_kernel_topk_reduce=True, + combine_format="bf16", + ) + resources = SimpleNamespace( + workspace=SimpleNamespace(symmetric=symmetric, local=local), + ) + weights = Mxfp8Weights(*(torch.empty(0) for _ in range(4))) + adapter = Mxfp8InputAdapter() + monkeypatch.setattr(adapter_module, "_as_mxfp8", lambda _: staged_activation) + monkeypatch.setattr(adapter, "_prepare_weights", lambda *_: weights) + + launch = adapter.stage( + request, + resources, + config, + local_workspace_zero_bytes=0, + shared_workspace_zero_bytes=0, + pre_reduced_activation_offset=None, + pre_reduced_activation_bytes_per_token=0, + pre_reduced_activation_sf_offset=None, + pre_reduced_activation_sf_bytes_per_token=0, + col_quant_data_rows=0, + col_quant_sf_elements=0, + ) + + assert launch.activation_sf.shape == (capacity, 16) + assert launch.activation_sf.data_ptr() == symmetric["activation_scale"].data_ptr() + + +@pytest.mark.L0 +def test_column_requant_workspace_is_allocated_only_when_enabled(): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend._workspace import WorkspaceRequirements + + with MoeEp(**_forward_config()) as op: + disabled = WorkspaceRequirements.for_mxfp8( + op._forward_config, + kernel_local_workspace_bytes=128, + kernel_shared_workspace_bytes=128, + ) + enabled = WorkspaceRequirements.for_mxfp8( + op._forward_config, + kernel_local_workspace_bytes=128, + kernel_shared_workspace_bytes=128, + col_quant_data_bytes=640, + col_quant_sf_bytes=80, + ) + + disabled_names = {region.name for region in disabled.local_regions} + assert "col_quant_data" not in disabled_names + assert "col_quant_sf" not in disabled_names + enabled_sizes = {region.name: region.nbytes for region in enabled.local_regions} + assert enabled_sizes["col_quant_data"] == 640 + assert enabled_sizes["col_quant_sf"] == 80 + + with pytest.raises(ValueError, match="must be enabled together"): + WorkspaceRequirements.for_mxfp8( + op._forward_config, + kernel_local_workspace_bytes=128, + kernel_shared_workspace_bytes=128, + col_quant_data_bytes=640, + ) + + +# Reference and quantization self-checks. + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "intermediate_format", + [None, MoeFormat.MXFP8], + ids=["fp32-intermediate", "mxfp8-intermediate"], +) +def test_reference_mxfp8_inputs_bf16_combine_matches_naive( + intermediate_format, +): + torch.manual_seed(19) + experts, tokens, hidden, intermediate = 2, 3, 128, 128 + activation = torch.randn(tokens, hidden) + fc1_weight = torch.randn(experts, hidden, 2 * intermediate) / 8 + fc2_weight = torch.randn(experts, intermediate, hidden) / 8 + q_activation = quantize_blockwise(activation, MoeFormat.MXFP8, axis=1) + q_fc1 = quantize_blockwise(fc1_weight, MoeFormat.MXFP8, axis=1) + q_fc2 = quantize_blockwise(fc2_weight, MoeFormat.MXFP8, axis=1) + topk_idx = torch.tensor([[0], [1], [0]], dtype=torch.int64) + topk_weights = torch.ones(tokens, 1) + op = MoeEpReference( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=1, + combine_format="bf16", + output_format="bf16", + intermediate_format=intermediate_format, + ) + + actual = op(q_activation, q_fc1, q_fc2, topk_idx, topk_weights) + expected = _naive_reference( + q_activation.dequantize(), + q_fc1.dequantize(), + q_fc2.dequantize(), + topk_idx, + topk_weights, + apply_topk_in_fc1=True, + combine_format=MoeFormat.BF16, + intermediate_format=intermediate_format, + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.L0 +def test_reference_interleaved_fc1_matches_logical_fc1(): + from moe_ep.moe_ep_reference import BlockScaledTensor + + torch.manual_seed(29) + experts, tokens, hidden, intermediate = 2, 3, 128, 128 + activation = quantize_blockwise( + torch.randn(tokens, hidden), + MoeFormat.MXFP8, + axis=1, + ) + logical_fc1 = quantize_blockwise( + torch.randn(experts, hidden, 2 * intermediate) / 8, + MoeFormat.MXFP8, + axis=1, + ) + fc2 = quantize_blockwise( + torch.randn(experts, intermediate, hidden) / 8, + MoeFormat.MXFP8, + axis=1, + ) + + def interleave_last(tensor): + shape = tensor.shape + return tensor.view(*shape[:-1], 2, intermediate // 32, 32).transpose(-3, -2).reshape(shape) + + interleaved_fc1 = BlockScaledTensor( + data=interleave_last(logical_fc1.data), + scale=interleave_last(logical_fc1.scale), + format=logical_fc1.format, + logical_shape=logical_fc1.logical_shape, + axis=logical_fc1.axis, + ) + topk_idx = torch.tensor([[0], [1], [0]], dtype=torch.int64) + topk_weights = torch.ones(tokens, 1) + kwargs = dict( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=1, + ) + + logical = MoeEpReference(**kwargs)( + activation, + logical_fc1, + fc2, + topk_idx, + topk_weights, + ) + interleaved = MoeEpReference(**kwargs, weight_interleave_size=32)( + activation, + interleaved_fc1, + fc2, + topk_idx, + topk_weights, + ) + + torch.testing.assert_close(interleaved, logical, atol=0, rtol=0) + + +# Host-side EP topology and runtime bootstrap. + + +@pytest.mark.L0 +def test_resolve_ep_topology_preserves_group_rank_order(monkeypatch): + from cudnn.moe_ep.api import _resolve_ep_topology + + group = object() + monkeypatch.setattr(dist, "is_available", lambda: True) + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda selected=None: 2) + monkeypatch.setattr( + dist, + "get_rank", + lambda selected=None: 1, + ) + monkeypatch.setattr( + dist, + "get_global_rank", + lambda selected, group_rank: (3, 1)[group_rank], + ) + + assert _resolve_ep_topology(group) == (2, 1, (3, 1)) + + +@pytest.mark.L0 +def test_resolve_ep_topology_rejects_nonmember(monkeypatch): + from cudnn.moe_ep.api import _resolve_ep_topology + + monkeypatch.setattr(dist, "is_available", lambda: True) + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda group: 2) + monkeypatch.setattr(dist, "get_rank", lambda group: -1) + + with pytest.raises(ValueError, match="must be a member"): + _resolve_ep_topology(object()) + + +@pytest.mark.L0 +def test_resolve_runtime_world_revalidates_ordered_membership(monkeypatch): + from cudnn.moe_ep._megamoe_backend._runtime import _resolve_world + + group = object() + monkeypatch.setattr(dist, "is_available", lambda: True) + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda selected: 2) + monkeypatch.setattr(dist, "get_rank", lambda selected: 1) + monkeypatch.setattr( + dist, + "get_global_rank", + lambda selected, group_rank: (3, 1)[group_rank], + ) + config = SimpleNamespace( + ep_group=group, + ep_size=2, + ep_rank=1, + ep_global_ranks=(3, 1), + ) + + world = _resolve_world(config) + assert world.identity == (1, 2, (3, 1)) + + config.ep_global_ranks = (1, 3) + with pytest.raises(RuntimeError, match="membership does not match"): + _resolve_world(config) + + +@pytest.mark.L0 +def test_megamoe_capability_and_kernel_config_accept_ep_above_16(): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend._capability import validate_config + from cudnn.moe_ep._megamoe_backend.mxfp8._config import ( + Mxfp8KernelConfig, + ) + + with MoeEp(**_forward_config()) as op: + config = replace( + op._forward_config, + num_experts=32, + experts_per_rank=1, + ep_size=32, + ep_rank=31, + ep_group=object(), + ep_global_ranks=tuple(range(32)), + ) + + validate_config(config) + kernel_config = Mxfp8KernelConfig.from_operator_config(config) + assert kernel_config.world_size == 32 + assert kernel_config.local_rank == 31 + + +@pytest.mark.L0 +def test_ep32_peer_mapping_selects_version_compatible_payload(): + import cutlass + from cutlass._mlir import ir + from packaging.version import Version + + from cudnn.moe_ep._megamoe_backend._comm import PeerMapping + from cudnn.moe_ep._megamoe_backend.cutedsl_src.communication.nvlink_domain.symmetric_buffer import ( + SymmetricBufferDevice, + ) + + offsets = tuple(index * 4096 for index in range(32)) + mapping = PeerMapping( + base_address=0x1000, + offsets=offsets, + rank=0, + ) + host = mapping.to_sym_buffer_host() + with ir.Context(): + device_type = SymmetricBufferDevice( + None, + host.max_ranks, + ).__get_mlir_types__()[0] + device_type_text = str(device_type) + + assert host.offsets == offsets + assert int(host.max_ranks) == 32 + dsl_release = Version(Version(cutlass.__version__).base_version) + grid_constant_width_is_free = dsl_release < Version("4.0.0") or dsl_release >= Version("4.7.0") + expected_type = "!llvm.ptr" if grid_constant_width_is_free else "vector<32xi64>" + assert device_type_text == expected_type + + +@pytest.fixture +def runtime_module(): + from cudnn.moe_ep._megamoe_backend import _runtime + + with _runtime._PROCESS_RUNTIME_REGISTRY.lock: + _runtime._PROCESS_RUNTIME_REGISTRY.active = None + yield _runtime + with _runtime._PROCESS_RUNTIME_REGISTRY.lock: + _runtime._PROCESS_RUNTIME_REGISTRY.active = None + + +class _FakeRuntimeProvider: + def __init__(self, runtime_module, state=None): + self._runtime_module = runtime_module + self._state = state or runtime_module.RuntimeInitState.NOT_INITIALIZED + self._world = None + self.initialize_count = 0 + self.finalize_count = 0 + + def initialization_state(self): + return self._state + + def initialize(self, device, world): + del device + self.initialize_count += 1 + self._world = world + self._state = self._runtime_module.RuntimeInitState.INITIALIZED + + def rank(self): + return self._world.rank + + def world_size(self): + return self._world.size + + def device(self): + return torch.device("cuda", 0) + + def finalize(self): + self.finalize_count += 1 + self._state = self._runtime_module.RuntimeInitState.NOT_INITIALIZED + + +@pytest.mark.L0 +def test_runtime_manager_shares_only_identical_subgroup(runtime_module): + world = runtime_module.RuntimeWorld( + rank=1, + size=2, + group=object(), + global_ranks=(1, 3), + ) + provider = _FakeRuntimeProvider(runtime_module) + manager = runtime_module.RuntimeManager( + provider_factory=lambda: provider, + world_resolver=lambda config: world, + ) + + first = manager.acquire(object(), torch.device("cuda", 0)) + second = manager.acquire(object(), torch.device("cuda", 0)) + assert manager.ref_count == 2 + assert second.global_ranks == (1, 3) + + second.close() + assert manager.ref_count == 1 + first.close() + assert manager.ref_count == 0 + assert provider.finalize_count == 1 + + +@pytest.mark.L0 +def test_runtime_manager_keep_alive_reuses_until_explicit_shutdown(runtime_module): + world = runtime_module.RuntimeWorld( + rank=0, + size=2, + group=object(), + global_ranks=(0, 1), + ) + provider = _FakeRuntimeProvider(runtime_module) + manager = runtime_module.RuntimeManager( + provider_factory=lambda: provider, + world_resolver=lambda config: world, + keep_alive=True, + ) + + first = manager.acquire(object(), torch.device("cuda", 0)) + first.close() + assert manager.ref_count == 0 + assert provider.initialize_count == 1 + assert provider.finalize_count == 0 + + second = manager.acquire(object(), torch.device("cuda", 0)) + assert manager.ref_count == 1 + assert provider.initialize_count == 1 + second.close() + + manager.shutdown() + assert manager.ref_count == 0 + assert provider.finalize_count == 1 + + +@pytest.mark.L0 +def test_runtime_manager_rejects_different_same_geometry_subgroup( + runtime_module, +): + first_world = runtime_module.RuntimeWorld( + rank=0, + size=2, + group=object(), + global_ranks=(0, 2), + ) + second_world = runtime_module.RuntimeWorld( + rank=0, + size=2, + group=object(), + global_ranks=(0, 3), + ) + provider = _FakeRuntimeProvider(runtime_module) + first_manager = runtime_module.RuntimeManager( + provider_factory=lambda: provider, + world_resolver=lambda config: first_world, + ) + second_manager = runtime_module.RuntimeManager( + provider_factory=lambda: provider, + world_resolver=lambda config: second_world, + ) + + handle = first_manager.acquire(object(), torch.device("cuda", 0)) + with pytest.raises(RuntimeError, match="different EP subgroup"): + second_manager.acquire(object(), torch.device("cuda", 0)) + handle.close() + + +@pytest.mark.L0 +def test_runtime_manager_rejects_unverifiable_external_subgroup( + runtime_module, + monkeypatch, +): + world = runtime_module.RuntimeWorld( + rank=0, + size=2, + group=object(), + global_ranks=(1, 3), + ) + provider = _FakeRuntimeProvider( + runtime_module, + runtime_module.RuntimeInitState.INITIALIZED, + ) + provider._world = world + manager = runtime_module.RuntimeManager( + provider_factory=lambda: provider, + world_resolver=lambda config: world, + ) + monkeypatch.setattr( + runtime_module, + "_spans_default_distributed_world", + lambda selected: False, + ) + + with pytest.raises(RuntimeError, match="cannot safely attach"): + manager.acquire(object(), torch.device("cuda", 0)) + + +@pytest.mark.L0 +def test_nvshmem_uid_broadcast_uses_subgroup_root_global_rank( + runtime_module, + monkeypatch, +): + class _FakeDevice: + def __init__(self, index): + self.index = index + + def set_current(self): + return None + + cuda_module = ModuleType("cuda") + cuda_core_module = ModuleType("cuda.core") + cuda_experimental_module = ModuleType("cuda.core.experimental") + cuda_experimental_module.Device = _FakeDevice + cuda_core_module.experimental = cuda_experimental_module + cuda_module.core = cuda_core_module + monkeypatch.setitem(sys.modules, "cuda", cuda_module) + monkeypatch.setitem(sys.modules, "cuda.core", cuda_core_module) + monkeypatch.setitem( + sys.modules, + "cuda.core.experimental", + cuda_experimental_module, + ) + + init_args = {} + + class _FakeUid: + def __init__(self): + self._data = np.arange(16, dtype=np.uint8) + + core = SimpleNamespace( + get_unique_id=lambda empty: _FakeUid(), + init=lambda **kwargs: init_args.update(kwargs), + ) + monkeypatch.setattr(runtime_module, "_load_nvshmem_core", lambda: core) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + + group = object() + broadcast_args = {} + monkeypatch.setattr(dist, "get_backend", lambda selected: "gloo") + monkeypatch.setattr( + dist, + "get_global_rank", + lambda selected, group_rank: (1, 3)[group_rank], + ) + + def _broadcast(tensor, *, src, group): + broadcast_args.update(tensor=tensor, src=src, group=group) + + monkeypatch.setattr(dist, "broadcast", _broadcast) + monkeypatch.setattr(dist, "barrier", lambda *, group: None) + + world = runtime_module.RuntimeWorld( + rank=0, + size=2, + group=group, + global_ranks=(1, 3), + ) + runtime_module._DefaultNvshmemRuntimeProvider().initialize( + torch.device("cuda", 0), + world, + ) + + assert broadcast_args["src"] == 1 + assert broadcast_args["group"] is group + assert broadcast_args["tensor"].device.type == "cpu" + assert init_args["rank"] == 0 + assert init_args["nranks"] == 2 diff --git a/test/python/moe_ep/test_moe_ep_multinode.py b/test/python/moe_ep/test_moe_ep_multinode.py new file mode 100644 index 000000000..568537af1 --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_multinode.py @@ -0,0 +1,331 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Torchrun-native multi-node MoE EP forward/backward acceptance tests.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist + +from moe_ep.moe_ep_distributed_workers import ( + _run_backward_reference_case, + _run_forward_output_case, +) +from moe_ep.moe_ep_test_support import make_distributed_forward_inputs + +pytestmark = [ + pytest.mark.L1, + pytest.mark.gpu_exclusive, + pytest.mark.moe_ep_multinode, +] + +_TORCHRUN_ENV = ("LOCAL_RANK", "LOCAL_WORLD_SIZE", "RANK", "WORLD_SIZE") +_PROCESS_GROUP_TIMEOUT = timedelta(minutes=10) + + +def _bind_torchrun_device_before_pytest_fixtures() -> None: + """Bind before the root conftest creates its session CUDA handle.""" + + value = os.environ.get("LOCAL_RANK") + if value is None or not torch.cuda.is_available(): + return + local_rank = int(value) + if 0 <= local_rank < torch.cuda.device_count(): + torch.cuda.set_device(local_rank) + + +_bind_torchrun_device_before_pytest_fixtures() + + +@dataclass(frozen=True) +class _TorchrunWorld: + rank: int + world_size: int + local_rank: int + local_world_size: int + device: torch.device + + +def _require_torchrun_environment() -> tuple[int, int, int, int]: + missing = [name for name in _TORCHRUN_ENV if name not in os.environ] + if missing: + pytest.skip("multi-node MoE EP tests require torchrun environment variables: " + ", ".join(missing)) + return ( + int(os.environ["RANK"]), + int(os.environ["WORLD_SIZE"]), + int(os.environ["LOCAL_RANK"]), + int(os.environ["LOCAL_WORLD_SIZE"]), + ) + + +@pytest.fixture(scope="session") +def torchrun_world(): + if not dist.is_available() or not dist.is_nccl_available(): + pytest.skip("multi-node Rubin MXFP8 tests require NCCL") + + rank, world_size, local_rank, local_world_size = _require_torchrun_environment() + if local_rank < 0 or local_rank >= torch.cuda.device_count(): + pytest.skip(f"torchrun LOCAL_RANK={local_rank} is not backed by a visible GPU") + + device = torch.device("cuda", local_rank) + if torch.cuda.get_device_capability(device) != (10, 7): + pytest.skip("multi-node Rubin MXFP8 tests require exactly SM107 " "(compute capability 10.7) on every rank") + try: + import nvshmem.core # noqa: F401 + except (ImportError, OSError): + pytest.skip("multi-node Rubin MXFP8 tests require NVSHMEM") + + os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") + torch.cuda.set_device(device) + if dist.is_initialized(): + if dist.get_rank() != rank or dist.get_world_size() != world_size: + raise RuntimeError("existing process group does not match torchrun RANK/WORLD_SIZE") + else: + dist.init_process_group( + backend="nccl", + init_method="env://", + device_id=device, + timeout=_PROCESS_GROUP_TIMEOUT, + ) + + context = _TorchrunWorld( + rank=rank, + world_size=world_size, + local_rank=local_rank, + local_world_size=local_world_size, + device=device, + ) + try: + yield context + finally: + if dist.is_initialized(): + try: + dist.barrier() + from cudnn.moe_ep._megamoe_backend._runtime import ( + get_runtime_manager, + ) + + get_runtime_manager().shutdown() + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + ( + "ep_size", + "required_world_size", + "required_local_world_size", + "ep_global_ranks", + ), + [ + pytest.param( + 4, + 8, + 4, + (0, 1, 4, 5), + id="forward-ep4-world8", + ), + pytest.param( + 6, + 8, + 4, + (0, 1, 2, 4, 5, 6), + id="forward-ep6-world8", + ), + pytest.param( + 12, + 12, + 4, + tuple(range(12)), + id="forward-ep12-world12", + ), + pytest.param( + 16, + 16, + 4, + tuple(range(16)), + id="forward-ep16-world16", + ), + ], +) +@pytest.mark.parametrize("combine_format", ["bf16", "mxfp8"]) +def test_mxfp8_forward_multinode_matches_reference( + torchrun_world, + ep_size, + required_world_size, + required_local_world_size, + ep_global_ranks, + combine_format, +): + world = torchrun_world + if world.world_size != required_world_size or world.local_world_size != required_local_world_size: + pytest.skip( + f"EP{ep_size} requires torchrun WORLD_SIZE={required_world_size}, " + f"LOCAL_WORLD_SIZE={required_local_world_size}; got " + f"WORLD_SIZE={world.world_size}, " + f"LOCAL_WORLD_SIZE={world.local_world_size}" + ) + + if ep_size == world.world_size: + ep_group = dist.group.WORLD + else: + # All WORLD ranks must create subgroups in the same order, including + # idle ranks that are not members of this balanced EP group. + ep_group = dist.new_group( + list(ep_global_ranks), + backend="nccl", + timeout=_PROCESS_GROUP_TIMEOUT, + ) + + is_ep_member = world.rank in ep_global_ranks + try: + if is_ep_member: + ep_rank = dist.get_rank(ep_group) + _run_forward_output_case( + device=world.device, + ep_group=ep_group, + ep_rank=ep_rank, + ep_size=ep_size, + combine_format=combine_format, + expected_global_ranks=ep_global_ranks, + ) + dist.barrier() + finally: + if ep_group is not dist.group.WORLD and is_ep_member: + # The runtime manager intentionally remains alive after op.close(). + # Finalize it while this exact subgroup is still valid so the next + # parametrized case may create a fresh ProcessGroup object with the + # same membership. + from cudnn.moe_ep._megamoe_backend._runtime import ( + get_runtime_manager, + ) + + get_runtime_manager().shutdown() + dist.barrier(group=ep_group) + dist.destroy_process_group(ep_group) + dist.barrier() + + +@pytest.mark.parametrize( + ("ep_size", "required_world_size", "combine_format"), + [ + pytest.param( + 8, + 8, + "bf16", + id="backward-ep8-world8-bf16", + ), + pytest.param( + 8, + 8, + "mxfp8", + id="backward-ep8-world8-mxfp8", + ), + pytest.param( + 16, + 16, + "bf16", + id="backward-ep16-world16-bf16", + ), + pytest.param( + 32, + 32, + "bf16", + id="backward-ep32-world32-bf16-minimal", + ), + pytest.param( + 32, + 32, + "mxfp8", + id="backward-ep32-world32-mxfp8", + ), + ], +) +def test_stateless_training_multinode_matches_independent_reference( + torchrun_world, + ep_size, + required_world_size, + combine_format, +): + world = torchrun_world + if world.world_size != required_world_size: + pytest.skip(f"EP{ep_size} requires torchrun WORLD_SIZE={required_world_size}; " f"got WORLD_SIZE={world.world_size}") + + _run_backward_reference_case( + device=world.device, + ep_group=dist.group.WORLD, + ep_rank=world.rank, + ep_size=ep_size, + combine_format=combine_format, + ) + + +@pytest.mark.parametrize( + ("rank_zero_lane_count", "other_lane_count"), + [ + pytest.param( + 2, + 1, + id="backward-ep8-world8-abi-mismatch", + ), + ], +) +def test_training_prepare_multinode_rejects_rank_abi_mismatch( + torchrun_world, + rank_zero_lane_count, + other_lane_count, +): + world = torchrun_world + if world.world_size != 8 or world.local_world_size != 4: + pytest.skip( + "EP8 ABI mismatch requires torchrun WORLD_SIZE=8, " + "LOCAL_WORLD_SIZE=4; got " + f"WORLD_SIZE={world.world_size}, " + f"LOCAL_WORLD_SIZE={world.local_world_size}" + ) + + from cudnn import MoeEp + + op = MoeEp( + num_experts=16, + hidden_size=128, + intermediate_size=256, + top_k=2, + ep_group=dist.group.WORLD, + max_tokens_per_rank=8, + max_recv_size_per_rank=3, + drop_on_overflow=True, + combine_format="bf16", + weight_interleave_size=32, + ) + caught_error = None + try: + lane_count = rank_zero_lane_count if world.rank == 0 else other_lane_count + try: + op.prepare_training( + lane_count=lane_count, + device=world.device, + ) + except Exception as error: + caught_error = error + + # Reachable only after the collective prepare path returns or raises on + # every rank. The NCCL barrier blocks the host and launches no MoE + # kernel or device-side assertion. + dist.barrier( + group=dist.group.WORLD, + device_ids=[world.local_rank], + ) + + assert isinstance(caught_error, RuntimeError), f"rank {world.rank} expected RuntimeError from collective prepare, " f"got {caught_error!r}" + message = str(caught_error) + assert "symmetric workspace region counts differ" in message or "ABI differs" in message, f"rank {world.rank} got unexpected prepare error: {message}" + finally: + op.close() diff --git a/test/python/pytest.ini b/test/python/pytest.ini index 412f6d38a..7ab68cea5 100644 --- a/test/python/pytest.ini +++ b/test/python/pytest.ini @@ -6,5 +6,6 @@ markers = L3: specifies L3 level (use -m L3) L4: specifies L4 level (use -m L4) gpu_exclusive: tests that require exclusive GPU access (no concurrent kernels from other processes) + moe_ep_multinode: torchrun-native multi-node MoE EP tests addopts = -m L0 --tb=short --no-header