MoE Operators¶

本页暂无中文版。以下为英文原文。

Every op on this page is used the same way: construct it once, then call it. The constructor takes what the kernel is compiled with; the call takes the tensors. Both are documented under each op — __init__ and forward, where forward is what runs when you call op(...).

A routed mixture-of-experts layer is available two ways here. FusedMoEFwdOp runs the whole FFN, and FusedMoESharedExpertFwdOp adds a shared expert beside the routed ones. The rest are its stages, callable on their own: the op that picks each token's experts, the ops that move tokens into an expert-contiguous layout and back, and the expert GEMMs that run on it. MoEGroupedGemmFwdOp is one grouped GEMM; MoEExpertMLPFwdOp is the pair of them with the gated activation fused into the first; FusedMoEExpertsFwdOp is that MLP with the permutes around it, on the tight (no-pad) layout, and IndexedExpertMLPFwdOp is the backend it picks instead when the routes are few enough to read the weights once per route rather than once per expert. SharedExpertMLPFwdOp is the dense gated MLP of the shared expert. The routing has to produce the layout the GEMM expects.

Fused forward¶

tileops.moe.FusedMoEFwdOp ¶

Routed Mixture-of-Experts (MoE) FFN.

Covers Qwen3 (softmax) style configurations where top-k comes straight from the gating scores, and Kimi K2 / DeepSeek-V3 style ones where a per-expert correction_bias is passed: top-k is then selected from sigmoid(score) + correction_bias while the final weights use the original (unbiased) scores, renormalized.

num_expert_group > 1 routes as DeepSeek-V3 does: the experts split into num_expert_group groups of consecutive experts, each scored by the sum of its two best biased scores when correction_bias is passed and by its best score otherwise, and top-k is selected inside the topk_group best groups. A group within 1e-5 (relative; also 1e-6 absolute) of the last kept group score may be kept in its place, as in FusedTopKFwdOp. The output is multiplied by routed_scaling_factor once.

gpt-oss is scoring_func="softmax", renormalize=True, activation="swigluoai" and both expert biases. activation, and the row order of w_gate_up and b_gate_up it reads, are as MoEGroupedGemmFwdOp defines them; the optional expert biases b_gate_up and b_down join each route as FusedMoEExpertsFwdOp states.

init ¶

__init__(
    top_k,
    scoring_func="softmax",
    renormalize=False,
    routed_scaling_factor=1.0,
    num_expert_group=1,
    topk_group=1,
    *,
    activation="silu_and_mul",
    target=None,
    prepare_finalize=None,
    experts=None
)

Build the op. Shapes and dtype are taken from each call.

Parameters:

top_k (int) –

K -- experts selected per token.
scoring_func (str, default: 'softmax' ) –

"softmax" (Qwen3) or "sigmoid" (Kimi K2 / DeepSeek-V3).
renormalize (bool, default: False ) –

Renormalize top-k weights to sum to 1.
routed_scaling_factor (float, default: 1.0 ) –

Multiplier on expert output (Kimi K2: 2.827, DeepSeek-V3: 2.5).
num_expert_group (int, default: 1 ) –

Groups of consecutive experts the routing is restricted by; must divide E. 1 selects over all experts. DeepSeek-V3: 8.
topk_group (int, default: 1 ) –

Groups kept per token, holding at least top_k experts. DeepSeek-V3: 4.
activation (str, default: 'silu_and_mul' ) –

Gated activation applied to gate_up: "silu_and_mul" or "gelu_and_mul" over a stacked w_gate_up, "swigluoai" over an interleaved one.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device. The sub-ops it builds are given the same one.
prepare_finalize (Optional[FusedMoEPrepareAndFinalize], default: None ) –

Override the PrepareAndFinalize implementation.
experts (Optional[FusedMoEExpertsModular], default: None ) –

Override the Experts implementation.

forward ¶

forward(
    hidden_states,
    gating_output,
    w_gate_up,
    w_down,
    correction_bias=None,
    b_gate_up=None,
    b_down=None,
)

Route each token to its top-k experts and return the [T, H] weighted sum.

Parameters:

hidden_states (Tensor) –

[T, H] tokens.
gating_output (Tensor) –

[T, E] float32 router logits.
w_gate_up (Tensor) –

[E, 2F, H] gate and up weights, stacked or interleaved as activation reads them.
w_down (Tensor) –

[E, H, F] down weights.
correction_bias (Optional[Tensor], default: None ) –

Optional [E] float32 bias added to the sigmoid scores for the selection only.
b_gate_up (Optional[Tensor], default: None ) –

Optional [E, 2F] gate/up bias in w_gate_up's row order.
b_down (Optional[Tensor], default: None ) –

Optional [E, H] down bias.

tileops.moe.FusedMoESharedExpertFwdOp ¶

FusedMoE with shared expert support, optionally TP-aware.

Extends FusedMoE to compute both shared and routed expert outputs. Passing the shared expert's weights enables it; the shared expert is the SharedExpertMLPFwdOp sub-op, which applies silu_and_mul.

The routed half routes as FusedMoEFwdOp: with num_expert_group > 1 the top-k is selected inside the topk_group best groups of consecutive experts, each scored by the sum of its two best biased scores when correction_bias is passed and by its best score otherwise; a group within 1e-5 (relative; also 1e-6 absolute) of the last kept group score may be kept in its place. routed_scaling_factor scales the routed output only, once.

TP support (shared expert only): When tp_size > 1, the op shards the shared expert weights internally: - shared_w_gate_up [2*F_s, H] is split along dim=0 (ColumnParallel) - shared_w_down [H, F_s] is split along dim=1 (RowParallel) The returned shared_out is a partial sum; the caller must all-reduce across TP ranks. The routed expert path is not affected.

Returns:

(shared_output, routed_output) –

tuple of [T, H] tensors. shared_output is None when the shared weights are not passed. shared_output is a partial sum when tp_size > 1.

init ¶

__init__(
    top_k,
    scoring_func="softmax",
    renormalize=False,
    routed_scaling_factor=1.0,
    num_expert_group=1,
    topk_group=1,
    tp_size=1,
    tp_rank=0,
    *,
    activation="silu_and_mul",
    target=None,
    prepare_finalize=None,
    experts=None
)

Build the op. Shapes and dtype are taken from each call.

Parameters:

top_k (int) –

K -- experts selected per token.
scoring_func (str, default: 'softmax' ) –

"softmax" (Qwen3) or "sigmoid" (Kimi K2 / DeepSeek-V3).
renormalize (bool, default: False ) –

Renormalize top-k weights to sum to 1.
routed_scaling_factor (float, default: 1.0 ) –

Multiplier on the routed expert output.
num_expert_group (int, default: 1 ) –

Groups of consecutive experts the routing is restricted by; must divide E. 1 selects over all experts. DeepSeek-V3: 8.
topk_group (int, default: 1 ) –

Groups kept per token, holding at least top_k experts. DeepSeek-V3: 4.
tp_size (int, default: 1 ) –

Tensor parallel world size. Default 1 (no TP).
tp_rank (int, default: 0 ) –

This rank's index in the TP group. Default 0.
activation (str, default: 'silu_and_mul' ) –

Gated activation of the routed experts; the shared expert supports silu_and_mul only.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.
prepare_finalize (Optional[FusedMoEPrepareAndFinalize], default: None ) –

Override the PrepareAndFinalize implementation.
experts (Optional[FusedMoEExpertsModular], default: None ) –

Override the Experts implementation.

forward ¶

forward(
    hidden_states,
    gating_output,
    w_gate_up,
    w_down,
    correction_bias=None,
    shared_w_gate_up=None,
    shared_w_down=None,
)

Run shared + routed MoE FFN.

Parameters:

hidden_states (Tensor) –

[T, H] input hidden states.
gating_output (Tensor) –

[T, E] gating logits.
w_gate_up (Tensor) –

[E, 2F, H] routed expert gate+up weights.
w_down (Tensor) –

[E, H, F] routed expert down weights.
correction_bias (Optional[Tensor], default: None ) –

Optional [E] bias for Kimi K2 / DeepSeek-V3 routing.
shared_w_gate_up (Optional[Tensor], default: None ) –

[2*F_s, H] shared expert gate+up weights (full); passing it enables the shared expert. When tp_size > 1, sharded along dim=0.
shared_w_down (Optional[Tensor], default: None ) –

[H, F_s] shared expert down weight (full), passed with shared_w_gate_up. When tp_size > 1, sharded along dim=1.

Returns:

(shared_output, routed_output) –

tuple of [T, H] tensors. shared_output is None when the shared weights are not passed, and a partial sum when tp_size > 1; the caller all-reduces across TP ranks.

Routing and layout¶

tileops.moe.FusedTopKFwdOp ¶

Mixture-of-Experts (MoE) top-k routing operator.

Applies scoring (softmax or sigmoid) to router logits and selects the top-k experts per token. Scoring and selection run in float32 whatever the logits dtype, with fast-math transcendentals, so a score can differ from torch's float32 score by a few ulp: an expert, or a group, whose selection score is within 1e-5 (relative; a group score also 1e-6 absolute) of the last one kept may be kept in its place; a selection that is not that close matches the float32 reference, ties to the lower index. The weights are within 1e-4 (relative) of the float32 scores of the experts returned.

With num_expert_group > 1 it is DeepSeek grouped routing, as vLLM's grouped_topk: the E experts split into num_expert_group groups of E / num_expert_group consecutive experts, each group is scored by the sum of its two best biased scores when correction_bias is passed and by its best score otherwise, the topk_group best groups are kept (ties to the lower group index), and the top-k experts are selected inside them. The returned weights are the unbiased scores of the selected experts, renormalized when renormalize, times routed_scaling_factor. num_expert_group=1 selects over all experts.

Example

op = FusedTopKFwdOp(top_k=8)
topk_weights, topk_ids = op(gating_output)
# topk_weights: [512, 8] float32
# topk_ids:     [512, 8] int32

init ¶

__init__(
    top_k,
    scoring_func="softmax",
    renormalize=False,
    num_expert_group=1,
    topk_group=1,
    routed_scaling_factor=1.0,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from each call.

Parameters:

top_k (int) –

Number of experts to select per token K.
scoring_func (str, default: 'softmax' ) –

"softmax" (Qwen3/Qwen2) or "sigmoid" (DeepSeek-V3/GLM-4/Kimi K2). Passing correction_bias to forward requires "sigmoid": the bias is added to sigmoid scores for selection only, and the output weights stay the original scores.
renormalize (bool, default: False ) –

If True, normalize top-k weights to sum to 1.
num_expert_group (int, default: 1 ) –

Groups of consecutive experts the selection is restricted by; must divide E. 1 (the default) selects over all experts. DeepSeek-V3: 8.
topk_group (int, default: 1 ) –

Groups kept per token, at most num_expert_group; they must hold at least top_k experts. DeepSeek-V3: 4.
routed_scaling_factor (float, default: 1.0 ) –

Positive multiplier on the returned weights, applied after renormalization. DeepSeek-V3: 2.5.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    gating_output,
    correction_bias=None,
)

Run top-k routing.

Parameters:

gating_output (Tensor) –

[T, E] router logits (bf16, fp16, or float32).
correction_bias (Optional[Tensor], default: None ) –

[E] float32 per-expert bias, or None. Passing it requires scoring_func="sigmoid".

Returns:

topk_weights ( Tensor ) –

[T, K] float32, times routed_scaling_factor.
topk_ids ( Tensor ) –

[T, K] int32.

tileops.moe.MoEPrePermuteFwdOp ¶

Materialize rank-grouped activations into a local expert layout.

local_expert_ids must already be in [0, num_local_experts). Global placement and communication belong to EPDispatch.

init ¶

__init__(
    layout,
    num_local_experts,
    *,
    target=None
)

Configure a pre-permute boundary for one layout and expert domain.

Parameters:

layout (MGroupedLayoutSpec) –

How the expert rows are materialized.
num_local_experts (int) –

Number of local experts the ids index.
target (object, default: None ) –

Which backend serves this instance; detected from the tensors when None.

forward ¶

forward(
    hidden_states,
    local_expert_ids,
)

Return (expert_input, layout_metadata, inverse_indices).

tileops.moe.MoEPermuteAlignFwdOp ¶

Route tokens to experts and pad each expert's token count to block_size.

Takes topk_ids and produces the three index arrays required by MoE grouped GEMM: sorted token indices, per-block expert ids, and the total padded token count.

With numel = T * K routes and count[e] of them on expert e, the result is:

num_tokens_post_pad: the sum over experts of count[e] rounded up to block_size.
sorted_token_ids[:num_tokens_post_pad]: one range per expert, in expert order, of that rounded-up length. The first count[e] slots of the range hold the indices of the routes on e, each once, in an order that is unspecified and may differ between calls on the same input; the rest of the range holds numel.
expert_ids[:num_tokens_post_pad // block_size]: the owner of each block_size block of those ranges.
Past those lengths, every entry of sorted_token_ids is numel and every entry of expert_ids lies outside [0, num_experts).

Example

op = MoEPermuteAlignFwdOp(num_experts=8, block_size=16)
sorted_ids, expert_ids, num_post_pad = op(topk_ids)

init ¶

__init__(
    num_experts,
    block_size=64,
    *,
    target=None
)

Build the op. The routed extents are taken from each call.

Parameters:

num_experts (int) –

Number of experts.
block_size (int, default: 64 ) –

GEMM tile size (M dimension); default 64.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    topk_ids,
)

Run permute-align.

Parameters:

topk_ids (Tensor) –

[T, K] int32 expert indices (0-indexed).

Returns:

sorted_token_ids ( Tensor ) –

[T * K + (num_experts + 1) * (block_size - 1)] int32
expert_ids ( Tensor ) –

[ceil_div(that, block_size)] int32
num_tokens_post_pad ( Tensor ) –

[1] int32

tileops.moe.MoEPostPermuteFwdOp ¶

Restore token order and apply the declared local routing epilogue.

A caller-provided out must not overlap expert_output, topk_weights or inverse_indices. Overlap raises ValueError before any write; disjoint views of one buffer are accepted.

init ¶

__init__(
    layout,
    epilogue=None,
    out_dtype=None,
    *,
    target=None
)

Configure inverse permutation and the exactly-once routing epilogue.

Parameters:

layout (MGroupedLayoutSpec) –

How the expert rows are materialized.
epilogue (RoutingEpilogueSpec | None, default: None ) –

The routing epilogue; None takes the default one.
out_dtype (dtype | None, default: None ) –

The dtype the reduced result is written in; None keeps the expert output's own.
target (object, default: None ) –

Which backend serves this instance; detected from the tensors when None.

forward ¶

forward(
    expert_output,
    topk_weights,
    inverse_indices,
    out=None,
)

Restore token order, apply routing weights, reduce top-k, and cast.

Expert GEMMs¶

tileops.moe.MoEGroupedGemmFwdOp ¶

M-grouped GEMM over expert-materialized rows: out[rows of g] = a[rows of g] @ b[g]^T.

layout fixes how the rows of a are organised by expert and what the layout_metadata tensor means; E is read off b, M/N/K off the operands. Accumulation is fp32; the output is written in the operand dtype unless out_dtype asks for fp32.

With activation set, b holds the gate and up projections along N ([E, 2 * ffn, K]) and out has ffn columns: the gated activation is fused into the GEMM's epilogue and the [.., 2 * ffn] intermediate is never written. The gated activations, which every MoE op's activation names, are defined here. silu_and_mul and gelu_and_mul read b stacked, the ffn gate rows then the ffn up rows, and keep act(gate) * up. swigluoai (gpt-oss) reads it interleaved, row 2j gate j and row 2j + 1 up j, and keeps g * sigmoid(1.702 * g) * (u + 1) with g = min(gate, 7) and u = clamp(up, -7, 7); the constants are fixed.

An optional bias ([E, N], or [E, 2 * ffn] in b's row order when an activation is fused) is added to each expert's fp32 product before the activation and the cast to the output dtype.

Per layout the operands are (trans-free, b is [E, N, K]):

layout	`a`	`layout_metadata`	`out`
contiguous · physical_psum	`[M, K]`	`[E]` segment end rows	`[M, N]`
contiguous · per_row	`[M, K]`	`[M]` expert id per row	`[M, N]`
masked	`[E, max_m, K]`	`[E]` valid rows per expert	`[E, max_m, N]`

Rows an aligned layout pads with, and a masked slab's rows past its valid count, hold unspecified values in out. The metadata's values must satisfy the layout (ordering, segment ends, ranges); the op does not check them, since doing so would synchronise.

init ¶

__init__(
    layout,
    *,
    activation=None,
    out_dtype=None,
    target=None
)

Fix the expert layout, the fused activation and the output dtype policy.

Parameters:

layout (MGroupedLayoutSpec) –

How a's rows are grouped by expert.
activation (str | None, default: None ) –

None for a plain GEMM, or a gated activation fused into the epilogue, which halves the output width: "silu_and_mul" or "gelu_and_mul" over a stacked gate||up b, "swigluoai" over an interleaved one.
out_dtype (dtype | None, default: None ) –

None writes the operand dtype, torch.float32 keeps the fp32 accumulator.
target (object, default: None ) –

Which backend serves this instance; detected from the tensors when None.

forward ¶

forward(
    a,
    b,
    layout_metadata,
    out=None,
    *,
    bias=None
)

Run one GEMM per expert over the rows the layout assigns it.

Parameters:

a (Tensor) –

Expert-materialized activations, [M, K] or [E, max_m, K] per the layout.
b (Tensor) –

Per-expert weights, [E, N, K].
layout_metadata (Tensor) –

int32 metadata whose shape and meaning the layout fixes.
out (Tensor | None, default: None ) –

Optional preallocated output in the resolved output dtype. It must not overlap a, b, layout_metadata or bias, or a ValueError is raised; disjoint views of one buffer are accepted.
bias (Tensor | None, default: None ) –

Optional per-expert bias in b's dtype and row order, [E, N]; keyword-only.

Returns:

Tensor –

[M, N] or [E, max_m, N] in the operand dtype, or fp32 when asked for.

tileops.moe.MoEGroupedGemmFP8FwdOp ¶

Block-scaled FP8 GEMM per expert over a masked slab: out[g, :m_g] = a[g, :m_g] @ b[g]^T.

a is [E, max_m, K] and b [E, N, K], both float8_e4m3fn. a_scale holds one float32 scale per row and 128 columns of a ([E, max_m, K / 128]), b_scale one per 128x128 block of b ([E, N / 128, K / 128]); K and N are multiples of 128. layout_metadata[g] is expert g's valid row count m_g.

Each 128-wide K block is contracted from the raw FP8 values, then multiplied by its two scales and added to a float32 accumulator, so the scales are applied once per 128 K, as deep_gemm.m_grouped_fp8_gemm_nt_masked does. The result is written in bfloat16. Rows at or past layout_metadata[g] hold undefined values. The counts are not checked, since checking would synchronise: a count past max_m is read as max_m and a negative one as 0, and nothing outside out is written.

Accuracy: on the valid rows the output matches an fp32 reference that dequantizes both operands under their scales to atol = 2e-2 * sqrt(max(1, K / 1024)) and rtol = 2e-2 (workloads.gemm.fp8_matmul_verification), a bound that grows with K as the fp8 accumulation noise does. The scales are applied in DeepGEMM's order; on the manifest's E = 8 rows the output is bit-identical to deep_gemm.m_grouped_fp8_gemm_nt_masked.

init ¶

__init__(
    layout,
    *,
    target=None
)

Fix the masked expert layout.

Parameters:

layout (MaskedLayoutSpec) –

The masked slab capacity max_m each expert's rows are padded to.
target (Target, default: None ) –

Which backend serves this instance; detected from the tensors when None.

forward ¶

forward(
    a,
    a_scale,
    b,
    b_scale,
    layout_metadata,
    out=None,
)

Run one block-scaled GEMM per expert over its valid rows.

Parameters:

a (Tensor) –

float8_e4m3fn activations, [E, max_m, K].
a_scale (Tensor) –

float32 per-row scales of a, [E, max_m, K / 128].
b (Tensor) –

float8_e4m3fn per-expert weights, [E, N, K].
b_scale (Tensor) –

float32 per-block scales of b, [E, N / 128, K / 128].
layout_metadata (Tensor) –

int32 valid row count of each expert, [E].
out (Optional[Tensor], default: None ) –

Optional preallocated contiguous bfloat16 output, [E, max_m, N], written in place and returned; None allocates one. Must not overlap any input, or a ValueError is raised before writing.

Returns:

Tensor –

[E, max_m, N] in bfloat16; rows at or past layout_metadata[g] are
Tensor –

undefined.

tileops.moe.MoEExpertMLPFwdOp ¶

Two grouped GEMMs on one expert layout, the gated activation fused into the first.

out = act(expert_input @ w_gate_up[g]^T + b_gate_up[g]) @ w_down[g]^T + b_down[g] per expert g, where act is the gated activation: w_gate_up holds the gate and up projections along N and the gate_up GEMM's epilogue applies the activation and halves it, so the [.., 2 * ffn] intermediate is never written. A composite: it registers no operator of its own, its graph is its two leaves'.

activation, and the order of w_gate_up's gate and up rows it reads, are as MoEGroupedGemmFwdOp defines them.

The biases are optional and independent of each other and of the activation: b_gate_up ([E, 2 * ffn], in w_gate_up's row order) is added to the fp32 gate/up product before the activation, b_down ([E, hidden]) to the fp32 down product before its cast.

init ¶

__init__(
    layout,
    activation="silu_and_mul",
    *,
    target=None
)

Configure two grouped GEMMs on layout, the first fusing the gated activation.

Parameters:

layout (MGroupedLayoutSpec) –

Shared by both GEMMs and the metadata.
activation (str, default: 'silu_and_mul' ) –

"silu_and_mul", "gelu_and_mul" (stacked w_gate_up) or "swigluoai" (interleaved w_gate_up), fused into the gate_up GEMM's epilogue.
target (object, default: None ) –

Which backend serves the delegates.

forward ¶

forward(
    expert_input,
    w_gate_up,
    w_down,
    layout_metadata,
    out=None,
    *,
    b_gate_up=None,
    b_down=None
)

Run the gate/up GEMM with its fused activation, then the down GEMM.

Parameters:

expert_input (Tensor) –

[M, hidden] or [E, max_m, hidden] per the layout.
w_gate_up (Tensor) –

[E, 2 * ffn, hidden], gate and up along the output axis, stacked or interleaved as activation reads them.
w_down (Tensor) –

[E, hidden, ffn].
layout_metadata (Tensor) –

int32 metadata the layout fixes; shared by both GEMMs.
out (Tensor | None, default: None ) –

Optional preallocated [.., hidden] output in the operand dtype. It must not overlap w_down, layout_metadata or b_down, which the down GEMM reads while it writes out, or a ValueError is raised; disjoint views of one buffer are accepted.
b_gate_up (Tensor | None, default: None ) –

Optional [E, 2 * ffn] bias in w_gate_up's row order; keyword-only.
b_down (Tensor | None, default: None ) –

Optional [E, hidden] bias; keyword-only.

Returns:

Tensor –

[M, hidden] or [E, max_m, hidden] in the operand dtype.

tileops.moe.FusedMoEExpertsFwdOp ¶

Expert MLP with indexed small-route and tight grouped backends.

output[t] = scale * sum_k w[t, k] * (down_k(act(gate_up_k(h[t]) + b_gate_up[k])) + b_down[k]) over each token's routed experts k, where scale is routed_scaling_factor. Each bias joins its GEMM's result before the activation or the cast; the routing weight then scales the biased down output, as gpt-oss and vLLM's fused experts apply it. The biases are optional and independent of each other and of the activation.

act is the gated activation activation names, and w_gate_up and b_gate_up hold gate and up rows in the order it reads, both as MoEGroupedGemmFwdOp defines them.

forward() writes (T, H): reduction is done internally by the PostPermute/Unpermute stage, so make_weighted_reduce() returns WeightedReduceNoOp. A call with few routes per expert (T * K <= 2 * E) on silu_and_mul with no bias, H % 128 == 0 and F % 256 == 0 runs the indexed small-route op; every other call runs the staged pipeline. The two moduli are the widths the indexed kernels are tiled for, not a limit of what they compute.

init ¶

__init__(
    routed_scaling_factor=1.0,
    *,
    activation="silu_and_mul",
    target=None
)

Build the op. Shapes and dtype are taken from each call.

Parameters:

routed_scaling_factor (float, default: 1.0 ) –

Scalar applied to the final reduced output.
activation (str, default: 'silu_and_mul' ) –

Gated activation applied to gate_up: 'silu_and_mul' or 'gelu_and_mul' over a stacked w_gate_up, 'swigluoai' over an interleaved one.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    output,
    hidden_states,
    w_gate_up,
    w_down,
    topk_weights,
    topk_ids,
    b_gate_up=None,
    b_down=None,
)

Run the local expert pipeline, writing the reduced result into output.

Parameters:

output (Tensor) –

[T, H] buffer the reduced result is written into.
hidden_states (Tensor) –

[T, H] tokens.
w_gate_up (Tensor) –

[E, 2F, H] gate and up weights, stacked or interleaved as activation reads them.
w_down (Tensor) –

[E, H, F] down weights.
topk_weights (Tensor) –

[T, K] float32 routing weights.
topk_ids (Tensor) –

[T, K] int32 expert of each route.
b_gate_up (Tensor | None, default: None ) –

Optional [E, 2F] gate/up bias in w_gate_up's row order.
b_down (Tensor | None, default: None ) –

Optional [E, H] down bias.

tileops.moe.IndexedExpertMLPFwdOp ¶

Route-major expert MLP with device-side reuse dispatch.

Each of the T * K routes is a row of its expert's GEMM, dispatched per route rather than per expert segment. Routes that share an expert are grouped, and only the leader copies that expert's weights, so the minimum DRAM traffic the roofline prices is one read per distinct expert. That pays off while the routes are few; :class:FusedMoEExpertsFwdOp picks this op over the staged pipeline on the shapes where it does.

init ¶

__init__(
    routed_scaling_factor=1.0,
    *,
    target=None
)

Fix the scalar applied to the reduced output; the route extents come per call.

Parameters:

routed_scaling_factor (float, default: 1.0 ) –

Scalar applied to the final reduced output.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    output,
    hidden_states,
    w_gate_up,
    w_down,
    topk_weights,
    topk_ids,
)

Write the weighted and reduced expert result into output.

tileops.moe.SharedExpertMLPFwdOp ¶

Dense gated MLP: down(silu(hidden @ gate.T) * (hidden @ up.T)).

w_gate_up stacks the gate rows over the up rows. Both GEMMs accumulate in float32; the gated activation is rounded to the input dtype before the down projection.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtype are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    hidden_states,
    w_gate_up,
    w_down,
)

Apply the shared expert to every token.

Parameters:

hidden_states (Tensor) –

\([T \times H]\), float16 or bfloat16.
w_gate_up (Tensor) –

\([2S \times H]\), the gate rows then the up rows.
w_down (Tensor) –

\([H \times S]\).

Returns:

Tensor –

\([T \times H]\) in the dtype of hidden_states.

MoE Operators¶

Fused forward¶

tileops.moe.FusedMoEFwdOp ¶

__init__ ¶

forward ¶

tileops.moe.FusedMoESharedExpertFwdOp ¶

__init__ ¶

forward ¶

Routing and layout¶

tileops.moe.FusedTopKFwdOp ¶

__init__ ¶

forward ¶

tileops.moe.MoEPrePermuteFwdOp ¶

__init__ ¶

forward ¶

tileops.moe.MoEPermuteAlignFwdOp ¶

__init__ ¶

forward ¶

tileops.moe.MoEPostPermuteFwdOp ¶

__init__ ¶

forward ¶

Expert GEMMs¶

tileops.moe.MoEGroupedGemmFwdOp ¶

__init__ ¶

forward ¶

tileops.moe.MoEGroupedGemmFP8FwdOp ¶

__init__ ¶

forward ¶

tileops.moe.MoEExpertMLPFwdOp ¶

__init__ ¶

forward ¶

tileops.moe.FusedMoEExpertsFwdOp ¶

__init__ ¶

forward ¶

tileops.moe.IndexedExpertMLPFwdOp ¶

__init__ ¶

forward ¶

tileops.moe.SharedExpertMLPFwdOp ¶

__init__ ¶

forward ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶